From fff3c11129275c81f26bbe2b9c8d5b1052919900 Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Wed, 23 Apr 2025 15:00:54 -0400 Subject: [PATCH 1/6] Prevent MCP ClientSession hang Per https://modelcontextprotocol.io/specification/draft/basic/lifecycle#timeouts "Implementations SHOULD establish timeouts for all sent requests, to prevent hung connections and resource exhaustion. When the request has not received a success or error response within the timeout period, the sender SHOULD issue a cancellation notification for that request and stop waiting for a response. SDKs and other middleware SHOULD allow these timeouts to be configured on a per-request basis." --- src/agents/mcp/server.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 9a137bbdd..18ee1c4c1 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -3,6 +3,7 @@ import abc import asyncio from contextlib import AbstractAsyncContextManager, AsyncExitStack +from datetime import timedelta from pathlib import Path from typing import Any, Literal @@ -69,6 +70,9 @@ def __init__(self, cache_tools_list: bool): self._cleanup_lock: asyncio.Lock = asyncio.Lock() self.cache_tools_list = cache_tools_list + self.timeout: float = 5 + """The timeout for the MCP ClientSession. Defaults to 5 seconds.""" + # The cache is always dirty at startup, so that we fetch tools at least once self._cache_dirty = True self._tools_list: list[MCPTool] | None = None @@ -101,7 +105,11 @@ async def connect(self): try: transport = await self.exit_stack.enter_async_context(self.create_streams()) read, write = transport - session = await self.exit_stack.enter_async_context(ClientSession(read, write)) + session = await self.exit_stack.enter_async_context( + ClientSession( + read, write, read_timeout_seconds=timedelta(seconds=self.timeout) + ) + ) await session.initialize() self.session = session except Exception as e: From c6de1b29c914f642d93fa3031bc1de1570808a8c Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Thu, 24 Apr 2025 08:25:50 -0400 Subject: [PATCH 2/6] Update server.py --- src/agents/mcp/server.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 18ee1c4c1..fd8fa66c5 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -55,7 +55,7 @@ async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> C class _MCPServerWithClientSession(MCPServer, abc.ABC): """Base class for MCP servers that use a `ClientSession` to communicate with the server.""" - def __init__(self, cache_tools_list: bool): + def __init__(self, cache_tools_list: bool, client_session_timeout_seconds: float | None): """ Args: cache_tools_list: Whether to cache the tools list. If `True`, the tools list will be @@ -70,8 +70,7 @@ def __init__(self, cache_tools_list: bool): self._cleanup_lock: asyncio.Lock = asyncio.Lock() self.cache_tools_list = cache_tools_list - self.timeout: float = 5 - """The timeout for the MCP ClientSession. Defaults to 5 seconds.""" + self.client_session_timeout_seconds = client_session_timeout_seconds # The cache is always dirty at startup, so that we fetch tools at least once self._cache_dirty = True @@ -107,7 +106,9 @@ async def connect(self): read, write = transport session = await self.exit_stack.enter_async_context( ClientSession( - read, write, read_timeout_seconds=timedelta(seconds=self.timeout) + read, + write, + read_timeout_seconds=timedelta(seconds=self.client_session_timeout_seconds), ) ) await session.initialize() @@ -191,6 +192,7 @@ def __init__( params: MCPServerStdioParams, cache_tools_list: bool = False, name: str | None = None, + client_session_timeout_seconds: float | None = 5, ): """Create a new MCP server based on the stdio transport. @@ -208,7 +210,7 @@ def __init__( name: A readable name for the server. If not provided, we'll create one from the command. """ - super().__init__(cache_tools_list) + super().__init__(cache_tools_list, client_session_timeout_seconds) self.params = StdioServerParameters( command=params["command"], @@ -265,6 +267,7 @@ def __init__( params: MCPServerSseParams, cache_tools_list: bool = False, name: str | None = None, + client_session_timeout_seconds: float | None = 5, ): """Create a new MCP server based on the HTTP with SSE transport. @@ -283,7 +286,7 @@ def __init__( name: A readable name for the server. If not provided, we'll create one from the URL. """ - super().__init__(cache_tools_list) + super().__init__(cache_tools_list, client_session_timeout_seconds) self.params = params self._name = name or f"sse: {self.params['url']}" From d209c1ccf67f3e76e1bd02a671e124c4b475660b Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Thu, 24 Apr 2025 08:59:38 -0400 Subject: [PATCH 3/6] Update server.py --- src/agents/mcp/server.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index fd8fa66c5..6f7d4cca4 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -209,6 +209,7 @@ def __init__( improve latency (by avoiding a round-trip to the server every time). name: A readable name for the server. If not provided, we'll create one from the command. + client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. """ super().__init__(cache_tools_list, client_session_timeout_seconds) @@ -285,6 +286,8 @@ def __init__( name: A readable name for the server. If not provided, we'll create one from the URL. + + client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. """ super().__init__(cache_tools_list, client_session_timeout_seconds) From 4287a454f8bb363d4704bf327558b3dc08867d73 Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Thu, 24 Apr 2025 09:00:35 -0400 Subject: [PATCH 4/6] Update server.py --- src/agents/mcp/server.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 6f7d4cca4..b4cfce0dd 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -64,6 +64,8 @@ def __init__(self, cache_tools_list: bool, client_session_timeout_seconds: float by calling `invalidate_tools_cache()`. You should set this to `True` if you know the server will not change its tools list, because it can drastically improve latency (by avoiding a round-trip to the server every time). + + client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. """ self.session: ClientSession | None = None self.exit_stack: AsyncExitStack = AsyncExitStack() From 215867d34885d52646d81f1e655e51178fb7ba70 Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Thu, 24 Apr 2025 09:08:03 -0400 Subject: [PATCH 5/6] Update server.py --- src/agents/mcp/server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index b4cfce0dd..3af4e59e5 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -110,7 +110,9 @@ async def connect(self): ClientSession( read, write, - read_timeout_seconds=timedelta(seconds=self.client_session_timeout_seconds), + timedelta(seconds=self.client_session_timeout_seconds) + if self.client_session_timeout_seconds + else None, ) ) await session.initialize() From 8ac5dffa89d8ad3446a10267991149f89fcf7aee Mon Sep 17 00:00:00 2001 From: Nathan Brake Date: Thu, 24 Apr 2025 11:51:49 -0400 Subject: [PATCH 6/6] lint and fix a test --- src/agents/mcp/server.py | 2 +- tests/mcp/test_server_errors.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 3af4e59e5..9916c92b0 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -64,7 +64,7 @@ def __init__(self, cache_tools_list: bool, client_session_timeout_seconds: float by calling `invalidate_tools_cache()`. You should set this to `True` if you know the server will not change its tools list, because it can drastically improve latency (by avoiding a round-trip to the server every time). - + client_session_timeout_seconds: the read timeout passed to the MCP ClientSession. """ self.session: ClientSession | None = None diff --git a/tests/mcp/test_server_errors.py b/tests/mcp/test_server_errors.py index bdca7ce62..fbd8db17d 100644 --- a/tests/mcp/test_server_errors.py +++ b/tests/mcp/test_server_errors.py @@ -6,7 +6,7 @@ class CrashingClientSessionServer(_MCPServerWithClientSession): def __init__(self): - super().__init__(cache_tools_list=False) + super().__init__(cache_tools_list=False, client_session_timeout_seconds=5) self.cleanup_called = False def create_streams(self):