From b65c4b9e74852b2c598553ba0eb14c6ffbcc6512 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 10:12:37 +0200 Subject: [PATCH 01/34] Fix streaming mode for CLI, MCP, TCP protocols and add SSE reconnection call_tool_streaming failed for several protocols instead of emitting the full result as a single chunk like the HTTP protocol does: - CLI raised NotImplementedError on every streaming call. - MCP yielded the un-awaited coroutine instead of the result. - TCP was a plain coroutine returning a generator, so the client's `async for` failed. UDP gets the same shape and type annotation. SSE improvements: - Implement `reconnect` / `retry_timeout`, which were accepted but never acted on. When an established stream drops, reconnect after `retry_timeout` (or the server's `retry:` value) with `Last-Event-ID`, capped at MAX_RECONNECT_ATTEMPTS per call. A clean end of stream completes the call; connection or HTTP errors on the initial request still fail immediately. Redirects stay refused on every attempt. - Handle CRLF line endings and a trailing unterminated event. Tests added for every fixed protocol and for the SSE reconnect paths. Co-Authored-By: Claude Fable 5.1 --- .../utcp_cli/cli_communication_protocol.py | 18 +- .../tests/test_cli_communication_protocol.py | 16 ++ .../utcp_http/sse_communication_protocol.py | 225 +++++++++++------- .../tests/test_sse_communication_protocol.py | 73 ++++++ .../utcp_mcp/mcp_communication_protocol.py | 3 +- .../mcp/tests/test_mcp_transport.py | 7 + .../utcp_socket/tcp_communication_protocol.py | 11 +- .../utcp_socket/udp_communication_protocol.py | 11 +- .../tests/test_tcp_communication_protocol.py | 28 ++- .../tests/test_udp_communication_protocol.py | 27 ++- 10 files changed, 315 insertions(+), 104 deletions(-) diff --git a/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py b/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py index bca07fc..422d70a 100644 --- a/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py +++ b/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py @@ -972,9 +972,19 @@ async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], too async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: """REQUIRED - Streaming calls are not supported for the CLI protocol. + Execute a tool call through the CLI transport streamingly. - Raises: - NotImplementedError: Always, as this functionality is not supported. + The CLI protocol does not natively support streaming, so the command is + executed to completion and the full result is yielded as a single chunk. + + Args: + caller: The UTCP client that is calling this method. + tool_name: Name of the tool to call. + tool_args: Dictionary of arguments to pass to the tool. + tool_call_template: Call template of the tool to call. + + Yields: + The complete tool result as a single item. """ - raise NotImplementedError("Streaming is not supported by the CLI communication protocol.") + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result diff --git a/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py b/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py index f96ffa7..0d98ee8 100644 --- a/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py +++ b/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py @@ -276,6 +276,22 @@ async def test_call_tool_json_output(transport: CliCommunicationProtocol, mock_c assert "Echo:" in result["result"] and "Hello" in result["result"] +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk(transport: CliCommunicationProtocol, mock_cli_script, python_executable): + """Streaming mode should emit the full result as one chunk instead of failing.""" + call_template = CliCallTemplate( + commands=[ + {"command": f"{python_executable} {mock_cli_script} --message UTCP_ARG_message_UTCP_END"} + ] + ) + + chunks = [chunk async for chunk in transport.call_tool_streaming(None, "echo", {"message": "Hello World"}, call_template)] + + assert len(chunks) == 1 + assert isinstance(chunks[0], dict) + assert "Echo:" in chunks[0]["result"] and "Hello" in chunks[0]["result"] + + @pytest.mark.asyncio async def test_call_tool_math_operation(transport: CliCommunicationProtocol, mock_cli_script, python_executable): """Test calling a math tool with numeric arguments.""" diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index 83afac2..ccc7486 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -35,6 +35,11 @@ class SseCommunicationProtocol(CommunicationProtocol): Handles Server-Sent Events based tool providers with streaming capabilities. """ + # Upper bound on reconnection attempts for a single tool call when the + # established stream drops and the call template has ``reconnect`` enabled. + # Keeps a tool call bounded even if the server keeps dropping the connection. + MAX_RECONNECT_ATTEMPTS: int = 5 + def __init__(self, logger: Optional[Callable[[str], None]] = None): self._oauth_tokens: Dict[str, Dict[str, Any]] = {} @@ -224,97 +229,143 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, token = await self._handle_oauth2(tool_call_template.auth) request_headers["Authorization"] = f"Bearer {token}" - session = aiohttp.ClientSession() - # Always close the session, success or failure. The previous - # version only closed on the except path, leaking the session - # on the (typical) success path. - try: - method = "POST" if body_content is not None else "GET" - data = body_content if "application/json" not in request_headers.get("Content-Type", "") else None - json_data = body_content if "application/json" in request_headers.get("Content-Type", "") else None - - # SSE handshake must not follow redirects: the streaming - # response has to stay open for the lifetime of the tool - # call, which is incompatible with the per-hop validator's - # release semantics, and SSE redirects are pathological in - # practice. Reject 3xx outright so an attacker-controlled - # endpoint cannot redirect the handshake into an internal - # service (GHSA-9qhg-99ww-9mqc). - response = await session.request( - method, url, params=query_params, headers=request_headers, - auth=auth, cookies=cookies, json=json_data, data=data, - timeout=None, allow_redirects=False, - ) - if 300 <= response.status < 400: - response.release() - raise RuntimeError( - f"SSE endpoint at {url!r} returned a {response.status} " - f"redirect. Redirects are not followed during SSE " - f"handshakes; update the call template to point at " - f"the final URL directly." - ) - response.raise_for_status() - async for event in self._process_sse_stream(response, tool_call_template.event_type): - yield event - except Exception as e: - logger.error(f"Error establishing SSE connection to '{tool_call_template.name}': {e}") - raise - finally: - await session.close() + method = "POST" if body_content is not None else "GET" + content_type = request_headers.get("Content-Type", "") + data = body_content if "application/json" not in content_type else None + json_data = body_content if "application/json" in content_type else None + + reconnect = bool(tool_call_template.reconnect) + retry_delay_ms = tool_call_template.retry_timeout + last_event_id: Optional[str] = None + reconnect_attempts = 0 + provider_name = tool_call_template.name + + while True: + attempt_headers = dict(request_headers) + if last_event_id is not None: + # Let the server resume from where we left off (SSE spec). + attempt_headers["Last-Event-ID"] = last_event_id + + session = aiohttp.ClientSession() + try: + try: + # SSE handshake must not follow redirects: the streaming + # response has to stay open for the lifetime of the tool + # call, which is incompatible with the per-hop validator's + # release semantics, and SSE redirects are pathological in + # practice. Reject 3xx outright so an attacker-controlled + # endpoint cannot redirect the handshake into an internal + # service (GHSA-9qhg-99ww-9mqc). + response = await session.request( + method, url, params=query_params, headers=attempt_headers, + auth=auth, cookies=cookies, json=json_data, data=data, + timeout=None, allow_redirects=False, + ) + if 300 <= response.status < 400: + response.release() + raise RuntimeError( + f"SSE endpoint at {url!r} returned a {response.status} " + f"redirect. Redirects are not followed during SSE " + f"handshakes; update the call template to point at " + f"the final URL directly." + ) + response.raise_for_status() + except Exception as e: + # Failing to establish the connection (or a non-2xx status) is a + # definitive answer, not a connection loss: fail fast, no retry. + logger.error(f"Error establishing SSE connection to '{provider_name}': {e}") + raise + + try: + async for event in self._iter_sse_events(response): + if event.get("id") is not None: + last_event_id = event["id"] + if event.get("retry") is not None: + retry_delay_ms = event["retry"] + if "data" not in event: + continue + if tool_call_template.event_type and event.get("event") != tool_call_template.event_type: + continue + yield self._parse_event_data(event["data"]) + # The server ended the stream cleanly: the tool call is complete. + return + except (aiohttp.ClientError, asyncio.TimeoutError) as e: + reconnect_attempts += 1 + if not reconnect or reconnect_attempts > self.MAX_RECONNECT_ATTEMPTS: + logger.error(f"SSE connection to '{provider_name}' lost and not reconnecting: {e}") + raise + logger.warning( + f"SSE connection to '{provider_name}' lost ({e}); reconnecting in {retry_delay_ms} ms " + f"(attempt {reconnect_attempts}/{self.MAX_RECONNECT_ATTEMPTS})" + ) + finally: + # Always release the connection, whether the stream completed, failed, + # or the consumer stopped iterating early. + if not session.closed: + await session.close() + + await asyncio.sleep(retry_delay_ms / 1000) - async def _process_sse_stream(self, response: aiohttp.ClientResponse, event_type=None): - """Process the SSE stream and yield events.""" + async def _iter_sse_events(self, response: aiohttp.ClientResponse) -> AsyncIterator[Dict[str, Any]]: + """Parse the SSE wire format and yield one dict per event block. + + Each dict may contain ``event``, ``id``, ``retry`` (int) and ``data`` (str, with + multi-line data joined by newlines). Blocks that only carry ``id``/``retry`` + are yielded too (without ``data``) so the caller can track reconnection + state; comment-only blocks are skipped. + """ buffer = "" - try: - async for chunk in response.content.iter_any(): - buffer += chunk.decode('utf-8') - while '\n\n' in buffer: - event_string, buffer = buffer.split('\n\n', 1) - - # Ignore empty event strings - if not event_string.strip(): - continue - - # Process the event string - lines = event_string.split('\n') - current_event = {} - data_lines = [] - for line in lines: - if line.startswith(':'): - continue # It's a comment - - if ':' in line: - field, value = line.split(':', 1) - value = value.lstrip() - if field == 'event': - current_event['event'] = value - elif field == 'data': - data_lines.append(value) - elif field == 'id': - current_event['id'] = value - elif field == 'retry': - try: - current_event['retry'] = int(value) - except ValueError: - pass - - if not data_lines: - continue - - current_event['data'] = '\n'.join(data_lines) - - if event_type and current_event.get('event') != event_type: - continue + def flush(event_string: str): + if not event_string.strip(): + return None + current_event: Dict[str, Any] = {} + data_lines: List[str] = [] + for line in event_string.split('\n'): + if line.startswith(':'): + continue # comment / keep-alive + if ':' in line: + field, value = line.split(':', 1) + if value.startswith(' '): + value = value[1:] + else: + field, value = line, '' + if field == 'event': + current_event['event'] = value + elif field == 'data': + data_lines.append(value) + elif field == 'id': + current_event['id'] = value + elif field == 'retry': try: - yield json.loads(current_event['data']) - except json.JSONDecodeError: - yield current_event['data'] - except Exception as e: - logger.error(f"Error processing SSE stream: {e}") - raise - finally: - pass # Session is managed and closed by deregister_tool_provider + current_event['retry'] = int(value) + except ValueError: + pass + if data_lines: + current_event['data'] = '\n'.join(data_lines) + return current_event or None + + async for chunk in response.content.iter_any(): + # Normalise CRLF / CR line endings so the event delimiter is always "\n\n". + buffer += chunk.decode('utf-8').replace('\r\n', '\n').replace('\r', '\n') + while '\n\n' in buffer: + event_string, buffer = buffer.split('\n\n', 1) + event = flush(event_string) + if event is not None: + yield event + + # Flush a trailing event that was not terminated by a blank line. + event = flush(buffer) + if event is not None: + yield event + + @staticmethod + def _parse_event_data(data: str) -> Any: + """Return the JSON-decoded payload when possible, otherwise the raw string.""" + try: + return json.loads(data) + except json.JSONDecodeError: + return data async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: """Handle OAuth2 client credentials flow, trying both body and diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index 76cb41e..c6a17e0 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -105,6 +105,27 @@ async def token_header_auth_handler(request): async def error_handler(request): return web.Response(status=500, text="Internal Server Error") +async def flaky_events_handler(request): + """Serves the first event then drops the TCP connection on the first connection + (or on every connection when ``always_drop`` is set). A reconnecting client is + served the remaining events and a clean end of stream.""" + state = request.app["flaky"] + state["connections"] += 1 + state["last_event_ids"].append(request.headers.get("Last-Event-ID")) + + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + + if state["always_drop"] or state["connections"] == 1: + await response.write(SAMPLE_SSE_EVENTS[0].encode('utf-8')) + await asyncio.sleep(0.01) + request.transport.close() + return response + + for event in SAMPLE_SSE_EVENTS[1:]: + await response.write(event.encode('utf-8')) + return response + # --- Pytest Fixtures --- @pytest_asyncio.fixture @@ -121,6 +142,8 @@ def app(): app.router.add_post("/token", token_handler) app.router.add_post("/token_header_auth", token_header_auth_handler) app.router.add_get("/error", error_handler) + app.router.add_get("/flaky_events", flaky_events_handler) + app["flaky"] = {"connections": 0, "last_event_ids": [], "always_drop": False} return app @pytest_asyncio.fixture @@ -377,3 +400,53 @@ async def test_call_tool_error_nonstream(sse_transport, aiohttp_client, app): with pytest.raises(aiohttp.ClientResponseError) as excinfo: await sse_transport.call_tool(None, "test_tool", {}, call_template) assert excinfo.value.status == 500 + + +# --- Reconnection --- + +@pytest.mark.asyncio +async def test_call_tool_reconnects_after_connection_loss(sse_transport, aiohttp_client, app): + """An established stream that drops is resumed with Last-Event-ID and yields every event once.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate( + name="test-sse", url=str(client.make_url("/flaky_events")), reconnect=True, retry_timeout=10 + ) + + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.test_tool", {}, call_template)] + + assert results == [{"message": "First part"}, {"message": "Second part"}, {"message": "End of stream"}] + assert app["flaky"]["connections"] == 2 + assert app["flaky"]["last_event_ids"] == [None, "1"] + + +@pytest.mark.asyncio +async def test_call_tool_connection_loss_without_reconnect_raises(sse_transport, aiohttp_client, app): + """With reconnect disabled a dropped stream surfaces as an error after the events received so far.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate( + name="test-sse", url=str(client.make_url("/flaky_events")), reconnect=False, retry_timeout=10 + ) + + received = [] + with pytest.raises(aiohttp.ClientError): + async for e in sse_transport.call_tool_streaming(None, "test-sse.test_tool", {}, call_template): + received.append(e) + + assert received == [{"message": "First part"}] + assert app["flaky"]["connections"] == 1 + + +@pytest.mark.asyncio +async def test_call_tool_reconnect_gives_up_after_max_attempts(sse_transport, aiohttp_client, app): + """A server that keeps dropping the stream cannot make a tool call hang forever.""" + app["flaky"]["always_drop"] = True + client = await aiohttp_client(app) + call_template = SseCallTemplate( + name="test-sse", url=str(client.make_url("/flaky_events")), reconnect=True, retry_timeout=1 + ) + + with pytest.raises(aiohttp.ClientError): + async for _ in sse_transport.call_tool_streaming(None, "test-sse.test_tool", {}, call_template): + pass + + assert app["flaky"]["connections"] == 1 + SseCommunicationProtocol.MAX_RECONNECT_ATTEMPTS diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 7204b43..2c6c5d1 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -379,7 +379,8 @@ async def _get_resource_server(self, resource_name: str, tool_call_template: Mcp async def call_tool_streaming(self, caller: 'UtcpClient', tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: """REQUIRED Streaming calls are not supported for MCP protocol, so we just call the tool and return the result as one item.""" - yield self.call_tool(caller, tool_name, tool_args, tool_call_template) + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result def _process_tool_result(self, result, tool_name: str) -> Any: self._log_info(f"Processing tool result for '{tool_name}', type: {type(result)}") diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index d127791..ce9872d 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -244,3 +244,10 @@ async def test_resource_tool_without_registration(transport: McpCommunicationPro # Should still work and return content assert isinstance(result, dict) assert "contents" in result + + +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + """Streaming mode should emit the awaited result as one chunk, not a coroutine.""" + chunks = [chunk async for chunk in transport.call_tool_streaming(None, f"{SERVER_NAME}.echo", {"message": "test"}, mcp_manual)] + assert chunks == [{"reply": "you said: test"}] diff --git a/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py b/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py index b2f08c3..4d9ff0a 100644 --- a/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py +++ b/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py @@ -8,7 +8,7 @@ import socket import struct import sys -from typing import Dict, Any, List, Optional, Callable, Union +from typing import Dict, Any, List, Optional, Callable, Union, AsyncGenerator from utcp.interfaces.communication_protocol import CommunicationProtocol from utcp_socket.tcp_call_template import TCPProvider, TCPProviderSerializer @@ -404,10 +404,11 @@ async def deregister_manual(self, caller, manual_call_template: CallTemplate) -> raise ValueError("TCPTransport can only be used with TCPProvider") self._log_info(f"Deregistering TCP provider '{manual_call_template.name}' (no-op)") - async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate): - async def _generator(): - yield await self.call_tool(caller, tool_name, tool_args, tool_call_template) - return _generator() + async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: + """REQUIRED + Streaming variant: the TCP protocol does not natively stream, so the full result is yielded as a single chunk.""" + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> Any: """Call a TCP tool.""" diff --git a/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py b/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py index 89ae3e3..fa1f98e 100644 --- a/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py +++ b/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py @@ -7,7 +7,7 @@ import json import socket import traceback -from typing import Dict, Any, List, Optional, Callable, Union +from typing import Dict, Any, List, Optional, Callable, Union, AsyncGenerator from utcp.interfaces.communication_protocol import CommunicationProtocol from utcp_socket.udp_call_template import UDPProvider, UDPProviderSerializer @@ -331,7 +331,8 @@ async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], too # While this works, it's inconsistent with the other implementation in tcp_communication_protocol.py (lines 384-387) which properly uses async def with an inner generator. # For consistency and clarity, this should also use async def directly: # - # async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate): - # yield await self.call_tool(caller, tool_name, tool_args, tool_call_template) - async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate): - yield await self.call_tool(caller, tool_name, tool_args, tool_call_template) + async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: + """REQUIRED + Streaming variant: the UDP protocol does not natively stream, so the full result is yielded as a single chunk.""" + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result diff --git a/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py b/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py index d359fd9..a82d14f 100644 --- a/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py +++ b/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py @@ -177,4 +177,30 @@ async def test_register_manual_fallbacks_to_manual_template_tcp(): assert tool.tool_call_template.name == provider.name finally: server.close() - await server.wait_closed() \ No newline at end of file + await server.wait_closed() + + +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk_tcp(): + """Streaming mode should be an async generator that yields the full result once.""" + server, port, set_response = await start_tcp_server() + set_response({"echo": "hello"}) + + try: + provider = TCPProvider( + name="tcp-provider", + host="127.0.0.1", + port=port, + request_data_format="json", + response_byte_format="utf-8", + framing_strategy="stream", + timeout=2000 + ) + transport_client = TCPTransport() + expected = await transport_client.call_tool(None, "tcp-provider.tcp_tool", {"x": 1}, provider) + chunks = [chunk async for chunk in transport_client.call_tool_streaming(None, "tcp-provider.tcp_tool", {"x": 1}, provider)] + + assert chunks == [expected] + finally: + server.close() + await server.wait_closed() diff --git a/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py b/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py index d6a770c..26fd402 100644 --- a/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py +++ b/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py @@ -173,4 +173,29 @@ async def test_register_manual_fallbacks_to_manual_template_udp(): assert tool.tool_call_template.port == provider.port assert tool.tool_call_template.name == provider.name finally: - transport.close() \ No newline at end of file + transport.close() + + +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk_udp(): + """Streaming mode should be an async generator that yields the full result once.""" + transport, port, set_response = await start_udp_server() + set_response({"echo": "hello"}) + + try: + provider = UDPProvider( + name="udp-provider", + host="127.0.0.1", + port=port, + number_of_response_datagrams=1, + request_data_format="json", + response_byte_format="utf-8", + timeout=2000 + ) + transport_client = UDPTransport() + expected = await transport_client.call_tool(None, "udp-provider.udp_tool", {"x": 1}, provider) + chunks = [chunk async for chunk in transport_client.call_tool_streaming(None, "udp-provider.udp_tool", {"x": 1}, provider)] + + assert chunks == [expected] + finally: + transport.close() From a4ac19033f637afae589cfdf9fbbf5d18d813f78 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 10:26:06 +0200 Subject: [PATCH 02/34] Pin mcp plugin to mcp 1.x mcp 2.x removed `mcp.server.fastmcp.FastMCP` and `mcp.shared.exceptions.McpError`, which the MCP test mocks import. CI installs the newest mcp, so every job failed at collection with mcp 2.1.1 (the last green run on dev predates the mcp 2 release). The plugin targets the 1.x API; a fresh install with the pin resolves to mcp 1.29.1 and the MCP suite passes. Migrating to mcp 2 is a separate task. Co-Authored-By: Claude Fable 5.1 --- plugins/communication_protocols/mcp/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/communication_protocols/mcp/pyproject.toml b/plugins/communication_protocols/mcp/pyproject.toml index 87461b7..ec66664 100644 --- a/plugins/communication_protocols/mcp/pyproject.toml +++ b/plugins/communication_protocols/mcp/pyproject.toml @@ -13,7 +13,7 @@ readme = "README.md" requires-python = ">=3.11" dependencies = [ "pydantic>=2.0", - "mcp>=1.12", + "mcp>=1.12,<2", "utcp>=1.1", "mcp-use>=1.3", "langchain>=0.3.27,<0.4.0", From d98ceaf2d52713b969a080d19354435e0f502c9e Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:21:43 +0200 Subject: [PATCH 03/34] mcp: quiet stdio child stderr by default, tighten structuredContent unwrap Python counterpart of typescript-utcp PR #33 plus its follow-ups, so both SDKs behave the same in the next release. - Stdio MCP children no longer inherit the host's stderr. mcp-use's MCPClient.from_dict offers no way to set the errlog that StdioConnector hands to the SDK's stdio_client, so a thin MCPClient subclass sets the connector's errlog to os.devnull between construction and initialization. UTCP_MCP_CHILD_STDERR=inherit restores the old behavior for debugging (same switch as the TypeScript SDK), and a stdio server that fails to start logs a hint pointing at it. - structuredContent is used when it is not None (the previous hasattr check was always true on CallToolResult). A FastMCP-style single-key {"result": value} wrapper is unwrapped; an object that merely has a "result" key among others is a genuine object return and now passes through untouched instead of losing its sibling keys. - Tests for both, plus a README section on child process stderr. Circular $ref handling from #33 needs no port: this plugin passes MCP schemas through without dereferencing them. Co-Authored-By: Claude Fable 5.1 --- plugins/communication_protocols/mcp/README.md | 10 +++ .../utcp_mcp/mcp_communication_protocol.py | 84 ++++++++++++++++--- .../mcp/tests/test_mcp_transport.py | 39 +++++++++ 3 files changed, 123 insertions(+), 10 deletions(-) diff --git a/plugins/communication_protocols/mcp/README.md b/plugins/communication_protocols/mcp/README.md index 0aa06f4..189296d 100644 --- a/plugins/communication_protocols/mcp/README.md +++ b/plugins/communication_protocols/mcp/README.md @@ -187,6 +187,16 @@ except TimeoutError: print("MCP server connection timed out") ``` +### Child Process stderr + +Stdio MCP servers often write banners, telemetry notices and auth chatter to stderr, multiplied by every server you federate. `utcp-mcp` therefore discards the child's stderr by default. To see it while debugging a server that fails to start, opt back in for the host process: + +```bash +UTCP_MCP_CHILD_STDERR=inherit python your_app.py +``` + +Any other value, or leaving the variable unset, keeps stderr suppressed. When a stdio server fails to connect, the error log reminds you of this switch. + ### List Available Tools ```python # Discover tools from MCP server diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 7204b43..78bd5b8 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -1,5 +1,6 @@ +import os import sys -from typing import Any, Dict, Optional, AsyncGenerator, TYPE_CHECKING, Tuple +from typing import Any, Dict, Optional, AsyncGenerator, TYPE_CHECKING, Tuple, TextIO import json from mcp_use import MCPClient @@ -23,6 +24,56 @@ logger = logging.getLogger(__name__) +# Environment variable that opts stdio MCP children back into writing to the +# host's stderr. Same name and semantics as the TypeScript SDK. +CHILD_STDERR_ENV_VAR = "UTCP_MCP_CHILD_STDERR" + +_devnull: Optional[TextIO] = None + + +def _child_stderr_target() -> TextIO: + """Return the stream stdio MCP children should write their stderr to. + + Defaults to ``os.devnull`` so a chatty server (banners, telemetry notices, + auth chatter, multiplied by every federated server) does not flood the host + terminal during discovery. Set ``UTCP_MCP_CHILD_STDERR=inherit`` to see it + while debugging. A file object rather than ``subprocess.DEVNULL`` because + the connector's contract is a text stream. + """ + if os.environ.get(CHILD_STDERR_ENV_VAR) == "inherit": + return sys.stderr + global _devnull + if _devnull is None or _devnull.closed: + _devnull = open(os.devnull, "w") + return _devnull + + +class _QuietStdioMCPClient(MCPClient): + """``MCPClient`` that routes stdio children's stderr per ``UTCP_MCP_CHILD_STDERR``. + + ``MCPClient.from_dict`` offers no way to set the ``errlog`` that + ``StdioConnector`` hands to the MCP SDK's ``stdio_client``, so every child + would inherit the host's stderr. The connector only reads ``errlog`` when it + connects, so it is enough to set it between construction and initialization. + """ + + async def create_session(self, server_name: str, auto_initialize: bool = True): + session = await super().create_session(server_name, auto_initialize=False) + if session is None: + return None + if hasattr(session.connector, "errlog"): + session.connector.errlog = _child_stderr_target() + if auto_initialize: + try: + await session.initialize() + except Exception: + # Mirror the base class: a session that failed to initialize must + # not stay cached, or the next lookup would hand back a dead one. + self.sessions.pop(server_name, None) + raise + return session + + class McpCommunicationProtocol(CommunicationProtocol): """REQUIRED MCP transport implementation that connects to MCP servers via stdio or HTTP. @@ -52,7 +103,7 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate'): if self._mcp_client is None or self._mcp_client.config != manual_call_template.config.mcpServers: # Create a new MCPClient with the server configuration config = {"mcpServers": manual_call_template.config.mcpServers} - self._mcp_client = MCPClient.from_dict(config) + self._mcp_client = _QuietStdioMCPClient.from_dict(config) async def _get_or_create_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): """Get an existing session or create a new one using MCPClient.""" @@ -66,7 +117,17 @@ async def _get_or_create_session(self, server_name: str, manual_call_template: ' except ValueError: # Session doesn't exist, create a new one self._log_info(f"Creating new session for server: {server_name}") - session = await self._mcp_client.create_session(server_name, auto_initialize=True) + try: + session = await self._mcp_client.create_session(server_name, auto_initialize=True) + except Exception as e: + server_config = manual_call_template.config.mcpServers.get(server_name) + is_stdio = isinstance(server_config, dict) and "command" in server_config + if is_stdio and os.environ.get(CHILD_STDERR_ENV_VAR) != "inherit": + self._log_error( + f"Failed to start stdio MCP server '{server_name}': {e}. The child's stderr was " + f"suppressed; re-run with {CHILD_STDERR_ENV_VAR}=inherit to see what it printed while starting." + ) + raise return session async def _cleanup_session(self, server_name: str): @@ -384,13 +445,16 @@ async def call_tool_streaming(self, caller: 'UtcpClient', tool_name: str, tool_a def _process_tool_result(self, result, tool_name: str) -> Any: self._log_info(f"Processing tool result for '{tool_name}', type: {type(result)}") - # Check for structured output first - this is the expected behavior - if hasattr(result, 'structuredContent'): - self._log_info(f"Found structuredContent: {result.structuredContent}") - # If structuredContent has a 'result' key, unwrap it - if isinstance(result.structuredContent, dict) and 'result' in result.structuredContent: - return result.structuredContent['result'] - return result.structuredContent + # Prefer structuredContent (MCP spec field) whenever the server sent it. + structured = getattr(result, 'structuredContent', None) + if structured is not None: + self._log_info(f"Found structuredContent: {structured}") + # FastMCP wraps non-object returns as {"result": value}; unwrap exactly + # that single-key shape. An object that merely has a "result" key among + # others is a genuine object return and passes through untouched. + if isinstance(structured, dict) and set(structured.keys()) == {"result"}: + return structured["result"] + return structured # Process content if available (fallback) if hasattr(result, 'content'): diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index d127791..15ed220 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -244,3 +244,42 @@ async def test_resource_tool_without_registration(transport: McpCommunicationPro # Should still work and return content assert isinstance(result, dict) assert "contents" in result + + +# --- Child stderr routing and structuredContent unwrapping --- + +@pytest.mark.asyncio +async def test_stdio_child_stderr_suppressed_by_default(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate, monkeypatch): + """Without the opt-in, stdio children write stderr to os.devnull, not the host's stderr.""" + monkeypatch.delenv("UTCP_MCP_CHILD_STDERR", raising=False) + session = await transport._get_or_create_session(SERVER_NAME, mcp_manual) + try: + assert session.connector.errlog is not sys.stderr + assert session.connector.errlog.name == os.devnull + finally: + await transport._cleanup_session(SERVER_NAME) + + +@pytest.mark.asyncio +async def test_stdio_child_stderr_inherit_opt_in(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate, monkeypatch): + """UTCP_MCP_CHILD_STDERR=inherit restores the host's stderr for debugging.""" + monkeypatch.setenv("UTCP_MCP_CHILD_STDERR", "inherit") + session = await transport._get_or_create_session(SERVER_NAME, mcp_manual) + try: + assert session.connector.errlog is sys.stderr + finally: + await transport._cleanup_session(SERVER_NAME) + + +@pytest.mark.asyncio +async def test_process_tool_result_unwraps_only_single_key_result_wrapper(transport: McpCommunicationProtocol): + """A FastMCP {"result": x} wrapper is unwrapped; a real object with a result key is not.""" + from types import SimpleNamespace + assert transport._process_tool_result(SimpleNamespace(structuredContent={"result": 42}, content=[]), "t") == 42 + assert transport._process_tool_result( + SimpleNamespace(structuredContent={"result": 1, "extra": 2}, content=[]), "t" + ) == {"result": 1, "extra": 2} + assert transport._process_tool_result(SimpleNamespace(structuredContent={"answer": 42}, content=[]), "t") == {"answer": 42} + # No structuredContent: fall back to text content. + text_only = SimpleNamespace(structuredContent=None, content=[SimpleNamespace(text="7")]) + assert transport._process_tool_result(text_only, "t") == 7 From 4f19b9d8d1b88ac843f6aff4024eaa6bab22a146 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 10:26:06 +0200 Subject: [PATCH 04/34] Pin mcp plugin to mcp 1.x mcp 2.x removed `mcp.server.fastmcp.FastMCP` and `mcp.shared.exceptions.McpError`, which the MCP test mocks import. CI installs the newest mcp, so every job failed at collection with mcp 2.1.1 (the last green run on dev predates the mcp 2 release). The plugin targets the 1.x API; a fresh install with the pin resolves to mcp 1.29.1 and the MCP suite passes. Migrating to mcp 2 is a separate task. Co-Authored-By: Claude Fable 5.1 --- plugins/communication_protocols/mcp/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugins/communication_protocols/mcp/pyproject.toml b/plugins/communication_protocols/mcp/pyproject.toml index 87461b7..ec66664 100644 --- a/plugins/communication_protocols/mcp/pyproject.toml +++ b/plugins/communication_protocols/mcp/pyproject.toml @@ -13,7 +13,7 @@ readme = "README.md" requires-python = ">=3.11" dependencies = [ "pydantic>=2.0", - "mcp>=1.12", + "mcp>=1.12,<2", "utcp>=1.1", "mcp-use>=1.3", "langchain>=0.3.27,<0.4.0", From 31a1dfb8a311bd38f492d34cf56c9cc661c614ac Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:32:30 +0200 Subject: [PATCH 05/34] mcp: only unwrap non-object {"result"} wrappers FastMCP wraps only non-object returns as {"result": value}, so a single-key {"result": {...}} is a genuine object return and must keep its shape. Unwrap only when the inner value is not a dict. Tests added for the list wrapper and the genuine single-key object return. Mirrors the cubic review fix on typescript-utcp#42. Co-Authored-By: Claude Fable 5.1 --- .../src/utcp_mcp/mcp_communication_protocol.py | 17 +++++++++++++---- .../mcp/tests/test_mcp_transport.py | 6 ++++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 78bd5b8..2bb6c14 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -449,10 +449,19 @@ def _process_tool_result(self, result, tool_name: str) -> Any: structured = getattr(result, 'structuredContent', None) if structured is not None: self._log_info(f"Found structuredContent: {structured}") - # FastMCP wraps non-object returns as {"result": value}; unwrap exactly - # that single-key shape. An object that merely has a "result" key among - # others is a genuine object return and passes through untouched. - if isinstance(structured, dict) and set(structured.keys()) == {"result"}: + # FastMCP wraps NON-OBJECT returns (primitives, lists, None) as + # {"result": value}; object returns are sent as-is. Unwrap exactly that + # shape: a single "result" key whose value is not a dict. A single-key + # {"result": {...}} is therefore a genuine object return and passes + # through untouched, as does any dict with other keys. A genuine + # {"result": } return is indistinguishable from the + # wrapper on the wire and is unwrapped too; that ambiguity is inherent + # to the FastMCP convention. + if ( + isinstance(structured, dict) + and set(structured.keys()) == {"result"} + and not isinstance(structured["result"], dict) + ): return structured["result"] return structured diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index 15ed220..b7fbced 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -280,6 +280,12 @@ async def test_process_tool_result_unwraps_only_single_key_result_wrapper(transp SimpleNamespace(structuredContent={"result": 1, "extra": 2}, content=[]), "t" ) == {"result": 1, "extra": 2} assert transport._process_tool_result(SimpleNamespace(structuredContent={"answer": 42}, content=[]), "t") == {"answer": 42} + assert transport._process_tool_result(SimpleNamespace(structuredContent={"result": ["a", "b"]}, content=[]), "t") == ["a", "b"] + # FastMCP only wraps non-object returns, so {"result": {...}} is a genuine + # object return from the tool and must keep its shape. + assert transport._process_tool_result( + SimpleNamespace(structuredContent={"result": {"nested": True}}, content=[]), "t" + ) == {"result": {"nested": True}} # No structuredContent: fall back to text content. text_only = SimpleNamespace(structuredContent=None, content=[SimpleNamespace(text="7")]) assert transport._process_tool_result(text_only, "t") == 7 From 08a2fd29dc41770f854674fa78778b8a748065f0 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:21:24 +0200 Subject: [PATCH 06/34] sse: bound the handshake, retry failed reconnect handshakes, cap delays, fix CRLF framing Addresses cubic review on #100: - The handshake (until response headers arrive) is bounded by HANDSHAKE_TIMEOUT_SECONDS (30 s) via asyncio.wait_for, so a server that accepts the connection but never answers cannot hang the call. Reading the body stays unbounded: an SSE stream may legitimately be quiet. - A reconnect handshake that fails (refused, 503, timeout) now counts as one attempt and is retried; only the initial handshake fails fast. - The reconnect delay is capped at MAX_RECONNECT_DELAY_MS (60 s) whatever retry_timeout or a server-sent retry: field asks for, so the attempt cap actually bounds the total wait. - A CRLF split across two chunks no longer becomes two LFs and ends the event early: a trailing CR is held until the next chunk. Decoding is now incremental too, so a multi-byte UTF-8 character straddling chunks no longer raises. - An event that exceeds MAX_EVENT_BUFFER_CHARS (16 Mi) without a blank-line delimiter raises SseProtocolError, which is never retried. Tests for each of the above. Co-Authored-By: Claude Fable 5.1 --- .../utcp_http/sse_communication_protocol.py | 89 +++++++++++-- .../tests/test_sse_communication_protocol.py | 126 ++++++++++++++++++ 2 files changed, 201 insertions(+), 14 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index ccc7486..a4d284f 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -3,6 +3,7 @@ import aiohttp import json import asyncio +import codecs import re from urllib.parse import quote import base64 @@ -28,6 +29,11 @@ logger = logging.getLogger(__name__) + +class SseProtocolError(RuntimeError): + """The server violated the SSE wire format. Not a connection loss, so never retried.""" + + class SseCommunicationProtocol(CommunicationProtocol): """REQUIRED SSE communication protocol implementation for UTCP client. @@ -39,6 +45,17 @@ class SseCommunicationProtocol(CommunicationProtocol): # established stream drops and the call template has ``reconnect`` enabled. # Keeps a tool call bounded even if the server keeps dropping the connection. MAX_RECONNECT_ATTEMPTS: int = 5 + # Cap on the delay before a reconnect, whatever ``retry_timeout`` or a + # server-sent ``retry:`` field asks for. Together with MAX_RECONNECT_ATTEMPTS + # this bounds the total time a call can spend waiting to reconnect. + MAX_RECONNECT_DELAY_MS: int = 60_000 + # Time allowed for the SSE handshake, i.e. until response headers arrive. + # Reading the body is unbounded: an SSE stream may legitimately stay quiet. + HANDSHAKE_TIMEOUT_SECONDS: float = 30.0 + # Largest partial event the parser buffers before declaring the stream + # malformed. Guards against a server that streams data without ever sending + # the blank-line event delimiter. + MAX_EVENT_BUFFER_CHARS: int = 16 * 1024 * 1024 def __init__(self, logger: Optional[Callable[[str], None]] = None): self._oauth_tokens: Dict[str, Dict[str, Any]] = {} @@ -256,10 +273,15 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, # practice. Reject 3xx outright so an attacker-controlled # endpoint cannot redirect the handshake into an internal # service (GHSA-9qhg-99ww-9mqc). - response = await session.request( - method, url, params=query_params, headers=attempt_headers, - auth=auth, cookies=cookies, json=json_data, data=data, - timeout=None, allow_redirects=False, + # Bound the handshake only (until response headers arrive); + # the body read stays unbounded because a stream may be quiet. + response = await asyncio.wait_for( + session.request( + method, url, params=query_params, headers=attempt_headers, + auth=auth, cookies=cookies, json=json_data, data=data, + timeout=None, allow_redirects=False, + ), + timeout=self.HANDSHAKE_TIMEOUT_SECONDS, ) if 300 <= response.status < 400: response.release() @@ -271,10 +293,24 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, ) response.raise_for_status() except Exception as e: - # Failing to establish the connection (or a non-2xx status) is a - # definitive answer, not a connection loss: fail fast, no retry. - logger.error(f"Error establishing SSE connection to '{provider_name}': {e}") - raise + if reconnect_attempts == 0: + # The initial handshake failing (refused, timed out, non-2xx) is a + # definitive answer about the endpoint: fail fast, no retry. + logger.error(f"Error establishing SSE connection to '{provider_name}': {e}") + raise + # A reconnect handshake failing is part of the outage we are riding + # out (the server may still be restarting): count it and try again. + reconnect_attempts += 1 + if reconnect_attempts > self.MAX_RECONNECT_ATTEMPTS: + logger.error(f"SSE reconnect to '{provider_name}' failed and attempts are exhausted: {e}") + raise + delay_ms = min(retry_delay_ms, self.MAX_RECONNECT_DELAY_MS) + logger.warning( + f"SSE reconnect to '{provider_name}' failed ({e}); retrying in {delay_ms} ms " + f"(attempt {reconnect_attempts}/{self.MAX_RECONNECT_ATTEMPTS})" + ) + await asyncio.sleep(delay_ms / 1000) + continue try: async for event in self._iter_sse_events(response): @@ -295,7 +331,8 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, logger.error(f"SSE connection to '{provider_name}' lost and not reconnecting: {e}") raise logger.warning( - f"SSE connection to '{provider_name}' lost ({e}); reconnecting in {retry_delay_ms} ms " + f"SSE connection to '{provider_name}' lost ({e}); reconnecting in " + f"{min(retry_delay_ms, self.MAX_RECONNECT_DELAY_MS)} ms " f"(attempt {reconnect_attempts}/{self.MAX_RECONNECT_ATTEMPTS})" ) finally: @@ -304,7 +341,7 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, if not session.closed: await session.close() - await asyncio.sleep(retry_delay_ms / 1000) + await asyncio.sleep(min(retry_delay_ms, self.MAX_RECONNECT_DELAY_MS) / 1000) async def _iter_sse_events(self, response: aiohttp.ClientResponse) -> AsyncIterator[Dict[str, Any]]: """Parse the SSE wire format and yield one dict per event block. @@ -345,16 +382,40 @@ def flush(event_string: str): current_event['data'] = '\n'.join(data_lines) return current_event or None - async for chunk in response.content.iter_any(): + # Incremental decoding: a multi-byte UTF-8 character may straddle two chunks. + decoder = codecs.getincrementaldecoder("utf-8")() + # A "\r" that ended the previous chunk is held back until the next chunk + # shows whether a "\n" follows; otherwise a CRLF split across two reads + # would become two LFs and dispatch an event early. + pending_cr = False + + def normalise(text: str) -> str: + nonlocal pending_cr + if pending_cr: + text = "\r" + text + pending_cr = False + if text.endswith("\r"): + text = text[:-1] + pending_cr = True # Normalise CRLF / CR line endings so the event delimiter is always "\n\n". - buffer += chunk.decode('utf-8').replace('\r\n', '\n').replace('\r', '\n') - while '\n\n' in buffer: - event_string, buffer = buffer.split('\n\n', 1) + return text.replace("\r\n", "\n").replace("\r", "\n") + + async for chunk in response.content.iter_any(): + buffer += normalise(decoder.decode(chunk)) + while "\n\n" in buffer: + event_string, buffer = buffer.split("\n\n", 1) event = flush(event_string) if event is not None: yield event + if len(buffer) > self.MAX_EVENT_BUFFER_CHARS: + raise SseProtocolError( + f"SSE event exceeded {self.MAX_EVENT_BUFFER_CHARS} characters without a blank-line delimiter" + ) # Flush a trailing event that was not terminated by a blank line. + buffer += normalise(decoder.decode(b"", final=True)) + if pending_cr: + buffer += "\n" event = flush(buffer) if event is not None: yield event diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index c6a17e0..fdd8ded 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -105,6 +105,65 @@ async def token_header_auth_handler(request): async def error_handler(request): return web.Response(status=500, text="Internal Server Error") + +async def crlf_split_events_handler(request): + """One multi-line CRLF event whose CRLF is split across two writes.""" + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + await response.write(b"data: line1\r") + await asyncio.sleep(0.05) + await response.write(b"\ndata: line2\r\n\r\n") + return response + + +async def no_delimiter_events_handler(request): + """Streams data lines without ever sending the blank-line event delimiter.""" + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + for _ in range(20): + await response.write(b"data: " + b"x" * 500 + b"\n") + return response + + +async def flaky_503_events_handler(request): + """Drops the stream after the first event, answers the first reconnect with a + 503, then serves the rest on the second reconnect.""" + state = request.app["flaky503"] + state["connections"] += 1 + if state["connections"] == 2: + return web.Response(status=503, text="restarting") + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + if state["connections"] == 1: + await response.write(SAMPLE_SSE_EVENTS[0].encode('utf-8')) + await asyncio.sleep(0.01) + request.transport.close() + return response + for event in SAMPLE_SSE_EVENTS[1:]: + await response.write(event.encode('utf-8')) + return response + + +async def slow_handshake_handler(request): + """Accepts the connection but does not send response headers for a long time.""" + await asyncio.sleep(5) + return web.Response(status=204) + + +async def huge_retry_events_handler(request): + """First connection asks for a very long retry delay, then drops.""" + state = request.app["huge_retry"] + state["connections"] += 1 + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + if state["connections"] == 1: + await response.write(b'id: 1\nretry: 100000\ndata: {"seq": 1}\n\n') + await asyncio.sleep(0.01) + request.transport.close() + return response + await response.write(b'id: 2\ndata: {"seq": 2}\n\n') + return response + async def flaky_events_handler(request): """Serves the first event then drops the TCP connection on the first connection (or on every connection when ``always_drop`` is set). A reconnecting client is @@ -144,6 +203,13 @@ def app(): app.router.add_get("/error", error_handler) app.router.add_get("/flaky_events", flaky_events_handler) app["flaky"] = {"connections": 0, "last_event_ids": [], "always_drop": False} + app.router.add_get("/crlf_split_events", crlf_split_events_handler) + app.router.add_get("/no_delimiter_events", no_delimiter_events_handler) + app.router.add_get("/flaky_503_events", flaky_503_events_handler) + app["flaky503"] = {"connections": 0} + app.router.add_get("/slow_handshake", slow_handshake_handler) + app.router.add_get("/huge_retry_events", huge_retry_events_handler) + app["huge_retry"] = {"connections": 0} return app @pytest_asyncio.fixture @@ -450,3 +516,63 @@ async def test_call_tool_reconnect_gives_up_after_max_attempts(sse_transport, ai pass assert app["flaky"]["connections"] == 1 + SseCommunicationProtocol.MAX_RECONNECT_ATTEMPTS + + +# --- Review follow-ups: framing robustness and bounded reconnects --- + +@pytest.mark.asyncio +async def test_crlf_split_across_chunks_is_one_event(sse_transport, aiohttp_client, app): + """A CRLF whose CR and LF arrive in different chunks must not end the event early.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/crlf_split_events"))) + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] + assert results == ["line1\nline2"] + + +@pytest.mark.asyncio +async def test_oversized_event_without_delimiter_raises_and_does_not_reconnect(sse_transport, aiohttp_client, app, monkeypatch): + """A stream that never sends the blank-line delimiter is rejected, not buffered forever.""" + from utcp_http.sse_communication_protocol import SseCommunicationProtocol, SseProtocolError + monkeypatch.setattr(SseCommunicationProtocol, "MAX_EVENT_BUFFER_CHARS", 1000) + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/no_delimiter_events")), reconnect=True, retry_timeout=1) + with pytest.raises(SseProtocolError): + async for _ in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template): + pass + + +@pytest.mark.asyncio +async def test_reconnect_handshake_failure_is_retried(sse_transport, aiohttp_client, app): + """A 503 on a reconnect handshake counts as one attempt and is retried, unlike the initial handshake.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/flaky_503_events")), reconnect=True, retry_timeout=10) + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] + assert results == [{"message": "First part"}, {"message": "Second part"}, {"message": "End of stream"}] + assert app["flaky503"]["connections"] == 3 + + +@pytest.mark.asyncio +async def test_initial_handshake_timeout_raises(sse_transport, aiohttp_client, app, monkeypatch): + """A server that accepts the connection but never sends headers cannot hang the call.""" + from utcp_http.sse_communication_protocol import SseCommunicationProtocol + monkeypatch.setattr(SseCommunicationProtocol, "HANDSHAKE_TIMEOUT_SECONDS", 0.3) + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/slow_handshake"))) + with pytest.raises((asyncio.TimeoutError, TimeoutError)): + async for _ in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template): + pass + + +@pytest.mark.asyncio +async def test_reconnect_delay_is_capped(sse_transport, aiohttp_client, app, monkeypatch): + """A server-sent retry of 100 s cannot stall the reconnect past MAX_RECONNECT_DELAY_MS.""" + import time + from utcp_http.sse_communication_protocol import SseCommunicationProtocol + monkeypatch.setattr(SseCommunicationProtocol, "MAX_RECONNECT_DELAY_MS", 50) + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/huge_retry_events")), reconnect=True, retry_timeout=10) + started = time.monotonic() + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] + assert results == [{"seq": 1}, {"seq": 2}] + assert app["huge_retry"]["connections"] == 2 + assert time.monotonic() - started < 3 From d86d312afcc0dc1b3bbe5361691d1b0bb0ee839a Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:19:42 +0200 Subject: [PATCH 07/34] sse test: assert the oversized-event failure does not reconnect Count connections in the no-delimiter handler and assert exactly one, so the test actually verifies that SseProtocolError bypasses the reconnect path instead of relying on the error being re-raised on a retry. Co-Authored-By: Claude Fable 5.1 --- .../http/tests/test_sse_communication_protocol.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index fdd8ded..48f73c6 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -118,6 +118,7 @@ async def crlf_split_events_handler(request): async def no_delimiter_events_handler(request): """Streams data lines without ever sending the blank-line event delimiter.""" + request.app["no_delimiter"]["connections"] += 1 response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) await response.prepare(request) for _ in range(20): @@ -205,6 +206,7 @@ def app(): app["flaky"] = {"connections": 0, "last_event_ids": [], "always_drop": False} app.router.add_get("/crlf_split_events", crlf_split_events_handler) app.router.add_get("/no_delimiter_events", no_delimiter_events_handler) + app["no_delimiter"] = {"connections": 0} app.router.add_get("/flaky_503_events", flaky_503_events_handler) app["flaky503"] = {"connections": 0} app.router.add_get("/slow_handshake", slow_handshake_handler) @@ -539,6 +541,8 @@ async def test_oversized_event_without_delimiter_raises_and_does_not_reconnect(s with pytest.raises(SseProtocolError): async for _ in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template): pass + # A protocol violation is not a connection loss: exactly one connection, no reconnect. + assert app["no_delimiter"]["connections"] == 1 @pytest.mark.asyncio From 90793cd16bf46833b788b53b950d88cb5c4c3360 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:55:47 +0200 Subject: [PATCH 08/34] http: surface the server's error body on failed calls and discovery Python counterpart of typescript-utcp #26 / #44. raise_for_status() raises a ClientResponseError whose message is only the reason phrase ("Forbidden"); the body, where servers put the real reason ({"error": "..."}), was discarded, so a refused call or discovery surfaced as nothing more than a status code. New utcp_http._errors.raise_for_status_with_body reads the body on a 4xx/5xx and raises a ClientResponseError of the same status and headers whose message is ": ", where detail is a string error / message / detail field when the body is JSON, otherwise the raw body (truncated). The raw text is attached as .body. Used by the HTTP protocol's tool calls and discovery and by SSE and Streamable HTTP discovery. The exception type is unchanged, so existing handlers keep working. Tests: body in the call error, object-valued error field shows its structure, and discovery errors[] carries the body for all three protocols. Co-Authored-By: Claude Fable 5.1 --- .../http/src/utcp_http/_errors.py | 65 +++++++++++++++++++ .../utcp_http/http_communication_protocol.py | 5 +- .../utcp_http/sse_communication_protocol.py | 3 +- .../streamable_http_communication_protocol.py | 3 +- .../tests/test_http_communication_protocol.py | 48 ++++++++++++++ .../tests/test_sse_communication_protocol.py | 16 +++++ ..._streamable_http_communication_protocol.py | 15 +++++ 7 files changed, 151 insertions(+), 4 deletions(-) create mode 100644 plugins/communication_protocols/http/src/utcp_http/_errors.py diff --git a/plugins/communication_protocols/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py new file mode 100644 index 0000000..f4df4ea --- /dev/null +++ b/plugins/communication_protocols/http/src/utcp_http/_errors.py @@ -0,0 +1,65 @@ +"""Surface the server's error body on failed HTTP calls. + +``aiohttp.ClientResponse.raise_for_status()`` raises a ``ClientResponseError`` +whose ``message`` is only the reason phrase ("Forbidden"). Servers put the +real reason in the response body, typically ``{"error": "..."}``, and that +was discarded, so a refused call or discovery surfaced as nothing more than a +status code. Mirrors the TypeScript SDK's ``_normalizeToolError``. +""" +import json +from typing import Optional + +import aiohttp + +# Bodies are folded into an exception message; keep pathological ones bounded. +MAX_DETAIL_CHARS = 2000 + + +def error_detail_from_body(text: str) -> Optional[str]: + """Extract a human-readable reason from an error response body. + + Prefers a string ``error`` / ``message`` / ``detail`` field when the body is + a JSON object; an object-valued field falls through to the raw JSON so the + real structure shows. Returns ``None`` for an empty body. + """ + body = text.strip() + if not body: + return None + try: + data = json.loads(body) + except ValueError: + return body[:MAX_DETAIL_CHARS] + if isinstance(data, dict): + for key in ("error", "message", "detail"): + value = data.get(key) + if isinstance(value, str) and value.strip(): + return value.strip()[:MAX_DETAIL_CHARS] + return body[:MAX_DETAIL_CHARS] + + +async def raise_for_status_with_body(response: aiohttp.ClientResponse) -> None: + """Like ``response.raise_for_status()``, but with the response body in the error. + + On a 4xx/5xx, reads the body and raises a ``ClientResponseError`` of the + same status and headers whose ``message`` is ``": "``. The + raw body text is attached as ``body`` for callers that want the structure. + Does nothing on a 2xx/3xx. + """ + if response.status < 400: + return + try: + text = await response.text() + except Exception: + text = "" + detail = error_detail_from_body(text) + reason = response.reason or "" + message = f"{reason}: {detail}" if detail else reason + error = aiohttp.ClientResponseError( + response.request_info, + response.history, + status=response.status, + message=message, + headers=response.headers, + ) + error.body = text # type: ignore[attr-defined] + raise error diff --git a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py index 8a3a187..99d5de2 100644 --- a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py @@ -34,6 +34,7 @@ from aiohttp import ClientSession, BasicAuth as AiohttpBasicAuth from utcp_http.openapi_converter import OpenApiConverter from utcp_http._security import ensure_secure_url, safe_request_with_redirects +from utcp_http._errors import raise_for_status_with_body import logging logging.basicConfig( @@ -210,7 +211,7 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R timeout=aiohttp.ClientTimeout(total=10.0), auth_header_names=auth_header_names, ) as response: - response.raise_for_status() # Raise exception for 4XX/5XX responses + await raise_for_status_with_body(response) # 4XX/5XX, with the server's body in the message # Check content type to determine how to parse the response content_type = response.headers.get('Content-Type', '') @@ -359,7 +360,7 @@ async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], too timeout=aiohttp.ClientTimeout(total=30.0), auth_header_names=auth_header_names, ) as response: - response.raise_for_status() + await raise_for_status_with_body(response) content_type = response.headers.get('Content-Type', '').lower() if 'application/json' in content_type: diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index a4d284f..fdb2ff5 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -18,6 +18,7 @@ from utcp.data.auth_implementations.oauth2_auth import OAuth2Auth from utcp_http.sse_call_template import SseCallTemplate from aiohttp import ClientSession, BasicAuth as AiohttpBasicAuth +from utcp_http._errors import raise_for_status_with_body from utcp_http._security import ensure_secure_url, safe_request_with_redirects import traceback import logging @@ -170,7 +171,7 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R timeout=aiohttp.ClientTimeout(total=10.0), auth_header_names=auth_header_names, ) as response: - response.raise_for_status() + await raise_for_status_with_body(response) response_data = await response.json() utcp_manual = UtcpManualSerializer().validate_dict(response_data) return RegisterManualResult( diff --git a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py index fde6cb5..df4b94f 100644 --- a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py @@ -15,6 +15,7 @@ from utcp.data.auth_implementations import OAuth2Auth from utcp_http.streamable_http_call_template import StreamableHttpCallTemplate from aiohttp import ClientSession, BasicAuth as AiohttpBasicAuth, ClientResponse +from utcp_http._errors import raise_for_status_with_body from utcp_http._security import ensure_secure_url, safe_request_with_redirects import logging @@ -149,7 +150,7 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R timeout=aiohttp.ClientTimeout(total=10.0), auth_header_names=auth_header_names, ) as response: - response.raise_for_status() + await raise_for_status_with_body(response) response_data = await response.json() utcp_manual = UtcpManualSerializer().validate_dict(response_data) return RegisterManualResult( diff --git a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py index 518b8df..f6de717 100644 --- a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py @@ -139,6 +139,17 @@ async def error_handler(request): app.router.add_post('/token', token_handler) app.router.add_post('/token_header_auth', token_header_auth_handler) app.router.add_get('/error', error_handler) + + # Non-2xx with a descriptive body, like a real API refusing a call. + async def forbidden_handler(request): + return web.json_response({"error": "You are not allowed to do that, and here is exactly why."}, status=403) + + # Some APIs nest an object under `error`; the message must show its JSON. + async def forbidden_object_handler(request): + return web.json_response({"error": {"code": "INVALID_FIELD", "reason": "value out of range"}}, status=422) + + app.router.add_route('*', '/forbidden', forbidden_handler) + app.router.add_route('*', '/forbidden-object', forbidden_object_handler) return app @@ -736,3 +747,40 @@ def test_auth_tools_integration(): serialized = serializer.to_dict(call_template) assert "auth_tools" in serialized assert serialized["auth_tools"]["auth_type"] == "api_key" + + +# --- Server error bodies are surfaced, not just status codes --- + +@pytest.mark.asyncio +async def test_call_tool_surfaces_server_error_body(http_transport, aiohttp_client, app): + """A refused call carries the server's reason, not only "403, message='Forbidden'".""" + client = await aiohttp_client(app) + call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}/forbidden", http_method="POST") + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + await http_transport.call_tool(None, "t.tool", {"param1": "value1"}, call_template) + assert excinfo.value.status == 403 + assert "You are not allowed to do that, and here is exactly why." in str(excinfo.value) + assert '"error"' in excinfo.value.body + + +@pytest.mark.asyncio +async def test_call_tool_surfaces_object_valued_error_field(http_transport, aiohttp_client, app): + """An object under `error` shows its JSON structure in the message.""" + client = await aiohttp_client(app) + call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}/forbidden-object", http_method="POST") + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + await http_transport.call_tool(None, "t.tool", {}, call_template) + assert excinfo.value.status == 422 + assert "INVALID_FIELD" in str(excinfo.value) + assert "value out of range" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_register_manual_surfaces_server_error_body(http_transport, aiohttp_client, app): + """A refused discovery reports the server's reason in errors[].""" + client = await aiohttp_client(app) + call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}/forbidden", http_method="GET") + result = await http_transport.register_manual(None, call_template) + assert result.success is False + assert "You are not allowed to do that, and here is exactly why." in result.errors[0] + assert "403" in result.errors[0] diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index 48f73c6..96d5351 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -106,6 +106,10 @@ async def error_handler(request): return web.Response(status=500, text="Internal Server Error") +async def forbidden_discovery_handler(request): + return web.Response(status=403, text="discovery refused: tenant is not provisioned for streaming") + + async def crlf_split_events_handler(request): """One multi-line CRLF event whose CRLF is split across two writes.""" response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) @@ -202,6 +206,7 @@ def app(): app.router.add_post("/token", token_handler) app.router.add_post("/token_header_auth", token_header_auth_handler) app.router.add_get("/error", error_handler) + app.router.add_get("/forbidden-discovery", forbidden_discovery_handler) app.router.add_get("/flaky_events", flaky_events_handler) app["flaky"] = {"connections": 0, "last_event_ids": [], "always_drop": False} app.router.add_get("/crlf_split_events", crlf_split_events_handler) @@ -580,3 +585,14 @@ async def test_reconnect_delay_is_capped(sse_transport, aiohttp_client, app, mon assert results == [{"seq": 1}, {"seq": 2}] assert app["huge_retry"]["connections"] == 2 assert time.monotonic() - started < 3 + + +@pytest.mark.asyncio +async def test_register_manual_surfaces_server_error_body(sse_transport, aiohttp_client, app): + """A refused discovery reports the server's body in errors[], not just the status.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/forbidden-discovery"))) + result = await sse_transport.register_manual(None, call_template) + assert result.success is False + assert "discovery refused: tenant is not provisioned for streaming" in result.errors[0] + assert "403" in result.errors[0] diff --git a/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py b/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py index d86a44c..e026f45 100644 --- a/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py @@ -109,6 +109,9 @@ async def check_oauth(request): async def error_endpoint(request): return web.Response(status=500, text="Internal Server Error") + async def forbidden_discovery(request): + return web.Response(status=403, text="discovery refused: tenant is not provisioned for streaming") + app = web.Application() app.add_routes([ web.get('/discover', discover), @@ -120,6 +123,7 @@ async def error_endpoint(request): web.post('/token', oauth_token_handler), web.post('/token-header', oauth_token_header_handler), web.get('/error', error_endpoint), + web.get('/forbidden-discovery', forbidden_discovery), ]) return app @@ -341,3 +345,14 @@ async def test_call_tool_with_oauth2_header_fallback_nonstream(streamable_http_t result = await streamable_http_transport.call_tool(None, "test_tool", {}, call_template) assert result == SAMPLE_NDJSON_RESPONSE + + +@pytest.mark.asyncio +async def test_register_manual_surfaces_server_error_body(streamable_http_transport, aiohttp_client, app): + """A refused discovery reports the server's body in errors[], not just the status.""" + client = await aiohttp_client(app) + call_template = StreamableHttpCallTemplate(name="test-provider", url=f"{client.make_url('/forbidden-discovery')}") + result = await streamable_http_transport.register_manual(None, call_template) + assert result.success is False + assert "discovery refused: tenant is not provisioned for streaming" in result.errors[0] + assert "403" in result.errors[0] From b539b399a3841821ec5c873d111b9c7e17eb1fcb Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:09:09 +0200 Subject: [PATCH 09/34] http errors: structured error fields win, and error bodies are read bounded Addresses cubic review on #102: - error_detail_from_body no longer skips an object-valued `error` to reach a lower-priority generic string: the first of error / message / detail that is present decides, and a structured value returns the raw JSON so its shape stays visible. Null and blank strings are still skipped. - The error body is read incrementally and capped at MAX_BODY_READ_BYTES (64 KiB) instead of buffered in full and truncated afterwards, so an arbitrarily large 4xx/5xx body from an untrusted endpoint cannot grow memory unbounded. Decoding is lenient and honours the response charset. Tests for precedence, null/blank skipping, and a 1 MiB error body. Co-Authored-By: Claude Fable 5.1 --- .../http/src/utcp_http/_errors.py | 55 +++++++++++++++---- .../tests/test_http_communication_protocol.py | 54 ++++++++++++++++++ 2 files changed, 97 insertions(+), 12 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py index f4df4ea..55dbac3 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_errors.py +++ b/plugins/communication_protocols/http/src/utcp_http/_errors.py @@ -14,13 +14,25 @@ # Bodies are folded into an exception message; keep pathological ones bounded. MAX_DETAIL_CHARS = 2000 +# How much of an error body is read at all. The response may come from an +# attacker-controlled endpoint (discovery URLs are exactly that trust +# surface), so the read is bounded up front rather than buffered in full and +# truncated afterwards. Comfortably larger than MAX_DETAIL_CHARS so a JSON +# body with a long ``error`` field still parses. +MAX_BODY_READ_BYTES = 64 * 1024 + +_DETAIL_KEYS = ("error", "message", "detail") + def error_detail_from_body(text: str) -> Optional[str]: """Extract a human-readable reason from an error response body. - Prefers a string ``error`` / ``message`` / ``detail`` field when the body is - a JSON object; an object-valued field falls through to the raw JSON so the - real structure shows. Returns ``None`` for an empty body. + When the body is a JSON object, the first of ``error`` / ``message`` / + ``detail`` that is present decides: a non-empty string is returned as the + reason; anything else (an object, a list, a number) means the server sent a + structured error, so the raw JSON is returned to keep that structure + visible rather than skipping ahead to a lower-priority generic string. + A non-JSON body is returned as-is. Returns ``None`` for an empty body. """ body = text.strip() if not body: @@ -30,25 +42,44 @@ def error_detail_from_body(text: str) -> Optional[str]: except ValueError: return body[:MAX_DETAIL_CHARS] if isinstance(data, dict): - for key in ("error", "message", "detail"): - value = data.get(key) - if isinstance(value, str) and value.strip(): - return value.strip()[:MAX_DETAIL_CHARS] + for key in _DETAIL_KEYS: + if key not in data or data[key] is None: + continue + value = data[key] + if isinstance(value, str): + if value.strip(): + return value.strip()[:MAX_DETAIL_CHARS] + continue + # Structured error: show it rather than a later generic string. + return body[:MAX_DETAIL_CHARS] return body[:MAX_DETAIL_CHARS] +async def _read_body_bounded(response: aiohttp.ClientResponse, limit: int) -> str: + """Read at most ``limit`` bytes of the body and decode them leniently.""" + chunks = [] + total = 0 + async for chunk in response.content.iter_chunked(8192): + chunks.append(chunk) + total += len(chunk) + if total >= limit: + break + raw = b"".join(chunks)[:limit] + return raw.decode(response.charset or "utf-8", errors="replace") + + async def raise_for_status_with_body(response: aiohttp.ClientResponse) -> None: """Like ``response.raise_for_status()``, but with the response body in the error. - On a 4xx/5xx, reads the body and raises a ``ClientResponseError`` of the - same status and headers whose ``message`` is ``": "``. The - raw body text is attached as ``body`` for callers that want the structure. - Does nothing on a 2xx/3xx. + On a 4xx/5xx, reads up to ``MAX_BODY_READ_BYTES`` of the body and raises a + ``ClientResponseError`` of the same status and headers whose ``message`` is + ``": "``. The text that was read is attached as ``body`` + for callers that want the structure. Does nothing on a 2xx/3xx. """ if response.status < 400: return try: - text = await response.text() + text = await _read_body_bounded(response, MAX_BODY_READ_BYTES) except Exception: text = "" detail = error_detail_from_body(text) diff --git a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py index f6de717..a6831fc 100644 --- a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py @@ -150,6 +150,20 @@ async def forbidden_object_handler(request): app.router.add_route('*', '/forbidden', forbidden_handler) app.router.add_route('*', '/forbidden-object', forbidden_object_handler) + + # A structured `error` next to a generic `message`: the structure must win. + async def forbidden_object_then_message_handler(request): + return web.json_response( + {"error": {"code": "INVALID_FIELD", "reason": "value out of range"}, "message": "Request failed"}, + status=422, + ) + + # A huge error body (think an HTML stack trace): the read itself is bounded. + async def forbidden_huge_handler(request): + return web.Response(status=403, text="x" * (1024 * 1024)) + + app.router.add_route('*', '/forbidden-object-then-message', forbidden_object_then_message_handler) + app.router.add_route('*', '/forbidden-huge', forbidden_huge_handler) return app @@ -784,3 +798,43 @@ async def test_register_manual_surfaces_server_error_body(http_transport, aiohtt assert result.success is False assert "You are not allowed to do that, and here is exactly why." in result.errors[0] assert "403" in result.errors[0] + + +@pytest.mark.asyncio +async def test_structured_error_wins_over_generic_message(http_transport, aiohttp_client, app): + """An object under `error` is shown even when a lower-priority string field exists.""" + client = await aiohttp_client(app) + call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}/forbidden-object-then-message", http_method="POST") + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + await http_transport.call_tool(None, "t.tool", {}, call_template) + assert "INVALID_FIELD" in str(excinfo.value) + assert "Request failed" not in excinfo.value.message.split(":", 1)[1].split("INVALID_FIELD")[0] + + +@pytest.mark.asyncio +async def test_huge_error_body_is_read_bounded(http_transport, aiohttp_client, app): + """A 1 MiB error body is neither buffered in full nor folded into the message in full.""" + from utcp_http._errors import MAX_BODY_READ_BYTES, MAX_DETAIL_CHARS + client = await aiohttp_client(app) + call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}/forbidden-huge", http_method="POST") + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + await http_transport.call_tool(None, "t.tool", {}, call_template) + assert excinfo.value.status == 403 + assert len(excinfo.value.body) <= MAX_BODY_READ_BYTES + assert len(excinfo.value.message) <= MAX_DETAIL_CHARS + 50 + + +def test_error_detail_from_body_precedence_and_fallbacks(): + from utcp_http._errors import error_detail_from_body + assert error_detail_from_body("") is None + assert error_detail_from_body(" ") is None + assert error_detail_from_body("plain text") == "plain text" + assert error_detail_from_body('{"error": "nope"}') == "nope" + assert error_detail_from_body('{"message": "nope"}') == "nope" + # Structured error beats a later generic string. + assert error_detail_from_body('{"error": {"code": "X"}, "message": "generic"}') == '{"error": {"code": "X"}, "message": "generic"}' + # An explicit null or blank string is skipped, not treated as structured. + assert error_detail_from_body('{"error": null, "message": "generic"}') == "generic" + assert error_detail_from_body('{"error": " ", "detail": "specific"}') == "specific" + # Non-object JSON falls back to the raw body. + assert error_detail_from_body('["a", "b"]') == '["a", "b"]' From badca3905ec59f90e0b3df0d64b4603a2404086f Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:15:47 +0200 Subject: [PATCH 10/34] http errors: guard the charset lookup; assert precedence on the parsed detail - Decoding a bounded error body now falls back to UTF-8 when the response declares an unknown charset (LookupError) or the lookup fails for any other reason, so the detail is never lost. Tests for a body without a Content-Type and one with an unknown charset. - The precedence test asserts on the parsed detail instead of slicing the message string. Co-Authored-By: Claude Fable 5.1 --- .../http/src/utcp_http/_errors.py | 10 +++++- .../tests/test_http_communication_protocol.py | 32 +++++++++++++++++-- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py index 55dbac3..b3a0e90 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_errors.py +++ b/plugins/communication_protocols/http/src/utcp_http/_errors.py @@ -65,7 +65,15 @@ async def _read_body_bounded(response: aiohttp.ClientResponse, limit: int) -> st if total >= limit: break raw = b"".join(chunks)[:limit] - return raw.decode(response.charset or "utf-8", errors="replace") + # ``charset`` only parses the Content-Type header, but stay defensive: the + # body was read directly, so aiohttp's own buffered-body machinery must not + # be relied on, and an unknown charset name must not lose the detail. + try: + encoding = response.charset or "utf-8" + raw.decode(encoding, errors="replace") + except (LookupError, RuntimeError, ValueError): + encoding = "utf-8" + return raw.decode(encoding, errors="replace") async def raise_for_status_with_body(response: aiohttp.ClientResponse) -> None: diff --git a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py index a6831fc..0dd25e5 100644 --- a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py @@ -164,6 +164,17 @@ async def forbidden_huge_handler(request): app.router.add_route('*', '/forbidden-object-then-message', forbidden_object_then_message_handler) app.router.add_route('*', '/forbidden-huge', forbidden_huge_handler) + + # No Content-Type at all, so no charset to decode with. + async def forbidden_no_charset_handler(request): + return web.Response(status=403, body=b'{"error": "no charset here"}') + + # A charset Python does not know. + async def forbidden_bad_charset_handler(request): + return web.Response(status=403, body=b'{"error": "odd charset"}', content_type="application/json", charset="x-unknown-charset") + + app.router.add_route('*', '/forbidden-no-charset', forbidden_no_charset_handler) + app.router.add_route('*', '/forbidden-bad-charset', forbidden_bad_charset_handler) return app @@ -807,8 +818,14 @@ async def test_structured_error_wins_over_generic_message(http_transport, aiohtt call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}/forbidden-object-then-message", http_method="POST") with pytest.raises(aiohttp.ClientResponseError) as excinfo: await http_transport.call_tool(None, "t.tool", {}, call_template) - assert "INVALID_FIELD" in str(excinfo.value) - assert "Request failed" not in excinfo.value.message.split(":", 1)[1].split("INVALID_FIELD")[0] + assert "INVALID_FIELD" in excinfo.value.message + # The detail is the whole structured body, not the generic "Request failed". + from utcp_http._errors import error_detail_from_body + import json as _json + assert _json.loads(error_detail_from_body(excinfo.value.body)) == { + "error": {"code": "INVALID_FIELD", "reason": "value out of range"}, + "message": "Request failed", + } @pytest.mark.asyncio @@ -838,3 +855,14 @@ def test_error_detail_from_body_precedence_and_fallbacks(): assert error_detail_from_body('{"error": " ", "detail": "specific"}') == "specific" # Non-object JSON falls back to the raw body. assert error_detail_from_body('["a", "b"]') == '["a", "b"]' + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path,expected", [("/forbidden-no-charset", "no charset here"), ("/forbidden-bad-charset", "odd charset")]) +async def test_error_body_is_surfaced_without_a_usable_charset(http_transport, aiohttp_client, app, path, expected): + """A missing or unknown charset must not lose the body; decode as UTF-8.""" + client = await aiohttp_client(app) + call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}{path}", http_method="POST") + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + await http_transport.call_tool(None, "t.tool", {}, call_template) + assert expected in excinfo.value.message From c151ebab7e193dfd14a6a1a7c8c4e49e9369d95f Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:00:38 +0200 Subject: [PATCH 11/34] Pre-release fixes: reuse the MCP client across calls; error bodies on streaming calls Pre-release review of dev: - mcp: _ensure_mcp_client compared the client's whole config dict with its mcpServers entry, which was always unequal, so every tool call built a new MCPClient and spawned a fresh server process that was never closed. Compare the mcpServers entry, and close the previous client's sessions when the configuration really changes. Test asserts one client and one session across two calls. - sse / streamable_http: the streaming call paths now raise with the server's error body like discovery and the TypeScript SDK already do. - _errors: decode the bounded body once. Co-Authored-By: Claude Fable 5.1 --- .../http/src/utcp_http/_errors.py | 6 ++---- .../src/utcp_http/sse_communication_protocol.py | 2 +- .../streamable_http_communication_protocol.py | 2 +- .../src/utcp_mcp/mcp_communication_protocol.py | 15 ++++++++++++++- .../mcp/tests/test_mcp_transport.py | 11 +++++++++++ 5 files changed, 29 insertions(+), 7 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py index b3a0e90..71ddb7e 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_errors.py +++ b/plugins/communication_protocols/http/src/utcp_http/_errors.py @@ -69,11 +69,9 @@ async def _read_body_bounded(response: aiohttp.ClientResponse, limit: int) -> st # body was read directly, so aiohttp's own buffered-body machinery must not # be relied on, and an unknown charset name must not lose the detail. try: - encoding = response.charset or "utf-8" - raw.decode(encoding, errors="replace") + return raw.decode(response.charset or "utf-8", errors="replace") except (LookupError, RuntimeError, ValueError): - encoding = "utf-8" - return raw.decode(encoding, errors="replace") + return raw.decode("utf-8", errors="replace") async def raise_for_status_with_body(response: aiohttp.ClientResponse) -> None: diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index fdb2ff5..f10c936 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -292,7 +292,7 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, f"handshakes; update the call template to point at " f"the final URL directly." ) - response.raise_for_status() + await raise_for_status_with_body(response) except Exception as e: if reconnect_attempts == 0: # The initial handshake failing (refused, timed out, non-2xx) is a diff --git a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py index df4b94f..0b93639 100644 --- a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py @@ -293,7 +293,7 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, f"followed during streaming handshakes; update the " f"call template to point at the final URL directly." ) - response.raise_for_status() + await raise_for_status_with_body(response) async for chunk in self._process_http_stream(response, tool_call_template.chunk_size, tool_call_template.name): yield chunk diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index f86f4ab..b488b2a 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -100,7 +100,20 @@ def _log_error(self, message: str): async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate'): """Ensure MCPClient is initialized with the current configuration.""" - if self._mcp_client is None or self._mcp_client.config != manual_call_template.config.mcpServers: + # ``MCPClient.config`` is the whole ``{"mcpServers": ...}`` dict, so it must + # be compared against its ``mcpServers`` entry. Comparing the whole dict to + # the servers mapping was always unequal, which rebuilt the client and + # spawned a fresh server process on every call without ever closing the + # previous ones. + current_servers = self._mcp_client.config.get("mcpServers") if self._mcp_client is not None else None + if self._mcp_client is None or current_servers != manual_call_template.config.mcpServers: + if self._mcp_client is not None: + # The configuration changed: release the previous client's sessions + # so their child processes do not outlive it. + try: + await self._mcp_client.close_all_sessions() + except Exception as e: + self._log_warning(f"Failed to close sessions of the previous MCP client: {e}") # Create a new MCPClient with the server configuration config = {"mcpServers": manual_call_template.config.mcpServers} self._mcp_client = _QuietStdioMCPClient.from_dict(config) diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index d2e2823..034c5bc 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -296,3 +296,14 @@ async def test_process_tool_result_unwraps_only_single_key_result_wrapper(transp # No structuredContent: fall back to text content. text_only = SimpleNamespace(structuredContent=None, content=[SimpleNamespace(text="7")]) assert transport._process_tool_result(text_only, "t") == 7 + + +@pytest.mark.asyncio +async def test_mcp_client_and_session_are_reused_across_calls(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + """Repeated calls with the same configuration reuse one client and one session + instead of spawning a new server process per call.""" + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "one"}, mcp_manual) + client_after_first = transport._mcp_client + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "two"}, mcp_manual) + assert transport._mcp_client is client_after_first + assert list(transport._mcp_client.sessions.keys()) == [SERVER_NAME] From 8d7142191bc182863f343ab664fa9fd9d7813c94 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:17:09 +0200 Subject: [PATCH 12/34] Pre-release review fixes, second batch - sse: an event block without an `event:` field has the type "message" (spec), so event_type="message" now matches it. An empty `id:` resets the last event ID and no Last-Event-ID header is sent for it; ids containing NUL are ignored. A 200 whose Content-Type is not text/event-stream raises SseProtocolError instead of parsing into zero events. Calls that send a request body are never reconnected: a re-issued POST could re-execute a non-idempotent tool. - _errors: a body nested deeper than the JSON parser's recursion limit raised RecursionError past the HTTP error handling; it is caught and treated as text. Control characters are collapsed so a server cannot forge log records or terminal escapes through an error message. - mcp: close() referenced a _session_locks attribute that was never assigned and always raised AttributeError after cleanup. A child that starts but fails the MCP handshake is now disconnected and removed from active_sessions instead of lingering. Tests for each. Co-Authored-By: Claude Fable 5.1 --- .../http/src/utcp_http/_errors.py | 22 ++++-- .../utcp_http/sse_communication_protocol.py | 28 ++++++-- .../tests/test_http_communication_protocol.py | 21 ++++++ .../tests/test_sse_communication_protocol.py | 72 +++++++++++++++++++ .../utcp_mcp/mcp_communication_protocol.py | 12 +++- .../mcp/tests/test_mcp_transport.py | 8 +++ 6 files changed, 150 insertions(+), 13 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py index 71ddb7e..a3ac17d 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_errors.py +++ b/plugins/communication_protocols/http/src/utcp_http/_errors.py @@ -7,6 +7,7 @@ status code. Mirrors the TypeScript SDK's ``_normalizeToolError``. """ import json +import re from typing import Optional import aiohttp @@ -23,6 +24,15 @@ _DETAIL_KEYS = ("error", "message", "detail") +# Control characters (newlines, ANSI escape introducers, NUL) are collapsed so +# server-controlled text folded into an exception message or a log line cannot +# forge extra log records or terminal escape sequences. +_CONTROL_CHARS = re.compile(r"[\x00-\x1f\x7f]+") + + +def _clean(text: str) -> str: + return _CONTROL_CHARS.sub(" ", text).strip()[:MAX_DETAIL_CHARS] + def error_detail_from_body(text: str) -> Optional[str]: """Extract a human-readable reason from an error response body. @@ -39,8 +49,10 @@ def error_detail_from_body(text: str) -> Optional[str]: return None try: data = json.loads(body) - except ValueError: - return body[:MAX_DETAIL_CHARS] + except (ValueError, RecursionError): + # RecursionError: a deeply nested body ("[[[[...") within the read cap + # can exceed the parser's recursion limit; it is still just text. + return _clean(body) if isinstance(data, dict): for key in _DETAIL_KEYS: if key not in data or data[key] is None: @@ -48,11 +60,11 @@ def error_detail_from_body(text: str) -> Optional[str]: value = data[key] if isinstance(value, str): if value.strip(): - return value.strip()[:MAX_DETAIL_CHARS] + return _clean(value) continue # Structured error: show it rather than a later generic string. - return body[:MAX_DETAIL_CHARS] - return body[:MAX_DETAIL_CHARS] + return _clean(body) + return _clean(body) async def _read_body_bounded(response: aiohttp.ClientResponse, limit: int) -> str: diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index f10c936..5ef191b 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -252,7 +252,11 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, data = body_content if "application/json" not in content_type else None json_data = body_content if "application/json" in content_type else None - reconnect = bool(tool_call_template.reconnect) + # Never re-send a request body: a reconnect re-issues the request, and for + # a POST that would re-execute a possibly non-idempotent tool. + reconnect = bool(tool_call_template.reconnect) and body_content is None + if tool_call_template.reconnect and body_content is not None: + logger.info(f"Reconnection is disabled for '{tool_call_template.name}' because the call sends a request body.") retry_delay_ms = tool_call_template.retry_timeout last_event_id: Optional[str] = None reconnect_attempts = 0 @@ -260,8 +264,9 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, while True: attempt_headers = dict(request_headers) - if last_event_id is not None: - # Let the server resume from where we left off (SSE spec). + if last_event_id: + # Let the server resume from where we left off (SSE spec). An empty + # last event ID means "none": the header is not sent. attempt_headers["Last-Event-ID"] = last_event_id session = aiohttp.ClientSession() @@ -293,6 +298,16 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, f"the final URL directly." ) await raise_for_status_with_body(response) + # Anything but an event stream would be parsed into silence: a + # JSON error document, say, yields zero events and a "successful" call. + content_type = response.headers.get("Content-Type", "") + if "text/event-stream" not in content_type.lower(): + response.release() + raise SseProtocolError( + f"Expected a text/event-stream response but got {content_type or 'no Content-Type'!r}" + ) + except SseProtocolError: + raise except Exception as e: if reconnect_attempts == 0: # The initial handshake failing (refused, timed out, non-2xx) is a @@ -315,13 +330,16 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, try: async for event in self._iter_sse_events(response): - if event.get("id") is not None: + # Per the SSE spec an id containing NUL is ignored, and an empty id + # resets the last event ID. + if event.get("id") is not None and "\x00" not in event["id"]: last_event_id = event["id"] if event.get("retry") is not None: retry_delay_ms = event["retry"] if "data" not in event: continue - if tool_call_template.event_type and event.get("event") != tool_call_template.event_type: + # An event block without an ``event:`` field has the type "message". + if tool_call_template.event_type and (event.get("event") or "message") != tool_call_template.event_type: continue yield self._parse_event_data(event["data"]) # The server ended the stream cleanly: the tool call is complete. diff --git a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py index 0dd25e5..4504e2c 100644 --- a/plugins/communication_protocols/http/tests/test_http_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_http_communication_protocol.py @@ -174,6 +174,12 @@ async def forbidden_bad_charset_handler(request): return web.Response(status=403, body=b'{"error": "odd charset"}', content_type="application/json", charset="x-unknown-charset") app.router.add_route('*', '/forbidden-no-charset', forbidden_no_charset_handler) + + # A body nested deeper than the JSON parser's recursion limit. + async def forbidden_deep_handler(request): + return web.Response(status=503, body=b"[" * 20000, content_type="application/json") + + app.router.add_route('*', '/forbidden-deep', forbidden_deep_handler) app.router.add_route('*', '/forbidden-bad-charset', forbidden_bad_charset_handler) return app @@ -866,3 +872,18 @@ async def test_error_body_is_surfaced_without_a_usable_charset(http_transport, a with pytest.raises(aiohttp.ClientResponseError) as excinfo: await http_transport.call_tool(None, "t.tool", {}, call_template) assert expected in excinfo.value.message + + +@pytest.mark.asyncio +async def test_deeply_nested_error_body_still_raises_client_response_error(http_transport, aiohttp_client, app): + """A body that overflows the JSON parser's recursion limit must not escape as RecursionError.""" + client = await aiohttp_client(app) + call_template = HttpCallTemplate(name="t", url=f"http://localhost:{client.port}/forbidden-deep", http_method="POST") + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + await http_transport.call_tool(None, "t.tool", {}, call_template) + assert excinfo.value.status == 503 + + +def test_error_detail_collapses_control_characters(): + from utcp_http._errors import error_detail_from_body + assert error_detail_from_body('{"error": "line one\\nline two\\u001b[31m"}') == "line one line two [31m" diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index 96d5351..2e00b11 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -169,6 +169,27 @@ async def huge_retry_events_handler(request): await response.write(b'id: 2\ndata: {"seq": 2}\n\n') return response +async def json_not_sse_handler(request): + """A 200 that is not an event stream at all.""" + return web.json_response({"error": "not a stream"}) + + +async def empty_id_events_handler(request): + """Sets an id, then resets it with an empty id, then drops the connection.""" + state = request.app["empty_id"] + state["connections"] += 1 + state["last_event_ids"].append(request.headers.get("Last-Event-ID")) + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + if state["connections"] == 1: + await response.write(b'id: 1\ndata: {"seq": 1}\n\nid\ndata: {"seq": 2}\n\n') + await asyncio.sleep(0.01) + request.transport.close() + return response + await response.write(b'data: {"seq": 3}\n\n') + return response + + async def flaky_events_handler(request): """Serves the first event then drops the TCP connection on the first connection (or on every connection when ``always_drop`` is set). A reconnecting client is @@ -203,6 +224,10 @@ def app(): app = web.Application() app.router.add_get("/tools", tools_handler) app.router.add_route('*', '/events', events_handler) + app.router.add_post("/flaky_events", flaky_events_handler) + app.router.add_get("/json_not_sse", json_not_sse_handler) + app.router.add_get("/empty_id_events", empty_id_events_handler) + app["empty_id"] = {"connections": 0, "last_event_ids": []} app.router.add_post("/token", token_handler) app.router.add_post("/token_header_auth", token_header_auth_handler) app.router.add_get("/error", error_handler) @@ -596,3 +621,50 @@ async def test_register_manual_surfaces_server_error_body(sse_transport, aiohttp assert result.success is False assert "discovery refused: tenant is not provisioned for streaming" in result.errors[0] assert "403" in result.errors[0] + + +# --- Spec conformance follow-ups --- + +@pytest.mark.asyncio +async def test_event_type_message_matches_events_without_an_event_field(sse_transport, aiohttp_client, app): + """Per the SSE spec an event block without `event:` has the type "message".""" + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/events")), event_type="message") + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] + assert results == [{"message": "First part"}] + + +@pytest.mark.asyncio +async def test_empty_id_resets_last_event_id(sse_transport, aiohttp_client, app): + """An empty `id` line resets the last event ID, so no Last-Event-ID header is sent on reconnect.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/empty_id_events")), reconnect=True, retry_timeout=10) + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] + assert results == [{"seq": 1}, {"seq": 2}, {"seq": 3}] + assert app["empty_id"]["last_event_ids"] == [None, None] + + +@pytest.mark.asyncio +async def test_non_event_stream_response_raises(sse_transport, aiohttp_client, app): + """A 200 that is not text/event-stream fails instead of yielding zero events.""" + from utcp_http.sse_communication_protocol import SseProtocolError + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/json_not_sse"))) + with pytest.raises(SseProtocolError): + async for _ in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template): + pass + + +@pytest.mark.asyncio +async def test_post_stream_is_not_reconnected(sse_transport, aiohttp_client, app): + """A dropped POST stream is not re-issued: that could re-execute a non-idempotent tool.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate( + name="test-sse", url=str(client.make_url("/flaky_events")), reconnect=True, retry_timeout=10, body_field="payload" + ) + received = [] + with pytest.raises(aiohttp.ClientError): + async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {"payload": {"n": 1}}, call_template): + received.append(e) + assert received == [{"message": "First part"}] + assert app["flaky"]["connections"] == 1 diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index b488b2a..e5ee9e5 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -67,9 +67,16 @@ async def create_session(self, server_name: str, auto_initialize: bool = True): try: await session.initialize() except Exception: - # Mirror the base class: a session that failed to initialize must - # not stay cached, or the next lookup would hand back a dead one. + # The base class only registers a session after a successful + # initialize; undo the early registration, and disconnect so a + # child that started but failed the MCP handshake does not linger. self.sessions.pop(server_name, None) + if server_name in self.active_sessions: + self.active_sessions.remove(server_name) + try: + await session.disconnect() + except Exception as disconnect_error: + logger.warning(f"Failed to disconnect '{server_name}' after a failed initialize: {disconnect_error}") raise return session @@ -569,7 +576,6 @@ async def close(self) -> None: """Close all active sessions and clean up resources.""" self._log_info("Closing MCP communication protocol and cleaning up all sessions") await self._cleanup_all_sessions() - self._session_locks.clear() self._log_info("MCP communication protocol closed successfully") async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index 034c5bc..38977d6 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -307,3 +307,11 @@ async def test_mcp_client_and_session_are_reused_across_calls(transport: McpComm await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "two"}, mcp_manual) assert transport._mcp_client is client_after_first assert list(transport._mcp_client.sessions.keys()) == [SERVER_NAME] + + +@pytest.mark.asyncio +async def test_close_after_use_does_not_raise(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + """close() used to raise AttributeError after cleaning up; it must complete.""" + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "one"}, mcp_manual) + await transport.close() + assert transport._mcp_client.sessions == {} From b83d77c7aa65155821675855b8af0193f0ed88dc Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:31:47 +0200 Subject: [PATCH 13/34] Pre-release review fixes, third batch - sse: `retry:` is honoured only when made of ASCII digits (spec); "-1" or "20ms" no longer change the reconnect delay. A stream that ends in the middle of an event no longer dispatches the incomplete event (spec: pending data is discarded at end of file). - streamable_http: Content-Type is matched case-insensitively. - udp: remove the stale comment block describing the bug this release fixed. - tests: the slow-handshake handler no longer outlives the test by seconds. Co-Authored-By: Claude Fable 5.1 --- .../utcp_http/sse_communication_protocol.py | 15 ++++-------- .../streamable_http_communication_protocol.py | 2 +- .../tests/test_sse_communication_protocol.py | 23 +++++++++++++++++-- .../utcp_socket/udp_communication_protocol.py | 4 ---- 4 files changed, 27 insertions(+), 17 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index 5ef191b..f7b1236 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -393,10 +393,9 @@ def flush(event_string: str): elif field == 'id': current_event['id'] = value elif field == 'retry': - try: + # Spec: only a value made of ASCII digits sets the reconnection time. + if value.isascii() and value.isdigit(): current_event['retry'] = int(value) - except ValueError: - pass if data_lines: current_event['data'] = '\n'.join(data_lines) return current_event or None @@ -431,13 +430,9 @@ def normalise(text: str) -> str: f"SSE event exceeded {self.MAX_EVENT_BUFFER_CHARS} characters without a blank-line delimiter" ) - # Flush a trailing event that was not terminated by a blank line. - buffer += normalise(decoder.decode(b"", final=True)) - if pending_cr: - buffer += "\n" - event = flush(buffer) - if event is not None: - yield event + # Spec: if the stream ends in the middle of an event, before the final + # blank line, the incomplete event is not dispatched. + decoder.decode(b"", final=True) @staticmethod def _parse_event_data(data: str) -> Any: diff --git a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py index 0b93639..0208605 100644 --- a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py @@ -310,7 +310,7 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, async def _process_http_stream(self, response: ClientResponse, chunk_size: Optional[int], provider_name: str) -> AsyncIterator[Any]: """Process the HTTP stream and yield chunks based on content type.""" try: - content_type = response.headers.get('Content-Type', '') + content_type = response.headers.get('Content-Type', '').lower() if 'application/x-ndjson' in content_type: async for line in response.content: diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index 2e00b11..2045fee 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -150,8 +150,9 @@ async def flaky_503_events_handler(request): async def slow_handshake_handler(request): - """Accepts the connection but does not send response headers for a long time.""" - await asyncio.sleep(5) + """Accepts the connection but does not send response headers until well past + the (patched) handshake timeout.""" + await asyncio.sleep(1) return web.Response(status=204) @@ -169,6 +170,15 @@ async def huge_retry_events_handler(request): await response.write(b'id: 2\ndata: {"seq": 2}\n\n') return response +async def bad_retry_events_handler(request): + """A retry field that is not made of digits must be ignored; the stream then + ends in the middle of an event, which must not be dispatched.""" + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + await response.write(b'retry: -1\ndata: {"seq": 1}\n\nretry: 20ms\n\ndata: {"seq": 2}') + return response + + async def json_not_sse_handler(request): """A 200 that is not an event stream at all.""" return web.json_response({"error": "not a stream"}) @@ -226,6 +236,7 @@ def app(): app.router.add_route('*', '/events', events_handler) app.router.add_post("/flaky_events", flaky_events_handler) app.router.add_get("/json_not_sse", json_not_sse_handler) + app.router.add_get("/bad_retry_events", bad_retry_events_handler) app.router.add_get("/empty_id_events", empty_id_events_handler) app["empty_id"] = {"connections": 0, "last_event_ids": []} app.router.add_post("/token", token_handler) @@ -668,3 +679,11 @@ async def test_post_stream_is_not_reconnected(sse_transport, aiohttp_client, app received.append(e) assert received == [{"message": "First part"}] assert app["flaky"]["connections"] == 1 + + +@pytest.mark.asyncio +async def test_malformed_retry_is_ignored_and_unterminated_trailing_event_is_dropped(sse_transport, aiohttp_client, app): + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/bad_retry_events"))) + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] + assert results == [{"seq": 1}] diff --git a/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py b/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py index fa1f98e..dec60c7 100644 --- a/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py +++ b/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py @@ -327,10 +327,6 @@ async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], too raise # Copilot AI (5 days ago): - # The call_tool_streaming method wraps a generator function but doesn't use the async def syntax for the method itself. - # While this works, it's inconsistent with the other implementation in tcp_communication_protocol.py (lines 384-387) which properly uses async def with an inner generator. - # For consistency and clarity, this should also use async def directly: - # async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: """REQUIRED Streaming variant: the UDP protocol does not natively stream, so the full result is yielded as a single chunk.""" From 2533abd9dc38fe7c13f0dc6ac501c465858a37ce Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:41:16 +0200 Subject: [PATCH 14/34] sse: compare the response media type exactly A substring check let a Content-Type such as text/event-stream-invalid through; compare the media-type portion exactly, parameters allowed. Co-Authored-By: Claude Fable 5.1 --- .../http/src/utcp_http/sse_communication_protocol.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index f7b1236..9abcaa8 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -301,7 +301,10 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, # Anything but an event stream would be parsed into silence: a # JSON error document, say, yields zero events and a "successful" call. content_type = response.headers.get("Content-Type", "") - if "text/event-stream" not in content_type.lower(): + # Compare the media type exactly (parameters such as charset allowed), + # so "text/event-stream-invalid" does not pass a substring check. + media_type = content_type.split(";", 1)[0].strip().lower() + if media_type != "text/event-stream": response.release() raise SseProtocolError( f"Expected a text/event-stream response but got {content_type or 'no Content-Type'!r}" From b0849f370002d936d3884d02f50dfaa26b724412 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:52:01 +0200 Subject: [PATCH 15/34] Pre-release review fixes, fourth batch - mcp: one MCPClient per distinct server configuration instead of a single shared client. The protocol object is registered once per process and shared by every manual, so a single client made manuals with different configurations evict each other's sessions, including ones still in use by a concurrent call. Clients are keyed by the canonical configuration, created under a lock so two concurrent first calls cannot each spawn a server, and never evicted by another manual's activity. close() closes every client's sessions and keeps a client whose shutdown failed so a later close() can retry it. - sse: the error-body read on a refused stream is bounded by the handshake timeout, so a server that answers 4xx/5xx and then stalls cannot hang the call. A final blank line ending in a lone CR still completes the last event; only genuinely incomplete events are discarded. Absurdly long retry digit strings are ignored. - _errors: control-character collapsing covers the C1 range too. - tests: streaming calls against a 5xx assert the body is surfaced for both SSE and Streamable HTTP; separate-clients-per-configuration test; the malformed-retry test is named for what it verifies. Co-Authored-By: Claude Fable 5.1 --- .../http/src/utcp_http/_errors.py | 3 +- .../utcp_http/sse_communication_protocol.py | 26 ++++-- .../tests/test_sse_communication_protocol.py | 31 ++++++- ..._streamable_http_communication_protocol.py | 12 +++ .../utcp_mcp/mcp_communication_protocol.py | 86 +++++++++++-------- .../mcp/tests/test_mcp_transport.py | 31 +++++-- 6 files changed, 142 insertions(+), 47 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py index a3ac17d..9d4fa96 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_errors.py +++ b/plugins/communication_protocols/http/src/utcp_http/_errors.py @@ -27,7 +27,8 @@ # Control characters (newlines, ANSI escape introducers, NUL) are collapsed so # server-controlled text folded into an exception message or a log line cannot # forge extra log records or terminal escape sequences. -_CONTROL_CHARS = re.compile(r"[\x00-\x1f\x7f]+") +# C0 and C1 control ranges: C1 (U+0080..U+009F) carries escape introducers too. +_CONTROL_CHARS = re.compile(r"[\x00-\x1f\x7f-\x9f]+") def _clean(text: str) -> str: diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index 9abcaa8..03e6d88 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -297,7 +297,9 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, f"handshakes; update the call template to point at " f"the final URL directly." ) - await raise_for_status_with_body(response) + # The error-body read is bounded like the handshake, so a server that + # answers 4xx/5xx and then stalls cannot hang the call either. + await asyncio.wait_for(raise_for_status_with_body(response), timeout=self.HANDSHAKE_TIMEOUT_SECONDS) # Anything but an event stream would be parsed into silence: a # JSON error document, say, yields zero events and a "successful" call. content_type = response.headers.get("Content-Type", "") @@ -397,8 +399,12 @@ def flush(event_string: str): current_event['id'] = value elif field == 'retry': # Spec: only a value made of ASCII digits sets the reconnection time. + # int() refuses absurdly long digit strings; ignore those too. if value.isascii() and value.isdigit(): - current_event['retry'] = int(value) + try: + current_event['retry'] = int(value) + except ValueError: + pass if data_lines: current_event['data'] = '\n'.join(data_lines) return current_event or None @@ -433,9 +439,19 @@ def normalise(text: str) -> str: f"SSE event exceeded {self.MAX_EVENT_BUFFER_CHARS} characters without a blank-line delimiter" ) - # Spec: if the stream ends in the middle of an event, before the final - # blank line, the incomplete event is not dispatched. - decoder.decode(b"", final=True) + # At end of stream, a held-back CR is a real line terminator and may + # complete the closing blank line of the last event. Dispatch whatever is + # fully delimited; per spec, an event still incomplete after that (no + # final blank line) is discarded. + buffer += normalise(decoder.decode(b"", final=True)) + if pending_cr: + buffer += "\n" + pending_cr = False + while "\n\n" in buffer: + event_string, buffer = buffer.split("\n\n", 1) + event = flush(event_string) + if event is not None: + yield event @staticmethod def _parse_event_data(data: str) -> Any: diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index 2045fee..c6cfedc 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -179,6 +179,14 @@ async def bad_retry_events_handler(request): return response +async def cr_eof_events_handler(request): + """A complete event whose closing blank line ends in a lone CR at end of stream.""" + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + await response.write(b'data: {"seq": 1}\n\r') + return response + + async def json_not_sse_handler(request): """A 200 that is not an event stream at all.""" return web.json_response({"error": "not a stream"}) @@ -237,6 +245,7 @@ def app(): app.router.add_post("/flaky_events", flaky_events_handler) app.router.add_get("/json_not_sse", json_not_sse_handler) app.router.add_get("/bad_retry_events", bad_retry_events_handler) + app.router.add_get("/cr_eof_events", cr_eof_events_handler) app.router.add_get("/empty_id_events", empty_id_events_handler) app["empty_id"] = {"connections": 0, "last_event_ids": []} app.router.add_post("/token", token_handler) @@ -682,8 +691,28 @@ async def test_post_stream_is_not_reconnected(sse_transport, aiohttp_client, app @pytest.mark.asyncio -async def test_malformed_retry_is_ignored_and_unterminated_trailing_event_is_dropped(sse_transport, aiohttp_client, app): +async def test_malformed_retry_does_not_abort_and_unterminated_trailing_event_is_dropped(sse_transport, aiohttp_client, app): client = await aiohttp_client(app) call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/bad_retry_events"))) results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] assert results == [{"seq": 1}] + + +@pytest.mark.asyncio +async def test_final_blank_line_ending_in_lone_cr_completes_last_event(sse_transport, aiohttp_client, app): + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/cr_eof_events"))) + results = [e async for e in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template)] + assert results == [{"seq": 1}] + + +@pytest.mark.asyncio +async def test_streaming_call_error_surfaces_server_body(sse_transport, aiohttp_client, app): + """A refused stream carries the server's body, like discovery does.""" + client = await aiohttp_client(app) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/error"))) + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + async for _ in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template): + pass + assert excinfo.value.status == 500 + assert "Internal Server Error" in excinfo.value.message diff --git a/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py b/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py index e026f45..4315952 100644 --- a/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py @@ -356,3 +356,15 @@ async def test_register_manual_surfaces_server_error_body(streamable_http_transp assert result.success is False assert "discovery refused: tenant is not provisioned for streaming" in result.errors[0] assert "403" in result.errors[0] + + +@pytest.mark.asyncio +async def test_streaming_call_error_surfaces_server_body(streamable_http_transport, aiohttp_client, app): + """A refused stream carries the server's body, like discovery does.""" + client = await aiohttp_client(app) + call_template = StreamableHttpCallTemplate(name="test-provider", url=f"{client.make_url('/error')}") + with pytest.raises(aiohttp.ClientResponseError) as excinfo: + async for _ in streamable_http_transport.call_tool_streaming(None, "test-provider.t", {}, call_template): + pass + assert excinfo.value.status == 500 + assert "Internal Server Error" in excinfo.value.message diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index e5ee9e5..c158119 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -1,3 +1,4 @@ +import asyncio import os import sys from typing import Any, Dict, Optional, AsyncGenerator, TYPE_CHECKING, Tuple, TextIO @@ -91,7 +92,12 @@ class McpCommunicationProtocol(CommunicationProtocol): def __init__(self): self._oauth_tokens: Dict[str, Dict[str, Any]] = {} - self._mcp_client: Optional[MCPClient] = None + # One MCPClient per distinct server configuration. This protocol object is + # registered once per process and shared by every manual, so a single + # client would make manuals with different configurations evict each + # other's sessions, including sessions still in use by a concurrent call. + self._mcp_clients: Dict[str, MCPClient] = {} + self._clients_lock = asyncio.Lock() def _log_info(self, message: str): """Log informational messages.""" @@ -105,40 +111,45 @@ def _log_error(self, message: str): """Log error messages.""" logger.error(f"[McpCommunicationProtocol] {message}") - async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate'): - """Ensure MCPClient is initialized with the current configuration.""" - # ``MCPClient.config`` is the whole ``{"mcpServers": ...}`` dict, so it must - # be compared against its ``mcpServers`` entry. Comparing the whole dict to - # the servers mapping was always unequal, which rebuilt the client and - # spawned a fresh server process on every call without ever closing the - # previous ones. - current_servers = self._mcp_client.config.get("mcpServers") if self._mcp_client is not None else None - if self._mcp_client is None or current_servers != manual_call_template.config.mcpServers: - if self._mcp_client is not None: - # The configuration changed: release the previous client's sessions - # so their child processes do not outlive it. - try: - await self._mcp_client.close_all_sessions() - except Exception as e: - self._log_warning(f"Failed to close sessions of the previous MCP client: {e}") - # Create a new MCPClient with the server configuration - config = {"mcpServers": manual_call_template.config.mcpServers} - self._mcp_client = _QuietStdioMCPClient.from_dict(config) + @staticmethod + def _config_key(manual_call_template: 'McpCallTemplate') -> str: + """Canonical key for a manual's server configuration.""" + return json.dumps(manual_call_template.config.mcpServers, sort_keys=True, default=str) + + async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> MCPClient: + """Return the MCPClient for this manual's configuration, creating it once. + + Clients are keyed by configuration and never evicted by another manual's + activity, so sessions are reused across calls and a call in flight on one + configuration is never torn down by a call on another. Creation is + serialised so two concurrent first calls cannot each spawn a client. + """ + key = self._config_key(manual_call_template) + client = self._mcp_clients.get(key) + if client is not None: + return client + async with self._clients_lock: + client = self._mcp_clients.get(key) + if client is None: + config = {"mcpServers": manual_call_template.config.mcpServers} + client = _QuietStdioMCPClient.from_dict(config) + self._mcp_clients[key] = client + return client async def _get_or_create_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): """Get an existing session or create a new one using MCPClient.""" - await self._ensure_mcp_client(manual_call_template) + client = await self._ensure_mcp_client(manual_call_template) try: # Try to get existing session - session = self._mcp_client.get_session(server_name) + session = client.get_session(server_name) self._log_info(f"Reusing existing session for server: {server_name}") return session except ValueError: # Session doesn't exist, create a new one self._log_info(f"Creating new session for server: {server_name}") try: - session = await self._mcp_client.create_session(server_name, auto_initialize=True) + session = await client.create_session(server_name, auto_initialize=True) except Exception as e: server_config = manual_call_template.config.mcpServers.get(server_name) is_stdio = isinstance(server_config, dict) and "command" in server_config @@ -150,16 +161,23 @@ async def _get_or_create_session(self, server_name: str, manual_call_template: ' raise return session - async def _cleanup_session(self, server_name: str): - """Clean up a specific session.""" - if self._mcp_client: - await self._mcp_client.close_session(server_name) + async def _cleanup_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): + """Clean up a specific session of the client serving this manual.""" + client = self._mcp_clients.get(self._config_key(manual_call_template)) + if client is not None and server_name in client.sessions: + await client.close_session(server_name) self._log_info(f"Cleaned up session for server: {server_name}") async def _cleanup_all_sessions(self): - """Clean up all active sessions.""" - if self._mcp_client: - await self._mcp_client.close_all_sessions() + """Close every session of every client. A client whose shutdown fails is + kept so a later close() can retry it instead of leaking its processes.""" + for key, client in list(self._mcp_clients.items()): + try: + await client.close_all_sessions() + del self._mcp_clients[key] + except Exception as e: + self._log_warning(f"Failed to close sessions of an MCP client: {e}") + if not self._mcp_clients: self._log_info("Cleaned up all sessions") def _add_server_to_tool_name(self, tools, server_name: str): @@ -192,7 +210,7 @@ async def _list_tools_with_session(self, server_name: str, manual_call_template: if is_session_error: # Only restart session for connection/transport level issues - await self._cleanup_session(server_name) + await self._cleanup_session(server_name, manual_call_template) self._log_warning(f"Session-level error for list_tools, retrying with fresh session: {e}") # Retry with a fresh session @@ -220,7 +238,7 @@ async def _list_resources_with_session(self, server_name: str, manual_call_templ return resources_response except Exception as e: # If there's an error, clean up the potentially bad session and try once more - await self._cleanup_session(server_name) + await self._cleanup_session(server_name, manual_call_template) self._log_warning(f"Session failed for list_resources, retrying: {e}") # Retry with a fresh session @@ -240,7 +258,7 @@ async def _read_resource_with_session(self, server_name: str, manual_call_templa return result except Exception as e: # If there's an error, clean up the potentially bad session and try once more - await self._cleanup_session(server_name) + await self._cleanup_session(server_name, manual_call_template) self._log_warning(f"Session failed for read_resource '{resource_uri}', retrying: {e}") # Retry with a fresh session @@ -569,7 +587,7 @@ async def deregister_manual(self, caller: 'UtcpClient', manual_call_template: Ca # Clean up sessions for all servers in this manual if manual_call_template.config and manual_call_template.config.mcpServers: for server_name, server_config in manual_call_template.config.mcpServers.items(): - await self._cleanup_session(server_name) + await self._cleanup_session(server_name, manual_call_template) self._log_info(f"Cleaned up session for server '{server_name}'") async def close(self) -> None: diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index 38977d6..65b4da9 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -264,7 +264,7 @@ async def test_stdio_child_stderr_suppressed_by_default(transport: McpCommunicat assert session.connector.errlog is not sys.stderr assert session.connector.errlog.name == os.devnull finally: - await transport._cleanup_session(SERVER_NAME) + await transport._cleanup_session(SERVER_NAME, mcp_manual) @pytest.mark.asyncio @@ -275,7 +275,7 @@ async def test_stdio_child_stderr_inherit_opt_in(transport: McpCommunicationProt try: assert session.connector.errlog is sys.stderr finally: - await transport._cleanup_session(SERVER_NAME) + await transport._cleanup_session(SERVER_NAME, mcp_manual) @pytest.mark.asyncio @@ -303,10 +303,29 @@ async def test_mcp_client_and_session_are_reused_across_calls(transport: McpComm """Repeated calls with the same configuration reuse one client and one session instead of spawning a new server process per call.""" await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "one"}, mcp_manual) - client_after_first = transport._mcp_client + assert len(transport._mcp_clients) == 1 + client_after_first = next(iter(transport._mcp_clients.values())) await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "two"}, mcp_manual) - assert transport._mcp_client is client_after_first - assert list(transport._mcp_client.sessions.keys()) == [SERVER_NAME] + assert len(transport._mcp_clients) == 1 + assert next(iter(transport._mcp_clients.values())) is client_after_first + assert list(client_after_first.sessions.keys()) == [SERVER_NAME] + + +@pytest.mark.asyncio +async def test_manuals_with_different_configurations_get_separate_clients(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + """The protocol object is shared by every manual; one manual's calls must not + evict another manual's sessions.""" + other_manual = McpCallTemplate( + name="other_manual", + call_template_type="mcp", + config=McpConfig(mcpServers={"other_server": dict(mcp_manual.config.mcpServers[SERVER_NAME])}), + ) + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "a"}, mcp_manual) + await transport.call_tool(None, "other_server.echo", {"message": "b"}, other_manual) + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "c"}, mcp_manual) + assert len(transport._mcp_clients) == 2 + sessions = sorted(name for c in transport._mcp_clients.values() for name in c.sessions) + assert sessions == sorted([SERVER_NAME, "other_server"]) @pytest.mark.asyncio @@ -314,4 +333,4 @@ async def test_close_after_use_does_not_raise(transport: McpCommunicationProtoco """close() used to raise AttributeError after cleaning up; it must complete.""" await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "one"}, mcp_manual) await transport.close() - assert transport._mcp_client.sessions == {} + assert transport._mcp_clients == {} From 7e00e6fbe24f9d5a6b2d551bba01d8ff1c3279a8 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:59:47 +0200 Subject: [PATCH 16/34] Pre-release review fixes, fifth batch - mcp: when a manual's configuration changes and no other manual uses the old one, the old client's sessions are closed instead of lingering until close(). Ownership is tracked per manual name. - sse: the residual buffer is checked against the event cap at end of stream too; retry digit strings are bounded to 18 digits before conversion. - tests: the connection-drop handlers wait 100 ms after writing the first event before closing the socket. On Windows CI the event and the close otherwise arrived together and aiohttp raised before delivering the event, failing six tests that pass elsewhere. Co-Authored-By: Claude Fable 5.1 --- .../utcp_http/sse_communication_protocol.py | 16 +++++++++------ .../tests/test_sse_communication_protocol.py | 20 +++++++++++++++---- .../utcp_mcp/mcp_communication_protocol.py | 20 ++++++++++++++++++- .../mcp/tests/test_mcp_transport.py | 15 ++++++++++++++ 4 files changed, 60 insertions(+), 11 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index 03e6d88..1c2ad6b 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -399,12 +399,10 @@ def flush(event_string: str): current_event['id'] = value elif field == 'retry': # Spec: only a value made of ASCII digits sets the reconnection time. - # int() refuses absurdly long digit strings; ignore those too. - if value.isascii() and value.isdigit(): - try: - current_event['retry'] = int(value) - except ValueError: - pass + # Anything longer than 18 digits is absurd (and would be capped anyway); + # bounding the length keeps the conversion cheap whatever the interpreter. + if value.isascii() and value.isdigit() and len(value) <= 18: + current_event['retry'] = int(value) if data_lines: current_event['data'] = '\n'.join(data_lines) return current_event or None @@ -452,6 +450,12 @@ def normalise(text: str) -> str: event = flush(event_string) if event is not None: yield event + # The residual (discarded) buffer is still subject to the cap, so an + # over-limit malformed stream fails the same way at end of stream. + if len(buffer) > self.MAX_EVENT_BUFFER_CHARS: + raise SseProtocolError( + f"SSE event exceeded {self.MAX_EVENT_BUFFER_CHARS} characters without a blank-line delimiter" + ) @staticmethod def _parse_event_data(data: str) -> Any: diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index c6cfedc..6c7df54 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -141,7 +141,10 @@ async def flaky_503_events_handler(request): await response.prepare(request) if state["connections"] == 1: await response.write(SAMPLE_SSE_EVENTS[0].encode('utf-8')) - await asyncio.sleep(0.01) + # Make sure the event has left the socket before dropping it: on Windows the + # data and the close otherwise arrive together and aiohttp raises before + # delivering the event. + await asyncio.sleep(0.1) request.transport.close() return response for event in SAMPLE_SSE_EVENTS[1:]: @@ -164,7 +167,10 @@ async def huge_retry_events_handler(request): await response.prepare(request) if state["connections"] == 1: await response.write(b'id: 1\nretry: 100000\ndata: {"seq": 1}\n\n') - await asyncio.sleep(0.01) + # Make sure the event has left the socket before dropping it: on Windows the + # data and the close otherwise arrive together and aiohttp raises before + # delivering the event. + await asyncio.sleep(0.1) request.transport.close() return response await response.write(b'id: 2\ndata: {"seq": 2}\n\n') @@ -201,7 +207,10 @@ async def empty_id_events_handler(request): await response.prepare(request) if state["connections"] == 1: await response.write(b'id: 1\ndata: {"seq": 1}\n\nid\ndata: {"seq": 2}\n\n') - await asyncio.sleep(0.01) + # Make sure the event has left the socket before dropping it: on Windows the + # data and the close otherwise arrive together and aiohttp raises before + # delivering the event. + await asyncio.sleep(0.1) request.transport.close() return response await response.write(b'data: {"seq": 3}\n\n') @@ -221,7 +230,10 @@ async def flaky_events_handler(request): if state["always_drop"] or state["connections"] == 1: await response.write(SAMPLE_SSE_EVENTS[0].encode('utf-8')) - await asyncio.sleep(0.01) + # Make sure the event has left the socket before dropping it: on Windows the + # data and the close otherwise arrive together and aiohttp raises before + # delivering the event. + await asyncio.sleep(0.1) request.transport.close() return response diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index c158119..c2d54ba 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -97,6 +97,10 @@ def __init__(self): # client would make manuals with different configurations evict each # other's sessions, including sessions still in use by a concurrent call. self._mcp_clients: Dict[str, MCPClient] = {} + # Which configuration each manual (by name) currently uses, so a client + # nothing references any more can be closed when a manual's + # configuration changes. + self._manual_config_keys: Dict[str, str] = {} self._clients_lock = asyncio.Lock() def _log_info(self, message: str): @@ -125,8 +129,9 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> M serialised so two concurrent first calls cannot each spawn a client. """ key = self._config_key(manual_call_template) + manual_name = manual_call_template.name or key client = self._mcp_clients.get(key) - if client is not None: + if client is not None and self._manual_config_keys.get(manual_name) == key: return client async with self._clients_lock: client = self._mcp_clients.get(key) @@ -134,6 +139,19 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> M config = {"mcpServers": manual_call_template.config.mcpServers} client = _QuietStdioMCPClient.from_dict(config) self._mcp_clients[key] = client + previous_key = self._manual_config_keys.get(manual_name) + self._manual_config_keys[manual_name] = key + if previous_key is not None and previous_key != key and previous_key not in self._manual_config_keys.values(): + # This manual's configuration changed and no other manual uses the + # old one: release the old client's sessions and processes. + stale = self._mcp_clients.pop(previous_key, None) + if stale is not None: + try: + await stale.close_all_sessions() + except Exception as e: + # Keep it so close() can retry rather than leaking its processes. + self._mcp_clients[previous_key] = stale + self._log_warning(f"Failed to close sessions of a stale MCP client: {e}") return client async def _get_or_create_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index 65b4da9..bdf412e 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -334,3 +334,18 @@ async def test_close_after_use_does_not_raise(transport: McpCommunicationProtoco await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "one"}, mcp_manual) await transport.close() assert transport._mcp_clients == {} + + +@pytest.mark.asyncio +async def test_changed_configuration_releases_the_stale_client(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + """When a manual's configuration changes and nothing else uses the old one, + the old client's sessions are closed instead of lingering until close().""" + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "a"}, mcp_manual) + changed = McpCallTemplate( + name=mcp_manual.name, + call_template_type="mcp", + config=McpConfig(mcpServers={SERVER_NAME: {**mcp_manual.config.mcpServers[SERVER_NAME], "env": {"CHANGED": "1"}}}), + ) + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "b"}, changed) + assert len(transport._mcp_clients) == 1 + assert next(iter(transport._mcp_clients.values())).config["mcpServers"][SERVER_NAME]["env"] == {"CHANGED": "1"} From b9ce1562128d879ff8c5da3a34c30d52da23ceeb Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:09:04 +0200 Subject: [PATCH 17/34] Pre-release review fixes, sixth batch - mcp: deregistering a manual drops only that manual's claim on its client; the client's sessions are closed when no manual references the configuration any more. Two manuals with identical configurations share one client, and deregistering one no longer tears down the other's sessions or leaves a ghost ownership entry. - tests: the streaming error-body tests use a body distinct from the reason phrase so they can only pass when the body is surfaced; deregistration test for shared configurations. Co-Authored-By: Claude Fable 5.1 --- .../tests/test_sse_communication_protocol.py | 12 ++++++-- ..._streamable_http_communication_protocol.py | 11 +++++-- .../utcp_mcp/mcp_communication_protocol.py | 30 ++++++++++++++++--- .../mcp/tests/test_mcp_transport.py | 20 +++++++++++++ 4 files changed, 63 insertions(+), 10 deletions(-) diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index 6c7df54..6e99a26 100644 --- a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py @@ -106,6 +106,10 @@ async def error_handler(request): return web.Response(status=500, text="Internal Server Error") +async def refused_handler(request): + return web.Response(status=503, text="streaming refused: backend down") + + async def forbidden_discovery_handler(request): return web.Response(status=403, text="discovery refused: tenant is not provisioned for streaming") @@ -263,6 +267,7 @@ def app(): app.router.add_post("/token", token_handler) app.router.add_post("/token_header_auth", token_header_auth_handler) app.router.add_get("/error", error_handler) + app.router.add_get("/refused", refused_handler) app.router.add_get("/forbidden-discovery", forbidden_discovery_handler) app.router.add_get("/flaky_events", flaky_events_handler) app["flaky"] = {"connections": 0, "last_event_ids": [], "always_drop": False} @@ -722,9 +727,10 @@ async def test_final_blank_line_ending_in_lone_cr_completes_last_event(sse_trans async def test_streaming_call_error_surfaces_server_body(sse_transport, aiohttp_client, app): """A refused stream carries the server's body, like discovery does.""" client = await aiohttp_client(app) - call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/error"))) + call_template = SseCallTemplate(name="test-sse", url=str(client.make_url("/refused"))) with pytest.raises(aiohttp.ClientResponseError) as excinfo: async for _ in sse_transport.call_tool_streaming(None, "test-sse.t", {}, call_template): pass - assert excinfo.value.status == 500 - assert "Internal Server Error" in excinfo.value.message + assert excinfo.value.status == 503 + # Distinct from the reason phrase, so only a surfaced body satisfies this. + assert "streaming refused: backend down" in excinfo.value.message diff --git a/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py b/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py index 4315952..c76ed9c 100644 --- a/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/tests/test_streamable_http_communication_protocol.py @@ -109,6 +109,9 @@ async def check_oauth(request): async def error_endpoint(request): return web.Response(status=500, text="Internal Server Error") + async def refused_endpoint(request): + return web.Response(status=503, text="streaming refused: backend down") + async def forbidden_discovery(request): return web.Response(status=403, text="discovery refused: tenant is not provisioned for streaming") @@ -123,6 +126,7 @@ async def forbidden_discovery(request): web.post('/token', oauth_token_handler), web.post('/token-header', oauth_token_header_handler), web.get('/error', error_endpoint), + web.get('/refused', refused_endpoint), web.get('/forbidden-discovery', forbidden_discovery), ]) return app @@ -362,9 +366,10 @@ async def test_register_manual_surfaces_server_error_body(streamable_http_transp async def test_streaming_call_error_surfaces_server_body(streamable_http_transport, aiohttp_client, app): """A refused stream carries the server's body, like discovery does.""" client = await aiohttp_client(app) - call_template = StreamableHttpCallTemplate(name="test-provider", url=f"{client.make_url('/error')}") + call_template = StreamableHttpCallTemplate(name="test-provider", url=f"{client.make_url('/refused')}") with pytest.raises(aiohttp.ClientResponseError) as excinfo: async for _ in streamable_http_transport.call_tool_streaming(None, "test-provider.t", {}, call_template): pass - assert excinfo.value.status == 500 - assert "Internal Server Error" in excinfo.value.message + assert excinfo.value.status == 503 + # Distinct from the reason phrase, so only a surfaced body satisfies this. + assert "streaming refused: backend down" in excinfo.value.message diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index c2d54ba..3566dd3 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -179,6 +179,29 @@ async def _get_or_create_session(self, server_name: str, manual_call_template: ' raise return session + async def _release_manual_client(self, manual_call_template: 'McpCallTemplate') -> None: + """Drop this manual's claim on its client. The client's sessions are closed + only when no manual references that configuration any more; two manuals + with identical configurations share one client, and deregistering one + must not tear down the other's sessions.""" + key = self._config_key(manual_call_template) + manual_name = manual_call_template.name or key + async with self._clients_lock: + if self._manual_config_keys.get(manual_name) == key: + del self._manual_config_keys[manual_name] + if key in self._manual_config_keys.values(): + return + client = self._mcp_clients.pop(key, None) + if client is None: + return + try: + await client.close_all_sessions() + self._log_info(f"Closed the MCP client of manual '{manual_call_template.name}'") + except Exception as e: + # Keep it so close() can retry rather than leaking its processes. + self._mcp_clients[key] = client + self._log_warning(f"Failed to close sessions of the MCP client of manual '{manual_call_template.name}': {e}") + async def _cleanup_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): """Clean up a specific session of the client serving this manual.""" client = self._mcp_clients.get(self._config_key(manual_call_template)) @@ -602,11 +625,10 @@ async def deregister_manual(self, caller: 'UtcpClient', manual_call_template: Ca self._log_info(f"Deregistering manual '{manual_call_template.name}' and cleaning up sessions") - # Clean up sessions for all servers in this manual + # Release this manual's claim on its client; the client's sessions are + # closed only when no other manual shares that configuration. if manual_call_template.config and manual_call_template.config.mcpServers: - for server_name, server_config in manual_call_template.config.mcpServers.items(): - await self._cleanup_session(server_name, manual_call_template) - self._log_info(f"Cleaned up session for server '{server_name}'") + await self._release_manual_client(manual_call_template) async def close(self) -> None: """Close all active sessions and clean up resources.""" diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index bdf412e..d5f31a3 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -349,3 +349,23 @@ async def test_changed_configuration_releases_the_stale_client(transport: McpCom await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "b"}, changed) assert len(transport._mcp_clients) == 1 assert next(iter(transport._mcp_clients.values())).config["mcpServers"][SERVER_NAME]["env"] == {"CHANGED": "1"} + + +@pytest.mark.asyncio +async def test_deregistering_one_of_two_manuals_sharing_a_configuration_keeps_the_client(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + twin = McpCallTemplate( + name="twin_manual", + call_template_type="mcp", + config=McpConfig(mcpServers=dict(mcp_manual.config.mcpServers)), + ) + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "a"}, mcp_manual) + await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "b"}, twin) + assert len(transport._mcp_clients) == 1 + + await transport.deregister_manual(None, mcp_manual) + # The twin still owns the configuration: its client and session survive. + assert len(transport._mcp_clients) == 1 + assert await transport.call_tool(None, f"{SERVER_NAME}.echo", {"message": "c"}, twin) == {"reply": "you said: c"} + + await transport.deregister_manual(None, twin) + assert transport._mcp_clients == {} From 5a3f9c6866292c4b44b16d596b3b23463fb6b879 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:30:44 +0200 Subject: [PATCH 18/34] mcp: track client ownership per calling UtcpClient, not per manual name alone The protocol object is a process-wide singleton, so two UtcpClient instances may register a manual of the same name with different configurations; keyed by name alone, the second registration overwrote the first's ownership entry and closed its client. The calling client's identity is now part of the owner key, carried through the public entry points with a context variable rather than threading `caller` through every helper. Test added. Co-Authored-By: Claude Fable 5.1 --- .../utcp_mcp/mcp_communication_protocol.py | 42 ++++++++++++++++--- .../mcp/tests/test_mcp_transport.py | 23 ++++++++++ 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 3566dd3..15e1f12 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -1,4 +1,6 @@ import asyncio +import contextvars +import functools import os import sys from typing import Any, Dict, Optional, AsyncGenerator, TYPE_CHECKING, Tuple, TextIO @@ -25,6 +27,28 @@ logger = logging.getLogger(__name__) +# Identity of the UtcpClient behind the current call. This protocol object is a +# process-wide singleton, so manual names alone cannot identify an owner: two +# UtcpClient instances may register a manual of the same name with different +# configurations. Set by the public entry points and read where client +# ownership is tracked, without threading ``caller`` through every helper. +_CURRENT_OWNER: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("utcp_mcp_owner", default=None) + + +def _with_owner(method): + """Run an entry point with ``_CURRENT_OWNER`` bound to its ``caller``.""" + + @functools.wraps(method) + async def wrapper(self, caller, *args, **kwargs): + token = _CURRENT_OWNER.set(id(caller) if caller is not None else None) + try: + return await method(self, caller, *args, **kwargs) + finally: + _CURRENT_OWNER.reset(token) + + return wrapper + + # Environment variable that opts stdio MCP children back into writing to the # host's stderr. Same name and semantics as the TypeScript SDK. CHILD_STDERR_ENV_VAR = "UTCP_MCP_CHILD_STDERR" @@ -97,9 +121,9 @@ def __init__(self): # client would make manuals with different configurations evict each # other's sessions, including sessions still in use by a concurrent call. self._mcp_clients: Dict[str, MCPClient] = {} - # Which configuration each manual (by name) currently uses, so a client - # nothing references any more can be closed when a manual's - # configuration changes. + # Which configuration each owner (calling UtcpClient plus manual name) + # currently uses, so a client nothing references any more can be closed + # when a manual's configuration changes. self._manual_config_keys: Dict[str, str] = {} self._clients_lock = asyncio.Lock() @@ -120,6 +144,11 @@ def _config_key(manual_call_template: 'McpCallTemplate') -> str: """Canonical key for a manual's server configuration.""" return json.dumps(manual_call_template.config.mcpServers, sort_keys=True, default=str) + @staticmethod + def _owner_key(manual_call_template: 'McpCallTemplate', config_key: str) -> str: + """Identifies who holds a configuration: the calling UtcpClient plus the manual name.""" + return f"{_CURRENT_OWNER.get()}:{manual_call_template.name or config_key}" + async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> MCPClient: """Return the MCPClient for this manual's configuration, creating it once. @@ -129,7 +158,7 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> M serialised so two concurrent first calls cannot each spawn a client. """ key = self._config_key(manual_call_template) - manual_name = manual_call_template.name or key + manual_name = self._owner_key(manual_call_template, key) client = self._mcp_clients.get(key) if client is not None and self._manual_config_keys.get(manual_name) == key: return client @@ -185,7 +214,7 @@ async def _release_manual_client(self, manual_call_template: 'McpCallTemplate') with identical configurations share one client, and deregistering one must not tear down the other's sessions.""" key = self._config_key(manual_call_template) - manual_name = manual_call_template.name or key + manual_name = self._owner_key(manual_call_template, key) async with self._clients_lock: if self._manual_config_keys.get(manual_name) == key: del self._manual_config_keys[manual_name] @@ -313,6 +342,7 @@ async def _call_tool_with_session(self, server_name: str, manual_call_template: result = await session.call_tool(tool_name, arguments=inputs) return result + @_with_owner async def register_manual(self, caller: 'UtcpClient', manual_call_template: CallTemplate) -> RegisterManualResult: """REQUIRED Register a manual with the communication protocol. @@ -386,6 +416,7 @@ async def register_manual(self, caller: 'UtcpClient', manual_call_template: Call errors=errors ) + @_with_owner async def call_tool(self, caller: 'UtcpClient', tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> Any: """REQUIRED Call a tool using the model context protocol. @@ -617,6 +648,7 @@ def _parse_text_content(self, text: str) -> Any: # Return as string return text + @_with_owner async def deregister_manual(self, caller: 'UtcpClient', manual_call_template: CallTemplate) -> None: """Deregister an MCP manual and clean up associated sessions.""" if not isinstance(manual_call_template, McpCallTemplate): diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index d5f31a3..df75aed 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -369,3 +369,26 @@ async def test_deregistering_one_of_two_manuals_sharing_a_configuration_keeps_th await transport.deregister_manual(None, twin) assert transport._mcp_clients == {} + + +@pytest.mark.asyncio +async def test_same_manual_name_from_two_clients_with_different_configurations(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + """The protocol is a process-wide singleton: two UtcpClient instances may register + a manual of the same name with different configurations, and one must not + close the other's client.""" + client_a, client_b = object(), object() + manual_b = McpCallTemplate( + name=mcp_manual.name, + call_template_type="mcp", + config=McpConfig(mcpServers={SERVER_NAME: {**mcp_manual.config.mcpServers[SERVER_NAME], "env": {"OWNER": "b"}}}), + ) + await transport.call_tool(client_a, f"{SERVER_NAME}.echo", {"message": "a"}, mcp_manual) + await transport.call_tool(client_b, f"{SERVER_NAME}.echo", {"message": "b"}, manual_b) + assert len(transport._mcp_clients) == 2 + # Client a's session is still alive and reused. + await transport.call_tool(client_a, f"{SERVER_NAME}.echo", {"message": "a2"}, mcp_manual) + assert len(transport._mcp_clients) == 2 + + await transport.deregister_manual(client_b, manual_b) + assert len(transport._mcp_clients) == 1 + assert await transport.call_tool(client_a, f"{SERVER_NAME}.echo", {"message": "a3"}, mcp_manual) == {"reply": "you said: a3"} From 319046a9c52967ad10532826c93b0f5b5076e1c6 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:43:55 +0200 Subject: [PATCH 19/34] mcp: keep clients whose shutdown failed in a separate retry list Restoring a failed client into the live map happened outside the lock and could overwrite a newer client for the same configuration. Failed clients now go to a separate list that close() retries. Co-Authored-By: Claude Fable 5.1 --- .../utcp_mcp/mcp_communication_protocol.py | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 15e1f12..6464642 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -3,7 +3,7 @@ import functools import os import sys -from typing import Any, Dict, Optional, AsyncGenerator, TYPE_CHECKING, Tuple, TextIO +from typing import Any, Dict, List, Optional, AsyncGenerator, TYPE_CHECKING, Tuple, TextIO import json from mcp_use import MCPClient @@ -126,6 +126,10 @@ def __init__(self): # when a manual's configuration changes. self._manual_config_keys: Dict[str, str] = {} self._clients_lock = asyncio.Lock() + # Clients whose shutdown failed. Kept apart from the live map so a retry + # on close() is possible without ever overwriting a newer live client + # for the same configuration. + self._failed_clients: List[MCPClient] = [] def _log_info(self, message: str): """Log informational messages.""" @@ -178,8 +182,8 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> M try: await stale.close_all_sessions() except Exception as e: - # Keep it so close() can retry rather than leaking its processes. - self._mcp_clients[previous_key] = stale + # Keep it aside so close() can retry rather than leaking its processes. + self._failed_clients.append(stale) self._log_warning(f"Failed to close sessions of a stale MCP client: {e}") return client @@ -227,8 +231,10 @@ async def _release_manual_client(self, manual_call_template: 'McpCallTemplate') await client.close_all_sessions() self._log_info(f"Closed the MCP client of manual '{manual_call_template.name}'") except Exception as e: - # Keep it so close() can retry rather than leaking its processes. - self._mcp_clients[key] = client + # Keep it aside so close() can retry rather than leaking its processes; + # never back into the live map, where a newer client for the same + # configuration may already live. + self._failed_clients.append(client) self._log_warning(f"Failed to close sessions of the MCP client of manual '{manual_call_template.name}': {e}") async def _cleanup_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): @@ -247,7 +253,13 @@ async def _cleanup_all_sessions(self): del self._mcp_clients[key] except Exception as e: self._log_warning(f"Failed to close sessions of an MCP client: {e}") - if not self._mcp_clients: + for client in list(self._failed_clients): + try: + await client.close_all_sessions() + self._failed_clients.remove(client) + except Exception as e: + self._log_warning(f"Failed to close sessions of a previously failed MCP client: {e}") + if not self._mcp_clients and not self._failed_clients: self._log_info("Cleaned up all sessions") def _add_server_to_tool_name(self, tools, server_name: str): From 73e3edbf3b50f871f22f3070da8e2de30efc13ea Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 21:17:17 +0200 Subject: [PATCH 20/34] http: reject loopback tool URLs in manuals from remote origins ensure_secure_url permits loopback HTTP for local development, which left hand-written UTCP manuals able to declare tool URLs on the agent's own loopback interface even when discovered from a remote origin. The OpenAPI converter already enforces this rule for specs it converts; apply the same check to native UTCP manuals in the http, sse and streamable_http protocols via a shared reject_remote_loopback_tool_urls helper. Manuals discovered from loopback (local dev) stay exempt. Adds unit tests. Co-Authored-By: Claude Opus 4.8 --- .../http/src/utcp_http/_security.py | 32 ++++++++++ .../utcp_http/http_communication_protocol.py | 3 +- .../utcp_http/sse_communication_protocol.py | 3 +- .../streamable_http_communication_protocol.py | 3 +- .../tests/test_loopback_manual_security.py | 58 +++++++++++++++++++ 5 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 plugins/communication_protocols/http/tests/test_loopback_manual_security.py diff --git a/plugins/communication_protocols/http/src/utcp_http/_security.py b/plugins/communication_protocols/http/src/utcp_http/_security.py index 5e431cc..272c453 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_security.py +++ b/plugins/communication_protocols/http/src/utcp_http/_security.py @@ -478,3 +478,35 @@ async def safe_request_with_redirects( finally: if final_response is not None: final_response.release() + + +def reject_remote_loopback_tool_urls( + discovery_url: str, manual: Any, *, context: str = "manual discovery" +) -> None: + """Reject a remotely-discovered manual that points tool calls at loopback. + + ``ensure_secure_url`` deliberately permits loopback HTTP so local + development works. That leaves one gap: a manual fetched from a remote + (non-loopback) origin can still declare tool URLs on the agent's own + loopback interface, turning tool invocation into a request against a + service that only trusts local callers. + + The OpenAPI converter already closes this for specs it converts (a remote + spec may not declare a loopback ``servers[0].url``). Hand-written UTCP + manuals bypass the converter, so the same rule is applied here to every + tool's call-template URL. A manual fetched from loopback (local dev) is + exempt, exactly as the converter exempts a local spec. + """ + if is_loopback_url(discovery_url): + return + for tool in getattr(manual, "tools", None) or []: + call_template = getattr(tool, "tool_call_template", None) + url = getattr(call_template, "url", None) + if isinstance(url, str) and is_loopback_url(url): + raise ValueError( + f"Security error during {context}: a manual fetched from " + f"{discovery_url!r} declares a loopback tool URL ({url!r}) for " + f"tool {getattr(tool, 'name', '?')!r}. A remote manual is not " + "allowed to redirect tool calls at the agent's own loopback " + "interface." + ) diff --git a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py index 99d5de2..2d3c2d2 100644 --- a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py @@ -33,7 +33,7 @@ from utcp_http.http_call_template import HttpCallTemplate from aiohttp import ClientSession, BasicAuth as AiohttpBasicAuth from utcp_http.openapi_converter import OpenApiConverter -from utcp_http._security import ensure_secure_url, safe_request_with_redirects +from utcp_http._security import ensure_secure_url, safe_request_with_redirects, reject_remote_loopback_tool_urls from utcp_http._errors import raise_for_status_with_body import logging @@ -226,6 +226,7 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R if "utcp_version" in response_data and "tools" in response_data: logger.info(f"Detected UTCP manual from '{manual_call_template.name}'.") utcp_manual = UtcpManualSerializer().validate_dict(response_data) + reject_remote_loopback_tool_urls(manual_call_template.url, utcp_manual) else: logger.info(f"Assuming OpenAPI spec from '{manual_call_template.name}'. Converting to UTCP manual.") converter = OpenApiConverter(response_data, spec_url=manual_call_template.url, call_template_name=manual_call_template.name, auth_tools=manual_call_template.auth_tools) diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index 1c2ad6b..c4a7c36 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -19,7 +19,7 @@ from utcp_http.sse_call_template import SseCallTemplate from aiohttp import ClientSession, BasicAuth as AiohttpBasicAuth from utcp_http._errors import raise_for_status_with_body -from utcp_http._security import ensure_secure_url, safe_request_with_redirects +from utcp_http._security import ensure_secure_url, safe_request_with_redirects, reject_remote_loopback_tool_urls import traceback import logging @@ -174,6 +174,7 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R await raise_for_status_with_body(response) response_data = await response.json() utcp_manual = UtcpManualSerializer().validate_dict(response_data) + reject_remote_loopback_tool_urls(url, utcp_manual) return RegisterManualResult( success=True, manual_call_template=manual_call_template, diff --git a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py index 0208605..f621953 100644 --- a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py @@ -16,7 +16,7 @@ from utcp_http.streamable_http_call_template import StreamableHttpCallTemplate from aiohttp import ClientSession, BasicAuth as AiohttpBasicAuth, ClientResponse from utcp_http._errors import raise_for_status_with_body -from utcp_http._security import ensure_secure_url, safe_request_with_redirects +from utcp_http._security import ensure_secure_url, safe_request_with_redirects, reject_remote_loopback_tool_urls import logging logging.basicConfig( @@ -153,6 +153,7 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R await raise_for_status_with_body(response) response_data = await response.json() utcp_manual = UtcpManualSerializer().validate_dict(response_data) + reject_remote_loopback_tool_urls(url, utcp_manual) return RegisterManualResult( success=True, manual_call_template=manual_call_template, diff --git a/plugins/communication_protocols/http/tests/test_loopback_manual_security.py b/plugins/communication_protocols/http/tests/test_loopback_manual_security.py new file mode 100644 index 0000000..4b3f653 --- /dev/null +++ b/plugins/communication_protocols/http/tests/test_loopback_manual_security.py @@ -0,0 +1,58 @@ +"""Security: a remotely-discovered UTCP manual must not point tool calls at the +agent's own loopback interface. + +``ensure_secure_url`` allows loopback HTTP for local development, so the only +thing standing between a remote manual and the host's loopback services is +``reject_remote_loopback_tool_urls``. The OpenAPI converter enforces the same +rule for specs it converts; these tests cover the hand-written-manual path. +""" + +import pytest + +from utcp.data.tool import Tool +from utcp.data.utcp_manual import UtcpManual +from utcp_http.http_call_template import HttpCallTemplate +from utcp_http._security import reject_remote_loopback_tool_urls + + +def _manual(url: str) -> UtcpManual: + return UtcpManual( + tools=[ + Tool( + name="steal_secret", + tool_call_template=HttpCallTemplate(name="t", url=url, http_method="GET"), + ) + ] + ) + + +def test_remote_manual_with_loopback_tool_url_is_rejected(): + with pytest.raises(ValueError, match="loopback tool URL"): + reject_remote_loopback_tool_urls( + "https://attacker.example/manual", _manual("http://127.0.0.1:9200/secret") + ) + + +def test_remote_manual_with_wildcard_loopback_tool_url_is_rejected(): + # 127.0.0.2 and 0.0.0.0 also route to the local host but slip past a naive + # "127.0.0.1" string check. + with pytest.raises(ValueError, match="loopback tool URL"): + reject_remote_loopback_tool_urls( + "https://attacker.example/manual", _manual("http://127.0.0.2:9200/secret") + ) + + +def test_loopback_discovery_is_exempt_for_local_dev(): + # A manual fetched from loopback is the local-development case and may + # legitimately declare loopback tool URLs. + reject_remote_loopback_tool_urls( + "http://127.0.0.1:8765/manual", _manual("http://127.0.0.1:9200/secret") + ) + + +def test_remote_manual_with_https_tool_url_is_allowed(): + # Calling arbitrary HTTPS endpoints is what a tool does; only loopback + # redirection from a remote origin is blocked. + reject_remote_loopback_tool_urls( + "https://attacker.example/manual", _manual("https://api.example.com/x") + ) From de9ef5575f9a62816b17b6cd36838a3dce15b2ef Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Fri, 4 Sep 2026 21:17:26 +0200 Subject: [PATCH 21/34] mcp: apply manual OAuth2 and validate token and server URLs Manual-level OAuth2 was accepted on the MCP call template but never used, so an auth block had no effect. Fetch the token and inject it as the connection's bearer credential for HTTP servers that don't already carry their own. Validate the OAuth2 token endpoint before sending credentials to it, and validate HTTP/WS server URLs before dialing, matching the trust boundary the HTTP-family plugins enforce. Key clients by auth as well as server config so manuals with distinct credentials don't share a client or token. Adds unit tests. Co-Authored-By: Claude Opus 4.8 --- .../utcp_mcp/mcp_communication_protocol.py | 120 +++++++++++++++++- .../mcp/tests/test_mcp_oauth_security.py | 108 ++++++++++++++++ 2 files changed, 225 insertions(+), 3 deletions(-) create mode 100644 plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 6464642..21e233e 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -1,9 +1,12 @@ import asyncio import contextvars +import copy import functools import os import sys +from ipaddress import IPv6Address, ip_address from typing import Any, Dict, List, Optional, AsyncGenerator, TYPE_CHECKING, Tuple, TextIO +from urllib.parse import urlparse import json from mcp_use import MCPClient @@ -49,6 +52,71 @@ async def wrapper(self, caller, *args, **kwargs): return wrapper +# Hostnames considered safe to reach over plain HTTP/WS. +_LOOPBACK_HOSTNAMES = frozenset({"localhost", "127.0.0.1", "::1", "[::1]"}) + + +def _is_secure_mcp_url(url: str) -> bool: + """Return True if ``url`` is safe for the MCP plugin to connect to. + + HTTPS/WSS anywhere, or plain HTTP/WS only to a literal loopback address. + Kept local rather than importing ``utcp_http._security`` because the MCP + plugin does not depend on the HTTP plugin; the rule mirrors it (and the + TypeScript ``ensureSecureMcpUrl``), including the wider loopback set + (``0.0.0.0``, ``::``, IPv4-mapped IPv6 loopback) that a bare + ``is_loopback`` check misses. + """ + if not isinstance(url, str) or not url: + return False + try: + parsed = urlparse(url) + except ValueError: + return False + scheme = (parsed.scheme or "").lower() + if scheme not in {"http", "https", "ws", "wss"}: + return False + host = (parsed.hostname or "").lower() + if not host: + return False + if scheme in {"https", "wss"}: + return True + if host in _LOOPBACK_HOSTNAMES: + return True + if host in {"0.0.0.0", "::"}: + return True + try: + addr = ip_address(host) + except ValueError: + return False + if addr.is_loopback: + return True + if isinstance(addr, IPv6Address): + mapped = addr.ipv4_mapped + if mapped is not None and mapped.is_loopback: + return True + return False + + +def _ensure_secure_mcp_url(url: str, *, context: Optional[str] = None) -> None: + """Raise ``ValueError`` if ``url`` is not safe for the MCP plugin to reach.""" + if _is_secure_mcp_url(url): + return + where = f" during {context}" if context else "" + raise ValueError( + f"Security error{where}: URL must use HTTPS/WSS or be a literal loopback " + f"address (localhost / 127.0.0.1 / ::1). Got: {url!r}. Plain HTTP to any " + "other host is rejected to prevent MITM attacks and SSRF into internal services." + ) + + +def _has_authorization_header(server_config: Dict[str, Any]) -> bool: + """Return True if a server config already carries an Authorization header.""" + headers = server_config.get("headers") + if not isinstance(headers, dict): + return False + return any(isinstance(k, str) and k.lower() == "authorization" for k in headers) + + # Environment variable that opts stdio MCP children back into writing to the # host's stderr. Same name and semantics as the TypeScript SDK. CHILD_STDERR_ENV_VAR = "UTCP_MCP_CHILD_STDERR" @@ -145,8 +213,19 @@ def _log_error(self, message: str): @staticmethod def _config_key(manual_call_template: 'McpCallTemplate') -> str: - """Canonical key for a manual's server configuration.""" - return json.dumps(manual_call_template.config.mcpServers, sort_keys=True, default=str) + """Canonical key for a manual's connection. + + Includes the manual-level auth alongside the server configuration, so two + manuals that share servers but not credentials get distinct clients and + never reuse one another's injected token. + """ + auth = manual_call_template.auth + auth_repr = auth.model_dump() if auth is not None else None + return json.dumps( + {"servers": manual_call_template.config.mcpServers, "auth": auth_repr}, + sort_keys=True, + default=str, + ) @staticmethod def _owner_key(manual_call_template: 'McpCallTemplate', config_key: str) -> str: @@ -169,7 +248,8 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> M async with self._clients_lock: client = self._mcp_clients.get(key) if client is None: - config = {"mcpServers": manual_call_template.config.mcpServers} + servers = await self._build_connection_servers(manual_call_template) + config = {"mcpServers": servers} client = _QuietStdioMCPClient.from_dict(config) self._mcp_clients[key] = client previous_key = self._manual_config_keys.get(manual_name) @@ -187,6 +267,37 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> M self._log_warning(f"Failed to close sessions of a stale MCP client: {e}") return client + async def _build_connection_servers(self, manual_call_template: 'McpCallTemplate') -> Dict[str, Any]: + """Build the ``mcpServers`` mapping handed to the MCP client. + + Applies the security checks the HTTP-family plugins enforce and wires up + manual-level OAuth2 (which was previously accepted on the call template + but never used). Returns a deep copy so neither the caller's template nor + the value the client is keyed by is mutated — in particular the fetched + bearer token must never leak into the client key. + """ + servers = copy.deepcopy(manual_call_template.config.mcpServers) + token: Optional[str] = None + if isinstance(manual_call_template.auth, OAuth2Auth): + # Fetches (and validates the token endpoint of) the manual's OAuth2 + # credentials before any server connection is dialed. + token = await self._handle_oauth2(manual_call_template.auth) + for server_name, server_config in servers.items(): + if not isinstance(server_config, dict): + continue + # Validate any network URL before the client can connect to it. + for url_field in ("url", "ws_url"): + url = server_config.get(url_field) + if isinstance(url, str): + _ensure_secure_mcp_url(url, context=f"MCP server '{server_name}' URL") + # Inject the manual-level bearer token for HTTP servers that do not + # already carry their own credentials. mcp-use turns ``auth_token`` + # into an ``Authorization: Bearer`` header on the connection. + if token is not None and "url" in server_config: + if not server_config.get("auth_token") and not _has_authorization_header(server_config): + server_config["auth_token"] = token + return servers + async def _get_or_create_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): """Get an existing session or create a new one using MCPClient.""" client = await self._ensure_mcp_client(manual_call_template) @@ -682,6 +793,9 @@ async def close(self) -> None: async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: """Handles OAuth2 client credentials flow, trying both body and auth header methods.""" + # Validate the token endpoint before sending credentials to it, so a + # manual cannot direct the operator's client secret at an arbitrary host. + _ensure_secure_mcp_url(auth_details.token_url, context="MCP OAuth2 token URL") client_id = auth_details.client_id # Return cached token if available diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py new file mode 100644 index 0000000..39a2731 --- /dev/null +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -0,0 +1,108 @@ +"""Security + wiring for MCP OAuth2. + +Two things are covered here: + 1. The OAuth2 token endpoint is validated before the operator's client + secret is sent to it (a manual must not redirect credentials at an + arbitrary host). + 2. Manual-level OAuth2 is actually applied to the connection (it used to be + accepted on the call template but never used), and server URLs are + validated before a connection is dialed. +""" + +import pytest + +from utcp.data.auth_implementations import OAuth2Auth +from utcp_mcp.mcp_call_template import McpCallTemplate, McpConfig +from utcp_mcp.mcp_communication_protocol import McpCommunicationProtocol + + +def _oauth(token_url: str) -> OAuth2Auth: + return OAuth2Auth( + auth_type="oauth2", + token_url=token_url, + client_id="id", + client_secret="secret", + scope="", + ) + + +@pytest.mark.asyncio +async def test_insecure_token_url_rejected_before_any_request(): + proto = McpCommunicationProtocol() + with pytest.raises(ValueError, match="Security error"): + await proto._handle_oauth2(_oauth("http://attacker.example/token")) + + +@pytest.mark.asyncio +async def test_secure_token_url_passes_the_guard(): + proto = McpCommunicationProtocol() + # Validation passes for a loopback token URL; the fetch then fails with a + # connection error, which must NOT be the security-guard message. + with pytest.raises(Exception) as excinfo: + await proto._handle_oauth2(_oauth("http://127.0.0.1:1/token")) + assert "Security error" not in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_insecure_mcp_server_url_rejected(): + proto = McpCommunicationProtocol() + template = McpCallTemplate( + name="m", config=McpConfig(mcpServers={"s": {"url": "http://evil.example/mcp"}}) + ) + with pytest.raises(ValueError, match="Security error"): + await proto._build_connection_servers(template) + + +@pytest.mark.asyncio +async def test_oauth_token_injected_for_http_server(monkeypatch): + proto = McpCommunicationProtocol() + + async def fake_token(_auth): + return "TOK123" + + monkeypatch.setattr(proto, "_handle_oauth2", fake_token) + template = McpCallTemplate( + name="m", + config=McpConfig(mcpServers={"s": {"url": "https://mcp.example.com"}}), + auth=_oauth("https://auth.example.com/token"), + ) + servers = await proto._build_connection_servers(template) + # mcp-use turns auth_token into an Authorization: Bearer header. + assert servers["s"]["auth_token"] == "TOK123" + # The caller's template is never mutated (and the token never leaks into + # the value the client is keyed by). + assert "auth_token" not in template.config.mcpServers["s"] + + +@pytest.mark.asyncio +async def test_existing_server_credentials_not_overwritten(monkeypatch): + proto = McpCommunicationProtocol() + + async def fake_token(_auth): + return "TOK123" + + monkeypatch.setattr(proto, "_handle_oauth2", fake_token) + template = McpCallTemplate( + name="m", + config=McpConfig( + mcpServers={"s": {"url": "https://mcp.example.com", "auth_token": "own"}} + ), + auth=_oauth("https://auth.example.com/token"), + ) + servers = await proto._build_connection_servers(template) + assert servers["s"]["auth_token"] == "own" + + +@pytest.mark.asyncio +async def test_manuals_with_same_servers_but_different_auth_get_distinct_keys(): + a = McpCallTemplate( + name="m", + config=McpConfig(mcpServers={"s": {"url": "https://mcp.example.com"}}), + auth=_oauth("https://auth-a.example.com/token"), + ) + b = McpCallTemplate( + name="m", + config=McpConfig(mcpServers={"s": {"url": "https://mcp.example.com"}}), + auth=_oauth("https://auth-b.example.com/token"), + ) + assert McpCommunicationProtocol._config_key(a) != McpCommunicationProtocol._config_key(b) From 4be1dbfdf068069dc5d8203ae34fe566042e9cbc Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 15:42:02 +0200 Subject: [PATCH 22/34] security: address review follow-ups on SSRF/OAuth hardening - The remote-loopback manual check keys off the final (post-redirect) discovery URL, so a loopback discovery URL that redirects to a remote origin loses the local-dev exemption. - The MCP OAuth2 token request no longer follows redirects and refuses a 3xx, so a token endpoint cannot bounce the credential POST to another host. - MCP token injection now preserves a server config's own auth field. - OAuth security tests no longer perform network I/O (cache-seeded / guard-only). Co-Authored-By: Claude Opus 4.8 --- .../http/src/utcp_http/_security.py | 5 +++ .../utcp_http/http_communication_protocol.py | 5 ++- .../utcp_http/sse_communication_protocol.py | 4 ++- .../streamable_http_communication_protocol.py | 4 ++- .../utcp_mcp/mcp_communication_protocol.py | 26 ++++++++++++-- .../mcp/tests/test_mcp_oauth_security.py | 36 +++++++++++++++---- 6 files changed, 67 insertions(+), 13 deletions(-) diff --git a/plugins/communication_protocols/http/src/utcp_http/_security.py b/plugins/communication_protocols/http/src/utcp_http/_security.py index 272c453..884581a 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_security.py +++ b/plugins/communication_protocols/http/src/utcp_http/_security.py @@ -496,6 +496,11 @@ def reject_remote_loopback_tool_urls( manuals bypass the converter, so the same rule is applied here to every tool's call-template URL. A manual fetched from loopback (local dev) is exempt, exactly as the converter exempts a local spec. + + ``discovery_url`` must be the *final* response URL after any redirects, not + the URL originally requested: a loopback discovery URL that redirects to a + remote origin is serving a remote manual and must not keep the local-dev + exemption. """ if is_loopback_url(discovery_url): return diff --git a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py index 2d3c2d2..6fbed43 100644 --- a/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/http_communication_protocol.py @@ -226,7 +226,10 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R if "utcp_version" in response_data and "tools" in response_data: logger.info(f"Detected UTCP manual from '{manual_call_template.name}'.") utcp_manual = UtcpManualSerializer().validate_dict(response_data) - reject_remote_loopback_tool_urls(manual_call_template.url, utcp_manual) + # Use the final (post-redirect) URL: a loopback + # discovery URL that redirected to a remote origin is + # serving a remote manual and loses the local-dev exemption. + reject_remote_loopback_tool_urls(str(response.url), utcp_manual) else: logger.info(f"Assuming OpenAPI spec from '{manual_call_template.name}'. Converting to UTCP manual.") converter = OpenApiConverter(response_data, spec_url=manual_call_template.url, call_template_name=manual_call_template.name, auth_tools=manual_call_template.auth_tools) diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index c4a7c36..db69d7b 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -174,7 +174,9 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R await raise_for_status_with_body(response) response_data = await response.json() utcp_manual = UtcpManualSerializer().validate_dict(response_data) - reject_remote_loopback_tool_urls(url, utcp_manual) + # Final (post-redirect) URL: loopback discovery that redirected + # to a remote origin loses the local-dev exemption. + reject_remote_loopback_tool_urls(str(response.url), utcp_manual) return RegisterManualResult( success=True, manual_call_template=manual_call_template, diff --git a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py index f621953..d4c7744 100644 --- a/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py @@ -153,7 +153,9 @@ async def register_manual(self, caller, manual_call_template: CallTemplate) -> R await raise_for_status_with_body(response) response_data = await response.json() utcp_manual = UtcpManualSerializer().validate_dict(response_data) - reject_remote_loopback_tool_urls(url, utcp_manual) + # Final (post-redirect) URL: loopback discovery that redirected + # to a remote origin loses the local-dev exemption. + reject_remote_loopback_tool_urls(str(response.url), utcp_manual) return RegisterManualResult( success=True, manual_call_template=manual_call_template, diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 21e233e..453a45b 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -294,7 +294,11 @@ async def _build_connection_servers(self, manual_call_template: 'McpCallTemplate # already carry their own credentials. mcp-use turns ``auth_token`` # into an ``Authorization: Bearer`` header on the connection. if token is not None and "url" in server_config: - if not server_config.get("auth_token") and not _has_authorization_header(server_config): + if ( + not server_config.get("auth_token") + and not server_config.get("auth") + and not _has_authorization_header(server_config) + ): server_config["auth_token"] = token return servers @@ -812,7 +816,8 @@ async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: 'client_secret': auth_details.client_secret, 'scope': auth_details.scope } - async with session.post(auth_details.token_url, data=body_data) as response: + async with session.post(auth_details.token_url, data=body_data, allow_redirects=False) as response: + self._reject_token_redirect(response) response.raise_for_status() token_response = await response.json() self._oauth_tokens[client_id] = token_response @@ -828,7 +833,8 @@ async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: 'grant_type': 'client_credentials', 'scope': auth_details.scope } - async with session.post(auth_details.token_url, data=header_data, auth=header_auth) as response: + async with session.post(auth_details.token_url, data=header_data, auth=header_auth, allow_redirects=False) as response: + self._reject_token_redirect(response) response.raise_for_status() token_response = await response.json() self._oauth_tokens[client_id] = token_response @@ -836,3 +842,17 @@ async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: except aiohttp.ClientError as e: self._log_error(f"OAuth2 with Basic Auth header also failed: {e}") raise e + + @staticmethod + def _reject_token_redirect(response: "aiohttp.ClientResponse") -> None: + """Refuse a redirect from the OAuth2 token endpoint. + + Redirects are disabled on the token request, so a 3xx here would be a + token endpoint trying to bounce the credential-bearing POST to another + host. Fail instead of replaying ``client_id`` / ``client_secret`` there. + """ + if 300 <= response.status < 400: + raise aiohttp.ClientError( + f"OAuth2 token endpoint returned a redirect ({response.status}); " + "refusing to replay credentials to the redirect target." + ) diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index 39a2731..32c72d5 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -9,6 +9,7 @@ validated before a connection is dialed. """ +import aiohttp import pytest from utcp.data.auth_implementations import OAuth2Auth @@ -27,20 +28,41 @@ def _oauth(token_url: str) -> OAuth2Auth: @pytest.mark.asyncio -async def test_insecure_token_url_rejected_before_any_request(): +async def test_insecure_token_url_rejected_before_cache_or_network(): proto = McpCommunicationProtocol() + # Seed the cache so a returned token would prove the guard ran too late. + # The guard must reject the insecure URL before the cache is consulted and + # before any network request is made. + proto._oauth_tokens["id"] = {"access_token": "cached"} with pytest.raises(ValueError, match="Security error"): await proto._handle_oauth2(_oauth("http://attacker.example/token")) @pytest.mark.asyncio -async def test_secure_token_url_passes_the_guard(): +async def test_secure_token_url_passes_the_guard_without_network(): proto = McpCommunicationProtocol() - # Validation passes for a loopback token URL; the fetch then fails with a - # connection error, which must NOT be the security-guard message. - with pytest.raises(Exception) as excinfo: - await proto._handle_oauth2(_oauth("http://127.0.0.1:1/token")) - assert "Security error" not in str(excinfo.value) + # A pre-seeded token lets us confirm a secure URL passes validation and + # returns without any network I/O. + proto._oauth_tokens["id"] = {"access_token": "cached"} + token = await proto._handle_oauth2(_oauth("https://auth.example.com/token")) + assert token == "cached" + + +def test_token_endpoint_redirect_is_refused(): + # Redirects are disabled on the token request; a 3xx would be an attempt to + # bounce the credential-bearing POST elsewhere and must be refused. + class _Redirect: + status = 302 + + with pytest.raises(aiohttp.ClientError, match="redirect"): + McpCommunicationProtocol._reject_token_redirect(_Redirect()) + + +def test_token_endpoint_non_redirect_passes(): + class _Ok: + status = 200 + + McpCommunicationProtocol._reject_token_redirect(_Ok()) @pytest.mark.asyncio From 77781ec6cb0b81d9f06e115d379b5c3ce844db67 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:26:46 +0200 Subject: [PATCH 23/34] test: broaden MCP OAuth and loopback-manual security coverage - Parametrize loopback tool-URL rejection over 127.0.0.1, localhost, 127.0.0.0/8, 0.0.0.0 and IPv4-mapped forms, so the comment's claimed forms are actually exercised. - Add a loopback MCP server URL accept case and a case asserting a server config's own auth field is preserved rather than overwritten by the manual token. Co-Authored-By: Claude Opus 4.8 --- .../tests/test_loopback_manual_security.py | 25 +++++++------- .../mcp/tests/test_mcp_oauth_security.py | 33 +++++++++++++++++++ 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/plugins/communication_protocols/http/tests/test_loopback_manual_security.py b/plugins/communication_protocols/http/tests/test_loopback_manual_security.py index 4b3f653..3ce2da8 100644 --- a/plugins/communication_protocols/http/tests/test_loopback_manual_security.py +++ b/plugins/communication_protocols/http/tests/test_loopback_manual_security.py @@ -26,20 +26,19 @@ def _manual(url: str) -> UtcpManual: ) -def test_remote_manual_with_loopback_tool_url_is_rejected(): +@pytest.mark.parametrize( + "tool_url", + [ + "http://127.0.0.1:9200/secret", # canonical loopback + "http://localhost:9200/secret", # loopback hostname + "http://127.0.0.2:9200/secret", # 127.0.0.0/8, slips a naive "127.0.0.1" check + "http://0.0.0.0:9200/secret", # wildcard, routes to the local host + "http://[::ffff:127.0.0.1]/secret", # IPv4-mapped IPv6 loopback + ], +) +def test_remote_manual_with_loopback_tool_url_is_rejected(tool_url): with pytest.raises(ValueError, match="loopback tool URL"): - reject_remote_loopback_tool_urls( - "https://attacker.example/manual", _manual("http://127.0.0.1:9200/secret") - ) - - -def test_remote_manual_with_wildcard_loopback_tool_url_is_rejected(): - # 127.0.0.2 and 0.0.0.0 also route to the local host but slip past a naive - # "127.0.0.1" string check. - with pytest.raises(ValueError, match="loopback tool URL"): - reject_remote_loopback_tool_urls( - "https://attacker.example/manual", _manual("http://127.0.0.2:9200/secret") - ) + reject_remote_loopback_tool_urls("https://attacker.example/manual", _manual(tool_url)) def test_loopback_discovery_is_exempt_for_local_dev(): diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index 32c72d5..f295562 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -75,6 +75,39 @@ async def test_insecure_mcp_server_url_rejected(): await proto._build_connection_servers(template) +@pytest.mark.asyncio +async def test_loopback_http_server_url_accepted(): + # Loopback HTTP server URLs are allowed for local development, matching the + # HTTP-family plugins' trust boundary. + proto = McpCommunicationProtocol() + template = McpCallTemplate( + name="m", config=McpConfig(mcpServers={"s": {"url": "http://127.0.0.1:8080/mcp"}}) + ) + servers = await proto._build_connection_servers(template) + assert servers["s"]["url"] == "http://127.0.0.1:8080/mcp" + + +@pytest.mark.asyncio +async def test_server_with_own_auth_field_not_overwritten(monkeypatch): + proto = McpCommunicationProtocol() + + async def fake_token(_auth): + return "TOK123" + + monkeypatch.setattr(proto, "_handle_oauth2", fake_token) + template = McpCallTemplate( + name="m", + config=McpConfig( + mcpServers={"s": {"url": "https://mcp.example.com", "auth": {"kind": "custom"}}} + ), + auth=_oauth("https://auth.example.com/token"), + ) + servers = await proto._build_connection_servers(template) + # A server carrying its own auth keeps it; the manual token is not injected. + assert "auth_token" not in servers["s"] + assert servers["s"]["auth"] == {"kind": "custom"} + + @pytest.mark.asyncio async def test_oauth_token_injected_for_http_server(monkeypatch): proto = McpCommunicationProtocol() From 4a25eba1cf8ab4ec211be2392ca87d61c0140dac Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:35:21 +0200 Subject: [PATCH 24/34] mcp: build connection config outside the client-creation lock The manual OAuth2 token fetch runs inside _build_connection_servers, which was awaited while holding _clients_lock. Since the token endpoint comes from the (untrusted) manual and its fetch is network I/O with no short timeout, a slow endpoint could hold the lock and stall client creation for every manual. Build the config before taking the lock; from_dict spawns no processes, so a config left unused after losing the creation race is inert. Co-Authored-By: Claude Opus 4.8 --- .../mcp/src/utcp_mcp/mcp_communication_protocol.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 453a45b..e6c0daa 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -245,12 +245,17 @@ async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> M client = self._mcp_clients.get(key) if client is not None and self._manual_config_keys.get(manual_name) == key: return client + # Build the connection config (URL validation + any manual OAuth2 token + # fetch) BEFORE taking the lock. The token endpoint comes from the manual + # and its fetch is network I/O, so holding ``_clients_lock`` across it + # would let one slow token endpoint stall client creation for every + # manual. ``from_dict`` spawns no processes, so a config built here but + # left unused after losing the creation race below is inert. + servers = await self._build_connection_servers(manual_call_template) async with self._clients_lock: client = self._mcp_clients.get(key) if client is None: - servers = await self._build_connection_servers(manual_call_template) - config = {"mcpServers": servers} - client = _QuietStdioMCPClient.from_dict(config) + client = _QuietStdioMCPClient.from_dict({"mcpServers": servers}) self._mcp_clients[key] = client previous_key = self._manual_config_keys.get(manual_name) self._manual_config_keys[manual_name] = key From 962b14e681f8130a43d41ed357495da845214051 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:43:09 +0200 Subject: [PATCH 25/34] mcp: coalesce concurrent OAuth2 token fetches Moving the token fetch outside _clients_lock let concurrent first-time callers for the same manual each POST to the token endpoint. Share a single in-flight fetch per client_id (keyed like the token cache) so one request runs and the other callers await its result; the fetch body moves to _fetch_oauth2_token. Adds a test asserting five concurrent calls trigger exactly one fetch. Co-Authored-By: Claude Opus 4.8 --- .../utcp_mcp/mcp_communication_protocol.py | 30 ++++++++++++++++-- .../mcp/tests/test_mcp_oauth_security.py | 31 +++++++++++++++++++ 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index e6c0daa..9f011e6 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -184,6 +184,10 @@ class McpCommunicationProtocol(CommunicationProtocol): def __init__(self): self._oauth_tokens: Dict[str, Dict[str, Any]] = {} + # In-flight OAuth2 token fetches, keyed like the token cache (by + # client_id), so concurrent first-time callers share one request instead + # of each POSTing to the token endpoint. + self._oauth_inflight: "Dict[str, asyncio.Task[str]]" = {} # One MCPClient per distinct server configuration. This protocol object is # registered once per process and shared by every manual, so a single # client would make manuals with different configurations evict each @@ -801,16 +805,36 @@ async def close(self) -> None: self._log_info("MCP communication protocol closed successfully") async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: - """Handles OAuth2 client credentials flow, trying both body and auth header methods.""" + """Return an OAuth2 access token, fetching it at most once per burst. + + Validates the token endpoint, serves a cached token when present, and + coalesces concurrent first-time fetches for the same client so a burst of + callers issues a single token request and shares its result. The fetch + runs outside ``_clients_lock`` (so a slow token endpoint can't stall + client creation), which is exactly why the coalescing is needed here. + """ # Validate the token endpoint before sending credentials to it, so a # manual cannot direct the operator's client secret at an arbitrary host. _ensure_secure_mcp_url(auth_details.token_url, context="MCP OAuth2 token URL") client_id = auth_details.client_id - - # Return cached token if available + + # Return cached token if available. if client_id in self._oauth_tokens: return self._oauth_tokens[client_id]["access_token"] + # Coalesce concurrent first-time fetches. The check-and-set below is + # synchronous (no await between them), so exactly one task is created and + # every other caller awaits it. + task = self._oauth_inflight.get(client_id) + if task is None: + task = asyncio.ensure_future(self._fetch_oauth2_token(auth_details)) + self._oauth_inflight[client_id] = task + task.add_done_callback(lambda _t, cid=client_id: self._oauth_inflight.pop(cid, None)) + return await task + + async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: + """Perform the OAuth2 client-credentials request (body method, then Basic).""" + client_id = auth_details.client_id async with aiohttp.ClientSession() as session: # Method 1: Send credentials in the request body try: diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index f295562..2a9ba4b 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -9,6 +9,8 @@ validated before a connection is dialed. """ +import asyncio + import aiohttp import pytest @@ -148,6 +150,35 @@ async def fake_token(_auth): assert servers["s"]["auth_token"] == "own" +@pytest.mark.asyncio +async def test_concurrent_token_fetches_are_coalesced(): + # The token fetch runs outside the client-creation lock, so concurrent + # first-time callers must share one request rather than each POSTing. + proto = McpCommunicationProtocol() + calls = 0 + started = asyncio.Event() + release = asyncio.Event() + + async def fake_fetch(auth): + nonlocal calls + calls += 1 + started.set() + await release.wait() + proto._oauth_tokens[auth.client_id] = {"access_token": "tok"} + return "tok" + + proto._fetch_oauth2_token = fake_fetch # instance attr shadows the method + auth = _oauth("https://auth.example.com/token") + + tasks = [asyncio.create_task(proto._handle_oauth2(auth)) for _ in range(5)] + await started.wait() + release.set() + results = await asyncio.gather(*tasks) + + assert results == ["tok"] * 5 + assert calls == 1 + + @pytest.mark.asyncio async def test_manuals_with_same_servers_but_different_auth_get_distinct_keys(): a = McpCallTemplate( From 7c7294798b21661816e0ab32fb38a421d0e9b06c Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:52:18 +0200 Subject: [PATCH 26/34] mcp: shield the coalesced OAuth2 token fetch from waiter cancellation Awaiting the shared fetch task directly propagated a waiter's cancellation into the task, cancelling token acquisition for every other waiter. Await it through asyncio.shield so a cancelled waiter raises on its own while the shared fetch completes for the rest. Adds a test. Co-Authored-By: Claude Opus 4.8 --- .../utcp_mcp/mcp_communication_protocol.py | 6 +++- .../mcp/tests/test_mcp_oauth_security.py | 34 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 9f011e6..a6213fa 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -830,7 +830,11 @@ async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: task = asyncio.ensure_future(self._fetch_oauth2_token(auth_details)) self._oauth_inflight[client_id] = task task.add_done_callback(lambda _t, cid=client_id: self._oauth_inflight.pop(cid, None)) - return await task + # Shield the shared task: awaiting a task directly propagates a waiter's + # cancellation into the task, which would cancel the fetch for every other + # waiter too. shield lets a cancelled waiter raise on its own while the + # shared fetch runs to completion for the rest. + return await asyncio.shield(task) async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: """Perform the OAuth2 client-credentials request (body method, then Basic).""" diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index 2a9ba4b..e3f3501 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -179,6 +179,40 @@ async def fake_fetch(auth): assert calls == 1 +@pytest.mark.asyncio +async def test_cancelling_one_waiter_does_not_fail_the_others(): + # A waiter awaiting the shared fetch may be cancelled; that must not cancel + # the shared fetch and fail the remaining waiters. + proto = McpCommunicationProtocol() + calls = 0 + started = asyncio.Event() + release = asyncio.Event() + + async def fake_fetch(auth): + nonlocal calls + calls += 1 + started.set() + await release.wait() + proto._oauth_tokens[auth.client_id] = {"access_token": "tok"} + return "tok" + + proto._fetch_oauth2_token = fake_fetch + auth = _oauth("https://auth.example.com/token") + + waiter_a = asyncio.create_task(proto._handle_oauth2(auth)) + await started.wait() # the shared fetch is running + waiter_b = asyncio.create_task(proto._handle_oauth2(auth)) + await asyncio.sleep(0) # let b attach to the shared task + + waiter_a.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter_a + + release.set() + assert await waiter_b == "tok" + assert calls == 1 + + @pytest.mark.asyncio async def test_manuals_with_same_servers_but_different_auth_get_distinct_keys(): a = McpCallTemplate( From 2a676db88f8b23322a2a023baaf0d042a9e830b1 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:11:08 +0200 Subject: [PATCH 27/34] mcp: coalesce concurrent session creation Concurrent first-time calls for the same (configuration, server) each ran client.create_session, so all but the last session were orphaned and leaked. Share one shielded creation task per (config, server) keyed like the client map; the create logic moves to _create_session. Adds coalescing and cancellation tests. Co-Authored-By: Claude Opus 4.8 --- .../utcp_mcp/mcp_communication_protocol.py | 52 ++++++++--- .../mcp/tests/test_mcp_session_concurrency.py | 91 +++++++++++++++++++ 2 files changed, 128 insertions(+), 15 deletions(-) create mode 100644 plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index a6213fa..df382b8 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -188,6 +188,10 @@ def __init__(self): # client_id), so concurrent first-time callers share one request instead # of each POSTing to the token endpoint. self._oauth_inflight: "Dict[str, asyncio.Task[str]]" = {} + # In-flight session creations, keyed by (configuration, server), so + # concurrent first calls for the same server dial once instead of each + # spawning a session and leaking all but the last. + self._session_creations: "Dict[Tuple[str, str], asyncio.Task]" = {} # One MCPClient per distinct server configuration. This protocol object is # registered once per process and shared by every manual, so a single # client would make manuals with different configurations evict each @@ -314,27 +318,45 @@ async def _build_connection_servers(self, manual_call_template: 'McpCallTemplate async def _get_or_create_session(self, server_name: str, manual_call_template: 'McpCallTemplate'): """Get an existing session or create a new one using MCPClient.""" client = await self._ensure_mcp_client(manual_call_template) - + try: # Try to get existing session session = client.get_session(server_name) self._log_info(f"Reusing existing session for server: {server_name}") return session except ValueError: - # Session doesn't exist, create a new one - self._log_info(f"Creating new session for server: {server_name}") - try: - session = await client.create_session(server_name, auto_initialize=True) - except Exception as e: - server_config = manual_call_template.config.mcpServers.get(server_name) - is_stdio = isinstance(server_config, dict) and "command" in server_config - if is_stdio and os.environ.get(CHILD_STDERR_ENV_VAR) != "inherit": - self._log_error( - f"Failed to start stdio MCP server '{server_name}': {e}. The child's stderr was " - f"suppressed; re-run with {CHILD_STDERR_ENV_VAR}=inherit to see what it printed while starting." - ) - raise - return session + pass + + # Coalesce concurrent creations for the same (configuration, server) so a + # burst of first calls dials once instead of each spawning a session and + # leaking all but the last. The check-and-set is synchronous, so exactly + # one task is created. + inflight_key = (self._config_key(manual_call_template), server_name) + task = self._session_creations.get(inflight_key) + if task is None: + task = asyncio.ensure_future( + self._create_session(server_name, client, manual_call_template) + ) + self._session_creations[inflight_key] = task + task.add_done_callback(lambda _t, k=inflight_key: self._session_creations.pop(k, None)) + # Shield so a cancelled waiter does not cancel the shared creation for + # the others (see _handle_oauth2 for the same reasoning). + return await asyncio.shield(task) + + async def _create_session(self, server_name: str, client: MCPClient, manual_call_template: 'McpCallTemplate'): + """Create (and initialize) a new session for ``server_name`` on ``client``.""" + self._log_info(f"Creating new session for server: {server_name}") + try: + return await client.create_session(server_name, auto_initialize=True) + except Exception as e: + server_config = manual_call_template.config.mcpServers.get(server_name) + is_stdio = isinstance(server_config, dict) and "command" in server_config + if is_stdio and os.environ.get(CHILD_STDERR_ENV_VAR) != "inherit": + self._log_error( + f"Failed to start stdio MCP server '{server_name}': {e}. The child's stderr was " + f"suppressed; re-run with {CHILD_STDERR_ENV_VAR}=inherit to see what it printed while starting." + ) + raise async def _release_manual_client(self, manual_call_template: 'McpCallTemplate') -> None: """Drop this manual's claim on its client. The client's sessions are closed diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py b/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py new file mode 100644 index 0000000..8afc872 --- /dev/null +++ b/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py @@ -0,0 +1,91 @@ +"""Concurrent session creation must dial once, not spawn (and leak) duplicates. + +``_get_or_create_session`` coalesces concurrent first-time creations for the +same (configuration, server) into a single shared task, shielded so one +waiter's cancellation can't cancel the creation for the others. +""" + +import asyncio + +import pytest + +from utcp_mcp.mcp_call_template import McpCallTemplate, McpConfig +from utcp_mcp.mcp_communication_protocol import McpCommunicationProtocol + + +def _template() -> McpCallTemplate: + return McpCallTemplate(name="m", config=McpConfig(mcpServers={"s": {"command": "true"}})) + + +class _NoSessionClient: + """Stands in for an MCPClient that has no session yet.""" + + def get_session(self, name): + raise ValueError("no session") + + +def _stub_client(proto: McpCommunicationProtocol, monkeypatch): + async def fake_ensure(_tmpl): + return _NoSessionClient() + + monkeypatch.setattr(proto, "_ensure_mcp_client", fake_ensure) + + +@pytest.mark.asyncio +async def test_concurrent_session_creation_is_coalesced(monkeypatch): + proto = McpCommunicationProtocol() + _stub_client(proto, monkeypatch) + + creations = 0 + release = asyncio.Event() + + async def fake_create(server_name, client, tmpl): + nonlocal creations + creations += 1 + await release.wait() + return f"session-{server_name}" + + monkeypatch.setattr(proto, "_create_session", fake_create) + + tmpl = _template() + tasks = [asyncio.create_task(proto._get_or_create_session("s", tmpl)) for _ in range(5)] + await asyncio.sleep(0) # let every caller attach to the shared creation + release.set() + results = await asyncio.gather(*tasks) + + assert results == ["session-s"] * 5 + assert creations == 1 + assert proto._session_creations == {} # slot cleared on settle + + +@pytest.mark.asyncio +async def test_cancelling_one_session_waiter_does_not_fail_the_others(monkeypatch): + proto = McpCommunicationProtocol() + _stub_client(proto, monkeypatch) + + creations = 0 + started = asyncio.Event() + release = asyncio.Event() + + async def fake_create(server_name, client, tmpl): + nonlocal creations + creations += 1 + started.set() + await release.wait() + return f"session-{server_name}" + + monkeypatch.setattr(proto, "_create_session", fake_create) + + tmpl = _template() + waiter_a = asyncio.create_task(proto._get_or_create_session("s", tmpl)) + await started.wait() + waiter_b = asyncio.create_task(proto._get_or_create_session("s", tmpl)) + await asyncio.sleep(0) + + waiter_a.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter_a + + release.set() + assert await waiter_b == "session-s" + assert creations == 1 From 2324f6945e2912dadf164f62c0b9dd6051f0fa8c Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:41:21 +0200 Subject: [PATCH 28/34] mcp: key OAuth state by full credential config; scope session creation to the client instance The token cache and in-flight fetch map were keyed by client_id alone, so two manuals sharing a client_id but differing in token URL, secret or scope received each other's tokens. Key both by the full OAuth configuration, matching the HTTP plugin and the TypeScript fix. In-flight session creations are now keyed by client identity rather than configuration, so a client retired while a creation is pending cannot hand that task (bound to the retired client) to a later same-config client. Tests updated. Co-Authored-By: Claude Fable 5.1 --- .../utcp_mcp/mcp_communication_protocol.py | 37 ++++++++++++++----- .../mcp/tests/test_mcp_oauth_security.py | 14 ++++--- .../mcp/tests/test_mcp_session_concurrency.py | 6 ++- 3 files changed, 40 insertions(+), 17 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index df382b8..79d152e 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -191,7 +191,7 @@ def __init__(self): # In-flight session creations, keyed by (configuration, server), so # concurrent first calls for the same server dial once instead of each # spawning a session and leaking all but the last. - self._session_creations: "Dict[Tuple[str, str], asyncio.Task]" = {} + self._session_creations: "Dict[Tuple[int, str], asyncio.Task]" = {} # One MCPClient per distinct server configuration. This protocol object is # registered once per process and shared by every manual, so a single # client would make manuals with different configurations evict each @@ -240,6 +240,17 @@ def _owner_key(manual_call_template: 'McpCallTemplate', config_key: str) -> str: """Identifies who holds a configuration: the calling UtcpClient plus the manual name.""" return f"{_CURRENT_OWNER.get()}:{manual_call_template.name or config_key}" + @staticmethod + def _oauth_cache_key(auth: OAuth2Auth) -> str: + """Key for the OAuth token cache and in-flight map. + + Keyed by the FULL configuration, not ``client_id`` alone: two manuals may + share a client_id but point at different issuers, scopes or secrets, and + must not receive each other's tokens. Matches the HTTP plugin. Carries + the secret, so it is used only as a dict key and never logged. + """ + return json.dumps([auth.token_url, auth.client_id, auth.client_secret, auth.scope or ""]) + async def _ensure_mcp_client(self, manual_call_template: 'McpCallTemplate') -> MCPClient: """Return the MCPClient for this manual's configuration, creating it once. @@ -331,7 +342,12 @@ async def _get_or_create_session(self, server_name: str, manual_call_template: ' # burst of first calls dials once instead of each spawning a session and # leaking all but the last. The check-and-set is synchronous, so exactly # one task is created. - inflight_key = (self._config_key(manual_call_template), server_name) + # Keyed by the CLIENT INSTANCE, not the configuration: a client can be + # retired (deregistered or drained) while a creation on it is pending, + # and a later client with the same configuration must not join that + # task, which is bound to the retired client. The task holds a reference + # to its client, so the id cannot be reused while the entry exists. + inflight_key = (id(client), server_name) task = self._session_creations.get(inflight_key) if task is None: task = asyncio.ensure_future( @@ -838,20 +854,20 @@ async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: # Validate the token endpoint before sending credentials to it, so a # manual cannot direct the operator's client secret at an arbitrary host. _ensure_secure_mcp_url(auth_details.token_url, context="MCP OAuth2 token URL") - client_id = auth_details.client_id + cache_key = self._oauth_cache_key(auth_details) # Return cached token if available. - if client_id in self._oauth_tokens: - return self._oauth_tokens[client_id]["access_token"] + if cache_key in self._oauth_tokens: + return self._oauth_tokens[cache_key]["access_token"] # Coalesce concurrent first-time fetches. The check-and-set below is # synchronous (no await between them), so exactly one task is created and # every other caller awaits it. - task = self._oauth_inflight.get(client_id) + task = self._oauth_inflight.get(cache_key) if task is None: task = asyncio.ensure_future(self._fetch_oauth2_token(auth_details)) - self._oauth_inflight[client_id] = task - task.add_done_callback(lambda _t, cid=client_id: self._oauth_inflight.pop(cid, None)) + self._oauth_inflight[cache_key] = task + task.add_done_callback(lambda _t, k=cache_key: self._oauth_inflight.pop(k, None)) # Shield the shared task: awaiting a task directly propagates a waiter's # cancellation into the task, which would cancel the fetch for every other # waiter too. shield lets a cancelled waiter raise on its own while the @@ -861,6 +877,7 @@ async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: """Perform the OAuth2 client-credentials request (body method, then Basic).""" client_id = auth_details.client_id + cache_key = self._oauth_cache_key(auth_details) async with aiohttp.ClientSession() as session: # Method 1: Send credentials in the request body try: @@ -875,7 +892,7 @@ async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: self._reject_token_redirect(response) response.raise_for_status() token_response = await response.json() - self._oauth_tokens[client_id] = token_response + self._oauth_tokens[cache_key] = token_response return token_response["access_token"] except aiohttp.ClientError as e: self._log_error(f"OAuth2 with credentials in body failed: {e}. Trying Basic Auth header.") @@ -892,7 +909,7 @@ async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: self._reject_token_redirect(response) response.raise_for_status() token_response = await response.json() - self._oauth_tokens[client_id] = token_response + self._oauth_tokens[cache_key] = token_response return token_response["access_token"] except aiohttp.ClientError as e: self._log_error(f"OAuth2 with Basic Auth header also failed: {e}") diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index e3f3501..d7126ee 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -35,9 +35,10 @@ async def test_insecure_token_url_rejected_before_cache_or_network(): # Seed the cache so a returned token would prove the guard ran too late. # The guard must reject the insecure URL before the cache is consulted and # before any network request is made. - proto._oauth_tokens["id"] = {"access_token": "cached"} + auth = _oauth("http://attacker.example/token") + proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "cached"} with pytest.raises(ValueError, match="Security error"): - await proto._handle_oauth2(_oauth("http://attacker.example/token")) + await proto._handle_oauth2(auth) @pytest.mark.asyncio @@ -45,8 +46,9 @@ async def test_secure_token_url_passes_the_guard_without_network(): proto = McpCommunicationProtocol() # A pre-seeded token lets us confirm a secure URL passes validation and # returns without any network I/O. - proto._oauth_tokens["id"] = {"access_token": "cached"} - token = await proto._handle_oauth2(_oauth("https://auth.example.com/token")) + auth = _oauth("https://auth.example.com/token") + proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "cached"} + token = await proto._handle_oauth2(auth) assert token == "cached" @@ -164,7 +166,7 @@ async def fake_fetch(auth): calls += 1 started.set() await release.wait() - proto._oauth_tokens[auth.client_id] = {"access_token": "tok"} + proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "tok"} return "tok" proto._fetch_oauth2_token = fake_fetch # instance attr shadows the method @@ -193,7 +195,7 @@ async def fake_fetch(auth): calls += 1 started.set() await release.wait() - proto._oauth_tokens[auth.client_id] = {"access_token": "tok"} + proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "tok"} return "tok" proto._fetch_oauth2_token = fake_fetch diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py b/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py index 8afc872..e933800 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py @@ -25,8 +25,12 @@ def get_session(self, name): def _stub_client(proto: McpCommunicationProtocol, monkeypatch): + # One shared instance, as the real _ensure_mcp_client returns the same client + # for the same configuration; in-flight creations are keyed by client identity. + client = _NoSessionClient() + async def fake_ensure(_tmpl): - return _NoSessionClient() + return client monkeypatch.setattr(proto, "_ensure_mcp_client", fake_ensure) From 3173aa11ddca2674b3348c53ba7db023941b5c8b Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:04:10 +0200 Subject: [PATCH 29/34] mcp: drop cached OAuth tokens on close() A drain left credentials cached on this shared instance, unlike the TypeScript plugin. Clear the token cache in close(); in-flight fetches are left to self-prune when they settle. Adds a test. Co-Authored-By: Claude Fable 5.1 --- .../mcp/src/utcp_mcp/mcp_communication_protocol.py | 4 ++++ .../mcp/tests/test_mcp_oauth_security.py | 10 ++++++++++ 2 files changed, 14 insertions(+) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 79d152e..1de551e 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -840,6 +840,10 @@ async def close(self) -> None: """Close all active sessions and clean up resources.""" self._log_info("Closing MCP communication protocol and cleaning up all sessions") await self._cleanup_all_sessions() + # Drop cached tokens on a drain so this shared instance does not hold + # credentials past it (matches the TypeScript plugin). In-flight fetches + # are left alone: they self-prune when they settle. + self._oauth_tokens.clear() self._log_info("MCP communication protocol closed successfully") async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index d7126ee..f5b8c4d 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -228,3 +228,13 @@ async def test_manuals_with_same_servers_but_different_auth_get_distinct_keys(): auth=_oauth("https://auth-b.example.com/token"), ) assert McpCommunicationProtocol._config_key(a) != McpCommunicationProtocol._config_key(b) + + +@pytest.mark.asyncio +async def test_close_drops_cached_tokens(): + # A drain must not leave credentials cached on this shared instance. + proto = McpCommunicationProtocol() + auth = _oauth("https://auth.example.com/token") + proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "tok"} + await proto.close() + assert proto._oauth_tokens == {} From 3322fa58d21c3c4d7c960229d48153a849f3125f Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:27:15 +0200 Subject: [PATCH 30/34] mcp: gate OAuth cache writes on in-flight identity; cancel and drop pending fetches on close() The fetch wrote the cache unconditionally, so a fetch still in flight when close() ran repopulated the cache afterwards and the drain did not actually leave the instance credential-free. Apply the same invariant as the TypeScript plugin: the in-flight entry is the sole authority for caching. The settle handler caches a result only if its task is still the current entry, and removes the entry only then; _fetch_oauth2_token is now a pure fetch returning the token response. close() cancels in-flight fetches and drops their entries, so a fetch that still lands is no longer current and does not cache. Adds a post-close regression test; the coalescing stubs now return the token response for the handler to cache. Co-Authored-By: Claude Fable 5.1 --- .../utcp_mcp/mcp_communication_protocol.py | 47 ++++++++++++++----- .../mcp/tests/test_mcp_oauth_security.py | 33 +++++++++++-- 2 files changed, 64 insertions(+), 16 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 1de551e..c9b37dd 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -840,9 +840,14 @@ async def close(self) -> None: """Close all active sessions and clean up resources.""" self._log_info("Closing MCP communication protocol and cleaning up all sessions") await self._cleanup_all_sessions() - # Drop cached tokens on a drain so this shared instance does not hold - # credentials past it (matches the TypeScript plugin). In-flight fetches - # are left alone: they self-prune when they settle. + # Drain the OAuth state so this shared instance holds no credentials past + # close(). Cancel in-flight fetches (asyncio can) and drop their entries; + # a fetch that still lands finds it is no longer the current entry and, + # by the caching rule in _on_oauth_fetch_done, does not repopulate the + # cache. + for task in list(self._oauth_inflight.values()): + task.cancel() + self._oauth_inflight.clear() self._oauth_tokens.clear() self._log_info("MCP communication protocol closed successfully") @@ -871,17 +876,37 @@ async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: if task is None: task = asyncio.ensure_future(self._fetch_oauth2_token(auth_details)) self._oauth_inflight[cache_key] = task - task.add_done_callback(lambda _t, k=cache_key: self._oauth_inflight.pop(k, None)) + task.add_done_callback(functools.partial(self._on_oauth_fetch_done, cache_key)) # Shield the shared task: awaiting a task directly propagates a waiter's # cancellation into the task, which would cancel the fetch for every other # waiter too. shield lets a cancelled waiter raise on its own while the # shared fetch runs to completion for the rest. - return await asyncio.shield(task) + return (await asyncio.shield(task))["access_token"] + + def _on_oauth_fetch_done(self, cache_key: str, task: "asyncio.Task") -> None: + """Settle handler for a shared token fetch. - async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: - """Perform the OAuth2 client-credentials request (body method, then Basic).""" + The in-flight entry is the sole authority for who may write the cache: + only a fetch that is STILL the current entry when it settles caches its + result, and only then does it remove itself. An entry dropped by + close() therefore never repopulates the cache and never clobbers a + successor — there is no version counter to coordinate or leak. + """ + if self._oauth_inflight.get(cache_key) is not task: + return + self._oauth_inflight.pop(cache_key, None) + if task.cancelled() or task.exception() is not None: + return + self._oauth_tokens[cache_key] = task.result() + + async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> Dict[str, Any]: + """Perform the OAuth2 client-credentials request (body method, then Basic). + + Pure fetch: returns the token response and never writes the cache — + whether the result may be cached is decided by ``_on_oauth_fetch_done``, + which knows whether this fetch is still the current in-flight entry. + """ client_id = auth_details.client_id - cache_key = self._oauth_cache_key(auth_details) async with aiohttp.ClientSession() as session: # Method 1: Send credentials in the request body try: @@ -896,8 +921,7 @@ async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: self._reject_token_redirect(response) response.raise_for_status() token_response = await response.json() - self._oauth_tokens[cache_key] = token_response - return token_response["access_token"] + return token_response except aiohttp.ClientError as e: self._log_error(f"OAuth2 with credentials in body failed: {e}. Trying Basic Auth header.") @@ -913,8 +937,7 @@ async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> str: self._reject_token_redirect(response) response.raise_for_status() token_response = await response.json() - self._oauth_tokens[cache_key] = token_response - return token_response["access_token"] + return token_response except aiohttp.ClientError as e: self._log_error(f"OAuth2 with Basic Auth header also failed: {e}") raise e diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index f5b8c4d..aa34fca 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -166,8 +166,7 @@ async def fake_fetch(auth): calls += 1 started.set() await release.wait() - proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "tok"} - return "tok" + return {"access_token": "tok"} # cached by the settle handler, if still current proto._fetch_oauth2_token = fake_fetch # instance attr shadows the method auth = _oauth("https://auth.example.com/token") @@ -195,8 +194,7 @@ async def fake_fetch(auth): calls += 1 started.set() await release.wait() - proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "tok"} - return "tok" + return {"access_token": "tok"} # cached by the settle handler, if still current proto._fetch_oauth2_token = fake_fetch auth = _oauth("https://auth.example.com/token") @@ -238,3 +236,30 @@ async def test_close_drops_cached_tokens(): proto._oauth_tokens[McpCommunicationProtocol._oauth_cache_key(auth)] = {"access_token": "tok"} await proto.close() assert proto._oauth_tokens == {} + + +@pytest.mark.asyncio +async def test_fetch_landing_after_close_does_not_repopulate_cache(): + # close() drops the in-flight entry; a fetch that still lands is no longer + # the current entry and must not write the cache (identity gate), so the + # drain leaves no credential behind. + proto = McpCommunicationProtocol() + release = asyncio.Event() + + async def fake_fetch(_auth): + await release.wait() + return {"access_token": "late"} + + proto._fetch_oauth2_token = fake_fetch + auth = _oauth("https://auth.example.com/token") + waiter = asyncio.create_task(proto._handle_oauth2(auth)) + await asyncio.sleep(0) # the shared fetch is registered and running + assert len(proto._oauth_inflight) == 1 + + await proto.close() # cancels the fetch and drops its entry + release.set() + with pytest.raises(asyncio.CancelledError): + await waiter + + assert proto._oauth_tokens == {} + assert proto._oauth_inflight == {} From f2f6d0567cd58d600578bb8ab1cbfb036e018546 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:40:58 +0200 Subject: [PATCH 31/34] mcp: treat a token response without access_token as a failed fetch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A 200 whose body lacks access_token was returned as a success, cached by the settle handler, and then failed on the read path — and since the cache is only ever cleared by a drain, one malformed reply became a persistent OAuth failure. Validate at the fetch boundary instead of guarding the cache write: _require_access_token turns a malformed body into a ClientError, so the cache only ever receives validated responses, the body-vs-Basic fallback proceeds exactly as it would on a transport error, and callers get a real error rather than a KeyError. Matches the TypeScript plugin. Also fixes the post-close test so it actually exercises the identity gate: the fake fetch now signals it is running (inside its try) before the drain, since cancelling a not-yet-started coroutine throws at entry and would only ever test cancellation. Co-Authored-By: Claude Fable 5.1 --- .../utcp_mcp/mcp_communication_protocol.py | 21 ++++++-- .../mcp/tests/test_mcp_oauth_security.py | 50 +++++++++++++++---- 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index c9b37dd..d588a4a 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -920,8 +920,7 @@ async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> Dict[str, Any]: async with session.post(auth_details.token_url, data=body_data, allow_redirects=False) as response: self._reject_token_redirect(response) response.raise_for_status() - token_response = await response.json() - return token_response + return self._require_access_token(await response.json()) except aiohttp.ClientError as e: self._log_error(f"OAuth2 with credentials in body failed: {e}. Trying Basic Auth header.") @@ -936,12 +935,26 @@ async def _fetch_oauth2_token(self, auth_details: OAuth2Auth) -> Dict[str, Any]: async with session.post(auth_details.token_url, data=header_data, auth=header_auth, allow_redirects=False) as response: self._reject_token_redirect(response) response.raise_for_status() - token_response = await response.json() - return token_response + return self._require_access_token(await response.json()) except aiohttp.ClientError as e: self._log_error(f"OAuth2 with Basic Auth header also failed: {e}") raise e + @staticmethod + def _require_access_token(token_response: Any) -> Dict[str, Any]: + """A successful HTTP response is not a successful token fetch unless it + carries an ``access_token``. + + Treating a malformed body as a fetch failure is what keeps it out of the + cache — the cache only ever receives validated responses, so a single + bad reply cannot become a persistent failure on the read path — and it + lets the body-vs-Basic fallback proceed the same way a transport error + would. Matches the TypeScript plugin. + """ + if not isinstance(token_response, dict) or not token_response.get("access_token"): + raise aiohttp.ClientError("OAuth2 token endpoint responded without an access_token") + return token_response + @staticmethod def _reject_token_redirect(response: "aiohttp.ClientResponse") -> None: """Refuse a redirect from the OAuth2 token endpoint. diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index aa34fca..ed38ef3 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -240,26 +240,58 @@ async def test_close_drops_cached_tokens(): @pytest.mark.asyncio async def test_fetch_landing_after_close_does_not_repopulate_cache(): - # close() drops the in-flight entry; a fetch that still lands is no longer - # the current entry and must not write the cache (identity gate), so the - # drain leaves no credential behind. + # close() drops the in-flight entry. The fake fetch deliberately survives the + # drain's cancel and still LANDS with a value, so the only thing standing + # between that value and the cache is the identity gate — which must hold. proto = McpCommunicationProtocol() + started = asyncio.Event() release = asyncio.Event() async def fake_fetch(_auth): - await release.wait() + started.set() + try: + await release.wait() + except asyncio.CancelledError: + pass # survive the drain's cancel so the fetch genuinely lands with a value return {"access_token": "late"} proto._fetch_oauth2_token = fake_fetch auth = _oauth("https://auth.example.com/token") waiter = asyncio.create_task(proto._handle_oauth2(auth)) - await asyncio.sleep(0) # the shared fetch is registered and running + # Wait until the fetch is genuinely RUNNING (inside its try), not merely + # scheduled: cancelling a coroutine that has not started throws at its entry + # and the except never runs, which would test cancellation, not the gate. + await started.wait() assert len(proto._oauth_inflight) == 1 - await proto.close() # cancels the fetch and drops its entry - release.set() - with pytest.raises(asyncio.CancelledError): - await waiter + await proto.close() # drops the entry; the fetch survives the cancel and lands + assert await waiter == "late" # the caller still receives its token... + assert proto._oauth_tokens == {} # ...but the fetch was no longer current, so nothing cached + assert proto._oauth_inflight == {} + + +def test_require_access_token_rejects_malformed_responses(): + # A 200 without an access_token is a failed fetch, not a cacheable result. + with pytest.raises(aiohttp.ClientError, match="access_token"): + McpCommunicationProtocol._require_access_token({"token_type": "bearer"}) + with pytest.raises(aiohttp.ClientError, match="access_token"): + McpCommunicationProtocol._require_access_token(["not", "a", "dict"]) + assert McpCommunicationProtocol._require_access_token({"access_token": "t"}) == {"access_token": "t"} + + +@pytest.mark.asyncio +async def test_failed_fetch_is_never_cached_and_can_be_retried(): + # A fetch that fails (including on a malformed body) must leave neither a + # cache entry nor an in-flight entry behind, so the next call retries. + proto = McpCommunicationProtocol() + + async def failing_fetch(_auth): + raise aiohttp.ClientError("OAuth2 token endpoint responded without an access_token") + + proto._fetch_oauth2_token = failing_fetch + auth = _oauth("https://auth.example.com/token") + with pytest.raises(aiohttp.ClientError, match="access_token"): + await proto._handle_oauth2(auth) assert proto._oauth_tokens == {} assert proto._oauth_inflight == {} From ae606891866aec77c9cfb3bd1b200e78277b631c Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:01:43 +0200 Subject: [PATCH 32/34] mcp: require the OAuth access_token to be a non-empty string Completes the token-response validation predicate at the fetch boundary. It accepted any truthy access_token, so a number, True or an object would pass, be cached, and be injected as an invalid bearer credential on every reuse. A usable token is a non-empty string; anything else is a failed fetch and never reaches the cache. The test now covers the non-string cases, which are exactly what fails if the string requirement is removed. Co-Authored-By: Claude Fable 5.1 --- .../src/utcp_mcp/mcp_communication_protocol.py | 8 ++++++-- .../mcp/tests/test_mcp_oauth_security.py | 18 +++++++++++++----- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index d588a4a..7b4fb86 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -951,8 +951,12 @@ def _require_access_token(token_response: Any) -> Dict[str, Any]: lets the body-vs-Basic fallback proceed the same way a transport error would. Matches the TypeScript plugin. """ - if not isinstance(token_response, dict) or not token_response.get("access_token"): - raise aiohttp.ClientError("OAuth2 token endpoint responded without an access_token") + # "Usable" means a non-empty string: mcp-use formats whatever it is given + # into ``Bearer ``, so a truthy non-string (a number, ``True``, a + # dict) would be injected as an invalid credential on every reuse. + token = token_response.get("access_token") if isinstance(token_response, dict) else None + if not isinstance(token, str) or not token: + raise aiohttp.ClientError("OAuth2 token endpoint responded without a usable access_token") return token_response @staticmethod diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index ed38ef3..5c32fe3 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -271,11 +271,19 @@ async def fake_fetch(_auth): def test_require_access_token_rejects_malformed_responses(): - # A 200 without an access_token is a failed fetch, not a cacheable result. - with pytest.raises(aiohttp.ClientError, match="access_token"): - McpCommunicationProtocol._require_access_token({"token_type": "bearer"}) - with pytest.raises(aiohttp.ClientError, match="access_token"): - McpCommunicationProtocol._require_access_token(["not", "a", "dict"]) + # A 200 without a USABLE access_token is a failed fetch, not a cacheable + # result. Usable means a non-empty string: a truthy non-string (12345, True) + # would be injected as an invalid bearer credential on every reuse. The + # non-string cases are what fail if the isinstance(str) clause is removed. + for bad in ( + {"token_type": "bearer"}, + ["not", "a", "dict"], + {"access_token": 12345}, + {"access_token": True}, + {"access_token": ""}, + ): + with pytest.raises(aiohttp.ClientError, match="access_token"): + McpCommunicationProtocol._require_access_token(bad) assert McpCommunicationProtocol._require_access_token({"access_token": "t"}) == {"access_token": "t"} From ea1fe9de8549bd045b097d488c2d72e1b48245ce Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:24:31 +0200 Subject: [PATCH 33/34] mcp: define a usable OAuth access_token positively as non-empty visible ASCII Replaces a growing denylist of bad token shapes with the rule derived from the contract the token must satisfy: mcp-use places it verbatim into an Authorization: Bearer header, and a header value may contain only visible ASCII (RFC 9110 VCHAR, 0x21-0x7E), with a space ending the token. One rule makes every unusable shape inexpressible at once, including CR/LF header injection, NUL and non-ASCII. RFC 6750's narrower b64token alphabet was deliberately not used: it would reject legitimate opaque tokens. Tests cover each unusable shape (each fails if the VCHAR clause is removed) and a printable-punctuation token that must be accepted. Parity with typescript-utcp. Co-Authored-By: Claude Fable 5.1 --- .../utcp_mcp/mcp_communication_protocol.py | 24 +++++++++++++++---- .../mcp/tests/test_mcp_oauth_security.py | 17 +++++++++---- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 7b4fb86..3347a94 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -3,6 +3,7 @@ import copy import functools import os +import re import sys from ipaddress import IPv6Address, ip_address from typing import Any, Dict, List, Optional, AsyncGenerator, TYPE_CHECKING, Tuple, TextIO @@ -55,6 +56,10 @@ async def wrapper(self, caller, *args, **kwargs): # Hostnames considered safe to reach over plain HTTP/WS. _LOOPBACK_HOSTNAMES = frozenset({"localhost", "127.0.0.1", "::1", "[::1]"}) +# What a usable bearer token looks like: a non-empty run of visible ASCII +# (RFC 9110 VCHAR, 0x21-0x7E). See ``_require_access_token``. +_VISIBLE_ASCII = re.compile(r"[\x21-\x7E]+") + def _is_secure_mcp_url(url: str) -> bool: """Return True if ``url`` is safe for the MCP plugin to connect to. @@ -951,12 +956,21 @@ def _require_access_token(token_response: Any) -> Dict[str, Any]: lets the body-vs-Basic fallback proceed the same way a transport error would. Matches the TypeScript plugin. """ - # "Usable" means a non-empty string: mcp-use formats whatever it is given - # into ``Bearer ``, so a truthy non-string (a number, ``True``, a - # dict) would be injected as an invalid credential on every reuse. + # Defined POSITIVELY from the contract the token must satisfy, not as a + # list of bad shapes: mcp-use places it verbatim into + # ``Authorization: Bearer ``, and an HTTP header value may contain + # only visible ASCII (RFC 9110 VCHAR, 0x21-0x7E), with a space ending the + # token. So a usable token is a non-empty string of VCHAR. That single + # rule makes every unusable shape inexpressible at once (non-string, + # empty, whitespace, CR/LF header injection, NUL/control characters, + # non-ASCII) without over-fitting to RFC 6750's narrower b64token + # alphabet, which would reject legitimate opaque tokens. token = token_response.get("access_token") if isinstance(token_response, dict) else None - if not isinstance(token, str) or not token: - raise aiohttp.ClientError("OAuth2 token endpoint responded without a usable access_token") + if not isinstance(token, str) or not _VISIBLE_ASCII.fullmatch(token): + raise aiohttp.ClientError( + "OAuth2 token endpoint responded without a usable access_token " + "(must be a non-empty string of visible ASCII)" + ) return token_response @staticmethod diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py index 5c32fe3..7bb8356 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -271,20 +271,27 @@ async def fake_fetch(_auth): def test_require_access_token_rejects_malformed_responses(): - # A 200 without a USABLE access_token is a failed fetch, not a cacheable - # result. Usable means a non-empty string: a truthy non-string (12345, True) - # would be injected as an invalid bearer credential on every reuse. The - # non-string cases are what fail if the isinstance(str) clause is removed. + # The usability rule is positive (non-empty visible ASCII), so every shape + # that cannot be a valid ``Authorization: Bearer `` header is rejected + # by one rule. The string cases below are exactly what fail if the VCHAR + # requirement is removed. for bad in ( {"token_type": "bearer"}, ["not", "a", "dict"], {"access_token": 12345}, {"access_token": True}, {"access_token": ""}, + {"access_token": "tok en"}, # embedded space ends the token + {"access_token": "tok\r\nen"}, # CR/LF: header injection + {"access_token": "tok\x00en"}, # NUL / control character + {"access_token": "tok\u00e9n"}, # non-ASCII ): with pytest.raises(aiohttp.ClientError, match="access_token"): McpCommunicationProtocol._require_access_token(bad) - assert McpCommunicationProtocol._require_access_token({"access_token": "t"}) == {"access_token": "t"} + # Printable punctuation is accepted: the rule must not over-reject real + # opaque tokens (tightening to RFC 6750's b64token alphabet would). + ok = {"access_token": "a.b-c_d~e+f/g=:h"} + assert McpCommunicationProtocol._require_access_token(ok) == ok @pytest.mark.asyncio From 961e246cd1d6256b7ffc4c9e6f80c537c480cb72 Mon Sep 17 00:00:00 2001 From: Razvan Radulescu <43811028+h3xxit@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:32:12 +0200 Subject: [PATCH 34/34] chore: bump versions of the packages changed in this release utcp-cli 1.1.4 -> 1.1.5, utcp-http 1.1.11 -> 1.1.12, utcp-mcp 1.1.2 -> 1.1.3, utcp-socket 1.1.0 -> 1.1.1. Patch bumps: every change is a backwards-compatible bug or security fix. Core is unchanged and not bumped. Co-Authored-By: Claude Fable 5.1 --- plugins/communication_protocols/cli/pyproject.toml | 2 +- plugins/communication_protocols/http/pyproject.toml | 2 +- plugins/communication_protocols/mcp/pyproject.toml | 2 +- plugins/communication_protocols/socket/pyproject.toml | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/plugins/communication_protocols/cli/pyproject.toml b/plugins/communication_protocols/cli/pyproject.toml index cd0aa8d..e8c913d 100644 --- a/plugins/communication_protocols/cli/pyproject.toml +++ b/plugins/communication_protocols/cli/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "utcp-cli" -version = "1.1.4" +version = "1.1.5" authors = [ { name = "UTCP Contributors" }, ] diff --git a/plugins/communication_protocols/http/pyproject.toml b/plugins/communication_protocols/http/pyproject.toml index d6c7221..f59334a 100644 --- a/plugins/communication_protocols/http/pyproject.toml +++ b/plugins/communication_protocols/http/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "utcp-http" -version = "1.1.11" +version = "1.1.12" authors = [ { name = "UTCP Contributors" }, ] diff --git a/plugins/communication_protocols/mcp/pyproject.toml b/plugins/communication_protocols/mcp/pyproject.toml index ec66664..e06e496 100644 --- a/plugins/communication_protocols/mcp/pyproject.toml +++ b/plugins/communication_protocols/mcp/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "utcp-mcp" -version = "1.1.2" +version = "1.1.3" authors = [ { name = "UTCP Contributors" }, ] diff --git a/plugins/communication_protocols/socket/pyproject.toml b/plugins/communication_protocols/socket/pyproject.toml index dbbc1b0..8b14e99 100644 --- a/plugins/communication_protocols/socket/pyproject.toml +++ b/plugins/communication_protocols/socket/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "utcp-socket" -version = "1.1.0" +version = "1.1.1" authors = [ { name = "UTCP Contributors" }, ]