diff --git a/docs/minimal-pipeline.md b/docs/minimal-pipeline.md index e58abbfa4e..a3e079312e 100644 --- a/docs/minimal-pipeline.md +++ b/docs/minimal-pipeline.md @@ -62,4 +62,9 @@ The script runs in Linux CI. Coverage and preprocessing share only source parsin `scripts/check_tasks.py` uses nine focal functions sharing one source file, a loopback model, and real Cargo commands. It measures generation and repair request peaks for N=1, N=2, and the default N=4, requires overlap for N>1, and checks that compiler invocations never overlap or observe changing source bytes. A paused compiler lets the N=1 generation queue fill: one candidate is being validated, one result is buffered, and one active job waits to send. Releasing compilation must drain every result. -The same check covers direct and input/prefix/oracle integration generation, removal of invalid imports from saved candidates, and preservation/removal of temporary compiler inputs. Integration coverage reports remain available until every function has been read, and all nine functions must have passing tests and nonzero covered lines. Missing usage in one fixture response exercises worker-panic propagation and slot release; a stale generation backup exercises queue draining after consumer failure, and blocking the diagnostic output file exercises failure while the repair lock is held. Sources must be restored, failed repair backups retained, unrelated backups untouched, existing recovery material rejected, and post-failure statistics skipped. These checks use byte comparisons without hashing and run on macOS/Linux. CLI tests separately verify the default and rejection of zero before configuration loading. +The same check covers direct and input/prefix/oracle integration generation, removal of invalid imports from saved candidates, and preservation/removal of temporary compiler inputs. Integration coverage reports remain available until every function has been read, and all nine functions must have passing tests and nonzero covered lines. Valid answers without usage must complete successfully and mark request statistics incomplete. HTTP 401 during generation, oracle generation, and repair must fail without outer-layer retries. A blocked prompt directory exercises worker-panic propagation and slot release; a stale generation backup exercises queue draining after consumer failure, and blocking the diagnostic output file exercises failure while the repair lock is held. Sources must be restored, failed repair backups retained, unrelated backups untouched, existing recovery material rejected, and post-failure statistics skipped. These checks use byte comparisons without hashing and run on macOS/Linux. CLI tests separately verify the default and rejection of zero before configuration loading. + + +## Model Request Regression + +Local HTTP tests exercise the shared generation/repair request path, including response usage omissions, empty choices or content, invalid JSON, permanent 4xx errors, transient 429/5xx errors, dropped connections, and delayed response bodies. They verify actual attempt counts and bounded retry exhaustion without calling a real service. Existing task checks also verify request summaries and source/backup handling after a final request failure. `--request-timeout` defaults to 180 seconds per attempt and rejects zero before configuration loading. Request formatting and compilation repair retries remain business-level operations, not additional transport retry layers. diff --git a/docs/palm-rust-unit-test-generation.md b/docs/palm-rust-unit-test-generation.md index debe553ae1..b1dab57c8f 100644 --- a/docs/palm-rust-unit-test-generation.md +++ b/docs/palm-rust-unit-test-generation.md @@ -283,7 +283,7 @@ utgen fix --project-dir [--work-dir ...] [--tasks ] 生成任务取得名额后才开始工作，将结果送入容量为 N 的队列后释放名额；串行消费者负责导入检查和候选编译。队列中的结果及正在验证的候选不占生成名额。修复任务则在该函数的修复结果保存后释放名额。源码插入、target 清理、编译或测试、诊断读取及源码恢复串行执行，模型请求可以重叠；该参数不改变 Cargo 内部的依赖编译并行度。同一工作副本不能供多个 PALM 进程同时使用，不同实验应使用独立副本和 target 目录。 -任务异常会在等待所有任务结束后报告，并跳过后续覆盖率统计。修复成功时只删除本次创建的备份；任务失败时恢复源码并保留这些备份，已有备份不会被覆盖。候选本身编译不通过仍按测试结果处理。当前没有增加请求超时、重试策略、跨进程锁或强制终止后的恢复机制。 +任务异常会在等待所有任务结束后报告，并跳过后续覆盖率统计。修复成功时只删除本次创建的备份；任务失败时恢复源码并保留这些备份，已有备份不会被覆盖。候选本身编译不通过仍按测试结果处理。模型请求现在统一执行：`--request-timeout` 默认每次尝试 180 秒，包含响应体读取；网络错误、429（余额不足除外）和 5xx 最多尝试 3 次，其他 HTTP 错误或空回答直接报错。格式纠正和编译修复轮次单独计算，不再叠加网络重试。缺少 `usage` 的有效回答可以继续，但请求统计标记为不完整；生成和修复分别写入 `utgen/generation/gen-requests.json`、`fix-requests.json`。该期限不涵盖 Cargo 子进程，也不提供跨进程锁或强制终止恢复。 ### 8.2 LLM 配置 diff --git a/scripts/check_tasks.py b/scripts/check_tasks.py index 987f2e2e10..7b41011f85 100644 --- a/scripts/check_tasks.py +++ b/scripts/check_tasks.py @@ -136,14 +136,21 @@ def do_POST(self): answer = f'ChangeLog:1@{file}\nFixDescription: Use an integer.\nOriginalCode@{number}-{number}:\n[{number}]{old}\nFixedCode@{number}-{number}:\n[{number}]{new}\n' response = {'id': 'fixture', 'object': 'chat.completion', 'created': 0, 'model': 'fixture', 'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': answer}, 'finish_reason': 'stop'}]} - # Exercise a real worker panic in the existing response parser. - if not (state['panic'] and serial == 1): + # Valid answers must remain usable when a service omits usage. + if not (state['missing_usage'] and serial == 1): response['usage'] = {'prompt_tokens': 1, 'completion_tokens': 1, 'total_tokens': 2} + status = 200 + if ((state['fault'] == 'request' and serial == 1) + or (state['fault'] == 'oracle' and 'generate accurate test oracles' in system + and not state['fault_sent'])): + response = {'error': {'message': 'fixture unauthorized'}} + status = 401 + state['fault_sent'] = True body = json.dumps(response).encode() # Count requests awaiting a response, excluding HTTP-handler teardown. with lock: state['active'] -= 1 - self.send_response(200) + self.send_response(status) self.send_header('Content-Type', 'application/json') self.send_header('Content-Length', str(len(body))) self.end_headers() @@ -157,7 +164,7 @@ def do_POST(self): threading.Thread(target=server.serve_forever, daemon=True).start() env.update(PALM_API_BASE=f'http://127.0.0.1:{server.server_port}/v1', PALM_API_KEY='fixture', PALM_MODEL='fixture') - def run(label, command, limit=1, integration=False, panic=False, fail=False, backpressure=False): + def run(label, command, limit=1, integration=False, missing_usage=False, fault=None, fail=False, backpressure=False): monitor = work / label monitor.mkdir() (monitor / 'compiler.json').write_text(json.dumps(dict(active=0, peak=0, total=0, source_changes=0))) @@ -165,7 +172,7 @@ def run(label, command, limit=1, integration=False, panic=False, fail=False, bac env.pop('PALM_HOLD_FIRST_BUILD', None) with lock: state.clear() - state.update(active=0, peak=0, total=0, errors=[], limit=limit, integration=integration, panic=panic) + state.update(active=0, peak=0, total=0, errors=[], limit=limit, integration=integration, missing_usage=missing_usage, fault=fault, fault_sent=False) ready.clear() watcher = None if backpressure: @@ -209,6 +216,11 @@ def check_queue(): assert (root / 'src/lib.rs').read_text() == original, label assert (root / 'src/unrelated.rs').read_text() == '// unrelated source file\n', label assert unrelated.read_bytes() == b'previous recovery material\n', label + if command[1] in ['gen', 'fix'] and state['total']: + report = json.loads((root / f'utgen/generation/{command[1]}-requests.json').read_text()) + assert report['attempts'] == state['total'], report + assert report['responses_without_usage'] == int(missing_usage), report + assert report['usage_complete'] == (not missing_usage and not state['fault_sent']), report return state.copy() def reset(seed=None): @@ -233,15 +245,19 @@ def results(directory, compiled): seed = work / 'broken-candidates' for n in [1, 2, 4]: reset() - measured = run(f'gen-{n}', command('gen', n), limit=n, backpressure=n == 1) + measured = run(f'gen-{n}', command('gen', n), limit=n, missing_usage=n == 1, backpressure=n == 1) assert measured['total'] == count and measured['peak'] == n, measured results('result', 0) assert not (root / 'src/lib.rs.bak').exists() if n == 1: shutil.copytree(root / 'utgen/generation/pre_fix', seed) + measured = run('gen-cached', command('gen', 1)) + assert measured['total'] == 0 + report = json.loads((root / 'utgen/generation/gen-requests.json').read_text()) + assert report['attempts'] == 0 and report['usage_complete'], report for n in [1, 2, 4]: reset(seed) - measured = run(f'fix-{n}', command('fix', n), limit=n) + measured = run(f'fix-{n}', command('fix', n), limit=n, missing_usage=n == 1) assert measured['total'] == count and measured['peak'] == n, measured results('fixed_result', 1) assert not (root / 'src/lib.rs.bak').exists() @@ -269,7 +285,7 @@ def results(directory, compiled): assert not path.exists(), path for mode in ['gen', 'fix']: reset(seed if mode == 'fix' else None) - measured = run(f'{mode}-panic', command(mode, 1), panic=True, fail=True) + measured = run(f'{mode}-request-error', command(mode, 1), fault='request', fail=True) assert measured['total'] == count, measured # Remaining workers still finish. assert not (root / 'utgen' / ('result' if mode == 'gen' else 'fixed_result')).exists() if mode == 'fix': @@ -281,6 +297,19 @@ def results(directory, compiled): assert measured['total'] == 0 assert backup.read_bytes() == b'old recovery material\n' backup.unlink() + reset() + measured = run('oracle-request-error', command('gen', 1, '--integration', '--oracle'), + integration=True, fault='oracle', fail=True) + assert measured['total'] == count * 3 and measured['fault_sent'] + assert not (root / 'utgen/result').exists() + reset() + encoded = next(iter(json.loads((root / 'brinfo/name_map.json').read_text()).values())) + blocker = root / 'utgen/generation/prompt' / encoded + blocker.parent.mkdir(parents=True) + blocker.write_text('not a directory') + measured = run('gen-worker-panic', command('gen', 1), fail=True) + assert measured['total'] == count - 1 + assert not (root / 'utgen/result').exists() # A consumer failure must still drain the queue and let every sender finish. reset() backup = root / 'src/lib.rs.bak' diff --git a/utgen/Cargo.toml b/utgen/Cargo.toml index 4bcac11f6a..1cd5c0e237 100644 --- a/utgen/Cargo.toml +++ b/utgen/Cargo.toml @@ -15,6 +15,7 @@ quick-xml = "0.37" quote = "1.0" rand = "0.9" regex = "1.11" +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" simplelog = "0.12" @@ -24,8 +25,5 @@ time = "0.3" tokio = { version = "1", features = ["full"] } walkdir = "2.5" -[dev-dependencies] -reqwest = { version = "0.12", default-features = false } - [build-dependencies] build-utils = { path = "../build-utils" } diff --git a/utgen/README.md b/utgen/README.md index 3d25411edd..d1c99a0e34 100644 --- a/utgen/README.md +++ b/utgen/README.md @@ -32,6 +32,12 @@ An absolute path continues to work after changing into the target crate. Relativ Changing configuration takes effect on the next invocation without rebuilding. There are no built-in model credentials. Configuration values are not printed in configuration diagnostics. +Generation and repair share one non-streaming request path. `--request-timeout ` defaults to 180 and covers each attempt, including its response body. Transport failures, HTTP 429 (except `insufficient_quota`), and 5xx responses get at most three attempts, with 1- and 2-second delays. Other HTTP errors and malformed/empty answers are returned without network retries. SDK retries are not layered underneath this policy; the existing request/response types are retained while the existing reqwest dependency executes HTTP requests. + +A valid answer without `usage` remains usable. Numeric token fields contain reported usage only; missing usage contributes no known tokens, which must not be interpreted as zero actual usage. The command writes `utgen/generation/gen-requests.json` or `fix-requests.json` with attempt/failure counts, missing-usage counts, reported token totals, and `usage_complete`. That flag is false after any failed attempt or response without usage. These files also cover requests whose answers failed later parsing. They are written after workers finish, including request failures; successful cached runs record zero attempts. Validation failures before any request leave the target untouched. + +Formatting retries and compiler-guided repair rounds remain separate from transport retries. A final request failure, including during oracle generation or repair, fails the command and skips subsequent statistics. The request deadline does not terminate Cargo processes or impose a whole-function deadline. + Install the toolchain and coverage tool described in the [project README](../README.md#prerequisites). On a fresh working copy, run `utgen pre-process -p ` and then `utgen analyze -p ` to produce: ```text @@ -139,6 +145,7 @@ utgen gen -p --requirement --context | `-o, --oracle` | Use separate input-range, test-prefix, and oracle generation. Default: off; otherwise generate complete tests directly. | | `-i, --integration` | Generate integration tests under `tests/`, using the analysis visibility flag to select functions and compilation checks to filter candidates. Default: off. | | `-t, --tasks` | Default: 4. Maximum active focal-function generation jobs; must be positive. | +| `--request-timeout` | Default: 180 seconds per model request attempt; also available for `fix`. Must be positive. | Generation validates the branch index, context index, context files, and source paths before modifying the target. Missing or inconsistent artifacts are errors. A failed generation task is reported instead of being silently lost before statistics. `--tasks 0` is rejected. @@ -156,7 +163,7 @@ Repair reads the generated candidates and uses compiler diagnostics to revise th Source insertion, target cleanup, compilation/test execution, diagnostic reading, and source restoration are serialized within each command. Model requests may overlap. Cargo retains its own dependency-build parallelism. Concurrent PALM commands or experiments require separate working copies and separate target directories. -All workers are joined before cleanup. A task panic or infrastructure error returns a nonzero status and skips subsequent coverage statistics; candidate compilation errors remain ordinary repair outcomes. Repair restores its source backups before returning, deletes only backups created by that invocation on success, and retains them after a worker failure. Existing source backups are rejected without overwriting them. Temporary import/candidate files are restored or removed after validation, including unwinding. This is not recovery from forced process termination, and request/subprocess deadlines remain unchanged. +All workers are joined before cleanup. A task panic or infrastructure error returns a nonzero status and skips subsequent coverage statistics; candidate compilation errors remain ordinary repair outcomes. Repair restores its source backups before returning, deletes only backups created by that invocation on success, and retains them after a worker failure. Existing source backups are rejected without overwriting them. Temporary import/candidate files are restored or removed after validation, including unwinding. This is not recovery from forced process termination. Model request deadlines are described above; Cargo subprocess deadlines remain separate work. There is no integration-mode option for repair. Candidates are inserted into source files as unit tests, and post-repair statistics use that same mode. Do not interpret `gen --integration` followed by `fix` as preserving an integration-only evaluation. @@ -169,6 +176,7 @@ Paths under `utgen/` below are relative to `--project-dir`: | `utgen/generation/prompt/` and `answer/` | Prompts and model responses. | | `utgen/generation/pre_fix/` | Candidate tests and their compilation status before repair. | | `utgen/generation/llm_fix/` | Candidates and their compilation status after repair. | +| `utgen/generation/gen-requests.json` and `fix-requests.json` | Request attempts, failures, reported token totals, and usage completeness for the invocation. | | `utgen/result/` | Pre-repair coverage and execution statistics per focal function. | | `utgen/fixed_result/` | Post-repair coverage and execution statistics per focal function. | | `utgen/original_result.json` | Comparison statistics when the original integration-test backup is available. | diff --git a/utgen/src/gene/cot/input_infer.rs b/utgen/src/gene/cot/input_infer.rs index accbceb3b6..52c664d602 100644 --- a/utgen/src/gene/cot/input_infer.rs +++ b/utgen/src/gene/cot/input_infer.rs @@ -1,8 +1,6 @@ -use super::Prompt; use super::LLM; -use log::{error, info}; -use rand::Rng; -use tokio::time::{sleep, Duration}; +use super::Prompt; +use log::error; fn postprocess(inputs: &mut Vec) { for input in inputs.iter_mut() { @@ -10,40 +8,23 @@ fn postprocess(inputs: &mut Vec) { } } -pub async fn gen_input_range(llm: &LLM, pt_info: &Prompt, conds: &Vec) -> Option<(String, u32, u32)> { +pub async fn gen_input_range( + llm: &LLM, + pt_info: &Prompt, + conds: &Vec, +) -> Option<(String, u32, u32)> { let system_pt = &pt_info.system_pt; let static_pt = &pt_info.static_pt; - let mut completion_tokens = 0; - let mut prompt_tokens = 0; let user_pt = static_pt.clone() + &conds.join(""); - let mut retry = 0; - let mut answers = None; - while retry < 3 { - if retry != 0 { - let random_secs = { - let mut rng = rand::rng(); - rng.random_range(10..=30) - }; - sleep(Duration::from_secs(random_secs)).await; - } - let result = llm.fetch_answer(Some(&system_pt), &user_pt, 1, false).await; - if result.is_ok() { - let (result_answers, usage_completion, usage_prompt) = result.unwrap(); - completion_tokens += usage_completion; - prompt_tokens += usage_prompt; - answers = Some(result_answers); - break; - } - error!("{}. Retrying...", result.unwrap_err()); - retry += 1; - } - if answers.is_none() { - error!("Failed to fetch answer."); - return None; - } - let mut answers = answers.unwrap(); - assert!(answers.len() == 1); + let (mut answers, completion_tokens, prompt_tokens) = + match llm.fetch_answer(Some(system_pt), &user_pt, 1, false).await { + Ok(answer) => answer, + Err(error) => { + error!("Model request failed: {error}"); + return None; + } + }; postprocess(&mut answers); // info!("Answers: {:?}", answers); diff --git a/utgen/src/gene/cot/oracle_gen.rs b/utgen/src/gene/cot/oracle_gen.rs index 1eed1565c4..120c91d0e6 100644 --- a/utgen/src/gene/cot/oracle_gen.rs +++ b/utgen/src/gene/cot/oracle_gen.rs @@ -1,8 +1,6 @@ -use super::Prompt; use super::LLM; -use log::{error, info}; -use rand::Rng; -use tokio::time::{sleep, Duration}; +use super::Prompt; +use log::error; fn postprocess(inputs: &mut Vec) { for input in inputs.iter_mut() { @@ -19,39 +17,18 @@ pub async fn gen_oracle( ) -> Option<(String, u32, u32)> { let system_pt = &pt_info.system_pt; let static_pt = &pt_info.static_pt; - let mut completion_tokens = 0; - let mut prompt_tokens = 0; let mut user_pt = static_pt.clone() + &conds.join(""); - user_pt += - "Here is the test function code. You should generate the corresponding test oracles for it:\n"; + user_pt += "Here is the test function code. You should generate the corresponding test oracles for it:\n"; user_pt += &code.join("\n"); - let mut retry = 0; - let mut answers = None; - while retry < 3 { - if retry != 0 { - let random_secs = { - let mut rng = rand::rng(); - rng.random_range(10..=30) - }; - sleep(Duration::from_secs(random_secs)).await; - } - let result = llm.fetch_answer(Some(&system_pt), &user_pt, 1, false).await; - if result.is_ok() { - let (result_answers, usage_completion, usage_prompt) = result.unwrap(); - completion_tokens += usage_completion; - prompt_tokens += usage_prompt; - answers = Some(result_answers); - break; - } - error!("{}. Retrying...", result.unwrap_err()); - retry += 1; - } - if answers.is_none() { - error!("Failed to fetch answer."); - return None; - } - let mut answers = answers.unwrap(); + let (mut answers, completion_tokens, prompt_tokens) = + match llm.fetch_answer(Some(system_pt), &user_pt, 1, false).await { + Ok(answer) => answer, + Err(error) => { + error!("Model request failed: {error}"); + return None; + } + }; postprocess(&mut answers); // info!("Answers: {:?}", answers); diff --git a/utgen/src/gene/cot/prefix_gen.rs b/utgen/src/gene/cot/prefix_gen.rs index 7db48e2ab4..9ef482f324 100644 --- a/utgen/src/gene/cot/prefix_gen.rs +++ b/utgen/src/gene/cot/prefix_gen.rs @@ -1,5 +1,5 @@ -use super::Prompt; use super::LLM; +use super::Prompt; use super::{extract_test_functions, try_parse}; use crate::types::{ChainTestAnswer, TestInfo}; use log::{error, info}; @@ -8,7 +8,7 @@ use std::collections::HashSet; use std::fs; use std::io::Write; use std::path::Path; -use tokio::time::{sleep, Duration}; +use tokio::time::{Duration, sleep}; fn postprocess(inputs: &mut Vec) { for input in inputs.iter_mut() { @@ -44,33 +44,16 @@ pub async fn gen_prefix( }; sleep(Duration::from_secs(random_secs)).await; } - // Net retry - let mut net_retry = 0; - let mut answers = None; - while net_retry < 3 { - if net_retry != 0 { - let random_secs = { - let mut rng = rand::rng(); - rng.random_range(10..=30) - }; - sleep(Duration::from_secs(random_secs)).await; - } - let result = llm.fetch_answer(Some(&system_pt), &user_pt, 1, false).await; - if result.is_ok() { - let (result_answers, usage_completion, usage_prompt) = result.unwrap(); - completion_tokens += usage_completion; - prompt_tokens += usage_prompt; - answers = Some(result_answers); - break; - } - error!("{}. Retrying...", result.unwrap_err()); - net_retry += 1; - } - if answers.is_none() { - error!("Failed to fetch answer."); - return None; - } - let mut answers = answers.unwrap(); + let (mut answers, usage_completion, usage_prompt) = + match llm.fetch_answer(Some(system_pt), &user_pt, 1, false).await { + Ok(answer) => answer, + Err(error) => { + error!("Model request failed: {error}"); + return None; + } + }; + completion_tokens += usage_completion; + prompt_tokens += usage_prompt; postprocess(&mut answers); diff --git a/utgen/src/gene/cot/test_gen.rs b/utgen/src/gene/cot/test_gen.rs index b8e75dc04b..b460a08ac3 100644 --- a/utgen/src/gene/cot/test_gen.rs +++ b/utgen/src/gene/cot/test_gen.rs @@ -1,5 +1,5 @@ -use super::Prompt; use super::LLM; +use super::Prompt; use super::{extract_test_functions, try_parse}; use crate::types::{ChainTestAnswer, TestInfo}; use log::{error, info}; @@ -8,7 +8,7 @@ use std::collections::HashSet; use std::fs; use std::io::Write; use std::path::Path; -use tokio::time::{sleep, Duration}; +use tokio::time::{Duration, sleep}; fn postprocess(inputs: &mut Vec) { for input in inputs.iter_mut() { @@ -40,33 +40,16 @@ pub async fn gen_test( }; sleep(Duration::from_secs(random_secs)).await; } - // Net retry - let mut net_retry = 0; - let mut answers = None; - while net_retry < 3 { - if net_retry != 0 { - let random_secs = { - let mut rng = rand::rng(); - rng.random_range(10..=30) - }; - sleep(Duration::from_secs(random_secs)).await; - } - let result = llm.fetch_answer(Some(&system_pt), &user_pt, 1, false).await; - if result.is_ok() { - let (result_answers, usage_completion, usage_prompt) = result.unwrap(); - completion_tokens += usage_completion; - prompt_tokens += usage_prompt; - answers = Some(result_answers); - break; - } - error!("{}. Retrying...", result.unwrap_err()); - net_retry += 1; - } - if answers.is_none() { - error!("Failed to fetch answer."); - return None; - } - let mut answers = answers.unwrap(); + let (mut answers, usage_completion, usage_prompt) = + match llm.fetch_answer(Some(system_pt), &user_pt, 1, false).await { + Ok(answer) => answer, + Err(error) => { + error!("Model request failed: {error}"); + return None; + } + }; + completion_tokens += usage_completion; + prompt_tokens += usage_prompt; postprocess(&mut answers); diff --git a/utgen/src/gene/generation.rs b/utgen/src/gene/generation.rs index 9926cf79f6..d419719a0f 100644 --- a/utgen/src/gene/generation.rs +++ b/utgen/src/gene/generation.rs @@ -183,9 +183,9 @@ async fn gen_tests_cot( for test_info in chain_test_answer.get_tests_mut() { let mut code = common.clone(); code.extend(test_info.prefix_func()); - let oracle = gen_oracle(llm, &oracle_pt_info, &cond_prompt, &code).await; - if oracle.is_some() { - let (oracles , usage_completion, usage_prompt)= oracle.unwrap(); + { + let (oracles, usage_completion, usage_prompt) = + gen_oracle(llm, &oracle_pt_info, &cond_prompt, &code).await?; completion_tokens += usage_completion; prompt_tokens += usage_prompt; let oracles = oracles @@ -248,9 +248,9 @@ async fn gen_tests_cot( for test_info in chain_test_answer.get_tests_mut() { let mut code = common.clone(); code.extend(test_info.prefix_func()); - let oracle = gen_oracle(llm, &oracle_pt_info, &vec![], &code).await; - if oracle.is_some() { - let (oracles , usage_completion, usage_prompt) = oracle.unwrap(); + { + let (oracles, usage_completion, usage_prompt) = + gen_oracle(llm, &oracle_pt_info, &vec![], &code).await?; completion_tokens += usage_completion; prompt_tokens += usage_prompt; let oracles = oracles diff --git a/utgen/src/gene/llm.rs b/utgen/src/gene/llm.rs index 69ee059a94..1d062963d6 100644 --- a/utgen/src/gene/llm.rs +++ b/utgen/src/gene/llm.rs @@ -1,29 +1,116 @@ use crate::LlmConfig; -use async_openai::{ - Client, - config::OpenAIConfig, - types::{ - ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs, - ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs, - }, +use async_openai::types::{ + ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs, + ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequest, + CreateChatCompletionRequestArgs, CreateChatCompletionResponse, }; -use futures::StreamExt; +use serde::Serialize; +use std::{ + fs, io, + path::Path, + sync::{Arc, Mutex}, + time::Duration, +}; +use tokio::time::{sleep, timeout}; + +const MAX_ATTEMPTS: u32 = 3; + +#[derive(Default, Serialize)] +struct RequestStats { + attempts: u64, + failed_attempts: u64, + responses_without_usage: u64, + reported_completion_tokens: u64, + reported_prompt_tokens: u64, +} + +struct RequestError { + message: String, + retryable: bool, +} + +impl From for RequestError { + fn from(error: reqwest::Error) -> Self { + Self { + retryable: error.is_connect() + || error.is_timeout() + || error.is_body() + || error.is_request(), + message: error.to_string(), + } + } +} #[derive(Clone)] pub struct LLM { config: LlmConfig, - client: Client, + client: reqwest::Client, + request_timeout: Duration, + stats: Arc>, } impl LLM { pub fn new(config: LlmConfig) -> Self { - let api_config = OpenAIConfig::new() - .with_api_base(&config.base) - .with_api_key(&config.key); Self { config, - client: Client::with_config(api_config), + client: reqwest::Client::new(), + request_timeout: Duration::from_secs(180), + stats: Arc::new(Mutex::new(RequestStats::default())), + } + } + + pub fn with_request_timeout(mut self, duration: Duration) -> Self { + self.request_timeout = duration; + self + } + + pub fn request_count(&self) -> u64 { + self.stats.lock().unwrap().attempts + } + + /// Persist reported usage, including an explicit incomplete-usage marker. + pub fn write_request_summary(&self, path: &Path) -> io::Result<()> { + let stats = self.stats.lock().unwrap(); + let mut summary = serde_json::to_value(&*stats)?; + summary["usage_complete"] = + (stats.failed_attempts == 0 && stats.responses_without_usage == 0).into(); + fs::create_dir_all(path.parent().unwrap())?; + fs::write(path, serde_json::to_vec_pretty(&summary)?) + } + + async fn send_once( + &self, + request: &CreateChatCompletionRequest, + ) -> Result { + let response = self + .client + .post(format!( + "{}/chat/completions", + self.config.base.trim_end_matches('/') + )) + .bearer_auth(&self.config.key) + .json(request) + .send() + .await?; + let status = response.status(); + let body = response.bytes().await?; + if !status.is_success() { + let detail: serde_json::Value = serde_json::from_slice(&body).unwrap_or_default(); + let error = &detail["error"]; + let quota_exhausted = + error["code"] == "insufficient_quota" || error["type"] == "insufficient_quota"; + return Err(RequestError { + message: format!( + "Model HTTP {status}: {}", + error["message"].as_str().unwrap_or("request failed") + ), + retryable: (status.as_u16() == 429 && !quota_exhausted) || status.is_server_error(), + }); } + serde_json::from_slice(&body).map_err(|error| RequestError { + message: format!("Invalid model response JSON: {error}"), + retryable: false, + }) } pub async fn fetch_answer( @@ -32,28 +119,27 @@ impl LLM { user_pt: &str, n: u8, stream: bool, - ) -> Result<(Vec, u32, u32), Box> { - let client = &self.client; - let system_msg = if system_pt.is_none() { - None - } else { - Some( + ) -> io::Result<(Vec, u32, u32)> { + if n != 1 || stream { + return Err(io::Error::other( + "PALM supports one non-streaming answer per request", + )); + } + let mut messages = Vec::new(); + if let Some(system) = system_pt { + messages.push(ChatCompletionRequestMessage::System( ChatCompletionRequestSystemMessageArgs::default() - .content(system_pt.unwrap()) - .build()?, - ) - }; - let user_msg = ChatCompletionRequestUserMessageArgs::default() - .content(user_pt) - .build()?; - let messages = if system_msg.is_none() { - vec![ChatCompletionRequestMessage::User(user_msg)] - } else { - vec![ - ChatCompletionRequestMessage::System(system_msg.unwrap()), - ChatCompletionRequestMessage::User(user_msg), - ] - }; + .content(system) + .build() + .map_err(io::Error::other)?, + )); + } + messages.push(ChatCompletionRequestMessage::User( + ChatCompletionRequestUserMessageArgs::default() + .content(user_pt) + .build() + .map_err(io::Error::other)?, + )); let request = CreateChatCompletionRequestArgs::default() .model(&self.config.model) .max_tokens(10000_u32) @@ -62,40 +148,69 @@ impl LLM { .n(1) .stream(false) .messages(messages) - .build()?; - let mut result; - let mut completion_tokens = 0; - let mut prompt_tokens = 0; - if !stream { - let response = client.chat().create(request).await?; - result = response - .choices - .into_iter() - .filter_map(|c| c.message.content) - .collect(); - let usage = response.usage.unwrap(); - completion_tokens += usage.completion_tokens; - prompt_tokens += usage.prompt_tokens; - } else { - result = vec!["".to_string(); n as usize]; - let mut stream = client.chat().create_stream(request).await?; - while let Some(response) = stream.next().await { - match response { - Ok(chunk) => { - for choice in chunk.choices.into_iter() { - if let Some(content) = choice.delta.content { - result[choice.index as usize] += &content; - } + .build() + .map_err(io::Error::other)?; + for attempt in 1..=MAX_ATTEMPTS { + self.stats.lock().unwrap().attempts += 1; + let response = match timeout(self.request_timeout, self.send_once(&request)).await { + Ok(result) => result, + Err(_) => Err(RequestError { + message: format!( + "Model request timed out after {}s", + self.request_timeout.as_secs_f64() + ), + retryable: true, + }), + }; + match response { + Ok(response) => { + let mut stats = self.stats.lock().unwrap(); + let (completion, prompt) = match response.usage { + Some(usage) => { + stats.reported_completion_tokens += u64::from(usage.completion_tokens); + stats.reported_prompt_tokens += u64::from(usage.prompt_tokens); + (usage.completion_tokens, usage.prompt_tokens) + } + None => { + stats.responses_without_usage += 1; + log::warn!( + "Model response has no usage; reported token totals are incomplete" + ); + (0, 0) } - let usage = chunk.usage.unwrap(); - completion_tokens += usage.completion_tokens; - prompt_tokens += usage.prompt_tokens; + }; + if response.choices.len() != 1 { + return Err(io::Error::other( + "Model response must contain exactly one choice", + )); + } + let answer = response + .choices + .into_iter() + .next() + .unwrap() + .message + .content + .filter(|content| !content.trim().is_empty()) + .ok_or_else(|| { + io::Error::other("Model response contains no answer text") + })?; + return Ok((vec![answer], completion, prompt)); + } + Err(error) => { + self.stats.lock().unwrap().failed_attempts += 1; + if !error.retryable || attempt == MAX_ATTEMPTS { + return Err(io::Error::other(format!( + "{} (attempt {attempt}/{MAX_ATTEMPTS})", + error.message + ))); } - Err(e) => return Err(Box::new(e)), + log::warn!("{}; retrying ({attempt}/{MAX_ATTEMPTS})", error.message); + sleep(Duration::from_secs(u64::from(attempt))).await; } } } - Ok((result, completion_tokens, prompt_tokens)) + unreachable!("the last attempt returns success or failure") } pub async fn get_answer( @@ -103,66 +218,8 @@ impl LLM { prompt: &str, n: u8, stream: bool, - ) -> Result<(Vec, u32, u32), Box> { - let client = &self.client; - // let msg = ChatCompletionRequestUserMessageArgs::default() - // .content(prompt) - // .build()?; - let msg = ChatCompletionRequestUserMessageArgs::default() - .content(prompt) - .build() - .map_err(|e| Box::new(e) as Box)?; - let mut completion_tokens = 0; - let mut prompt_tokens = 0; - let request = CreateChatCompletionRequestArgs::default() - .model(&self.config.model) - .max_tokens(10000_u32) - .temperature(1.0) - .top_p(0_f32) - .n(1) - .stream(false) - .messages(vec![ChatCompletionRequestMessage::User(msg)]) - .build() - .map_err(|e| Box::new(e) as Box)?; - let mut result; - if !stream { - let response = client - .chat() - .create(request) - .await - .map_err(|e| Box::new(e) as Box)?; - result = response - .choices - .into_iter() - .filter_map(|c| c.message.content) - .collect(); - let usage = response.usage.unwrap(); - completion_tokens += usage.completion_tokens; - prompt_tokens += usage.prompt_tokens; - } else { - result = vec!["".to_string(); n as usize]; - let mut stream = client - .chat() - .create_stream(request) - .await - .map_err(|e| Box::new(e) as Box)?; - while let Some(response) = stream.next().await { - match response { - Ok(chunk) => { - for choice in chunk.choices.into_iter() { - if let Some(content) = choice.delta.content { - result[choice.index as usize] += &content; - } - } - let usage = chunk.usage.unwrap(); - completion_tokens += usage.completion_tokens; - prompt_tokens += usage.prompt_tokens; - } - Err(e) => return Err(Box::new(e)), - } - } - } - Ok((result, completion_tokens, prompt_tokens)) + ) -> io::Result<(Vec, u32, u32)> { + self.fetch_answer(None, prompt, n, stream).await } } @@ -175,44 +232,66 @@ mod tests { time::{Duration, timeout}, }; - // Exercise the actual HTTP client and request shape without a model service. - async fn mock_model() -> (LLM, tokio::task::JoinHandle) { + fn completion() -> serde_json::Value { + serde_json::json!({ + "id": "local-response", "object": "chat.completion", "created": 0, "model": "local-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "local answer"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5} + }) + } + + fn reply(status: u16) -> (u16, String, Duration) { + let body = if status == 200 { + completion() + } else { + serde_json::json!({"error": {"message": "fixture error"}}) + }; + (status, body.to_string(), Duration::ZERO) + } + + // Real HTTP, with counts and delayed bodies to verify the entire attempt deadline. + async fn mock_replies( + replies: Vec<(u16, String, Duration)>, + ) -> (LLM, tokio::task::JoinHandle>) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let server = tokio::spawn(async move { - timeout(Duration::from_secs(10), async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut data = Vec::new(); - let (headers_end, length) = loop { - let mut buffer = [0; 2048]; - let count = socket.read(&mut buffer).await.unwrap(); - assert!(count > 0, "request ended before headers"); - data.extend_from_slice(&buffer[..count]); - if let Some(end) = data.windows(4).position(|s| s == b"\r\n\r\n") { - let headers = String::from_utf8_lossy(&data[..end]).to_lowercase(); - assert!(headers.starts_with("post /v1/chat/completions http/1.1")); - assert!(headers.contains("authorization: bearer local-test-key")); - let length = headers.lines().find_map(|line| { - line.strip_prefix("content-length:").map(|n| n.trim().parse::().unwrap()) - }).unwrap(); - break (end + 4, length); + timeout(Duration::from_secs(15), async move { + let mut requests = Vec::new(); + for (status, body, delay) in replies { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut data = Vec::new(); + let (headers_end, length) = loop { + let mut buffer = [0; 2048]; + let count = socket.read(&mut buffer).await.unwrap(); + assert!(count > 0, "request ended before headers"); + data.extend_from_slice(&buffer[..count]); + if let Some(end) = data.windows(4).position(|s| s == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&data[..end]).to_lowercase(); + assert!(headers.starts_with("post /v1/chat/completions http/1.1")); + assert!(headers.contains("authorization: bearer local-test-key")); + let length = headers.lines().find_map(|line| { + line.strip_prefix("content-length:").map(|n| n.trim().parse::().unwrap()) + }).unwrap(); + break (end + 4, length); + } + }; + while data.len() < headers_end + length { + let mut buffer = [0; 2048]; + let count = socket.read(&mut buffer).await.unwrap(); + assert!(count > 0, "request ended before body"); + data.extend_from_slice(&buffer[..count]); } - }; - while data.len() < headers_end + length { - let mut buffer = [0; 2048]; - let count = socket.read(&mut buffer).await.unwrap(); - assert!(count > 0, "request ended before body"); - data.extend_from_slice(&buffer[..count]); + requests.push(serde_json::from_slice(&data[headers_end..headers_end + length]).unwrap()); + if status == 0 { continue; } // Deliberately drop the response connection. + let header = format!("HTTP/1.1 {status} Fixture\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()); + socket.write_all(header.as_bytes()).await.unwrap(); + sleep(delay).await; + // A timeout deliberately disconnects before its delayed body arrives. + let result = socket.write_all(body.as_bytes()).await; + if delay.is_zero() { result.unwrap(); } } - let request = serde_json::from_slice(&data[headers_end..headers_end + length]).unwrap(); - let body = serde_json::json!({ - "id": "local-response", "object": "chat.completion", "created": 0, "model": "local-model", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "local answer"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5} - }).to_string(); - let response = format!("HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", body.len(), body); - socket.write_all(response.as_bytes()).await.unwrap(); - request + requests }).await.expect("local mock server timed out") }); let mut llm = LLM::new(LlmConfig { @@ -220,14 +299,15 @@ mod tests { key: "local-test-key".into(), model: "local-model".into(), }); - // Only the local fixture bypasses proxies. LLM::new keeps the SDK's - // default proxy behavior for production and the opt-in live test. - llm.client = llm - .client - .with_http_client(reqwest::Client::builder().no_proxy().build().unwrap()); + // Only fixtures bypass proxies; normal clients retain proxy discovery. + llm.client = reqwest::Client::builder().no_proxy().build().unwrap(); (llm, server) } + async fn mock_model() -> (LLM, tokio::task::JoinHandle>) { + mock_replies(vec![reply(200)]).await + } + fn check_request(request: &serde_json::Value) { assert_eq!(request["model"], "local-model"); assert_eq!(request["max_tokens"], 10000); @@ -280,7 +360,7 @@ mod tests { .unwrap(); assert_eq!(answers, vec!["local answer"]); assert_eq!((completion, prompt), (3, 2)); - let request = server.await.unwrap(); + let request = server.await.unwrap().remove(0); check_request(&request); assert_eq!(request["messages"][0]["role"], "system"); assert_eq!(request["messages"][1]["content"], "generate a test"); @@ -298,12 +378,112 @@ mod tests { .unwrap(); assert_eq!(answers, vec!["local answer"]); assert_eq!((completion, prompt), (3, 2)); - let request = server.await.unwrap(); + let request = server.await.unwrap().remove(0); check_request(&request); assert_eq!(request["messages"][0]["role"], "user"); assert_eq!(request["messages"][0]["content"], "repair a test"); } + #[tokio::test] + async fn missing_usage_keeps_the_answer_and_marks_the_report_incomplete() { + let mut body = completion(); + body.as_object_mut().unwrap().remove("usage"); + let (llm, server) = mock_replies(vec![(200, body.to_string(), Duration::ZERO)]).await; + let answer = llm.fetch_answer(None, "test", 1, false).await.unwrap(); + assert_eq!(answer, (vec!["local answer".into()], 0, 0)); + assert_eq!(server.await.unwrap().len(), 1); + let path = + std::env::temp_dir().join(format!("palm-request-usage-{}.json", std::process::id())); + llm.write_request_summary(&path).unwrap(); + let summary: serde_json::Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap(); + fs::remove_file(path).unwrap(); + assert_eq!(summary["responses_without_usage"], 1); + assert_eq!(summary["usage_complete"], false); + } + + #[tokio::test] + async fn rejects_empty_or_malformed_answers_without_retrying() { + let mut bodies = Vec::new(); + let mut empty = completion(); + empty["choices"] = serde_json::json!([]); + bodies.push(empty.to_string()); + for content in [serde_json::Value::Null, serde_json::json!(" \n ")] { + let mut body = completion(); + body["choices"][0]["message"]["content"] = content; + bodies.push(body.to_string()); + } + bodies.push("{invalid JSON".into()); + for body in bodies { + let (llm, server) = mock_replies(vec![(200, body, Duration::ZERO)]).await; + assert!(llm.get_answer("test", 1, false).await.is_err()); + assert_eq!(server.await.unwrap().len(), 1); + assert_eq!(llm.stats.lock().unwrap().attempts, 1); + } + } + + #[tokio::test] + async fn permanent_http_errors_do_not_retry() { + for status in [400, 401, 403, 404, 422] { + let (llm, server) = mock_replies(vec![reply(status)]).await; + let error = llm.get_answer("test", 1, false).await.unwrap_err(); + assert!(error.to_string().contains(&status.to_string())); + assert_eq!(server.await.unwrap().len(), 1); + assert_eq!(llm.stats.lock().unwrap().attempts, 1); + } + let body = + serde_json::json!({"error": {"message": "no quota", "code": "insufficient_quota"}}); + let (llm, server) = mock_replies(vec![(429, body.to_string(), Duration::ZERO)]).await; + assert!(llm.get_answer("test", 1, false).await.is_err()); + assert_eq!(server.await.unwrap().len(), 1); + } + + #[tokio::test] + async fn transient_errors_retry_and_recover_within_three_attempts() { + let (llm, server) = mock_replies(vec![reply(429), reply(503), reply(200)]).await; + assert_eq!( + llm.get_answer("test", 1, false).await.unwrap().0, + vec!["local answer"] + ); + assert_eq!(server.await.unwrap().len(), 3); + let stats = llm.stats.lock().unwrap(); + assert_eq!((stats.attempts, stats.failed_attempts), (3, 2)); + } + + #[tokio::test] + async fn a_dropped_connection_is_retried() { + let (llm, server) = mock_replies(vec![reply(0), reply(200)]).await; + assert!(llm.get_answer("test", 1, false).await.is_ok()); + assert_eq!(server.await.unwrap().len(), 2); + } + + #[tokio::test] + async fn transient_errors_stop_after_three_attempts() { + for status in [429, 500] { + let (llm, server) = + mock_replies(vec![reply(status), reply(status), reply(status)]).await; + assert!( + llm.get_answer("test", 1, false) + .await + .unwrap_err() + .to_string() + .contains("attempt 3/3") + ); + assert_eq!(server.await.unwrap().len(), 3); + assert_eq!(llm.stats.lock().unwrap().attempts, 3); + } + } + + #[tokio::test] + async fn attempt_timeout_includes_reading_the_body_and_stops_after_three() { + let delayed = (200, completion().to_string(), Duration::from_millis(150)); + let (llm, server) = mock_replies(vec![delayed.clone(), delayed.clone(), delayed]).await; + let llm = llm.with_request_timeout(Duration::from_millis(50)); + let error = llm.get_answer("test", 1, false).await.unwrap_err(); + assert!(error.to_string().contains("timed out")); + assert_eq!(server.await.unwrap().len(), 3); + assert_eq!(llm.stats.lock().unwrap().failed_attempts, 3); + } + #[tokio::test] #[ignore = "requires explicitly configured model credentials and makes a real API request"] async fn test_llm() { diff --git a/utgen/src/main.rs b/utgen/src/main.rs index 46ffdfa02a..58db9416c4 100644 --- a/utgen/src/main.rs +++ b/utgen/src/main.rs @@ -2,7 +2,11 @@ use clap::{Args, Parser, Subcommand}; use log::info; use simplelog::{ColorChoice, ConfigBuilder, LevelFilter, TermLogger, TerminalMode}; use std::path::PathBuf; -use std::{env, num::NonZeroUsize}; +use std::{ + env, + num::{NonZeroU64, NonZeroUsize}, + time::Duration, +}; use utgen::{LLM, LlmConfig, analyze_project, collect_coverage, validate_repair}; use utgen::{comment_out_tests, gen_test_rate, gen_tests_project, llm_fix, rename_tests_to_bak}; @@ -41,6 +45,9 @@ enum Command { /// Maximum active function-generation tasks; compilation stays serial #[arg(short, long, default_value = "4")] tasks: NonZeroUsize, + /// Timeout in seconds for each model request attempt (at most three attempts) + #[arg(long, default_value = "180", value_name = "SECONDS")] + request_timeout: NonZeroU64, /// Whether to generate integration tests #[arg(short, long)] integration: bool, @@ -61,6 +68,9 @@ enum Command { /// Maximum active function-repair tasks; compilation stays serial #[arg(short, long, default_value = "4")] tasks: NonZeroUsize, + /// Timeout in seconds for each model request attempt (at most three attempts) + #[arg(long, default_value = "180", value_name = "SECONDS")] + request_timeout: NonZeroU64, }, } @@ -161,12 +171,14 @@ async fn main() -> Result<(), Box> { Command::Gen { options, tasks, + request_timeout, integration, requirement, context, oracle, } => { - let llm = LLM::new(LlmConfig::load(cli.config.as_deref())?); + let llm = LLM::new(LlmConfig::load(cli.config.as_deref())?) + .with_request_timeout(Duration::from_secs(request_timeout.get())); let (project_dir, work_dirs) = get_dirs(options)?; info!( "Generating tests for project at {}, with work directory(s) {:?}", @@ -174,38 +186,64 @@ async fn main() -> Result<(), Box> { work_dirs ); - // Generate tests for each work directory - for work_dir in work_dirs.iter() { - gen_tests_project( - &llm, - &project_dir, - &work_dir, - tasks.get(), - integration, - requirement, - context, - oracle, - ) - .await?; + let generated = async { + for work_dir in work_dirs.iter() { + gen_tests_project( + &llm, + &project_dir, + &work_dir, + tasks.get(), + integration, + requirement, + context, + oracle, + ) + .await?; + } + Ok::<_, std::io::Error>(()) } + .await; + let recorded = if generated.is_ok() || llm.request_count() > 0 { + llm.write_request_summary(&project_dir.join("utgen/generation/gen-requests.json")) + } else { + Ok(()) + }; + generated?; + recorded?; // Collect coverage rate and pass rate for each work directory for work_dir in work_dirs.iter() { gen_test_rate(&project_dir, &work_dir, integration, true); } } // fix unit tests - Command::Fix { options, tasks } => { - let llm = LLM::new(LlmConfig::load(cli.config.as_deref())?); + Command::Fix { + options, + tasks, + request_timeout, + } => { + let llm = LLM::new(LlmConfig::load(cli.config.as_deref())?) + .with_request_timeout(Duration::from_secs(request_timeout.get())); let (project_dir, work_dirs) = get_dirs(options)?; info!( "Fixing tests for project at {}, with work directory(s) {:?}", project_dir.display(), work_dirs ); - for work_dir in work_dirs.iter() { - validate_repair(&project_dir, work_dir)?; - llm_fix(&llm, project_dir.clone(), work_dir.clone(), tasks.get()).await?; + let repaired = async { + for work_dir in work_dirs.iter() { + validate_repair(&project_dir, work_dir)?; + llm_fix(&llm, project_dir.clone(), work_dir.clone(), tasks.get()).await?; + } + Ok::<_, std::io::Error>(()) } + .await; + let recorded = if repaired.is_ok() || llm.request_count() > 0 { + llm.write_request_summary(&project_dir.join("utgen/generation/fix-requests.json")) + } else { + Ok(()) + }; + repaired?; + recorded?; //gen_test_rate_aggregated(&project_dir, false); for work_dir in work_dirs.iter() { gen_test_rate(&project_dir, &work_dir, false, false); diff --git a/utgen/src/run/llm_fix.rs b/utgen/src/run/llm_fix.rs index 79eff03005..abf0474518 100644 --- a/utgen/src/run/llm_fix.rs +++ b/utgen/src/run/llm_fix.rs @@ -17,7 +17,6 @@ use serde::Deserialize; use std::{ cmp::min, collections::{BTreeSet, HashMap}, - error::Error, fs::{self, create_dir_all, exists, read_to_string, File}, i32, io::{self, Write}, @@ -27,7 +26,7 @@ use std::{ Arc, }, }; -use tokio::time::{sleep, Duration}; +use tokio::time::Duration; use tokio::{ sync::{Mutex, OnceCell, Semaphore}, time::Instant, @@ -36,11 +35,6 @@ use tokio::{ static FIX_LOCK: OnceCell> = OnceCell::const_new(); static FIX_TIMEOUT: u64 = 28800; -async fn get_fixes_from_llm(llm: &LLM, prompt: &String, n: i32) -> Result<(Vec, u32, u32), Box> { - let answers = llm.get_answer(&prompt, n as u8, false).await; - answers -} - fn compiler_error_parser_from_json(work_path: &Path) -> Vec { let file_path = work_path.join("error_output.json"); let file_content = read_to_string(&file_path).unwrap(); @@ -144,7 +138,7 @@ async fn compilation_fix_assistant_for_an_error( test_name: String, insert_kind: InsertKind, common: &Vec, -) -> Result<(Vec, TestCode, u32, u32), ()> { +) -> io::Result<(Vec, TestCode, u32, u32)> { // let mut local_error_group: Vec = Vec::new(); // local_error_group.push(error_message.clone()); let mut iterative_time = 0; @@ -176,35 +170,15 @@ async fn compilation_fix_assistant_for_an_error( // file.write_all(final_prompt.as_bytes()).unwrap(); let lock = FIX_LOCK.get_or_init(|| async { Mutex::new(()) }).await; - let mut retry = 0; let mut completion_tokens = 0; let mut prompt_tokens = 0; while iterative_time < max_time_to_iterative { iterative_time += 1; - let request_choices_result = get_fixes_from_llm(llm, &final_prompt, 1).await; - // let request_choices_result = Err(()); - // let request_choices_result = read_to_string(work_path.join("fix_request.txt")); - let mut request_choices: Vec = Vec::new(); - if let Ok((s, usage_completion, usage_prompt)) = request_choices_result { - completion_tokens += usage_completion; - prompt_tokens += usage_prompt; - request_choices = s; - // request_choices = s.split("ChangeLog:").map(|s| s.to_string()).collect(); - // request_choices.remove(0); - } else { - if retry < 3 { - retry += 1; - max_time_to_iterative += 1; - let random_secs = { - let mut rng = rand::rng(); - rng.random_range(10..=30) - }; - sleep(Duration::from_secs(random_secs)).await; - } else { - return Err(()); - } - } + let (request_choices, usage_completion, usage_prompt) = + llm.get_answer(&final_prompt, 1, false).await?; + completion_tokens += usage_completion; + prompt_tokens += usage_prompt; // for request_choice in request_choices.iter_mut() { // *request_choice = "ChangeLog:".to_string() + request_choice; // } @@ -323,9 +297,9 @@ async fn compilation_fix_assistant_for_one_fn( project_dir: PathBuf, work_path: PathBuf, test_gen_info: TestGenInfo, -) -> TestGenInfo { +) -> io::Result { let start_time = Instant::now(); - let timeout = Duration::from_secs(FIX_TIMEOUT); // 设置超时时间为 30 秒 + let timeout = Duration::from_secs(FIX_TIMEOUT); // Existing between-round limit; not a subprocess deadline. let mut test_gen_info = test_gen_info; let file_rela = test_gen_info.get_file(); @@ -348,8 +322,6 @@ async fn compilation_fix_assistant_for_one_fn( common.push("".to_string()); for chain_test in answer.get_tests_mut().iter_mut() { for (num, test_code) in chain_test.codes.iter_mut().enumerate() { - let mut network_error = false; - let mut network_error_times = 0; if !chain_test.can_compile[num].is_ok() && !chain_test.repaired[num] { let sig = format!("fn test_{}_{:02}()", fn_name, id); let mut fn_code = vec!["#[test]".to_string(), TIMEOUT_DERIVE.to_string()]; @@ -403,43 +375,21 @@ async fn compilation_fix_assistant_for_one_fn( // } // already_rng.push(random_num); let random_error = compile_error_set.get(random_num).unwrap(); - let once_fix_result = compilation_fix_assistant_for_an_error( - llm, - &random_error, - &compile_error_set, - &project_dir, - &work_path, - &file_path, - &test_file_content, - sig.clone(), - test_name.clone(), - insert_kind, - &common, - ) - .await; - if let Ok((new_error_set, new_error_content, usage_completion, usage_prompt)) = once_fix_result { - compile_error_set = new_error_set; - test_file_content = new_error_content; - completion_tokens += usage_completion; - prompt_tokens += usage_prompt; - // } - if compile_error_set.len() == 0 { - break; - } - } else { - if network_error_times < 3 { - network_error_times += 1; - initial_error_num += 1; - } else { - network_error = true; - break; - } + let (new_error_set, new_error_content, usage_completion, usage_prompt) = + compilation_fix_assistant_for_an_error( + llm, random_error, &compile_error_set, &project_dir, &work_path, + &file_path, &test_file_content, sig.clone(), test_name.clone(), + insert_kind, &common, + ).await.map_err(|error| io::Error::other(format!( + "Repair model request failed for {name}: {error}" + )))?; + compile_error_set = new_error_set; + test_file_content = new_error_content; + completion_tokens += usage_completion; + prompt_tokens += usage_prompt; + if compile_error_set.is_empty() { + break; } - // if new_error_set.len() < compile_error_set.len() { - } - if network_error { - error!("Fix failed for {}: network error", fn_name); - continue; } let sig = format!("fn test_{}_{:02}()", fn_name, id); let mut fn_code = vec!["#[test]".to_string(), TIMEOUT_DERIVE.to_string()]; @@ -479,7 +429,7 @@ async fn compilation_fix_assistant_for_one_fn( drop(guard); if start_time.elapsed() >= timeout { warn!("Fix time out for {}", fn_name); - return test_gen_info; + return Ok(test_gen_info); } } } @@ -488,7 +438,7 @@ async fn compilation_fix_assistant_for_one_fn( answer.set_prompt_tokens(prompt_tokens); } } - test_gen_info + Ok(test_gen_info) } fn merge_common_in_code(test_gen_infos: &mut Vec) { @@ -622,7 +572,7 @@ pub async fn llm_fix( work_path, test_gen_info, ) - .await; + .await?; let json_path = parent_dir_clone.join(encoded_name + ".json"); test_gen_info.dump_json(&json_path); let completed = counter_clone.fetch_add(1, Ordering::Relaxed) + 1; diff --git a/utgen/tests/cli_config.rs b/utgen/tests/cli_config.rs index 9b0995ba8d..d185bd2ccc 100644 --- a/utgen/tests/cli_config.rs +++ b/utgen/tests/cli_config.rs @@ -93,6 +93,30 @@ fn tasks_default_to_four_and_reject_zero_before_loading_configuration() { } } +#[test] +fn request_timeout_rejects_zero_before_loading_configuration() { + let fixture = Fixture::new(); + for command in ["gen", "fix"] { + let help = fixture + .command() + .args([command, "--help"]) + .output() + .unwrap(); + assert!(String::from_utf8_lossy(&help.stdout).contains("[default: 180]")); + let output = fixture + .command() + .args([command, "-p", ".", "--request-timeout", "0"]) + .env("PALM_CONFIG", "missing.json") + .output() + .unwrap(); + assert!(!output.status.success()); + let error = String::from_utf8_lossy(&output.stderr); + assert!(error.contains("--request-timeout"), "{error}"); + assert!(!error.contains("Cannot read"), "{error}"); + assert!(!fixture.0.join("utgen").exists()); + } +} + #[test] fn analyze_validates_the_crate_without_loading_model_configuration() { let fixture = Fixture::new();