From a4bf8b631d9928ce8516b5151cd7d329a73104f3 Mon Sep 17 00:00:00 2001 From: cppbear Date: Thu, 8 Oct 2026 13:23:17 +0800 Subject: [PATCH] feat: scope model runs by function and shared request budget --- .gitignore | 1 + docs/minimal-pipeline.md | 6 ++ docs/palm-rust-unit-test-generation.md | 6 +- examples/README.md | 18 ++++++ examples/bytes-smoke-2.txt | 2 + examples/bytes-smoke-8.txt | 8 +++ scripts/check_bytes.py | 8 +++ scripts/check_tasks.py | 83 +++++++++++++++++++++++++- utgen/README.md | 14 ++++- utgen/src/gene/llm.rs | 74 ++++++++++++++++++++++- utgen/src/gene/mod.rs | 6 +- utgen/src/lib.rs | 2 + utgen/src/main.rs | 68 +++++++++++++++++---- utgen/src/run/llm_fix.rs | 79 +++++++++--------------- utgen/src/run/run.rs | 17 +++--- utgen/src/selection.rs | 52 ++++++++++++++++ utgen/tests/cli_config.rs | 51 ++++++++++++++++ 17 files changed, 420 insertions(+), 75 deletions(-) create mode 100644 examples/bytes-smoke-2.txt create mode 100644 examples/bytes-smoke-8.txt create mode 100644 utgen/src/selection.rs diff --git a/.gitignore b/.gitignore index 3827a9d94c..fd59d011d5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ .vscode/ target/ *.txt +!/examples/bytes-smoke-*.txt api.json diff --git a/docs/minimal-pipeline.md b/docs/minimal-pipeline.md index a3e079312e..acfb59b919 100644 --- a/docs/minimal-pipeline.md +++ b/docs/minimal-pipeline.md @@ -68,3 +68,9 @@ The same check covers direct and input/prefix/oracle integration generation, rem ## 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. + +## Function Selection and Request Budgets + +`scripts/check_tasks.py` generates and repairs just two selected functions while other saved candidates remain in the same directories. It checks actual request counts, selected-only statistics, untouched unselected candidate files, and invocation records. Duplicate names and blank lines are accepted. Separate CLI tests reject empty/unknown lists and a zero request limit before configuration or target changes. + +With four function workers and a two-attempt limit, both generation and repair must send exactly two requests, return failure when more work needs a request, restore sources, and skip statistics. Local HTTP unit tests additionally cover retries consuming the limit, concurrent clones sharing it, and success on the final allowed attempt. None of these checks uses real model credentials. The bytes analysis check validates the checked-in two- and eight-function lists against the complete analysis index for a later model trial. diff --git a/docs/palm-rust-unit-test-generation.md b/docs/palm-rust-unit-test-generation.md index b1dab57c8f..4aa4ee9e8e 100644 --- a/docs/palm-rust-unit-test-generation.md +++ b/docs/palm-rust-unit-test-generation.md @@ -274,6 +274,8 @@ utgen fix --project-dir [--work-dir ...] [--tasks ] | `-p, --project-dir` | 项目根目录。路径会相对于当前目录解析并 canonicalize。 | | `-w, --work-dir` | 工作目录，可重复传入或用逗号分隔，相对路径基于当前命令执行目录，默认等于 project dir。 | | `-t, --tasks` | 默认 `4`，必须大于 `0`。生成和修复均限制同时活跃的被测函数任务数；同一函数内的条件链、生成阶段或修复轮次依次执行。 | +| `--functions-file` | 可选 UTF-8 函数清单，每行一个 `brinfo/name_map.json` 中的完整名称。生成、修复及其统计均按清单筛选；空行忽略、重复项去重，不支持通配符。 | +| `--max-requests` | 可选的整次命令请求尝试次数上限，必须大于 `0`；默认不设总上限。所有并发任务和重试共享该额度。 | | `-i, --integration` | 生成集成测试，放到 `tests/` 目录，按分析结果的 `visible` 标志筛选函数并进行编译检查。 | | `-r, --requirement` | 在提示词中加入条件链约束。 | | `-c, --context` | 在提示词中加入 `focxt/.rs` 上下文。 | @@ -285,9 +287,11 @@ utgen fix --project-dir [--work-dir ...] [--tasks ] 任务异常会在等待所有任务结束后报告，并跳过后续覆盖率统计。修复成功时只删除本次创建的备份；任务失败时恢复源码并保留这些备份，已有备份不会被覆盖。候选本身编译不通过仍按测试结果处理。模型请求现在统一执行：`--request-timeout` 默认每次尝试 180 秒，包含响应体读取；网络错误、429（余额不足除外）和 5xx 最多尝试 3 次，其他 HTTP 错误或空回答直接报错。格式纠正和编译修复轮次单独计算，不再叠加网络重试。缺少 `usage` 的有效回答可以继续，但请求统计标记为不完整；生成和修复分别写入 `utgen/generation/gen-requests.json`、`fix-requests.json`。该期限不涵盖 Cargo 子进程，也不提供跨进程锁或强制终止恢复。 +小规模试跑可在 `gen`、`fix` 中分别传入同一份 `--functions-file`，不会自动沿用上次清单。空清单和未知名称在修改目标前报错；修复时缺少所选函数的候选结果也会报错。`--max-requests` 包含所有并发任务、生成阶段、格式纠正、修复轮次及网络重试的请求尝试，计数与额度检查使用同一个共享计数器。额度用完且仍需请求时，命令失败，等待已有请求收尾并恢复源码；恰好用完额度且工作全部完成则成功。请求报告记录模型、选项、去重后的函数清单、候选处理状态和用量；候选状态不等于测试通过或覆盖率统计完成。未选函数的已有文件保留，跨实验应使用新副本。准备好的 [2 函数和 8 函数清单及操作说明](../examples/README.md#prepare-a-small-model-trial)用于后续真实模型验收，目前不代表已取得真实模型结果。 + ### 8.2 LLM 配置 -`utgen/src/gene/llm.rs` 使用 `async-openai` 调用兼容 OpenAI Chat Completions 的接口。模型配置由 `utgen/src/config.rs` 在运行时加载；可以复制 `utgen/res/api.example.json` 并填写为： +`utgen/src/gene/llm.rs` 保留 `async-openai` 的请求／响应类型，通过 `reqwest` 调用兼容 OpenAI Chat Completions 的接口，统一管理重试。模型配置由 `utgen/src/config.rs` 在运行时加载；可以复制 `utgen/res/api.example.json` 并填写为： ```text utgen/res/api.json diff --git a/examples/README.md b/examples/README.md index 2ed4a2a3d0..96059a13e5 100644 --- a/examples/README.md +++ b/examples/README.md @@ -44,6 +44,23 @@ utgen fix -p "$palm_example_dir" `--requirement` enables generation for representative condition chains; `--context` includes focal context in the prompt. Both flags are off by default. In unit mode, candidates are stored in JSON and temporarily inserted into the target source for checks and execution. +### Prepare a Small Model Trial + +The checked-in [two-function list](bytes-smoke-2.txt) selects `Bytes::len` and `BytesMut::len`. The [eight-function list](bytes-smoke-8.txt) adds `is_empty`, `split_off`, and `truncate` for both types. They cover 2 and 16 representative condition chains respectively. The repeated namespace segments are exact compiler-index names. The analysis regression checks both lists without calling a model. + +For a future trial, agree on the model configuration and separate generation/repair attempt limits first. Set `palm_gen_limit` and `palm_fix_limit` to those positive integers, then use the prepared working copy: + +```sh +utgen gen -p "$palm_example_dir" --requirement --context \ + --functions-file "$palm_repo/examples/bytes-smoke-2.txt" \ + --max-requests "$palm_gen_limit" +utgen fix -p "$palm_example_dir" \ + --functions-file "$palm_repo/examples/bytes-smoke-2.txt" \ + --max-requests "$palm_fix_limit" +``` + +Both commands default to four active function tasks and a 180-second deadline per request attempt. Every retry consumes the command's attempt limit. Review the request reports, candidate compilation/pass counts, coverage, and restored sources before expanding to eight functions on another fresh copy. Keep configuration and generation mode fixed when comparing these runs. No real-model result is claimed by the deterministic checks; the lists and limits prepare that later validation. + ### Integration Test Mode Use a fresh working copy and repeat extraction and preprocessing before choosing this mode, since existing generation results are reused independently of these flags: @@ -64,6 +81,7 @@ Paths below are relative to the working copy: | `utgen/generation/prompt/` and `answer/` | Generation prompts and model responses. | | `utgen/generation/pre_fix/` | Candidate tests and compilation results before repair. | | `utgen/generation/llm_fix/` | Candidate tests and compilation results after repair. | +| `utgen/generation/gen-requests.json` and `fix-requests.json` | Model, invocation selection/options, attempts, limits, token reporting, and candidate-stage status. | | `utgen/result/` and `utgen/fixed_result/` | Per-function coverage and execution statistics. | `coverage.xml` and `coverage.json` are intermediate files and may be removed after parsing. The current implementation writes JSON statistics; it does not generate the `result.html` report mentioned in earlier instructions. See [utgen](../utgen/README.md) for CLI limitations. diff --git a/examples/bytes-smoke-2.txt b/examples/bytes-smoke-2.txt new file mode 100644 index 0000000000..e506b0fcfb --- /dev/null +++ b/examples/bytes-smoke-2.txt @@ -0,0 +1,2 @@ +bytes::bytes::bytes::Bytes::len +bytes::bytes_mut::bytes_mut::BytesMut::len diff --git a/examples/bytes-smoke-8.txt b/examples/bytes-smoke-8.txt new file mode 100644 index 0000000000..9c4febd52a --- /dev/null +++ b/examples/bytes-smoke-8.txt @@ -0,0 +1,8 @@ +bytes::bytes::bytes::Bytes::len +bytes::bytes::bytes::Bytes::is_empty +bytes::bytes::bytes::Bytes::split_off +bytes::bytes::bytes::Bytes::truncate +bytes::bytes_mut::bytes_mut::BytesMut::len +bytes::bytes_mut::bytes_mut::BytesMut::is_empty +bytes::bytes_mut::bytes_mut::BytesMut::split_off +bytes::bytes_mut::bytes_mut::BytesMut::truncate diff --git a/scripts/check_bytes.py b/scripts/check_bytes.py index 50f1726dc9..9c6d5a42b3 100644 --- a/scripts/check_bytes.py +++ b/scripts/check_bytes.py @@ -119,6 +119,14 @@ def run(root, label, command): elapsed, output = run(root, 'analyze', ['utgen', 'analyze', '-p', str(root)]) assert snapshot(root) == prepared, 'analysis modified prepared sources' names, infos = check_artifacts(root, 663) + # Keep the future model-smoke lists aligned with real compiler identities. + for count, chains in [(2, 2), (8, 16)]: + selected = (REPO / f'examples/bytes-smoke-{count}.txt').read_text().splitlines() + assert len(selected) == len(set(selected)) == count + assert set(selected) <= names.keys() + actual_chains = sum(json.loads((root / f'brinfo/brdata/{names[name]}.json').read_text()) + ['size']['min_set'] for name in selected) + assert actual_chains == chains, (count, actual_chains) # Verify representative default, generic-impl, and macro-generated methods. for name, info in infos.items(): context = (root / f"focxt/{info['encoded_name']}.rs").read_text() diff --git a/scripts/check_tasks.py b/scripts/check_tasks.py index 7b41011f85..1640a7c5a5 100644 --- a/scripts/check_tasks.py +++ b/scripts/check_tasks.py @@ -221,6 +221,9 @@ def check_queue(): 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 + assert report['model'] == 'fixture', report + assert report['invocation']['command'] == command[1], report + assert report['invocation']['candidate_status'] == ('failed' if fail else 'completed'), report return state.copy() def reset(seed=None): @@ -232,9 +235,11 @@ def reset(seed=None): def command(mode, n, *extra): return ['utgen', mode, '-p', str(root), *([] if n == 4 else ['--tasks', str(n)]), *extra] - def results(directory, compiled): + def results(directory, compiled, names=None): data = [json.loads(p.read_text()) for p in (root / 'utgen' / directory).glob('*.json')] - assert len(data) == count, (directory, len(data)) + assert len(data) == (len(names) if names is not None else count), (directory, len(data)) + if names is not None: + assert {item['function_name'] for item in data} == set(names), data assert all(item['tests_compiled'] == compiled for item in data), data if compiled: assert all(item['tests_run'] == 1 and item['tests_passed'] == 1 @@ -261,6 +266,80 @@ def results(directory, compiled): assert measured['total'] == count and measured['peak'] == n, measured results('fixed_result', 1) assert not (root / 'src/lib.rs.bak').exists() + + # Existing results for other functions must not enter this invocation. + nmap = json.loads((root / 'brinfo/name_map.json').read_text()) + selected = sorted(nmap)[:2] + functions_file = work / 'functions.txt' + functions_file.write_text('\n' + '\n'.join(selected + [selected[0]]) + '\n') + selection_args = ['--functions-file', str(functions_file), '--max-requests', '2'] + reset(seed) + pre_dir = root / 'utgen/generation/pre_fix' + for name in selected: + (pre_dir / (nmap[name] + '.json')).unlink() + untouched = {p: p.read_bytes() for p in pre_dir.glob('*.json')} + measured = run('gen-selected', command('gen', 2, *selection_args), limit=2) + assert measured['total'] == 2 + results('result', 0, selected) + for path, content in untouched.items(): + assert path.read_bytes() == content, path + report = json.loads((root / 'utgen/generation/gen-requests.json').read_text()) + assert report['invocation']['functions'] == selected, report + assert report['max_requests'] == 2 and not report['budget_exhausted'], report + + missing = pre_dir / (nmap[selected[0]] + '.json') + candidate = missing.read_bytes() + missing.unlink() + measured = run('fix-missing-selected', command('fix', 2, *selection_args), fail=True) + assert measured['total'] == 0 + assert not (root / 'utgen/generation/llm_fix').exists() + assert not (root / 'src/lib.rs.bak').exists() + missing.write_bytes(candidate) + + fixed_dir = root / 'utgen/generation/llm_fix' + fixed_dir.mkdir() + other_name = sorted(nmap)[2] + other = fixed_dir / (nmap[other_name] + '.json') + shutil.copyfile(pre_dir / other.name, other) + old_other = other.read_bytes() + before_fix = {p: p.read_bytes() for p in pre_dir.glob('*.json')} + measured = run('fix-selected', command('fix', 2, *selection_args), limit=2) + assert measured['total'] == 2 + results('fixed_result', 1, selected) + assert len(list(fixed_dir.glob('*.json'))) == 3 + assert other.read_bytes() == old_other + for path, content in before_fix.items(): + assert path.read_bytes() == content, path + report = json.loads((root / 'utgen/generation/fix-requests.json').read_text()) + assert report['invocation']['functions'] == selected, report + assert report['max_requests'] == 2 and not report['budget_exhausted'], report + measured = run('fix-selected-cached', command('fix', 2, *selection_args)) + assert measured['total'] == 0 + report = json.loads((root / 'utgen/generation/fix-requests.json').read_text()) + assert report['attempts'] == 0 and report['invocation']['functions'] == selected, report + assert other.read_bytes() == old_other + + reset(seed) + for name in selected: + (pre_dir / (nmap[name] + '.json')).unlink() + measured = run('integration-selected', command('gen', 2, '--integration', *selection_args), + limit=2, integration=True) + assert measured['total'] == 2 + results('result', 1, selected) + + # More workers than requests: every clone shares the same hard cap. + for mode in ['gen', 'fix']: + reset(seed if mode == 'fix' else None) + measured = run(f'{mode}-budget', command(mode, 4, '--max-requests', '2'), limit=2, fail=True) + assert measured['total'] == 2, measured + report = json.loads((root / f'utgen/generation/{mode}-requests.json').read_text()) + assert report['attempts'] == 2 and report['budget_exhausted'], report + assert not (root / 'utgen' / ('result' if mode == 'gen' else 'fixed_result')).exists() + if mode == 'fix': + backup = root / 'src/lib.rs.bak' + assert backup.read_text() == original + backup.unlink() + for oracle in [False, True]: reset() preserved = {} diff --git a/utgen/README.md b/utgen/README.md index d1c99a0e34..6fe3ccb12c 100644 --- a/utgen/README.md +++ b/utgen/README.md @@ -146,6 +146,8 @@ utgen gen -p --requirement --context | `-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. | +| `--functions-file` | Optional path to exact full analysis-index names, one per line; also available for `fix`. | +| `--max-requests` | Optional positive request-attempt limit for the whole command, including retries; also available for `fix`. Default: no overall limit. | 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. @@ -153,6 +155,16 @@ Each generation job handles one focal function, including its condition chains a Generation may append an `ntest` dependency to the target's Cargo.toml. Existing `utgen/generation/pre_fix/.json` results are skipped, so use a fresh target copy for a different model, prompt, or generation mode. +### Limited function runs + +Copy names exactly from the keys of `brinfo/name_map.json` into a UTF-8 text file, one per line. `--functions-file` resolves against the current shell directory. Surrounding whitespace and blank lines are ignored, duplicate names are processed once, and there is no wildcard or comment syntax. Missing files, empty lists, and unknown names fail before configuration loading or target changes. Omitting the option preserves the normal all-function behavior; integration generation still applies its visibility filter. + +The same selection applies to generation, repair preparation, and per-function coverage/execution statistics. Pass it to both `gen` and `fix`; it is not inherited from an earlier command. Repair reports a selected function with no saved candidate as an error. Unselected candidate and result files remain unchanged. Output directories can therefore contain older results outside the current selection: the invocation record identifies this run's selection. Use separate fresh copies when comparing experiments, including different integration-test selections. + +`--max-requests N` counts every model request attempt across all workers, chains, generation stages, format retries, and repair rounds. Network retries consume the same limit. Checking and incrementing the count is one shared operation, so concurrent workers cannot overshoot it. Once another attempt is needed after the limit, the command fails, joins all workers, restores sources, and skips subsequent statistics; already sent requests may finish and successful candidates may be saved. Reaching exactly N attempts while completing all work is successful. A later command has its own budget. This is an attempt cap, not a token, monetary, or wall-clock budget. + +The request report also records `model`, `request_timeout_seconds`, `max_requests`, `budget_exhausted`, and an `invocation` object with command, directories, options, deduplicated function names (`null` means no explicit selection), `candidate_status`, and an error when present. `candidate_status` describes generation/repair worker completion before statistics, not passing tests or completed coverage. The report does not serialize the API key or endpoint. Each command replaces its own previous report, including zero-request cached runs; preflight failures before work do not create a new report. Existing candidate and statistics JSON provide the per-function outcomes. + ### Repair ```sh @@ -176,7 +188,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/generation/gen-requests.json` and `fix-requests.json` | Model/options/function selection, request limits and counts, reported token totals, usage completeness, and candidate-stage status. | | `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/llm.rs b/utgen/src/gene/llm.rs index 1d062963d6..f88dfdec4c 100644 --- a/utgen/src/gene/llm.rs +++ b/utgen/src/gene/llm.rs @@ -7,6 +7,7 @@ use async_openai::types::{ use serde::Serialize; use std::{ fs, io, + num::NonZeroU64, path::Path, sync::{Arc, Mutex}, time::Duration, @@ -22,6 +23,7 @@ struct RequestStats { responses_without_usage: u64, reported_completion_tokens: u64, reported_prompt_tokens: u64, + budget_exhausted: bool, } struct RequestError { @@ -46,6 +48,7 @@ pub struct LLM { config: LlmConfig, client: reqwest::Client, request_timeout: Duration, + max_requests: Option, stats: Arc>, } @@ -55,6 +58,7 @@ impl LLM { config, client: reqwest::Client::new(), request_timeout: Duration::from_secs(180), + max_requests: None, stats: Arc::new(Mutex::new(RequestStats::default())), } } @@ -68,12 +72,21 @@ impl LLM { self.stats.lock().unwrap().attempts } + pub fn with_max_requests(mut self, max_requests: Option) -> Self { + self.max_requests = max_requests; + self + } + /// Persist reported usage, including an explicit incomplete-usage marker. - pub fn write_request_summary(&self, path: &Path) -> io::Result<()> { + pub fn write_request_summary(&self, path: &Path, invocation: serde_json::Value) -> 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(); + summary["model"] = self.config.model.clone().into(); + summary["request_timeout_seconds"] = self.request_timeout.as_secs_f64().into(); + summary["max_requests"] = serde_json::to_value(self.max_requests)?; + summary["invocation"] = invocation; fs::create_dir_all(path.parent().unwrap())?; fs::write(path, serde_json::to_vec_pretty(&summary)?) } @@ -151,7 +164,18 @@ impl LLM { .build() .map_err(io::Error::other)?; for attempt in 1..=MAX_ATTEMPTS { - self.stats.lock().unwrap().attempts += 1; + { + let mut stats = self.stats.lock().unwrap(); + if let Some(max) = self.max_requests { + if stats.attempts >= max.get() { + stats.budget_exhausted = true; + return Err(io::Error::other(format!( + "Model request budget exhausted (--max-requests {max})" + ))); + } + } + stats.attempts += 1; + } let response = match timeout(self.request_timeout, self.send_once(&request)).await { Ok(result) => result, Err(_) => Err(RequestError { @@ -394,7 +418,7 @@ mod tests { 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(); + llm.write_request_summary(&path, serde_json::json!({})).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); @@ -473,6 +497,50 @@ mod tests { } } + #[tokio::test] + async fn request_budget_is_shared_across_concurrent_workers_and_retries() { + let (llm, server) = mock_replies(vec![reply(503), reply(200), reply(200)]).await; + let llm = llm.with_max_requests(NonZeroU64::new(3)); + let mut workers = tokio::task::JoinSet::new(); + for _ in 0..9 { + let llm = llm.clone(); + workers.spawn(async move { llm.get_answer("test", 1, false).await }); + } + let mut successes = 0; + while let Some(result) = workers.join_next().await { + match result.unwrap() { + Ok(_) => successes += 1, + Err(error) => assert!(error.to_string().contains("budget exhausted")), + } + } + assert_eq!(successes, 2); + assert_eq!(server.await.unwrap().len(), 3); + let stats = llm.stats.lock().unwrap(); + assert_eq!((stats.attempts, stats.failed_attempts), (3, 1)); + assert!(stats.budget_exhausted); + } + + #[tokio::test] + async fn request_budget_can_stop_a_single_workers_retries() { + let (llm, server) = mock_replies(vec![reply(503), reply(503)]).await; + let llm = llm.with_max_requests(NonZeroU64::new(2)); + let error = llm.get_answer("test", 1, false).await.unwrap_err(); + assert!(error.to_string().contains("budget exhausted")); + assert_eq!(server.await.unwrap().len(), 2); + let stats = llm.stats.lock().unwrap(); + assert_eq!((stats.attempts, stats.failed_attempts), (2, 2)); + assert!(stats.budget_exhausted); + } + + #[tokio::test] + async fn completing_on_the_last_allowed_attempt_is_successful() { + let (llm, server) = mock_replies(vec![reply(503), reply(200)]).await; + let llm = llm.with_max_requests(NonZeroU64::new(2)); + assert!(llm.get_answer("test", 1, false).await.is_ok()); + assert_eq!(server.await.unwrap().len(), 2); + assert!(!llm.stats.lock().unwrap().budget_exhausted); + } + #[tokio::test] async fn attempt_timeout_includes_reading_the_body_and_stops_after_three() { let delayed = (200, completion().to_string(), Duration::from_millis(150)); diff --git a/utgen/src/gene/mod.rs b/utgen/src/gene/mod.rs index 9919cb8f0c..45eefab175 100644 --- a/utgen/src/gene/mod.rs +++ b/utgen/src/gene/mod.rs @@ -6,7 +6,7 @@ mod test_ext; use crate::run::add_ntest_dependency; use crate::types::{BrData, RfocxtNameInformation, TestGenInfo}; -use crate::validate_analysis; +use crate::{FunctionSelection, validate_analysis}; use cot::{gen_input_range, gen_oracle, gen_prefix, gen_test}; use generation::{check_integration, check_unit, generation_tests}; use log::{error, info, warn}; @@ -60,6 +60,7 @@ pub async fn gen_tests_project( llm: &LLM, project_dir: &Path, work_dir: &Path, + functions: &FunctionSelection, tasks: usize, integration: bool, requirement: bool, @@ -91,6 +92,9 @@ pub async fn gen_tests_project( if brdata_path.is_file() { let brdata: BrData = serde_json::from_str(&fs::read_to_string(&brdata_path).unwrap()).unwrap(); + if !functions.contains(&brdata.name) { + continue; + } // if brdata.size.min_set < 2 { // info!("{} has less than 2 condition chains in min_set", brdata.name); // continue; diff --git a/utgen/src/lib.rs b/utgen/src/lib.rs index 0cd0e43779..2cc1d799a8 100644 --- a/utgen/src/lib.rs +++ b/utgen/src/lib.rs @@ -2,12 +2,14 @@ mod analysis; mod config; mod gene; mod run; +mod selection; mod types; mod utils; pub use gene::gen_tests_project; pub use config::{ConfigError, LlmConfig}; pub use gene::LLM; +pub use selection::FunctionSelection; pub use analysis::{analyze_project, validate_analysis, validate_repair}; pub use run::{collect_coverage, gen_test_rate, llm_fix, gen_test_rate_aggregated}; pub use utils::{comment_out_tests, rename_tests_to_bak}; diff --git a/utgen/src/main.rs b/utgen/src/main.rs index 58db9416c4..1e200b83ae 100644 --- a/utgen/src/main.rs +++ b/utgen/src/main.rs @@ -7,7 +7,9 @@ use std::{ num::{NonZeroU64, NonZeroUsize}, time::Duration, }; -use utgen::{LLM, LlmConfig, analyze_project, collect_coverage, validate_repair}; +use utgen::{ + FunctionSelection, 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}; /// Generate unit tests for a project @@ -48,6 +50,12 @@ enum Command { /// Timeout in seconds for each model request attempt (at most three attempts) #[arg(long, default_value = "180", value_name = "SECONDS")] request_timeout: NonZeroU64, + /// Exact full function names, one per line (blank lines ignored) + #[arg(long, value_name = "PATH")] + functions_file: Option, + /// Maximum model request attempts for the whole invocation, including retries + #[arg(long, value_name = "N")] + max_requests: Option, /// Whether to generate integration tests #[arg(short, long)] integration: bool, @@ -71,6 +79,12 @@ enum Command { /// Timeout in seconds for each model request attempt (at most three attempts) #[arg(long, default_value = "180", value_name = "SECONDS")] request_timeout: NonZeroU64, + /// Exact full function names, one per line (blank lines ignored) + #[arg(long, value_name = "PATH")] + functions_file: Option, + /// Maximum model request attempts for the whole invocation, including retries + #[arg(long, value_name = "N")] + max_requests: Option, }, } @@ -172,14 +186,18 @@ async fn main() -> Result<(), Box> { options, tasks, request_timeout, + functions_file, + max_requests, integration, requirement, context, oracle, } => { - 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)?; + let functions = FunctionSelection::load(functions_file.as_deref(), &work_dirs)?; + let llm = LLM::new(LlmConfig::load(cli.config.as_deref())?) + .with_request_timeout(Duration::from_secs(request_timeout.get())) + .with_max_requests(max_requests); info!( "Generating tests for project at {}, with work directory(s) {:?}", project_dir.display(), @@ -192,6 +210,7 @@ async fn main() -> Result<(), Box> { &llm, &project_dir, &work_dir, + &functions, tasks.get(), integration, requirement, @@ -204,7 +223,17 @@ async fn main() -> Result<(), Box> { } .await; let recorded = if generated.is_ok() || llm.request_count() > 0 { - llm.write_request_summary(&project_dir.join("utgen/generation/gen-requests.json")) + llm.write_request_summary( + &project_dir.join("utgen/generation/gen-requests.json"), + serde_json::json!({ + "command": "gen", "functions": functions.names(), + "project_dir": project_dir, "work_dirs": work_dirs, + "tasks": tasks, "integration": integration, "requirement": requirement, + "context": context, "oracle": oracle, + "candidate_status": if generated.is_ok() { "completed" } else { "failed" }, + "error": generated.as_ref().err().map(ToString::to_string), + }), + ) } else { Ok(()) }; @@ -212,7 +241,7 @@ async fn main() -> Result<(), Box> { 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); + gen_test_rate(&project_dir, &work_dir, integration, true, &functions); } } // fix unit tests @@ -220,10 +249,14 @@ async fn main() -> Result<(), Box> { options, tasks, request_timeout, + functions_file, + max_requests, } => { - 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)?; + let functions = FunctionSelection::load(functions_file.as_deref(), &work_dirs)?; + let llm = LLM::new(LlmConfig::load(cli.config.as_deref())?) + .with_request_timeout(Duration::from_secs(request_timeout.get())) + .with_max_requests(max_requests); info!( "Fixing tests for project at {}, with work directory(s) {:?}", project_dir.display(), @@ -232,13 +265,28 @@ async fn main() -> Result<(), Box> { 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?; + llm_fix( + &llm, + project_dir.clone(), + work_dir.clone(), + &functions, + 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")) + llm.write_request_summary( + &project_dir.join("utgen/generation/fix-requests.json"), + serde_json::json!({ + "command": "fix", "functions": functions.names(), + "project_dir": project_dir, "work_dirs": work_dirs, "tasks": tasks, + "candidate_status": if repaired.is_ok() { "completed" } else { "failed" }, + "error": repaired.as_ref().err().map(ToString::to_string), + }), + ) } else { Ok(()) }; @@ -246,7 +294,7 @@ async fn main() -> Result<(), Box> { recorded?; //gen_test_rate_aggregated(&project_dir, false); for work_dir in work_dirs.iter() { - gen_test_rate(&project_dir, &work_dir, false, false); + gen_test_rate(&project_dir, &work_dir, false, false, &functions); } } } diff --git a/utgen/src/run/llm_fix.rs b/utgen/src/run/llm_fix.rs index abf0474518..6a8f868a6a 100644 --- a/utgen/src/run/llm_fix.rs +++ b/utgen/src/run/llm_fix.rs @@ -5,6 +5,7 @@ use super::{ TIMEOUT_DERIVE, }; use crate::{ + FunctionSelection, gene::LLM, types::{InsertKind, TestGenInfo}, utils::{ @@ -17,7 +18,7 @@ use serde::Deserialize; use std::{ cmp::min, collections::{BTreeSet, HashMap}, - fs::{self, create_dir_all, exists, read_to_string, File}, + fs::{self, create_dir_all, read_to_string, File}, i32, io::{self, Write}, path::{Path, PathBuf}, @@ -463,6 +464,7 @@ pub async fn llm_fix( llm: &LLM, project_path: PathBuf, work_path: PathBuf, + functions: &FunctionSelection, tasks: usize, ) -> io::Result<()> { if tasks == 0 || tasks > Semaphore::MAX_PERMITS { @@ -470,66 +472,43 @@ pub async fn llm_fix( "--tasks is outside the supported positive range", )); } - let mut test_gen_infos; - let pre_dir = project_path.join("utgen/generation/pre_fix"); let parent_dir = project_path.join("utgen/generation/llm_fix"); let map_path = work_path.join("brinfo/name_map.json"); let nmap: HashMap = serde_json::from_str(&fs::read_to_string(&map_path).unwrap()).unwrap(); - if exists(&parent_dir).unwrap() { - for entry in fs::read_dir(pre_dir).unwrap() { - let entry = entry.unwrap(); - let source_path = entry.path(); - if source_path.is_file() { - if let Some(file_name) = source_path.file_name() { - let target_path = parent_dir.join(file_name); - if !target_path.exists() { - fs::copy(&source_path, &target_path).unwrap(); - } - } - } - } - test_gen_infos = get_test_gen_infos(&project_path, false); - merge_common_in_code(&mut test_gen_infos); - for test_gen_info in test_gen_infos.iter() { - let file_rela = test_gen_info.get_file(); - let file_path = project_path.join(&file_rela); - if (project_path == work_path && !file_rela.starts_with("src")) - || (project_path != work_path && !file_path.starts_with(&work_path)) - { - continue; - } - let json_path = - parent_dir.join(nmap.get(test_gen_info.get_name()).unwrap().to_owned() + ".json"); - test_gen_info.dump_json(&json_path); - } - } else { - create_dir_all(&parent_dir).unwrap(); - test_gen_infos = get_test_gen_infos(&project_path, true); - merge_common_in_code(&mut test_gen_infos); - for test_gen_info in test_gen_infos.iter() { - let file_rela = test_gen_info.get_file(); - let file_path = project_path.join(&file_rela); - if (project_path == work_path && !file_rela.starts_with("src")) - || (project_path != work_path && !file_path.starts_with(&work_path)) - { - continue; - } - let json_path = - parent_dir.join(nmap.get(test_gen_info.get_name()).unwrap().to_owned() + ".json"); - test_gen_info.dump_json(&json_path); - } - } - + // Prefer saved repair progress, then add candidates not yet copied from generation. + let mut test_gen_infos = get_test_gen_infos(&project_path, false); + let saved: BTreeSet<_> = test_gen_infos.iter() + .map(|info| info.get_name().to_owned()).collect(); + test_gen_infos.extend( + get_test_gen_infos(&project_path, true).into_iter() + .filter(|info| !saved.contains(info.get_name())), + ); test_gen_infos.retain(|info| { let relative = info.get_file(); - if project_path == work_path { + let in_work_dir = if project_path == work_path { relative.starts_with("src") } else { project_path.join(relative).starts_with(&work_path) - } + }; + in_work_dir && functions.contains(info.get_name()) }); + if let Some(names) = functions.names() { + for name in names.iter().filter(|name| nmap.contains_key(*name)) { + if !test_gen_infos.iter().any(|info| info.get_name() == name) { + return Err(io::Error::other(format!( + "No generated candidate for selected function {name}; run gen first" + ))); + } + } + } + merge_common_in_code(&mut test_gen_infos); + create_dir_all(&parent_dir)?; + for info in &test_gen_infos { + let json_path = parent_dir.join(nmap.get(info.get_name()).unwrap().to_owned() + ".json"); + info.dump_json(&json_path); + } // Back up each source once before starting workers. Reject stale backups // rather than reusing them, and only clean up files owned by this command. let sources: BTreeSet<_> = test_gen_infos diff --git a/utgen/src/run/run.rs b/utgen/src/run/run.rs index 221482adc9..5077290176 100644 --- a/utgen/src/run/run.rs +++ b/utgen/src/run/run.rs @@ -4,6 +4,7 @@ use super::{ prepare::get_test_gen_infos, }; use crate::{ + FunctionSelection, run::TIMEOUT_DERIVE, types::TestGenInfo, utils::{backup_file, delete_backup, insert_test, restore_file, target_clean}, @@ -288,13 +289,16 @@ pub fn gen_test_rate_aggregated(project_dir: &Path, is_pre: bool) { // keep `touched_files` somewhere, serialize if needed } -pub fn gen_test_rate(project_dir: &Path, work_dir: &Path, integration: bool, is_pre: bool) { - // add_ntest_dependency(work_dir); - let map_path = work_dir.join("brinfo/name_map.json"); - let nmap: HashMap = - serde_json::from_str(&fs::read_to_string(&map_path).unwrap()).unwrap(); +pub fn gen_test_rate( + project_dir: &Path, + work_dir: &Path, + integration: bool, + is_pre: bool, + functions: &FunctionSelection, +) { + let mut test_gen_infos = get_test_gen_infos(project_dir, is_pre); + test_gen_infos.retain(|info| functions.contains(info.get_name())); if !integration { - let mut test_gen_infos = get_test_gen_infos(project_dir, is_pre); let mut test_rate_infos: Vec = Vec::new(); gen_coverage_and_pass_rate(project_dir, work_dir, &test_gen_infos, &mut test_rate_infos); dump_result_for_test_rate_infos( @@ -309,7 +313,6 @@ pub fn gen_test_rate(project_dir: &Path, work_dir: &Path, integration: bool, is_ gen_coverage_rate_for_original_tests(project_dir, work_dir, &test_gen_infos); } } else { - let mut test_gen_infos = get_test_gen_infos(project_dir, is_pre); let integration_infos = gen_integration(&test_gen_infos, project_dir, work_dir); let mut test_rate_infos: Vec = Vec::new(); let test_type = TestType::CoverageRate; diff --git a/utgen/src/selection.rs b/utgen/src/selection.rs new file mode 100644 index 0000000000..865a13a5b5 --- /dev/null +++ b/utgen/src/selection.rs @@ -0,0 +1,52 @@ +use std::{ + collections::{BTreeSet, HashMap}, + fs, io, + path::{Path, PathBuf}, +}; + +/// Exact analysis-index names. No file means all functions. +#[derive(Default)] +pub struct FunctionSelection { + names: Option>, +} + +impl FunctionSelection { + pub fn load(path: Option<&Path>, work_dirs: &[PathBuf]) -> io::Result { + let Some(path) = path else { + return Ok(Self::default()); + }; + let names: BTreeSet = fs::read_to_string(path)? + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .map(str::to_owned) + .collect(); + if names.is_empty() { + return Err(io::Error::other( + "--functions-file contains no function names", + )); + } + let mut available = BTreeSet::new(); + for work_dir in work_dirs { + let map: HashMap = + serde_json::from_slice(&fs::read(work_dir.join("brinfo/name_map.json"))?)?; + available.extend(map.into_keys()); + } + let unknown: Vec<_> = names.difference(&available).cloned().collect(); + if !unknown.is_empty() { + return Err(io::Error::other(format!( + "Unknown functions in --functions-file: {}", + unknown.join(", ") + ))); + } + Ok(Self { names: Some(names) }) + } + + pub fn contains(&self, name: &str) -> bool { + self.names.as_ref().is_none_or(|names| names.contains(name)) + } + + pub fn names(&self) -> Option<&BTreeSet> { + self.names.as_ref() + } +} diff --git a/utgen/tests/cli_config.rs b/utgen/tests/cli_config.rs index d185bd2ccc..04bdff25d6 100644 --- a/utgen/tests/cli_config.rs +++ b/utgen/tests/cli_config.rs @@ -117,6 +117,57 @@ fn request_timeout_rejects_zero_before_loading_configuration() { } } +#[test] +fn request_budget_rejects_zero_before_loading_configuration() { + let fixture = Fixture::new(); + for command in ["gen", "fix"] { + let output = fixture + .command() + .args([command, "-p", ".", "--max-requests", "0"]) + .env("PALM_CONFIG", "missing.json") + .output() + .unwrap(); + assert!(!output.status.success()); + let error = String::from_utf8_lossy(&output.stderr); + assert!(error.contains("--max-requests"), "{error}"); + assert!(!error.contains("Cannot read"), "{error}"); + assert!(!fixture.0.join("utgen").exists()); + } +} + +#[test] +fn invalid_function_lists_fail_before_configuration_or_target_changes() { + let fixture = Fixture::new(); + fs::create_dir(fixture.0.join("brinfo")).unwrap(); + fs::write( + fixture.0.join("brinfo/name_map.json"), + r#"{"fixture::known": "known"}"#, + ) + .unwrap(); + let manifest = fs::read(fixture.0.join("Cargo.toml")).unwrap(); + for (contents, message) in [ + ("\n \n", "contains no function names"), + ("fixture::known\nfixture::unknown\n", "Unknown functions"), + ("fixture::*\n", "Unknown functions"), + ] { + fs::write(fixture.0.join("functions.txt"), contents).unwrap(); + for command in ["gen", "fix"] { + let output = fixture + .command() + .args([command, "-p", ".", "--functions-file", "functions.txt"]) + .env("PALM_CONFIG", "missing.json") + .output() + .unwrap(); + assert!(!output.status.success()); + let error = String::from_utf8_lossy(&output.stderr); + assert!(error.contains(message), "{error}"); + assert!(!error.contains("Cannot read"), "{error}"); + assert_eq!(manifest, fs::read(fixture.0.join("Cargo.toml")).unwrap()); + assert!(!fixture.0.join("utgen").exists()); + } + } +} + #[test] fn analyze_validates_the_crate_without_loading_model_configuration() { let fixture = Fixture::new();