Raise error if MCPServer.__aexit__ is called when _running_count is already 0 (#2696)

Author: federicociner

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -300,6 +300,8 @@ async def __aenter__(self) -> Self:
         return self
 
     async def __aexit__(self, *args: Any) -> bool | None:
+        if self._running_count == 0:
+            raise ValueError('MCPServer.__aexit__ called more times than __aenter__')
         async with self._enter_lock:
             self._running_count -= 1
             if self._running_count == 0 and self._exit_stack is not None:

tests/test_mcp.py

@@ -107,6 +107,19 @@ async def test_context_manager_initialization_error() -> None:
     assert server._write_stream._closed  # pyright: ignore[reportPrivateUsage]
 
 
+async def test_aexit_called_more_times_than_aenter():
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+
+    with pytest.raises(ValueError, match='MCPServer.__aexit__ called more times than __aenter__'):
+        await server.__aexit__(None, None, None)
+
+    async with server:
+        pass  # This will call __aenter__ and __aexit__ once each
+
+    with pytest.raises(ValueError, match='MCPServer.__aexit__ called more times than __aenter__'):
+        await server.__aexit__(None, None, None)
+
+
 async def test_stdio_server_with_tool_prefix(run_context: RunContext[int]):
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], tool_prefix='foo')
     async with server: