From 6631970daeb56438106b71e5b0351e29d33e2704 Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Thu, 24 Apr 2025 08:49:07 -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/mcpadapt/core.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/mcpadapt/core.py b/src/mcpadapt/core.py index 39c767c..30199a9 100644 --- a/src/mcpadapt/core.py +++ b/src/mcpadapt/core.py @@ -8,6 +8,7 @@ import threading from abc import ABC, abstractmethod from contextlib import AsyncExitStack, asynccontextmanager +from datetime import timedelta from functools import partial from typing import Any, AsyncGenerator, Callable, Coroutine @@ -71,6 +72,7 @@ def async_adapt( @asynccontextmanager async def mcptools( serverparams: StdioServerParameters | dict[str, Any], + client_session_timeout_seconds: float | None = 5, ) -> AsyncGenerator[tuple[ClientSession, list[mcp.types.Tool]], None]: """Async context manager that yields tools from an MCP server. @@ -81,7 +83,7 @@ async def mcptools( serverparams: Parameters passed to either the stdio client or sse client. * if StdioServerParameters, run the MCP server using the stdio protocol. * if dict, assume the dict corresponds to parameters to an sse MCP server. - + client_session_timeout_seconds: Timeout for MCP ClientSession HTTP calls Yields: A tuple of (MCP Client Session, list of MCP tools) available on the MCP server. @@ -99,7 +101,7 @@ async def mcptools( ) async with client as (read, write): - async with ClientSession(read, write) as session: + async with ClientSession(read, write, timedelta(seconds=client_session_timeout_seconds)) as session: # Initialize the connection and get the tools from the mcp server await session.initialize() tools = await session.list_tools() From 52b1faa360110151a73a48d854f113c8d32087f0 Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Thu, 24 Apr 2025 09:09:52 -0400 Subject: [PATCH 2/6] Update core.py --- src/mcpadapt/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcpadapt/core.py b/src/mcpadapt/core.py index 30199a9..596093e 100644 --- a/src/mcpadapt/core.py +++ b/src/mcpadapt/core.py @@ -101,7 +101,7 @@ async def mcptools( ) async with client as (read, write): - async with ClientSession(read, write, timedelta(seconds=client_session_timeout_seconds)) as session: + async with ClientSession(read, write, timedelta(seconds=client_session_timeout_seconds) if client_session_timeout_seconds else None) as session: # Initialize the connection and get the tools from the mcp server await session.initialize() tools = await session.list_tools() From 56c7218f76b3dfe1fdddf13b532d801fa96fa745 Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Thu, 24 Apr 2025 09:10:13 -0400 Subject: [PATCH 3/6] Update core.py --- src/mcpadapt/core.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mcpadapt/core.py b/src/mcpadapt/core.py index 596093e..8715517 100644 --- a/src/mcpadapt/core.py +++ b/src/mcpadapt/core.py @@ -84,6 +84,7 @@ async def mcptools( * if StdioServerParameters, run the MCP server using the stdio protocol. * if dict, assume the dict corresponds to parameters to an sse MCP server. client_session_timeout_seconds: Timeout for MCP ClientSession HTTP calls + Yields: A tuple of (MCP Client Session, list of MCP tools) available on the MCP server. From 89a2f758318d1cf0e1f7f598c1a71e5fee5f30cd Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Tue, 29 Apr 2025 08:10:57 -0400 Subject: [PATCH 4/6] Update core.py --- src/mcpadapt/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/mcpadapt/core.py b/src/mcpadapt/core.py index 8715517..8cfac88 100644 --- a/src/mcpadapt/core.py +++ b/src/mcpadapt/core.py @@ -83,7 +83,7 @@ async def mcptools( serverparams: Parameters passed to either the stdio client or sse client. * if StdioServerParameters, run the MCP server using the stdio protocol. * if dict, assume the dict corresponds to parameters to an sse MCP server. - client_session_timeout_seconds: Timeout for MCP ClientSession HTTP calls + client_session_timeout_seconds: Timeout for MCP ClientSession calls Yields: A tuple of (MCP Client Session, list of MCP tools) available on the MCP server. From 632155f3472a364d70b4be492514175ec658d42a Mon Sep 17 00:00:00 2001 From: Nathan Brake Date: Tue, 29 Apr 2025 09:51:21 -0400 Subject: [PATCH 5/6] Update based on PR comments --- src/mcpadapt/core.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/mcpadapt/core.py b/src/mcpadapt/core.py index 8cfac88..c23a694 100644 --- a/src/mcpadapt/core.py +++ b/src/mcpadapt/core.py @@ -72,7 +72,7 @@ def async_adapt( @asynccontextmanager async def mcptools( serverparams: StdioServerParameters | dict[str, Any], - client_session_timeout_seconds: float | None = 5, + client_session_timeout_seconds: float | timedelta | None = 5, ) -> AsyncGenerator[tuple[ClientSession, list[mcp.types.Tool]], None]: """Async context manager that yields tools from an MCP server. @@ -101,8 +101,19 @@ async def mcptools( f"Invalid serverparams, expected StdioServerParameters or dict found `{type(serverparams)}`" ) + timeout = None + if client_session_timeout_seconds is not None: + if isinstance(client_session_timeout_seconds, float): + timeout = timedelta(seconds=client_session_timeout_seconds) + elif isinstance(client_session_timeout_seconds, timedelta): + timeout = client_session_timeout_seconds + async with client as (read, write): - async with ClientSession(read, write, timedelta(seconds=client_session_timeout_seconds) if client_session_timeout_seconds else None) as session: + async with ClientSession( + read, + write, + timeout, + ) as session: # Initialize the connection and get the tools from the mcp server await session.initialize() tools = await session.list_tools() From 32820cc1df96763abf56a654c1b1c72c56b4eb8e Mon Sep 17 00:00:00 2001 From: Nathan Brake <33383515+njbrake@users.noreply.github.com> Date: Tue, 29 Apr 2025 12:05:46 -0400 Subject: [PATCH 6/6] Update core.py Co-authored-by: Guillaume Raille --- src/mcpadapt/core.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/mcpadapt/core.py b/src/mcpadapt/core.py index c23a694..5842a18 100644 --- a/src/mcpadapt/core.py +++ b/src/mcpadapt/core.py @@ -102,11 +102,10 @@ async def mcptools( ) timeout = None - if client_session_timeout_seconds is not None: - if isinstance(client_session_timeout_seconds, float): - timeout = timedelta(seconds=client_session_timeout_seconds) - elif isinstance(client_session_timeout_seconds, timedelta): - timeout = client_session_timeout_seconds + if isinstance(client_session_timeout_seconds, float): + timeout = timedelta(seconds=client_session_timeout_seconds) + elif isinstance(client_session_timeout_seconds, timedelta): + timeout = client_session_timeout_seconds async with client as (read, write): async with ClientSession(