diff --git a/plugins/communication_protocols/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py index b3a0e90..9d4fa96 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,16 @@ _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. +# 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: + 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 +50,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 +61,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: @@ -69,11 +82,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..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 @@ -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() @@ -292,7 +297,22 @@ 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() + # 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", "") + # 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}" + ) + except SseProtocolError: + raise except Exception as e: if reconnect_attempts == 0: # The initial handshake failing (refused, timed out, non-2xx) is a @@ -315,13 +335,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. @@ -375,10 +398,11 @@ 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. + # 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) - except ValueError: - pass if data_lines: current_event['data'] = '\n'.join(data_lines) return current_event or None @@ -413,13 +437,25 @@ 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. + # 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" - event = flush(buffer) - if event is not None: - yield event + 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 + # 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/src/utcp_http/streamable_http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py index df4b94f..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 @@ -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 @@ -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_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..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") @@ -141,7 +145,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:]: @@ -150,8 +157,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) @@ -163,12 +171,56 @@ 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') 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 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"}) + + +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') + # 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') + 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 @@ -182,7 +234,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 @@ -203,9 +258,16 @@ 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("/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) 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} @@ -596,3 +658,79 @@ 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 + + +@pytest.mark.asyncio +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("/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 == 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 e026f45..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 @@ -356,3 +360,16 @@ 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('/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 == 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 f86f4ab..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 @@ -1,6 +1,9 @@ +import asyncio +import contextvars +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 @@ -24,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" @@ -67,9 +92,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 @@ -84,7 +116,20 @@ 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] = {} + # 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() + # 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.""" @@ -98,27 +143,64 @@ 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.""" - 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 = _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) + + @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. + + 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) + 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 + 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 + 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 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 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 @@ -130,16 +212,54 @@ 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 _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 = 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] + 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 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'): + """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}") + 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): @@ -172,7 +292,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 @@ -200,7 +320,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 @@ -220,7 +340,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 @@ -234,6 +354,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. @@ -307,6 +428,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. @@ -538,6 +660,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): @@ -546,17 +669,15 @@ 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) - 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.""" 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 d2e2823..df75aed 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 @@ -296,3 +296,99 @@ 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) + 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 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 +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_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"} + + +@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 == {} + + +@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"} 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."""