Make Code Mode resource acquisition transactional

pydantic_ai_harness/code_mode/_toolset.py

@@ -7,7 +7,7 @@ import keyword
 import re
 import warnings
 from collections.abc import Callable, Generator, Sequence
-from contextlib import contextmanager
+from contextlib import AsyncExitStack, contextmanager
 from dataclasses import dataclass, field, replace
 from typing import Annotated, Any
 
@@ -297,6 +297,10 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
     # copy creates its own; `for_run_step` copies borrow the entered instance's pool by reference.
     _monty_pool: Monty | None = field(default=None, init=False, repr=False, compare=False)
 
+    # Owns both the synchronous Monty pool and asynchronous wrapped toolset. It is populated
+    # only after every resource enters successfully, then unwound in reverse entry order.
+    _exit_stack: AsyncExitStack | None = field(default=None, init=False, repr=False, compare=False)
+
     # Catalog string stashed during `get_tools` (when `dynamic_catalog`) and read back by
     # `get_instructions` in the same step. Empty when there's nothing to surface.
     _last_catalog: str = field(default='', init=False, repr=False)
@@ -329,19 +333,20 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
         (e.g. inside a durable-execution sandbox that forbids subprocesses) fails fast without
         leaving the wrapped toolset half-entered.
         """
-        self._monty_pool = Monty().__enter__()
-        await self.wrapped.__aenter__()
+        async with AsyncExitStack() as stack:
+            monty_pool = stack.enter_context(Monty())
+            await stack.enter_async_context(self.wrapped)
+            self._monty_pool = monty_pool
+            self._exit_stack = stack.pop_all()
         return self
 
     async def __aexit__(self, *args: Any) -> bool | None:
         """Exit the wrapped toolset, then tear down the worker pool."""
-        assert self._monty_pool is not None
-        monty_pool = self._monty_pool
+        exit_stack = self._exit_stack
+        assert exit_stack is not None
+        self._exit_stack = None
         self._monty_pool = None
-        try:
-            return await self.wrapped.__aexit__(*args)
-        finally:
-            monty_pool.__exit__(*args)
+        return await exit_stack.__aexit__(*args)
 
     @contextmanager
     def _acquire_pool(self) -> Generator[Monty]:
@@ -567,18 +572,11 @@ class CodeModeToolset(WrapperToolset[AgentDepsT]):
             # Serialize to JSON-compatible form so Monty receives only plain data.
             return _TOOL_RETURN_CONTENT_TA.dump_python(result)
 
-        # Static type checking on fresh REPL sessions (first call or after
-        # restart). Monty type-checks each fed snippet before executing it when
-        # the session is checked out with `type_check=True`, so a type error
-        # surfaces as `MontyTypingError` from `feed_start` before any tool runs.
-        # Skipped on subsequent (loaded) sessions: accumulated REPL state
-        # (variables from prior snippets) is invisible to the stateless checker,
-        # and `load` does not restore the checker's accumulated context anyway.
-        # `skip_type_check` is passed to `feed_start` too (not just `type_check`
-        # at checkout) so a loaded session never type-checks even if it lands on
-        # a pooled worker that a prior checkout left in type-check mode. Because
-        # `feed_start` raises before completing, `_repl_state` is left None on a
-        # type error, so the next retry is type-checked again.
+        # Type-check only the first snippet of a fresh REPL. Monty snapshots restore some checker
+        # context, but it can diverge from runtime state across snippets: imports are not available
+        # to the next check, and incrementally constructed dictionaries can be rejected where the
+        # runtime tool validator accepts them as TypedDict inputs. Keep runtime REPL state without
+        # introducing those false positives on later calls.
         type_check = fresh_repl and bool(callable_defs)
         type_check_stubs = self._build_type_check_stubs(callable_defs) if type_check else None
 

tests/code_mode/test_code_mode.py

@@ -435,6 +435,109 @@ class TestCodeMode:
     # `for_run` / `for_run_step` lifecycle
     # ---------------------------------------------------------------------------
 
+    async def test_enter_cleans_up_monty_if_wrapped_enter_fails(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """A later resource acquisition failure must close the already-entered Monty pool."""
+        events: list[str] = []
+
+        class TrackingMonty:
+            def __enter__(self) -> TrackingMonty:
+                events.append('monty enter')
+                return self
+
+            def __exit__(self, *args: Any) -> None:
+                events.append('monty exit')
+
+        class FailingToolset(FunctionToolset[object]):
+            async def __aenter__(self) -> FailingToolset:
+                events.append('wrapped enter')
+                raise RuntimeError('wrapped enter failed')
+
+        monkeypatch.setattr('pydantic_ai_harness.code_mode._toolset.Monty', TrackingMonty)
+        wrapper = CodeMode[object]().get_wrapper_toolset(FailingToolset())
+        assert isinstance(wrapper, CodeModeToolset)
+
+        with pytest.raises(RuntimeError, match='wrapped enter failed'):
+            await wrapper.__aenter__()
+
+        assert events == ['monty enter', 'wrapped enter', 'monty exit']
+
+    async def test_exit_releases_resources_in_reverse_entry_order(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """The wrapped toolset exits before the Monty pool it may depend on."""
+        events: list[str] = []
+
+        class TrackingMonty:
+            def __enter__(self) -> TrackingMonty:
+                events.append('monty enter')
+                return self
+
+            def __exit__(self, *args: Any) -> None:
+                events.append('monty exit')
+
+        class TrackingToolset(FunctionToolset[object]):
+            async def __aenter__(self) -> TrackingToolset:
+                events.append('wrapped enter')
+                return self
+
+            async def __aexit__(self, *args: Any) -> bool | None:
+                events.append('wrapped exit')
+                return None
+
+        monkeypatch.setattr('pydantic_ai_harness.code_mode._toolset.Monty', TrackingMonty)
+        wrapper = CodeMode[object]().get_wrapper_toolset(TrackingToolset())
+        assert isinstance(wrapper, CodeModeToolset)
+
+        async with wrapper:
+            assert events == ['monty enter', 'wrapped enter']
+
+        assert events == ['monty enter', 'wrapped enter', 'wrapped exit', 'monty exit']
+
+    async def test_agent_run_reuses_one_pool_with_separate_checkouts(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """The agent lifecycle owns one pool while each `run_code` call gets a session."""
+        from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
+        from pydantic_ai.models.function import AgentInfo, FunctionModel
+        from pydantic_monty import Monty as RealMonty
+
+        events: list[str] = []
+
+        class TrackingMonty:
+            def __init__(self, *args: Any, **kwargs: Any) -> None:
+                events.append('pool init')
+                self._pool = RealMonty(*args, **kwargs)
+
+            def __enter__(self) -> TrackingMonty:
+                events.append('pool enter')
+                self._pool.__enter__()
+                return self
+
+            def __exit__(self, *args: Any) -> bool | None:
+                events.append('pool exit')
+                return self._pool.__exit__(*args)
+
+            def checkout(self, *args: Any, **kwargs: Any) -> Any:
+                events.append('checkout')
+                return self._pool.checkout(*args, **kwargs)
+
+        def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+            response_count = sum(isinstance(message, ModelResponse) for message in messages)
+            if response_count == 0:
+                return ModelResponse(parts=[ToolCallPart('run_code', {'code': 'x = await add(a=1, b=2)'})])
+            if response_count == 1:
+                return ModelResponse(parts=[ToolCallPart('run_code', {'code': 'x * 10'})])
+            return ModelResponse(parts=[TextPart('done')])
+
+        monkeypatch.setattr('pydantic_ai_harness.code_mode._toolset.Monty', TrackingMonty)
+        agent: Agent[object, str] = Agent(FunctionModel(model_fn), capabilities=[CodeMode[object]()])
+
+        @agent.tool_plain
+        def add(a: int, b: int) -> int:  # pyright: ignore[reportUnusedFunction]
+            """Add two numbers."""
+            return a + b
+
+        result = await agent.run('use code mode twice')
+
+        assert result.output == 'done'
+        assert events == ['pool init', 'pool enter', 'checkout', 'checkout', 'pool exit']
+
     async def test_for_run_returns_fresh_instance_with_cleared_repl(self) -> None:
         """`for_run` must hand back a new toolset instance -- concurrent runs cannot share REPL state."""
         wrapper = CodeMode[object]().get_wrapper_toolset(_build_function_toolset(add))
@@ -1309,14 +1412,13 @@ class TestCodeMode:
         On a fresh REPL, the static type checker catches this before execution.
         """
 
-        # TODO: It is not being caught, why?
         wrapper = CodeMode[object]().get_wrapper_toolset(_build_function_toolset(add))
         assert isinstance(wrapper, CodeModeToolset)
         ctx = await build_ctx(None, wrapper)
         tools = await wrapper.get_tools(ctx)
 
         # Pass a string where int is expected -- type checker catches this.
-        with pytest.raises(ModelRetry, match='error in code'):
+        with pytest.raises(ModelRetry, match='Type error in code'):
             await wrapper.call_tool(
                 'run_code',
                 {'code': "await add(a='not_a_number', b=3)"},