Add tests

Author: DouweM

pydantic_ai_slim/pydantic_ai/_output.py

@@ -114,7 +114,7 @@ async def validate(
                     content=r.message,
                     tool_name=run_context.tool_name,
                 )
-                if run_context.tool_call_id:
+                if run_context.tool_call_id:  # pragma: no cover
                     m.tool_call_id = run_context.tool_call_id
                 raise ToolRetryError(m) from r
             else:

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -218,6 +218,11 @@ async def __aexit__(
         if self._running_count <= 0:
             await self._exit_stack.aclose()
 
+    @property
+    def is_running(self) -> bool:
+        """Check if the MCP server is running."""
+        return bool(self._running_count)
+
     async def _sampling_callback(
         self, context: RequestContext[ClientSession, Any], params: mcp_types.CreateMessageRequestParams
     ) -> mcp_types.CreateMessageResult | mcp_types.ErrorData:

pydantic_ai_slim/pydantic_ai/toolsets/combined.py

@@ -24,11 +24,11 @@ class CombinedToolset(AbstractToolset[AgentDepsT]):
     toolsets: list[AbstractToolset[AgentDepsT]]
     _toolset_per_tool_name: dict[str, AbstractToolset[AgentDepsT]]
     _exit_stack: AsyncExitStack | None
-    _running_count: int
+    _entered_count: int
 
     def __init__(self, toolsets: Sequence[AbstractToolset[AgentDepsT]]):
         self._exit_stack = None
-        self._running_count = 0
+        self._entered_count = 0
         self.toolsets = list(toolsets)
 
         self._toolset_per_tool_name = {}
@@ -44,18 +44,18 @@ def __init__(self, toolsets: Sequence[AbstractToolset[AgentDepsT]]):
                 self._toolset_per_tool_name[name] = toolset
 
     async def __aenter__(self) -> Self:
-        if self._running_count == 0:
+        if self._entered_count == 0:
             self._exit_stack = AsyncExitStack()
             for toolset in self.toolsets:
                 await self._exit_stack.enter_async_context(toolset)
-        self._running_count += 1
+        self._entered_count += 1
         return self
 
     async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
     ) -> bool | None:
-        self._running_count -= 1
-        if self._running_count <= 0 and self._exit_stack is not None:
+        self._entered_count -= 1
+        if self._entered_count <= 0 and self._exit_stack is not None:
             await self._exit_stack.aclose()
             self._exit_stack = None
         return None

tests/test_agent.py

@@ -3600,11 +3600,24 @@ async def only_if_plan_presented(
     )
 
 
-async def test_reentrant_context_manager():
-    agent = Agent('test')
+async def test_context_manager():
+    try:
+        from pydantic_ai.mcp import MCPServerStdio
+    except ImportError:
+        return
+
+    server1 = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    server2 = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    toolset = CombinedToolset([server1, PrefixedToolset(server2, 'prefix')])
+    agent = Agent('test', toolsets=[toolset])
+
     async with agent:
+        assert server1.is_running
+        assert server2.is_running
+
         async with agent:
-            pass
+            assert server1.is_running
+            assert server2.is_running
 
 
 def test_set_mcp_sampling_model():
@@ -3616,7 +3629,7 @@ def test_set_mcp_sampling_model():
     test_model = TestModel()
     server1 = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     server2 = MCPServerStdio('python', ['-m', 'tests.mcp_server'], sampling_model=test_model)
-    toolset = CombinedToolset([server1, PrefixedToolset(server2, 'prefix_')])
+    toolset = CombinedToolset([server1, PrefixedToolset(server2, 'prefix')])
     agent = Agent(None, toolsets=[toolset])
 
     with pytest.raises(UserError, match='No sampling model provided and no model set on the agent.'):

tests/test_logfire.py

@@ -554,7 +554,7 @@ async def add_numbers(x: int, y: int) -> int:
         assert result.output == snapshot('{"add_numbers":84}')
     except UnexpectedModelBehavior:
         if not tool_error:
-            raise
+            raise  # pragma: no cover
 
     summary = get_logfire_summary()