Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 31 additions & 9 deletions src/nooa/mcp/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 (
Expand All @@ -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:
Expand All @@ -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.

Expand All @@ -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
Expand All @@ -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'"
Expand Down
6 changes: 6 additions & 0 deletions src/nooa/mcp/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.

Expand All @@ -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).
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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():
Expand Down
80 changes: 79 additions & 1 deletion tests/test_mcp/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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():
Expand Down