fix: serialize streamable-http MCP requests per session (#2682)

Author: elainegan-openai

src/agents/mcp/server.py

@@ -336,6 +336,7 @@ def __init__(
         self.session: ClientSession | None = None
         self.exit_stack: AsyncExitStack = AsyncExitStack()
         self._cleanup_lock: asyncio.Lock = asyncio.Lock()
+        self._request_lock: asyncio.Lock = asyncio.Lock()
         self.cache_tools_list = cache_tools_list
         self.server_initialize_result: InitializeResult | None = None
 
@@ -349,6 +350,13 @@ def __init__(
         self._tools_list: list[MCPTool] | None = None
 
         self.tool_filter = tool_filter
+        self._serialize_session_requests = False
+
+    async def _maybe_serialize_request(self, func: Callable[[], Awaitable[T]]) -> T:
+        if not self._serialize_session_requests:
+            return await func()
+        async with self._request_lock:
+            return await func()
 
     async def _apply_tool_filter(
         self,
@@ -573,7 +581,9 @@ async def list_tools(
                 tools = self._tools_list
             else:
                 # Fetch the tools from the server
-                result = await self._run_with_retries(lambda: session.list_tools())
+                result = await self._run_with_retries(
+                    lambda: self._maybe_serialize_request(lambda: session.list_tools())
+                )
                 self._tools_list = result.tools
                 self._cache_dirty = False
                 tools = self._tools_list
@@ -609,9 +619,15 @@ async def call_tool(
         try:
             self._validate_required_parameters(tool_name=tool_name, arguments=arguments)
             if meta is None:
-                return await self._run_with_retries(lambda: session.call_tool(tool_name, arguments))
+                return await self._run_with_retries(
+                    lambda: self._maybe_serialize_request(
+                        lambda: session.call_tool(tool_name, arguments)
+                    )
+                )
             return await self._run_with_retries(
-                lambda: session.call_tool(tool_name, arguments, meta=meta)
+                lambda: self._maybe_serialize_request(
+                    lambda: session.call_tool(tool_name, arguments, meta=meta)
+                )
             )
         except httpx.HTTPStatusError as e:
             status_code = e.response.status_code
@@ -665,17 +681,19 @@ async def list_prompts(
         """List the prompts available on the server."""
         if not self.session:
             raise UserError("Server not initialized. Make sure you call `connect()` first.")
-
-        return await self.session.list_prompts()
+        session = self.session
+        assert session is not None
+        return await self._maybe_serialize_request(lambda: session.list_prompts())
 
     async def get_prompt(
         self, name: str, arguments: dict[str, Any] | None = None
     ) -> GetPromptResult:
         """Get a specific prompt from the server."""
         if not self.session:
             raise UserError("Server not initialized. Make sure you call `connect()` first.")
-
-        return await self.session.get_prompt(name, arguments)
+        session = self.session
+        assert session is not None
+        return await self._maybe_serialize_request(lambda: session.get_prompt(name, arguments))
 
     async def cleanup(self):
         """Cleanup the server."""
@@ -1084,6 +1102,7 @@ def __init__(
 
         self.params = params
         self._name = name or f"streamable_http: {self.params['url']}"
+        self._serialize_session_requests = True
 
     def create_streams(
         self,

tests/mcp/test_client_session_retries.py

@@ -1,3 +1,4 @@
+import asyncio
 from typing import cast
 
 import pytest
@@ -29,14 +30,15 @@ async def list_tools(self):
 
 
 class DummyServer(_MCPServerWithClientSession):
-    def __init__(self, session: DummySession, retries: int):
+    def __init__(self, session: DummySession, retries: int, *, serialize_requests: bool = False):
         super().__init__(
             cache_tools_list=False,
             client_session_timeout_seconds=None,
             max_retry_attempts=retries,
             retry_backoff_seconds_base=0,
         )
         self.session = cast(ClientSession, session)
+        self._serialize_session_requests = serialize_requests
 
     def create_streams(self):
         raise NotImplementedError
@@ -148,3 +150,70 @@ async def test_call_tool_rejects_non_object_arguments_before_remote_call():
         await server.call_tool("tool", cast(dict[str, object] | None, ["bad"]))
 
     assert session.call_tool_attempts == 0
+
+
+class ConcurrentCancellationSession:
+    def __init__(self):
+        self._slow_task: asyncio.Task[CallToolResult] | None = None
+        self._slow_started = asyncio.Event()
+
+    async def call_tool(self, tool_name, arguments, meta=None):
+        if tool_name == "slow":
+            self._slow_task = cast(asyncio.Task[CallToolResult], asyncio.current_task())
+            self._slow_started.set()
+            await asyncio.sleep(0.1)
+            return CallToolResult(content=[])
+
+        await self._slow_started.wait()
+        assert self._slow_task is not None
+        self._slow_task.cancel()
+        raise RuntimeError("synthetic request failure")
+
+    async def list_tools(self):
+        return ListToolsResult(tools=[MCPTool(name="tool", inputSchema={})])
+
+    async def list_prompts(self):
+        await self._slow_started.wait()
+        assert self._slow_task is not None
+        self._slow_task.cancel()
+        raise RuntimeError("synthetic request failure")
+
+    async def get_prompt(self, name, arguments=None):
+        await self._slow_started.wait()
+        assert self._slow_task is not None
+        self._slow_task.cancel()
+        raise RuntimeError("synthetic request failure")
+
+
+@pytest.mark.asyncio
+async def test_serialized_session_requests_prevent_sibling_cancellation():
+    session = ConcurrentCancellationSession()
+    server = DummyServer(session=cast(DummySession, session), retries=0, serialize_requests=True)
+
+    results = await asyncio.gather(
+        server.call_tool("slow", None),
+        server.call_tool("fail", None),
+        return_exceptions=True,
+    )
+
+    assert isinstance(results[0], CallToolResult)
+    assert isinstance(results[1], RuntimeError)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("prompt_method", ["list_prompts", "get_prompt"])
+async def test_serialized_prompt_requests_prevent_tool_cancellation(prompt_method: str):
+    session = ConcurrentCancellationSession()
+    server = DummyServer(session=cast(DummySession, session), retries=0, serialize_requests=True)
+
+    prompt_request = (
+        server.list_prompts() if prompt_method == "list_prompts" else server.get_prompt("prompt")
+    )
+    results = await asyncio.gather(
+        server.call_tool("slow", None),
+        prompt_request,
+        return_exceptions=True,
+    )
+
+    assert isinstance(results[0], CallToolResult)
+    assert isinstance(results[1], RuntimeError)