From 0bf4ff3211d38aae31cf12905d6f73067fe3fda2 Mon Sep 17 00:00:00 2001 From: hyeonsang010716 Date: Thu, 17 Sep 2026 10:44:43 +0900 Subject: [PATCH] fix(mcp): keep the Streamable HTTP session usable after a 5xx With MCP Python SDK v2, the response hook that `MCPServerStreamableHttp` installs on its HTTP client called `raise_for_status()` for every 5xx. The hook runs inside the MCP transport task group, so a single transient 5xx on any request tore down the shared transport. Every later `call_tool()` and `list_tools()` then failed with `MCPError: Connection closed` until `cleanup()` and `connect()`. A 5xx on the `notifications/initialized` notification to a legacy-protocol server did the same while `connect()` still reported success. The hook now skips the raise only for a post-handshake MCP transport message, meaning a request body that is a JSON-RPC message other than `server/discover` and `initialize`. MCP v2 fails such a message on its own, so the session stays usable and `max_retry_attempts` retries on it. Every other response keeps the existing HTTP error mapping. A handshake 5xx still fails `connect()` with the HTTP status, so a failed discovery probe is not mistaken for a legacy server. OAuth auth-flow sub-requests such as token and registration requests are not transport messages, so their failures also stay on that path instead of reaching MCP's OAuth exceptions, which carry the authorization server's response body. --- src/agents/mcp/server.py | 31 ++++- tests/mcp/test_mcp_v2_http.py | 245 +++++++++++++++++++++++++++++++--- 2 files changed, 255 insertions(+), 21 deletions(-) diff --git a/src/agents/mcp/server.py b/src/agents/mcp/server.py index 763b6a6fb8..304a9181e5 100644 --- a/src/agents/mcp/server.py +++ b/src/agents/mcp/server.py @@ -355,15 +355,29 @@ def create_client( return create_client -def _jsonrpc_request_method(request: Any) -> str | None: +def _jsonrpc_message_payload(request: Any) -> dict[str, Any] | None: + """Return the JSON-RPC message this request carries, or `None` for any other request.""" try: payload = json.loads(request.content) except (TypeError, ValueError, UnicodeDecodeError): return None - if not isinstance(payload, dict): + if not isinstance(payload, dict) or "jsonrpc" not in payload: return None + return payload + + +_V2_HANDSHAKE_REQUEST_METHODS = frozenset({"server/discover", "initialize"}) + + +def _is_v2_post_handshake_message(payload: dict[str, Any] | None) -> bool: + """Whether MCP v2 owns the failure handling for this transport message.""" + if payload is None: + return False method = payload.get("method") - return method if isinstance(method, str) else None + if method is None: + # A response the client sends for a server-initiated request. + return True + return isinstance(method, str) and method not in _V2_HANDSHAKE_REQUEST_METHODS def _configure_v2_session_id_hook( @@ -372,12 +386,17 @@ def _configure_v2_session_id_hook( on_session_id: Callable[[str], None] | None, ) -> None: async def handle_response(response: Any) -> None: - if response.status_code >= 500: + payload = _jsonrpc_message_payload(response.request) + if response.status_code >= 500 and not _is_v2_post_handshake_message(payload): + # MCP v2 fails a post-handshake transport message on its own, so raising here would + # tear down the transport that every later request shares. Every other response keeps + # the existing HTTP error mapping: a handshake 5xx must fail the connection instead of + # looking like a legacy server, and OAuth sub-request failures stay on that path too. response.raise_for_status() - method = _jsonrpc_request_method(response.request) if ( on_session_id is not None - and method == "initialize" + and payload is not None + and payload.get("method") == "initialize" and 200 <= response.status_code < 300 ): session_id = response.headers.get("mcp-session-id") diff --git a/tests/mcp/test_mcp_v2_http.py b/tests/mcp/test_mcp_v2_http.py index cf66823a12..0f12994a1a 100644 --- a/tests/mcp/test_mcp_v2_http.py +++ b/tests/mcp/test_mcp_v2_http.py @@ -3,6 +3,7 @@ import asyncio import json import socket +import traceback from typing import Any import httpx @@ -14,7 +15,7 @@ from agents.exceptions import UserError from agents.mcp import MCPServerStreamableHttp -from agents.mcp._compat import MCP_V2, create_v2_client +from agents.mcp._compat import MCP_V2, MCPError, create_v2_client from agents.mcp.server import ( _configure_v2_session_id_hook, _create_default_streamable_http_client, @@ -103,6 +104,30 @@ def handle_request(request): await client.aclose() +@pytest.mark.asyncio +async def test_v2_response_hook_raises_5xx_for_a_request_that_is_not_a_transport_message(): + def handle_request(request): + return httpx2.Response(503, request=request) + + client = httpx2.AsyncClient(transport=httpx2.MockTransport(handle_request)) + _configure_v2_session_id_hook(client, on_session_id=None) + + # An OAuth dynamic client registration body is JSON, but it is not a transport message, so + # its failures stay on the HTTP error path instead of MCP's body-bearing OAuth exceptions. + with pytest.raises(httpx2.HTTPStatusError): + await client.post( + "https://example.test/register", + content=json.dumps( + { + "client_name": "example", + "redirect_uris": ["https://example.test/callback"], + } + ), + ) + + await client.aclose() + + def test_v2_rejects_initialized_notification_tolerance_before_connecting(): server = MCPServerStreamableHttp( params={ @@ -261,34 +286,68 @@ async def handler(request): assert all(client.is_closed for client in clients) -@pytest.mark.asyncio -async def test_v2_streamable_http_retries_5xx_on_isolated_session(): - clients: list[Any] = [] - observed_statuses: list[int] = [] +def _first_tool_call_returns_503_factory(clients: list[Any], tool_call_statuses: list[int]): + """Build clients whose server answers only the first `tools/call` with HTTP 503.""" def factory(headers=None, timeout=None, auth=None): - tool_status_code = 503 if not clients else None - async def handler(request): - return _v2_response_for_request(request, tool_status_code=tool_status_code) - - async def observe_response(response): - observed_statuses.append(response.status_code) + payload = json.loads(request.content) if request.content else {} + if payload.get("method") == "tools/call": + status_code = 503 if not tool_call_statuses else 200 + tool_call_statuses.append(status_code) + if status_code == 503: + return _v2_response_for_request(request, tool_status_code=503) + return _v2_response_for_request(request) client = httpx2.AsyncClient( transport=httpx2.MockTransport(handler), headers=headers, timeout=timeout, auth=auth, - event_hooks={"response": [observe_response]}, ) clients.append(client) return client + return factory + + +@pytest.mark.asyncio +async def test_v2_streamable_http_5xx_fails_only_that_request(): + clients: list[Any] = [] + tool_call_statuses: list[int] = [] server = MCPServerStreamableHttp( params={ "url": "https://example.test/mcp", - "httpx_client_factory": factory, + "httpx_client_factory": _first_tool_call_returns_503_factory( + clients, tool_call_statuses + ), + }, + ) + + async with server: + with pytest.raises(MCPError): + await asyncio.wait_for(server.call_tool("test", {}), timeout=2) + result = await asyncio.wait_for(server.call_tool("test", {}), timeout=2) + tools = await asyncio.wait_for(server.list_tools(), timeout=2) + + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "ok" + assert [tool.name for tool in tools] == ["test"] + assert tool_call_statuses == [503, 200] + assert len(clients) == 1 + assert all(client.is_closed for client in clients) + + +@pytest.mark.asyncio +async def test_v2_streamable_http_retries_5xx_on_shared_session(): + clients: list[Any] = [] + tool_call_statuses: list[int] = [] + server = MCPServerStreamableHttp( + params={ + "url": "https://example.test/mcp", + "httpx_client_factory": _first_tool_call_returns_503_factory( + clients, tool_call_statuses + ), }, max_retry_attempts=1, retry_backoff_seconds_base=0, @@ -299,11 +358,167 @@ async def observe_response(response): assert isinstance(result.content[0], TextContent) assert result.content[0].text == "ok" - assert len(clients) == 2 - assert 503 in observed_statuses + assert tool_call_statuses == [503, 200] + assert len(clients) == 1 assert all(client.is_closed for client in clients) +@pytest.mark.asyncio +async def test_v2_streamable_http_initialized_notification_5xx_keeps_session_usable(): + clients: list[Any] = [] + + def factory(headers=None, timeout=None, auth=None): + async def handler(request): + payload = json.loads(request.content) if request.content else {} + if payload.get("method") == "notifications/initialized": + return httpx2.Response(503, request=request) + return _v2_response_for_request(request) + + client = httpx2.AsyncClient( + transport=httpx2.MockTransport(handler), + headers=headers, + timeout=timeout, + auth=auth, + ) + clients.append(client) + return client + + server = MCPServerStreamableHttp( + params={ + "url": "https://example.test/mcp", + "httpx_client_factory": factory, + }, + ) + + async with server: + result = await asyncio.wait_for(server.call_tool("test", {}), timeout=2) + tools = await asyncio.wait_for(server.list_tools(), timeout=2) + + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == "ok" + assert [tool.name for tool in tools] == ["test"] + assert len(clients) == 1 + + +@pytest.mark.asyncio +async def test_v2_streamable_http_handshake_5xx_fails_connect_without_legacy_fallback(): + methods: list[str | None] = [] + + def factory(headers=None, timeout=None, auth=None): + async def handler(request): + payload = json.loads(request.content) if request.content else {} + methods.append(payload.get("method")) + if payload.get("method") == "server/discover": + return httpx2.Response(503, request=request) + return _v2_response_for_request(request) + + return httpx2.AsyncClient( + transport=httpx2.MockTransport(handler), + headers=headers, + timeout=timeout, + auth=auth, + ) + + server = MCPServerStreamableHttp( + params={ + "url": "https://example.test/mcp", + "httpx_client_factory": factory, + }, + ) + + with pytest.raises(UserError, match="HTTP error 503"): + await server.connect() + + assert methods == ["server/discover"] + assert server.session is None + + +@pytest.mark.asyncio +async def test_v2_streamable_http_oauth_subrequest_5xx_keeps_http_error_mapping(): + from mcp.client.auth.extensions.client_credentials import ClientCredentialsOAuthProvider + + response_body_marker = "synthetic-authorization-server-body" + + class _UnauthenticatedStorage: + async def get_tokens(self): + return None + + async def set_tokens(self, tokens): + return None + + async def get_client_info(self): + return None + + async def set_client_info(self, client_information): + return None + + def factory(headers=None, timeout=None, auth=None): + async def handler(request): + path = request.url.path + if path == "/token": + return httpx2.Response(503, text=response_body_marker, request=request) + if path.startswith("/.well-known/oauth-protected-resource"): + return httpx2.Response( + 200, + json={ + "resource": "https://example.test/mcp", + "authorization_servers": ["https://example.test"], + }, + request=request, + ) + if path.startswith("/.well-known/oauth-authorization-server"): + return httpx2.Response( + 200, + json={ + "issuer": "https://example.test", + "authorization_endpoint": "https://example.test/authorize", + "token_endpoint": "https://example.test/token", + "response_types_supported": ["code"], + }, + request=request, + ) + if "authorization" not in request.headers: + return httpx2.Response( + 401, + headers={ + "www-authenticate": ( + "Bearer resource_metadata=" + '"https://example.test/.well-known/oauth-protected-resource/mcp"' + ) + }, + request=request, + ) + return _v2_response_for_request(request) + + return httpx2.AsyncClient( + transport=httpx2.MockTransport(handler), + headers=headers, + timeout=timeout, + auth=auth, + ) + + server = MCPServerStreamableHttp( + params={ + "url": "https://example.test/mcp", + "auth": ClientCredentialsOAuthProvider( + server_url="https://example.test/mcp", + storage=_UnauthenticatedStorage(), + client_id="placeholder-client-id", + client_secret="placeholder-client-secret", + ), + "httpx_client_factory": factory, + }, + ) + + with pytest.raises(UserError, match="HTTP error 503") as exc_info: + await server.connect() + + error = exc_info.value + rendered = "".join(traceback.format_exception(type(error), error, error.__traceback__)) + assert response_body_marker not in rendered + assert server.session is None + + @pytest.mark.asyncio async def test_v2_connect_cancellation_stops_pending_client_owner(monkeypatch): client_entered = asyncio.Event()