Merge pull request #46 from callingmedic911/fix/mcp-tool-call-timeout

fix: honor tool_call_timeout so slow MCP tool calls stop hanging

Author: sklinglernv

src/nooa/mcp/client.py

@@ -17,6 +17,10 @@
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import streamable_http_client
 
+# Establishing the connection is not a tool call, so it keeps its own short budget.
+# Matches the connect timeout the MCP SDK's own SSE transport defaults to.
+CONNECT_TIMEOUT_SECONDS = 5.0
+
 
 class MCPBaseClient(ABC):
     """Base client for creating an MCP transport session and connecting to an MCP server.
@@ -120,7 +124,7 @@ async def connect_to_server(self):
                 url=self._url,
                 headers=self._headers if self._headers else None,
             ) as (read, write),
-            ClientSession(read, write) as session,
+            ClientSession(read, write, read_timeout_seconds=self._tool_call_timeout) as session,
         ):
             await session.initialize()
             yield session
@@ -197,7 +201,7 @@ async def connect_to_server(self):
         )
         async with (
             stdio_client(server_params) as (read, write),
-            ClientSession(read, write) as session,
+            ClientSession(read, write, read_timeout_seconds=self._tool_call_timeout) as session,
         ):
             await session.initialize()
             yield session
@@ -273,9 +277,19 @@ async def connect_to_server(self):
             httpx.HTTPStatusError: If server returns HTTP error (e.g., 401 Unauthorized, 500)
             RuntimeError: If session initialization fails (MCP protocol error)
         """
-        # Create httpx client with custom headers
-        # streamable_http_client expects a pre-configured httpx.AsyncClient for headers
-        http_client = httpx.AsyncClient(headers=self._headers if self._headers else None)
+        # Create httpx client with custom headers.
+        # streamable_http_client expects a pre-configured httpx.AsyncClient, so this
+        # client's timeouts are the only ones that apply: the transport has no timeout
+        # arguments of its own to fall back on. Reading the response has to be allowed
+        # to take as long as a tool call may take, or a slow tool's reply arrives on a
+        # stream httpx already abandoned and the caller waits forever.
+        http_client = httpx.AsyncClient(
+            headers=self._headers if self._headers else None,
+            timeout=httpx.Timeout(
+                self._tool_call_timeout.total_seconds(),
+                connect=CONNECT_TIMEOUT_SECONDS,
+            ),
+        )
 
         try:
             async with (
@@ -288,7 +302,9 @@ async def connect_to_server(self):
             ):
                 # Store the session ID callback for later retrieval
                 self._get_mcp_session_id = get_session_id
-                async with ClientSession(read, write) as session:
+                async with ClientSession(
+                    read, write, read_timeout_seconds=self._tool_call_timeout
+                ) as session:
                     await session.initialize()
                     yield session
         finally:
@@ -303,6 +319,7 @@ def create_mcp_client(
     args: list[str] | None = None,
     env: dict[str, str] | None = None,
     headers: dict[str, str] | None = None,
+    tool_call_timeout: timedelta = timedelta(seconds=60),
 ) -> MCPBaseClient:
     """Create an MCP client based on the transport type and configuration.
 
@@ -313,6 +330,7 @@ def create_mcp_client(
         args: Command arguments (optional, for stdio transport)
         env: Environment variables for the server process (optional, for stdio transport)
         headers: Optional custom HTTP headers to include in requests (for HTTP transports)
+        tool_call_timeout: How long one tool call may take before it fails
 
     Returns:
         An MCPBaseClient instance configured for the specified transport
@@ -328,15 +346,19 @@ def create_mcp_client(
         case "stdio":
             if command is None:
                 raise ValueError("command must be provided for stdio transport")
-            return MCPStdioClient(command=command, args=args, env=env)
+            return MCPStdioClient(
+                command=command, args=args, env=env, tool_call_timeout=tool_call_timeout
+            )
         case "sse":
             if url is None:
                 raise ValueError("url must be provided for sse transport")
-            return MCPSSEClient(url=url, headers=headers)
+            return MCPSSEClient(url=url, headers=headers, tool_call_timeout=tool_call_timeout)
         case "streamable-http":
             if url is None:
                 raise ValueError("url must be provided for streamable-http transport")
-            return MCPStreamableHTTPClient(url=url, headers=headers)
+            return MCPStreamableHTTPClient(
+                url=url, headers=headers, tool_call_timeout=tool_call_timeout
+            )
         case _:
             raise ValueError(
                 f"Unsupported transport type: {transport}. Use 'stdio', 'sse', or 'streamable-http'"

src/nooa/mcp/tool.py

@@ -17,6 +17,7 @@
 import types
 from collections.abc import Awaitable, Callable, Sequence
 from dataclasses import dataclass, field
+from datetime import timedelta
 from pathlib import Path
 from typing import Any, Literal
 
@@ -767,6 +768,7 @@ def create_from_server(
         oauth_timeout: float | None = None,
         mcp_file: Path | None = None,
         servers: dict[str, dict[str, Any]] | None = None,
+        tool_call_timeout: timedelta = timedelta(seconds=60),
     ) -> MCPTool:
         """Create a per-server tool instance; connects to the MCP server.
 
@@ -789,6 +791,8 @@ def create_from_server(
             oauth_browser_open: Async hook to open the auth URL in a reachable browser (host handoff).
             mcp_file: Path to .mcp.json file (default: .mcp.json in cwd)
             servers: Optional inline server config from the TUI config.toml.
+            tool_call_timeout: How long one tool call may take before it fails.
+                Raise it for servers whose tools wrap slow work such as an LLM call.
 
         Returns:
             An MCPTool instance (dynamically generated class with methods for each tool).
@@ -845,6 +849,7 @@ def _run_sync(coro):
             args=args,
             env=env,
             headers=headers,
+            tool_call_timeout=tool_call_timeout,
         )
 
         # Connect and list tools (with OAuth retry if needed)
@@ -890,6 +895,7 @@ async def _connect_and_list():
                         args=args,
                         env=env,
                         headers=headers,
+                        tool_call_timeout=tool_call_timeout,
                     )
 
                     async def _connect_and_list_retry():

tests/test_mcp/test_client.py

@@ -371,6 +371,83 @@ async def test_sse_headers_passed_to_transport(
     )
 
 
+@pytest.mark.asyncio
+async def test_streamable_http_applies_tool_call_timeout(
+    streamable_http_client: MCPStreamableHTTPClient,
+    mock_client_session: AsyncMock,
+):
+    """streamable-http gives httpx and the session the caller's tool_call_timeout.
+
+    The transport has no timeout arguments of its own and uses whatever client it is
+    handed, so an httpx client built without one caps every tool call at httpx's 5s
+    default no matter what tool_call_timeout says.
+    """
+    with (
+        patch("nooa.mcp.client.httpx.AsyncClient") as mock_async_client,
+        patch("nooa.mcp.client.streamable_http_client") as mock_http,
+        patch("nooa.mcp.client.ClientSession") as mock_session_class,
+    ):
+        mock_http.return_value.__aenter__.return_value = (MagicMock(), MagicMock(), MagicMock())
+        mock_session_class.return_value.__aenter__.return_value = mock_client_session
+
+        async with streamable_http_client.connect_to_server():
+            pass
+
+    timeout = mock_async_client.call_args.kwargs["timeout"]
+    assert timeout.read == 90
+    assert timeout.write == 90
+    # Opening the connection is not a tool call and keeps its own short budget.
+    assert timeout.connect == 5.0
+    assert mock_session_class.call_args.kwargs["read_timeout_seconds"] == timedelta(seconds=90)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "client_fixture, transport_patch, expected_timeout",
+    [
+        ("sse_client", "nooa.mcp.client.sse_client", timedelta(seconds=45)),
+        ("stdio_client", "nooa.mcp.client.stdio_client", timedelta(seconds=30)),
+    ],
+)
+async def test_session_enforces_tool_call_timeout(
+    client_fixture: str,
+    transport_patch: str,
+    expected_timeout: timedelta,
+    request: pytest.FixtureRequest,
+    mock_mcp_transport: tuple[MagicMock, MagicMock],
+    mock_client_session: AsyncMock,
+):
+    """Every transport hands tool_call_timeout to the session that enforces it."""
+    client: MCPBaseClient = request.getfixturevalue(client_fixture)
+
+    with (
+        patch(transport_patch) as mock_transport,
+        patch("nooa.mcp.client.ClientSession") as mock_session_class,
+    ):
+        mock_transport.return_value.__aenter__.return_value = mock_mcp_transport
+        mock_session_class.return_value.__aenter__.return_value = mock_client_session
+
+        async with client.connect_to_server():
+            pass
+
+    assert mock_session_class.call_args.kwargs["read_timeout_seconds"] == expected_timeout
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"transport": "stdio", "command": "python"},
+        {"transport": "sse", "url": "https://example.test/sse"},
+        {"transport": "streamable-http", "url": "https://example.test/mcp"},
+    ],
+)
+def test_create_mcp_client_forwards_tool_call_timeout(kwargs: dict[str, str]):
+    """create_mcp_client is the documented entry point and must pass the timeout on."""
+    client = create_mcp_client(tool_call_timeout=timedelta(seconds=7), **kwargs)
+
+    assert client.tool_call_timeout == timedelta(seconds=7)
+
+
 @pytest.mark.asyncio
 async def test_streamable_http_connect_context_manager(
     streamable_http_client: MCPStreamableHTTPClient,
@@ -439,7 +516,8 @@ async def test_streamable_http_headers_passed_to_httpx_client(
                     pass
 
                 # Verify httpx.AsyncClient was created with expected headers
-                mock_httpx_client.assert_called_once_with(headers=expected_headers)
+                mock_httpx_client.assert_called_once()
+                assert mock_httpx_client.call_args.kwargs["headers"] == expected_headers
 
 
 def test_dynamic_method_supports_json_container_defaults():