From 13b67580332dbaa58607b0e7ba6420e318b15b75 Mon Sep 17 00:00:00 2001 From: Per Larsen Date: Sun, 6 Sep 2026 18:04:20 -0700 Subject: [PATCH 1/3] postprocess: await model calls throughout the transform pipeline Use the native async OpenAI and Gemini clients and await trimming, response validation/retries, and comment transfer without changing prompts or cache keys. Keep file processing serial in this commit. Run the CLI in one event loop and close model transports before it exits, including on failure. Adapt existing tests with asyncio.run and exercise SDK request forwarding with fake clients; no live model calls are needed. --- c2rust-postprocess/postprocess/__init__.py | 52 ++++++----- c2rust-postprocess/postprocess/models/base.py | 6 +- .../postprocess/models/gemini.py | 10 +- c2rust-postprocess/postprocess/models/gpt.py | 11 ++- c2rust-postprocess/postprocess/models/mock.py | 2 +- .../postprocess/transforms/base.py | 18 ++-- .../postprocess/transforms/comments.py | 6 +- .../postprocess/transforms/trim.py | 4 +- c2rust-postprocess/tests/test_apply_file.py | 77 +++++++++------- .../tests/test_comments_transform.py | 87 ++++++++++-------- c2rust-postprocess/tests/test_models.py | 92 +++++++++++++++++++ .../tests/test_trim_transform.py | 15 +-- 12 files changed, 261 insertions(+), 119 deletions(-) create mode 100644 c2rust-postprocess/tests/test_models.py diff --git a/c2rust-postprocess/postprocess/__init__.py b/c2rust-postprocess/postprocess/__init__.py index 1f6753af5b..a21f5c0cb2 100644 --- a/c2rust-postprocess/postprocess/__init__.py +++ b/c2rust-postprocess/postprocess/__init__.py @@ -3,6 +3,7 @@ """ import argparse +import asyncio import logging import os from argparse import BooleanOptionalAction @@ -170,22 +171,36 @@ def get_model(model_id: str) -> AbstractGenerativeModel: def main(argv: Sequence[str] | None = None): try: - parser = build_arg_parser() - args = parser.parse_args(argv) + return asyncio.run(_main(argv)) + except BaselineError as error: + logging.error(error) + return 1 + except TransformError as error: + logging.exception(f"Aborting at first transform failure: {error}") + return 1 + except KeyboardInterrupt: + logging.warning("Interrupted by user, terminating...") + return 130 # 128 + SIGINT(2) + + +async def _main(argv: Sequence[str] | None = None) -> int: + parser = build_arg_parser() + args = parser.parse_args(argv) - logging.basicConfig(level=logging.getLevelName(args.log_level.upper())) + logging.basicConfig(level=logging.getLevelName(args.log_level.upper())) - if args.cache_dir is not None: - cache = DirectoryCache(args.cache_dir) - else: - cache = getattr(DirectoryCache, args.cache_scope)() - if args.update_cache and args.prune_cache_days > 0: - cache.prune(args.prune_cache_days) - if not args.update_cache: - cache = FrozenCache(cache) + if args.cache_dir is not None: + cache = DirectoryCache(args.cache_dir) + else: + cache = getattr(DirectoryCache, args.cache_scope)() + if args.update_cache and args.prune_cache_days > 0: + cache.prune(args.prune_cache_days) + if not args.update_cache: + cache = FrozenCache(cache) - model = get_model(args.llm_model) + model = get_model(args.llm_model) + try: # sort transform IDs to transforms always run in the same order to # maximize cache hits even if the user passed them in a different order transform_ids = sorted( @@ -214,7 +229,7 @@ def main(argv: Sequence[str] | None = None): ) for transform in transforms: result.extend( - transform.apply_dir( + await transform.apply_dir( root_rust_source_file=args.root_rust_source_file, exclude_list=IdentifierExcludeList(src_path=args.exclude_file), ident_filter=args.ident_filter, @@ -237,12 +252,5 @@ def main(argv: Sequence[str] | None = None): return 1 return 0 - except BaselineError as error: - logging.error(error) - return 1 - except TransformError as error: - logging.exception(f"Aborting at first transform failure: {error}") - return 1 - except KeyboardInterrupt: - logging.warning("Interrupted by user, terminating...") - return 130 # 128 + SIGINT(2) + finally: + await model.aclose() diff --git a/c2rust-postprocess/postprocess/models/base.py b/c2rust-postprocess/postprocess/models/base.py index 4473476b24..f4e63e4745 100644 --- a/c2rust-postprocess/postprocess/models/base.py +++ b/c2rust-postprocess/postprocess/models/base.py @@ -16,10 +16,14 @@ def id(self) -> str: return self._id @abstractmethod - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), max_tool_loops: int = 5, ) -> str | None: pass + + async def aclose(self) -> None: + """Release any resources held by the model client.""" + return None diff --git a/c2rust-postprocess/postprocess/models/gemini.py b/c2rust-postprocess/postprocess/models/gemini.py index 092e8d22bb..ebf908ba02 100644 --- a/c2rust-postprocess/postprocess/models/gemini.py +++ b/c2rust-postprocess/postprocess/models/gemini.py @@ -21,7 +21,7 @@ def __init__( super().__init__(id) self.client = genai.Client(api_key=api_key) - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), @@ -38,12 +38,18 @@ def generate_with_tools( ), ) - response = self.client.models.generate_content( + response = await self.client.aio.models.generate_content( model=self._id, contents=contents, config=config ) return response.text + async def aclose(self) -> None: + try: + await self.client.aio.aclose() + finally: + self.client.close() + def _convert_messages(self, messages: list[dict[str, Any]]) -> list[types.Content]: """ Converts standard list of dicts to Google GenAI 'Content' objects. diff --git a/c2rust-postprocess/postprocess/models/gpt.py b/c2rust-postprocess/postprocess/models/gpt.py index eb0c8904a4..4e8bfe0074 100644 --- a/c2rust-postprocess/postprocess/models/gpt.py +++ b/c2rust-postprocess/postprocess/models/gpt.py @@ -1,7 +1,7 @@ from collections.abc import Callable, Iterable from typing import Any -from openai import OpenAI +from openai import AsyncOpenAI from postprocess.models import AbstractGenerativeModel @@ -14,9 +14,9 @@ def __init__( base_url: str | None = None, ): super().__init__(id) - self.client = OpenAI(api_key=api_key, base_url=base_url) + self.client = AsyncOpenAI(api_key=api_key, base_url=base_url) - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), @@ -25,10 +25,13 @@ def generate_with_tools( # TODO: implement tool calling support assert not tools, "Tool calling not yet implemented for GPTModel" - response = self.client.responses.create( + response = await self.client.responses.create( model=self.id, input=messages[0]["content"], max_tool_calls=max_tool_loops, ) return response.output_text + + async def aclose(self) -> None: + await self.client.close() diff --git a/c2rust-postprocess/postprocess/models/mock.py b/c2rust-postprocess/postprocess/models/mock.py index 9095ccd8c2..1cc7b172c7 100644 --- a/c2rust-postprocess/postprocess/models/mock.py +++ b/c2rust-postprocess/postprocess/models/mock.py @@ -13,7 +13,7 @@ class MockGenerativeModel(AbstractGenerativeModel): def __init__(self): super().__init__(id="mock-llm") - def generate_with_tools( + async def generate_with_tools( self, messages: list[dict[str, Any]], tools: Iterable[Callable[..., Any]] = (), diff --git a/c2rust-postprocess/postprocess/transforms/base.py b/c2rust-postprocess/postprocess/transforms/base.py index b6e21ecd8a..cd9af3a5ec 100644 --- a/c2rust-postprocess/postprocess/transforms/base.py +++ b/c2rust-postprocess/postprocess/transforms/base.py @@ -62,7 +62,7 @@ def __init__( def system_instruction(self) -> str: return self._system_instruction - def apply_ident( + async def apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -74,7 +74,7 @@ def apply_ident( Apply the transform to one Rust definition and commit it through merge_rust. """ - new_definition = self.try_apply_ident( + new_definition = await self.try_apply_ident( rust_source_file=rust_source_file, rust_definition=rust_definition, c_definition=c_definition, @@ -99,7 +99,7 @@ def apply_ident( logging.info(f"{self.__class__.__name__}: Transformed Rust fn {identifier}") return new_definition - def try_apply_ident( + async def try_apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -115,7 +115,7 @@ def try_apply_ident( "or override apply_ident directly" ) - def generate( + async def generate( self, identifier: str, messages: list[dict[str, Any]], @@ -155,7 +155,7 @@ def generate( for attempt in range(self.max_attempts): try: - response = self.model.generate_with_tools(messages) + response = await self.model.generate_with_tools(messages) if response is None: raise TransformError(f"model returned no response for {identifier}") result = validate(response) @@ -181,7 +181,7 @@ def generate( f"for {identifier}: {error}" ) from error - def apply_dir( + async def apply_dir( self, root_rust_source_file: Path, exclude_list: IdentifierExcludeList, @@ -206,7 +206,7 @@ def apply_dir( ) assert rs_path.exists() result.extend( - self.apply_file( + await self.apply_file( rust_source_file=rs_path, exclude_list=exclude_list, ident_filter=ident_filter, @@ -218,7 +218,7 @@ def apply_dir( ) return result - def apply_file( + async def apply_file( self, rust_source_file: Path, exclude_list: IdentifierExcludeList, @@ -267,7 +267,7 @@ def apply_file( ) try: - new_definition = self.apply_ident( + new_definition = await self.apply_ident( rust_source_file=rust_source_file, rust_definition=rust_definition, c_definition=c_definition, diff --git a/c2rust-postprocess/postprocess/transforms/comments.py b/c2rust-postprocess/postprocess/transforms/comments.py index 8e6f1030f4..564e92e551 100644 --- a/c2rust-postprocess/postprocess/transforms/comments.py +++ b/c2rust-postprocess/postprocess/transforms/comments.py @@ -104,7 +104,7 @@ def __init__(self, cache: AbstractCache, model: AbstractGenerativeModel): super().__init__(SYSTEM_INSTRUCTION, cache, model) self.trim_transform = TrimTransform(cache, model) - def try_apply_ident( + async def try_apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -124,7 +124,7 @@ def try_apply_ident( logging.info(f"Skipping C function without comments: {identifier}") return None - match self.trim_transform.apply_ident( + match await self.trim_transform.apply_ident( rust_source_file=rust_source_file, rust_definition=rust_definition, c_definition=c_definition, @@ -205,7 +205,7 @@ def validate(response: str) -> str: return rust_fn - rust_fn = self.generate(identifier, messages, validate) + rust_fn = await self.generate(identifier, messages, validate) if rust_fn is None: return None diff --git a/c2rust-postprocess/postprocess/transforms/trim.py b/c2rust-postprocess/postprocess/transforms/trim.py index 64226c27f2..0c70344f05 100644 --- a/c2rust-postprocess/postprocess/transforms/trim.py +++ b/c2rust-postprocess/postprocess/transforms/trim.py @@ -61,7 +61,7 @@ class TrimTransform(AbstractTransform): def __init__(self, cache: AbstractCache, model: AbstractGenerativeModel): super().__init__(SYSTEM_INSTRUCTION, cache, model) - def apply_ident( + async def apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -122,7 +122,7 @@ def validate(response: str) -> str: return trimmed try: - return self.generate(identifier, messages, validate) + return await self.generate(identifier, messages, validate) except TransformError as error: # Trimming is best-effort; callers fall back to the untrimmed input. logging.warning(f"{self.__class__.__name__}: {error}") diff --git a/c2rust-postprocess/tests/test_apply_file.py b/c2rust-postprocess/tests/test_apply_file.py index eaa985a983..8f52f684f6 100644 --- a/c2rust-postprocess/tests/test_apply_file.py +++ b/c2rust-postprocess/tests/test_apply_file.py @@ -1,3 +1,4 @@ +import asyncio import json from pathlib import Path from typing import Any @@ -115,11 +116,13 @@ def fake_update(*, root_rust_source_file, identifier, new_definition): def test_accepted_batch_is_applied(patched_io, rust_file: Path) -> None: cache, transform = make_comments_transform() - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=BatchValidator(lambda: None), + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=BatchValidator(lambda: None), + ) ) assert result.failed == [] @@ -132,11 +135,13 @@ def test_rejection_counts_failure_and_invalidates_final_response( ) -> None: cache, transform = make_comments_transform() - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=BatchValidator(lambda: "type error"), + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=BatchValidator(lambda: "type error"), + ) ) assert result.failed == [(rust_file, "enabled", "rejected by cargo check")] @@ -150,11 +155,13 @@ def test_rejection_raises_without_keep_going(patched_io, rust_file: Path) -> Non cache, transform = make_comments_transform() with pytest.raises(TransformError, match="cargo check rejected"): - transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=False, - validator=BatchValidator(lambda: "type error"), + asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=False, + validator=BatchValidator(lambda: "type error"), + ) ) assert rust_file.read_text() == "original\n" @@ -167,12 +174,14 @@ def test_no_update_rust_is_purely_generative(patched_io, rust_file: Path) -> Non def exploding_check() -> str | None: raise AssertionError("validator must not run") - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - update_rust=False, - keep_going=True, - validator=BatchValidator(exploding_check), + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + update_rust=False, + keep_going=True, + validator=BatchValidator(exploding_check), + ) ) assert result.failures == 0 @@ -182,11 +191,13 @@ def exploding_check() -> str | None: def test_without_validator_candidates_are_applied(patched_io, rust_file: Path) -> None: cache, transform = make_comments_transform() - result = transform.apply_file( - rust_source_file=rust_file, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=None, + result = asyncio.run( + transform.apply_file( + rust_source_file=rust_file, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=None, + ) ) assert result.failures == 0 @@ -215,7 +226,7 @@ def __init__(self, rewrites: dict[str, str]): super().__init__("", CannedCache({}), MockGenerativeModel()) self.rewrites = rewrites - def try_apply_ident( + async def try_apply_ident( self, rust_source_file: Path, rust_definition: str, @@ -268,11 +279,13 @@ def test_end_to_end_type_invalid_rewrite_is_isolated(tmp_path: Path) -> None: checker = CargoChecker(manifest) assert checker() is None # baseline - result = transform.apply_file( - rust_source_file=lib_rs, - exclude_list=IdentifierExcludeList(None), - keep_going=True, - validator=BatchValidator(checker), + result = asyncio.run( + transform.apply_file( + rust_source_file=lib_rs, + exclude_list=IdentifierExcludeList(None), + keep_going=True, + validator=BatchValidator(checker), + ) ) content = lib_rs.read_text() diff --git a/c2rust-postprocess/tests/test_comments_transform.py b/c2rust-postprocess/tests/test_comments_transform.py index ee594a0e71..d522980394 100644 --- a/c2rust-postprocess/tests/test_comments_transform.py +++ b/c2rust-postprocess/tests/test_comments_transform.py @@ -1,3 +1,4 @@ +import asyncio from pathlib import Path from typing import Any @@ -79,12 +80,14 @@ def test_directive_line_comment_survives_preprocessed_check() -> None: cache = StaticCache(response) transform = CommentsTransform(cache=cache, model=MockGenerativeModel()) - transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=rust_definition, - c_definition=c_definition, - identifier="enabled", - update_rust=False, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=rust_definition, + c_definition=c_definition, + identifier="enabled", + update_rust=False, + ) ) transforms = [transform for transform, _ in cache.lookups] @@ -128,12 +131,14 @@ def fake_update(*, root_rust_source_file, identifier, new_definition): monkeypatch.setattr(base, "update_rust_definition", fake_update) - transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=True, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=True, + ) ) assert "/// @note" not in merged["f"] @@ -150,12 +155,14 @@ def test_syntactically_invalid_response_is_rejected() -> None: transform = CommentsTransform(cache=cache, model=MockGenerativeModel()) with pytest.raises(TransformError, match="not syntactically valid"): - transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) @@ -169,12 +176,14 @@ def test_response_that_changes_non_comment_code_is_rejected() -> None: transform = CommentsTransform(cache=cache, model=MockGenerativeModel()) with pytest.raises(TransformError, match="non-comment Rust code changed"): - transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) @@ -191,12 +200,14 @@ def test_response_may_reformat_non_comment_code() -> None: cache = StaticCache(response) transform = CommentsTransform(cache=cache, model=MockGenerativeModel()) - result = transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + result = asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) assert result is not None @@ -246,7 +257,7 @@ def __init__(self, responses: list[str]): self.responses = responses self.calls = 0 - def generate_with_tools(self, messages, tools=(), max_tool_loops=5): + async def generate_with_tools(self, messages, tools=(), max_tool_loops=5): self.calls += 1 return self.responses.pop(0) @@ -273,12 +284,14 @@ def apply_to_body_comment_fn( cache: AbstractCache, model: AbstractGenerativeModel ) -> str | None: transform = CommentsTransform(cache=cache, model=model) - return transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition=RUST_DEFINITION_NO_COMMENTS, - c_definition=C_DEFINITION_BODY_COMMENT, - identifier="f", - update_rust=False, + return asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition=RUST_DEFINITION_NO_COMMENTS, + c_definition=C_DEFINITION_BODY_COMMENT, + identifier="f", + update_rust=False, + ) ) diff --git a/c2rust-postprocess/tests/test_models.py b/c2rust-postprocess/tests/test_models.py new file mode 100644 index 0000000000..2416f7f3c1 --- /dev/null +++ b/c2rust-postprocess/tests/test_models.py @@ -0,0 +1,92 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from google.genai import types + +from postprocess.models import gemini, gpt + + +def test_gpt_awaits_response_and_closes_client(monkeypatch) -> None: + create = AsyncMock(return_value=SimpleNamespace(output_text="rewritten function")) + close = AsyncMock() + client = SimpleNamespace(responses=SimpleNamespace(create=create), close=close) + constructor = Mock(return_value=client) + monkeypatch.setattr(gpt, "AsyncOpenAI", constructor) + + async def run() -> str: + model = gpt.GPTModel( + id="test/model", api_key="test-key", base_url="https://example.test/v1" + ) + try: + return await model.generate_with_tools( + [{"role": "user", "content": "function prompt"}], max_tool_loops=2 + ) + finally: + await model.aclose() + + assert asyncio.run(run()) == "rewritten function" + constructor.assert_called_once_with( + api_key="test-key", base_url="https://example.test/v1" + ) + create.assert_awaited_once_with( + model="test/model", input="function prompt", max_tool_calls=2 + ) + close.assert_awaited_once_with() + + +@pytest.mark.parametrize("response_text", ["rewritten function", None]) +def test_gemini_awaits_response_and_closes_both_clients( + monkeypatch, response_text: str | None +) -> None: + generate = AsyncMock(return_value=SimpleNamespace(text=response_text)) + async_close = AsyncMock() + sync_close = Mock() + client = SimpleNamespace( + aio=SimpleNamespace( + models=SimpleNamespace(generate_content=generate), aclose=async_close + ), + close=sync_close, + ) + constructor = Mock(return_value=client) + monkeypatch.setattr(gemini.genai, "Client", constructor) + + def validate(response: str) -> str: + return response + + async def run() -> str | None: + model = gemini.GoogleGenerativeModel(id="gemini-test", api_key="test-key") + try: + return await model.generate_with_tools( + [ + {"role": "user", "content": "function prompt"}, + {"role": "assistant", "content": "previous response"}, + ], + tools=[validate], + max_tool_loops=2, + ) + finally: + await model.aclose() + + assert asyncio.run(run()) == response_text + constructor.assert_called_once_with(api_key="test-key") + generate.assert_awaited_once_with( + model="gemini-test", + contents=[ + types.Content( + role="user", parts=[types.Part.from_text(text="function prompt")] + ), + types.Content( + role="model", parts=[types.Part.from_text(text="previous response")] + ), + ], + config=types.GenerateContentConfig( + tools=[validate], + automatic_function_calling=types.AutomaticFunctionCallingConfig( + disable=False, maximum_remote_calls=2 + ), + ), + ) + async_close.assert_awaited_once_with() + sync_close.assert_called_once_with() diff --git a/c2rust-postprocess/tests/test_trim_transform.py b/c2rust-postprocess/tests/test_trim_transform.py index 22e4f1439b..114e93832d 100644 --- a/c2rust-postprocess/tests/test_trim_transform.py +++ b/c2rust-postprocess/tests/test_trim_transform.py @@ -1,3 +1,4 @@ +import asyncio from pathlib import Path from typing import Any @@ -109,12 +110,14 @@ def invalidate(self, *, transform: str, identifier: str) -> None: def apply_trim(c_definition: CDefinition, cache: AbstractCache) -> str | None: transform = TrimTransform(cache=cache, model=MockGenerativeModel()) - return transform.apply_ident( - rust_source_file=Path("unused.rs"), - rust_definition="fn f() {}", - c_definition=c_definition, - identifier="f", - update_rust=False, + return asyncio.run( + transform.apply_ident( + rust_source_file=Path("unused.rs"), + rust_definition="fn f() {}", + c_definition=c_definition, + identifier="f", + update_rust=False, + ) ) From a83ca914e21eda7ab3e2926264b076bc363944e3 Mon Sep 17 00:00:00 2001 From: Per Larsen Date: Sun, 6 Sep 2026 18:10:36 -0700 Subject: [PATCH 2/3] postprocess: observe cancellation before committing rewrite batches asyncio.run handles Ctrl-C by cancelling the main task. A synchronous merge or cargo check can return with that cancellation still pending, bypassing the existing rollback boundary. Await cancellation checkpoints inside the transaction before checking and before accepting the batch. Keep bisection sequential and cover cancellation during both merge and check, including restoration of the original file. --- .../postprocess/transforms/base.py | 2 +- c2rust-postprocess/postprocess/validate.py | 11 +++- c2rust-postprocess/tests/test_validate.py | 55 ++++++++++++++++--- 3 files changed, 55 insertions(+), 13 deletions(-) diff --git a/c2rust-postprocess/postprocess/transforms/base.py b/c2rust-postprocess/postprocess/transforms/base.py index cd9af3a5ec..fd99c636ff 100644 --- a/c2rust-postprocess/postprocess/transforms/base.py +++ b/c2rust-postprocess/postprocess/transforms/base.py @@ -319,7 +319,7 @@ async def apply_file( f"Validating {len(candidates)} rewrite(s) in {rust_source_file} " "with cargo check" ) - _, rejected = validator.validate(candidates) + _, rejected = await validator.validate(candidates) for candidate, error in rejected: candidate.invalidate() result.failed.append( diff --git a/c2rust-postprocess/postprocess/validate.py b/c2rust-postprocess/postprocess/validate.py index 61192b9e45..238afaae2b 100644 --- a/c2rust-postprocess/postprocess/validate.py +++ b/c2rust-postprocess/postprocess/validate.py @@ -1,5 +1,6 @@ """Transactional validation of applied rewrites via `cargo check`.""" +import asyncio import json import logging import subprocess @@ -35,7 +36,7 @@ class BatchValidator: def __init__(self, check: Callable[[], str | None]): self._check = check - def validate( + async def validate( self, candidates: Sequence[Candidate] ) -> tuple[list[Candidate], list[tuple[Candidate, str]]]: """ @@ -56,7 +57,11 @@ def validate( try: for candidate in candidates: candidate.apply() + # SIGINT cancels the main asyncio task. Observe cancellation from + # synchronous merges/checks while rollback is still possible. + await asyncio.sleep(0) error = self._check() + await asyncio.sleep(0) ok = error is None finally: # `finally` rather than `except Exception` so KeyboardInterrupt @@ -75,8 +80,8 @@ def validate( # candidates, so interacting candidates are isolated correctly. logging.info(f"Check failed for batch of {len(candidates)}; bisecting") mid = len(candidates) // 2 - accepted, rejected = self.validate(candidates[:mid]) - right_accepted, right_rejected = self.validate(candidates[mid:]) + accepted, rejected = await self.validate(candidates[:mid]) + right_accepted, right_rejected = await self.validate(candidates[mid:]) return accepted + right_accepted, rejected + right_rejected diff --git a/c2rust-postprocess/tests/test_validate.py b/c2rust-postprocess/tests/test_validate.py index e7653bc109..2fa8839770 100644 --- a/c2rust-postprocess/tests/test_validate.py +++ b/c2rust-postprocess/tests/test_validate.py @@ -1,3 +1,4 @@ +import asyncio import shutil import textwrap from pathlib import Path @@ -47,7 +48,7 @@ def test_valid_batch_is_retained_with_one_check(source: Path) -> None: check = ContentCheck(source) candidates = [append_candidate(source, "a"), append_candidate(source, "b")] - accepted, rejected = BatchValidator(check).validate(candidates) + accepted, rejected = asyncio.run(BatchValidator(check).validate(candidates)) assert accepted == candidates assert rejected == [] @@ -64,7 +65,7 @@ def test_mixed_batch_retains_only_valid_candidates(source: Path) -> None: append_candidate(source, "d"), ] - accepted, rejected = BatchValidator(check).validate(candidates) + accepted, rejected = asyncio.run(BatchValidator(check).validate(candidates)) assert [c.identifier for c in accepted] == ["a", "c", "d"] assert [(c.identifier, error) for c, error in rejected] == [("BAD", "found BAD")] @@ -72,8 +73,8 @@ def test_mixed_batch_retains_only_valid_candidates(source: Path) -> None: def test_single_rejected_candidate_restores_baseline_exactly(source: Path) -> None: - accepted, rejected = BatchValidator(ContentCheck(source)).validate( - [append_candidate(source, "BAD")] + accepted, rejected = asyncio.run( + BatchValidator(ContentCheck(source)).validate([append_candidate(source, "BAD")]) ) assert accepted == [] @@ -90,8 +91,10 @@ def check() -> str | None: return "x and y conflict" return None - accepted, rejected = BatchValidator(check).validate( - [append_candidate(source, "x"), append_candidate(source, "y")] + accepted, rejected = asyncio.run( + BatchValidator(check).validate( + [append_candidate(source, "x"), append_candidate(source, "y")] + ) ) assert [c.identifier for c in accepted] == ["x"] @@ -104,7 +107,7 @@ def check() -> str | None: raise RuntimeError("cargo crashed") with pytest.raises(RuntimeError, match="cargo crashed"): - BatchValidator(check).validate([append_candidate(source, "a")]) + asyncio.run(BatchValidator(check).validate([append_candidate(source, "a")])) assert source.read_text() == "baseline\n" @@ -117,13 +120,47 @@ def broken_apply() -> None: identifier="a", files=(source,), apply=broken_apply, invalidate=lambda: None ) with pytest.raises(RuntimeError, match="merge failed"): - BatchValidator(ContentCheck(source)).validate([candidate]) + asyncio.run(BatchValidator(ContentCheck(source)).validate([candidate])) assert source.read_text() == "baseline\n" +@pytest.mark.parametrize("cancel_during", ["apply", "check"]) +def test_pending_cancellation_restores_file(source: Path, cancel_during: str) -> None: + check_calls = 0 + + async def run() -> None: + task = asyncio.current_task() + assert task is not None + + def apply() -> None: + source.write_text("rewritten\n") + if cancel_during == "apply": + task.cancel() + + def check() -> str | None: + nonlocal check_calls + check_calls += 1 + assert source.read_text() == "rewritten\n" + if cancel_during == "check": + task.cancel() + return None + + candidate = Candidate( + identifier="f", files=(source,), apply=apply, invalidate=lambda: None + ) + await BatchValidator(check).validate([candidate]) + + # asyncio.run handles SIGINT by cancelling the main task; synchronous + # subprocess calls can return normally with that cancellation pending. + with pytest.raises(asyncio.CancelledError): + asyncio.run(run()) + assert source.read_text() == "baseline\n" + assert check_calls == (1 if cancel_during == "check" else 0) + + def test_empty_batch_runs_no_check(source: Path) -> None: check = ContentCheck(source) - assert BatchValidator(check).validate([]) == ([], []) + assert asyncio.run(BatchValidator(check).validate([])) == ([], []) assert check.calls == 0 From 83a10606ed083a37594ec3f919680cd92c7c5198 Mon Sep 17 00:00:00 2001 From: Per Larsen Date: Sun, 6 Sep 2026 18:10:42 -0700 Subject: [PATCH 3/3] postprocess: process functions concurrently with bounded jobs Use a small pool of async workers to overlap independent functions within each Rust file. Add -j/--jobs with a default of four and require a positive limit; -j 1 preserves serial processing. Keep trimming and retries within each worker, and process files and transform passes sequentially. Collect responses before merging, apply candidates in source order, and cancel and await outstanding workers on abort. Existing prompts and cache entries remain reusable without locks or new dependencies. Test the concurrency bound, ordered writes, failure handling, cancellation, serial replay of a parallel-populated frozen cache, CLI validation, and client cleanup. All 79 local tests passed; four c2rust-dependent tests skipped. Ruff and ty checks passed without live model calls. --- c2rust-postprocess/README.md | 6 + c2rust-postprocess/postprocess/__init__.py | 11 +- .../postprocess/transforms/base.py | 143 ++++++---- c2rust-postprocess/postprocess/utils.py | 7 + c2rust-postprocess/tests/test_cli.py | 60 +++++ c2rust-postprocess/tests/test_parallel.py | 252 ++++++++++++++++++ 6 files changed, 428 insertions(+), 51 deletions(-) create mode 100644 c2rust-postprocess/tests/test_cli.py create mode 100644 c2rust-postprocess/tests/test_parallel.py diff --git a/c2rust-postprocess/README.md b/c2rust-postprocess/README.md index 1b90e42e08..2819f56eb2 100644 --- a/c2rust-postprocess/README.md +++ b/c2rust-postprocess/README.md @@ -22,6 +22,12 @@ attempts to produce fully safe and idiomatic Rust; not a replacement for these e - `c2rust-postprocess path/to/transpiled_rust.rs`, or - `uv run postprocess path/to/transpiled_rust.rs` +Up to four functions in each Rust file are processed concurrently by default. +Use `-j 8` / `--jobs 8` to change the limit, or `-j 1` to process serially. +Each function's trimming and comment transfer still run in order. Files and +transform passes also run in order, and rewrites are applied and checked only +after all functions in a file finish. Existing cached responses are reused. + ## Excluding/Filtering `c2rust-postprocess` has a few ways to filter/exclude the function identifiers that are processed. diff --git a/c2rust-postprocess/postprocess/__init__.py b/c2rust-postprocess/postprocess/__init__.py index a21f5c0cb2..13ad3dd608 100644 --- a/c2rust-postprocess/postprocess/__init__.py +++ b/c2rust-postprocess/postprocess/__init__.py @@ -21,7 +21,7 @@ from postprocess.models.mock import MockGenerativeModel from postprocess.transforms import get_transform_by_id from postprocess.transforms.base import TransformError, TransformResult -from postprocess.utils import existing_file +from postprocess.utils import existing_file, positive_int from postprocess.validate import BaselineError, make_validator DEFAULT_LLM_MODEL = "gpt-5.6-luna" @@ -38,6 +38,14 @@ def build_arg_parser() -> argparse.ArgumentParser: help="Path to Rust source file referenced by Cargo.toml", ) + parser.add_argument( + "-j", + "--jobs", + type=positive_int, + default=4, + help="Maximum concurrent function transformations per file (default: 4)", + ) + parser.add_argument( "--log-level", type=str, @@ -237,6 +245,7 @@ async def _main(argv: Sequence[str] | None = None) -> int: keep_going=args.on_error != "abort", failure_log_level=failure_log_level, validator=validator, + jobs=args.jobs, ) ) diff --git a/c2rust-postprocess/postprocess/transforms/base.py b/c2rust-postprocess/postprocess/transforms/base.py index fd99c636ff..91139d5d08 100644 --- a/c2rust-postprocess/postprocess/transforms/base.py +++ b/c2rust-postprocess/postprocess/transforms/base.py @@ -1,3 +1,4 @@ +import asyncio import logging import re from collections.abc import Callable @@ -190,6 +191,7 @@ async def apply_dir( keep_going: bool = False, failure_log_level: int = logging.ERROR, validator: BatchValidator | None = None, + jobs: int = 1, ) -> TransformResult: """ Run `self.apply_file` on each `*.rs` in `dir` @@ -200,7 +202,7 @@ async def apply_dir( result = TransformResult() root_dir = root_rust_source_file.parent c_decls_json_suffix = ".c_decls.json" - for c_decls_path in root_dir.glob(f"**/*{c_decls_json_suffix}"): + for c_decls_path in sorted(root_dir.glob(f"**/*{c_decls_json_suffix}")): rs_path = c_decls_path.with_name( c_decls_path.name.removesuffix(c_decls_json_suffix) + ".rs" ) @@ -214,6 +216,7 @@ async def apply_dir( keep_going=keep_going, failure_log_level=failure_log_level, validator=validator, + jobs=jobs, ) ) return result @@ -227,7 +230,10 @@ async def apply_file( keep_going: bool = False, failure_log_level: int = logging.ERROR, validator: BatchValidator | None = None, + jobs: int = 1, ) -> TransformResult: + if jobs < 1: + raise ValueError("jobs must be at least 1") ident_regex = re.compile(ident_filter) if ident_filter else None result = TransformResult() @@ -237,58 +243,95 @@ async def apply_file( logging.info(f"Loaded {len(rust_definitions)} Rust definitions") logging.info(f"Loaded {len(c_definitions)} C definitions") - # Collect candidates without touching the file; application and - # validation happen per batch below. - candidates: list[Candidate] = [] - for identifier, rust_definition in rust_definitions.items(): - if exclude_list.contains(path=rust_source_file, identifier=identifier): - logging.info( - f"Skipping Rust fn {identifier} in {rust_source_file} " - f"due to exclude file {exclude_list.src_path}" - ) - continue - - if ident_regex and not ident_regex.search(identifier): - logging.info( - f"Skipping Rust fn {identifier} in {rust_source_file} " - f"due to ident filter {ident_filter}" - ) - continue - - if identifier not in c_definitions: - logging.warning(f"No corresponding C definition found for {identifier}") - continue - - c_definition = c_definitions[identifier] - - highlighted_c_definition = get_highlighted_c(c_definition.effective) - logging.debug( - f"C function {identifier} definition:\n{highlighted_c_definition}\n" - ) - - try: - new_definition = await self.apply_ident( - rust_source_file=rust_source_file, - rust_definition=rust_definition, - c_definition=c_definition, - identifier=identifier, - update_rust=False, - ) - except TransformError as error: - if not keep_going: - raise - logging.log( - failure_log_level, - f"Transform failed for {identifier} in {rust_source_file}: {error}", + # Each worker completes a function's trim/retry sequence before taking + # another one. Cache operations stay on this event loop; no file writes + # or cargo checks run until all workers finish. + definitions = iter(rust_definitions.items()) + new_definitions: dict[str, str] = {} + failures: dict[str, TransformFailure] = {} + + async def worker() -> None: + for identifier, rust_definition in definitions: + if exclude_list.contains(path=rust_source_file, identifier=identifier): + logging.info( + f"Skipping Rust fn {identifier} in {rust_source_file} " + f"due to exclude file {exclude_list.src_path}" + ) + continue + + if ident_regex and not ident_regex.search(identifier): + logging.info( + f"Skipping Rust fn {identifier} in {rust_source_file} " + f"due to ident filter {ident_filter}" + ) + continue + + if identifier not in c_definitions: + logging.warning( + f"No corresponding C definition found for {identifier}" + ) + continue + + c_definition = c_definitions[identifier] + + highlighted_c_definition = get_highlighted_c(c_definition.effective) + logging.debug( + f"C function {identifier} definition:\n{highlighted_c_definition}\n" ) - result.failed.append( - (rust_source_file, identifier, "failed to transform") - ) - continue - if new_definition is None: + try: + new_definition = await self.apply_ident( + rust_source_file=rust_source_file, + rust_definition=rust_definition, + c_definition=c_definition, + identifier=identifier, + update_rust=False, + ) + except TransformError as error: + if not keep_going: + raise + logging.log( + failure_log_level, + f"Transform failed for {identifier} " + f"in {rust_source_file}: {error}", + ) + failures[identifier] = ( + rust_source_file, + identifier, + "failed to transform", + ) + continue + + if new_definition is None: + continue + + new_definitions[identifier] = new_definition + + workers = [ + asyncio.create_task(worker()) + for _ in range(min(jobs, len(rust_definitions))) + ] + try: + await asyncio.gather(*workers) + except BaseException: + # Abort and interruption must cancel pending model calls before + # returning or closing the shared client. + for task in workers: + task.cancel() + await asyncio.gather(*workers, return_exceptions=True) + raise + + # Apply and report in source order, independent of response timing. + result.failed = [ + failures[identifier] + for identifier in rust_definitions + if identifier in failures + ] + candidates: list[Candidate] = [] + for identifier in rust_definitions: + if identifier not in new_definitions: continue - + new_definition = new_definitions[identifier] candidates.append( Candidate( identifier=identifier, diff --git a/c2rust-postprocess/postprocess/utils.py b/c2rust-postprocess/postprocess/utils.py index b577fb187d..73555be7cb 100644 --- a/c2rust-postprocess/postprocess/utils.py +++ b/c2rust-postprocess/postprocess/utils.py @@ -48,6 +48,13 @@ def existing_file(value: str) -> Path: raise argparse.ArgumentTypeError(f"{value!r} is not a readable file") +def positive_int(value: str) -> int: + number = int(value) + if number < 1: + raise argparse.ArgumentTypeError("must be at least 1") + return number + + # TODO: test def get_compile_commands(compile_commands_path: Path) -> list[dict[str, Any]]: try: diff --git a/c2rust-postprocess/tests/test_cli.py b/c2rust-postprocess/tests/test_cli.py new file mode 100644 index 0000000000..f6e257fe4c --- /dev/null +++ b/c2rust-postprocess/tests/test_cli.py @@ -0,0 +1,60 @@ +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest + +import postprocess +from postprocess.transforms.base import TransformError, TransformResult +from postprocess.validate import BaselineError + + +@pytest.mark.parametrize("value", ["0", "-1", "many"]) +def test_jobs_rejects_invalid_values(tmp_path: Path, value: str) -> None: + source = tmp_path / "lib.rs" + source.touch() + with pytest.raises(SystemExit) as error: + postprocess.build_arg_parser().parse_args([str(source), "-j", value]) + assert error.value.code == 2 + + +@pytest.mark.parametrize("option", ["-j", "--jobs"]) +def test_jobs_option_and_default(tmp_path: Path, option: str) -> None: + source = tmp_path / "lib.rs" + source.touch() + parser = postprocess.build_arg_parser() + assert parser.parse_args([str(source)]).jobs == 4 + assert parser.parse_args([str(source), option, "8"]).jobs == 8 + + +@pytest.mark.parametrize("failure", [None, "baseline", "transform"]) +def test_main_forwards_jobs_and_closes_model(monkeypatch, tmp_path, failure) -> None: + source = tmp_path / "lib.rs" + source.touch() + close = AsyncMock() + model = SimpleNamespace(aclose=close) + monkeypatch.setattr(postprocess, "get_model", lambda _: model) + + apply = AsyncMock(return_value=TransformResult()) + if failure == "transform": + apply.side_effect = TransformError("rejected") + transform = SimpleNamespace(apply_dir=apply) + monkeypatch.setattr(postprocess, "get_transform_by_id", lambda *a, **kw: transform) + + validator = Mock(return_value=None) + if failure == "baseline": + validator.side_effect = BaselineError("broken baseline") + monkeypatch.setattr(postprocess, "make_validator", validator) + + result = postprocess.main( + [str(source), "--cache-dir", str(tmp_path / "cache"), "-j", "8"] + ) + + assert result == (1 if failure else 0) + if failure == "baseline": + apply.assert_not_awaited() + else: + apply.assert_awaited_once() + assert apply.await_args is not None + assert apply.await_args.kwargs["jobs"] == 8 + close.assert_awaited_once_with() diff --git a/c2rust-postprocess/tests/test_parallel.py b/c2rust-postprocess/tests/test_parallel.py new file mode 100644 index 0000000000..baa0c01ad2 --- /dev/null +++ b/c2rust-postprocess/tests/test_parallel.py @@ -0,0 +1,252 @@ +import asyncio +from collections.abc import Awaitable, Callable +from pathlib import Path + +import pytest +from test_apply_file import CannedCache + +from postprocess.cache import DirectoryCache, FrozenCache +from postprocess.definitions import CDefinition +from postprocess.exclude_list import IdentifierExcludeList +from postprocess.models.base import AbstractGenerativeModel +from postprocess.models.mock import MockGenerativeModel +from postprocess.transforms import base +from postprocess.transforms.base import AbstractTransform, TransformError + +IDENTIFIERS = ["first", "second", "third", "fourth"] + + +class AsyncTransform(AbstractTransform): + def __init__(self, generate: Callable[[str], Awaitable[str | None]]): + super().__init__("", CannedCache({}), MockGenerativeModel()) + self.generate_definition = generate + + async def try_apply_ident( + self, + rust_source_file: Path, + rust_definition: str, + c_definition: CDefinition, + identifier: str, + ) -> str | None: + return await self.generate_definition(identifier) + + +@pytest.fixture +def source(monkeypatch, tmp_path: Path) -> tuple[Path, list[str]]: + path = tmp_path / "lib.rs" + path.write_text("original\n") + writes: list[str] = [] + monkeypatch.setattr( + base, + "get_rust_definitions", + lambda path: {identifier: "original" for identifier in IDENTIFIERS}, + ) + monkeypatch.setattr( + base, + "get_c_definitions", + lambda path: { + identifier: CDefinition( + definition="void f(void) {}", preprocessed_definition=None + ) + for identifier in IDENTIFIERS + }, + ) + + def update(*, root_rust_source_file, identifier, new_definition): + writes.append(identifier) + root_rust_source_file.write_text(new_definition) + + monkeypatch.setattr(base, "update_rust_definition", update) + return path, writes + + +@pytest.mark.parametrize("jobs", [1, 2]) +def test_jobs_bounds_concurrent_transformations(source, jobs: int) -> None: + path, writes = source + + async def run() -> None: + started: asyncio.Queue[str] = asyncio.Queue() + release: asyncio.Queue[None] = asyncio.Queue() + active = 0 + peak = 0 + + async def generate(identifier: str) -> str: + nonlocal active, peak + active += 1 + peak = max(peak, active) + started.put_nowait(identifier) + try: + await release.get() + return identifier + finally: + active -= 1 + + task = asyncio.create_task( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=jobs + ) + ) + for _ in range(jobs): + await asyncio.wait_for(started.get(), timeout=2) + assert active == jobs + assert writes == [] + for _ in IDENTIFIERS: + release.put_nowait(None) + result = await asyncio.wait_for(task, timeout=2) + assert result.failures == 0 + assert peak == jobs + assert active == 0 + + asyncio.run(run()) + assert writes == IDENTIFIERS + + +def test_out_of_order_responses_are_applied_in_source_order(source) -> None: + path, writes = source + + async def run() -> None: + started: asyncio.Queue[str] = asyncio.Queue() + completed: asyncio.Queue[str] = asyncio.Queue() + release = {identifier: asyncio.Event() for identifier in IDENTIFIERS} + + async def generate(identifier: str) -> str: + started.put_nowait(identifier) + await release[identifier].wait() + completed.put_nowait(identifier) + return identifier + + task = asyncio.create_task( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=4 + ) + ) + for _ in IDENTIFIERS: + await asyncio.wait_for(started.get(), timeout=2) + for identifier in reversed(IDENTIFIERS): + assert writes == [] + release[identifier].set() + assert await asyncio.wait_for(completed.get(), timeout=2) == identifier + await asyncio.wait_for(task, timeout=2) + + asyncio.run(run()) + assert writes == IDENTIFIERS + + +def test_keep_going_applies_successes_and_reports_failure(source) -> None: + path, writes = source + + async def generate(identifier: str) -> str: + await asyncio.sleep(0) + if identifier == "second": + raise TransformError("rejected") + return identifier + + async def run(): + return await asyncio.wait_for( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=2, keep_going=True + ), + timeout=2, + ) + + result = asyncio.run(run()) + assert result.failed == [(path, "second", "failed to transform")] + assert writes == ["first", "third", "fourth"] + + +def test_abort_cancels_pending_transformations_before_writing(source) -> None: + path, writes = source + + async def run() -> None: + first_started = asyncio.Event() + cancelled = asyncio.Event() + pending = asyncio.Event() + active: set[str] = set() + + async def generate(identifier: str) -> str: + active.add(identifier) + try: + if identifier == "second": + await first_started.wait() + raise TransformError("rejected") + if identifier == "first": + first_started.set() + await pending.wait() + return identifier + except asyncio.CancelledError: + if identifier == "first": + cancelled.set() + raise + finally: + active.remove(identifier) + + with pytest.raises(TransformError, match="rejected"): + await asyncio.wait_for( + AsyncTransform(generate).apply_file( + path, IdentifierExcludeList(None), jobs=2 + ), + timeout=2, + ) + assert cancelled.is_set() + assert not active + assert writes == [] + assert path.read_text() == "original\n" + + asyncio.run(run()) + + +def test_parallel_responses_are_reused_by_serial_frozen_run( + source, tmp_path: Path, monkeypatch +) -> None: + path, writes = source + + class CountingModel(AbstractGenerativeModel): + def __init__(self): + super().__init__("test-model") + self.calls = 0 + + async def generate_with_tools(self, messages, tools=(), max_tool_loops=5): + self.calls += 1 + await asyncio.sleep(0) + return messages[0]["content"] + + class CachedTransform(AbstractTransform): + async def try_apply_ident( + self, rust_source_file, rust_definition, c_definition, identifier + ): + return await self.generate( + identifier, + [{"role": "user", "content": identifier}], + lambda response: response, + ) + + cache = DirectoryCache(tmp_path / "cache") + model = CountingModel() + result = asyncio.run( + CachedTransform("", cache, model).apply_file( + path, IdentifierExcludeList(None), jobs=4 + ) + ) + assert result.failures == 0 + assert writes == IDENTIFIERS + assert model.calls == len(IDENTIFIERS) + assert len(list(cache.path.glob("*/*/*/metadata.toml"))) == len(IDENTIFIERS) + first_output = path.read_text() + + path.write_text("original\n") + writes.clear() + frozen = FrozenCache(DirectoryCache(cache.path)) + + def unexpected_update(**kwargs): + raise AssertionError("cache hits must not update responses") + + monkeypatch.setattr(frozen, "update", unexpected_update) + result = asyncio.run( + CachedTransform("", frozen, model).apply_file( + path, IdentifierExcludeList(None), jobs=1 + ) + ) + assert result.failures == 0 + assert writes == IDENTIFIERS + assert path.read_text() == first_output + assert model.calls == len(IDENTIFIERS)