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/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/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/http/src/utcp_http/_errors.py b/plugins/communication_protocols/http/src/utcp_http/_errors.py new file mode 100644 index 0000000..9d4fa96 --- /dev/null +++ b/plugins/communication_protocols/http/src/utcp_http/_errors.py @@ -0,0 +1,115 @@ +"""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 +import re +from typing import Optional + +import aiohttp + +# 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") + +# 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. + + 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: + return None + try: + data = json.loads(body) + 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: + continue + value = data[key] + if isinstance(value, str): + if value.strip(): + return _clean(value) + continue + # Structured error: show it rather than a later generic string. + return _clean(body) + return _clean(body) + + +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] + # ``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: + return raw.decode(response.charset or "utf-8", errors="replace") + except (LookupError, RuntimeError, ValueError): + return raw.decode("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 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 _read_body_bounded(response, MAX_BODY_READ_BYTES) + 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/_security.py b/plugins/communication_protocols/http/src/utcp_http/_security.py index 5e431cc..884581a 100644 --- a/plugins/communication_protocols/http/src/utcp_http/_security.py +++ b/plugins/communication_protocols/http/src/utcp_http/_security.py @@ -478,3 +478,40 @@ 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. + + ``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 + 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 8a3a187..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 @@ -33,7 +33,8 @@ 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 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', '') @@ -225,6 +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) + # 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) @@ -359,7 +364,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 83afac2..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 @@ -3,6 +3,7 @@ import aiohttp import json import asyncio +import codecs import re from urllib.parse import quote import base64 @@ -17,7 +18,8 @@ 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._security import ensure_secure_url, safe_request_with_redirects +from utcp_http._errors import raise_for_status_with_body +from utcp_http._security import ensure_secure_url, safe_request_with_redirects, reject_remote_loopback_tool_urls import traceback import logging @@ -28,6 +30,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. @@ -35,6 +42,22 @@ 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 + # 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]] = {} @@ -148,9 +171,12 @@ 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) + # 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, @@ -224,97 +250,223 @@ 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." + 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 + + # 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 + provider_name = tool_call_template.name + + while True: + attempt_headers = dict(request_headers) + 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() + 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). + # 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() + 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." + ) + # 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 + # 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): + # 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 + # 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. + 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 " + f"{min(retry_delay_ms, self.MAX_RECONNECT_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(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. + + 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 = "" + + 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': + # 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) + if data_lines: + current_event['data'] = '\n'.join(data_lines) + return current_event or None + + # 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". + 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" ) - response.raise_for_status() - async for event in self._process_sse_stream(response, tool_call_template.event_type): + + # 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 - except Exception as e: - logger.error(f"Error establishing SSE connection to '{tool_call_template.name}': {e}") - raise - finally: - await session.close() + # 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" + ) - async def _process_sse_stream(self, response: aiohttp.ClientResponse, event_type=None): - """Process the SSE stream and yield events.""" - buffer = "" + @staticmethod + def _parse_event_data(data: str) -> Any: + """Return the JSON-decoded payload when possible, otherwise the raw string.""" 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 - - 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 + 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/src/utcp_http/streamable_http_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/streamable_http_communication_protocol.py index fde6cb5..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 @@ -15,7 +15,8 @@ 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._security import ensure_secure_url, safe_request_with_redirects +from utcp_http._errors import raise_for_status_with_body +from utcp_http._security import ensure_secure_url, safe_request_with_redirects, reject_remote_loopback_tool_urls import logging logging.basicConfig( @@ -149,9 +150,12 @@ 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) + # 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, @@ -292,7 +296,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 @@ -309,7 +313,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 518b8df..4504e2c 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,48 @@ 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) + + # 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) + + # 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) + + # 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 @@ -736,3 +778,112 @@ 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] + + +@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 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 +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"]' + + +@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 + + +@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_loopback_manual_security.py b/plugins/communication_protocols/http/tests/test_loopback_manual_security.py new file mode 100644 index 0000000..3ce2da8 --- /dev/null +++ b/plugins/communication_protocols/http/tests/test_loopback_manual_security.py @@ -0,0 +1,57 @@ +"""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"), + ) + ] + ) + + +@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(tool_url)) + + +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") + ) 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..6e99a26 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,146 @@ async def token_header_auth_handler(request): 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") + + +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.""" + 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): + 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')) + # 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:]: + await response.write(event.encode('utf-8')) + return response + + +async def slow_handshake_handler(request): + """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) + + +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') + # 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 + 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')) + # 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:]: + await response.write(event.encode('utf-8')) + return response + # --- Pytest Fixtures --- @pytest_asyncio.fixture @@ -118,9 +258,27 @@ 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} + 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) + app.router.add_get("/huge_retry_events", huge_retry_events_handler) + app["huge_retry"] = {"connections": 0} return app @pytest_asyncio.fixture @@ -377,3 +535,202 @@ 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 + + +# --- 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 + # A protocol violation is not a connection loss: exactly one connection, no reconnect. + assert app["no_delimiter"]["connections"] == 1 + + +@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 + + +@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] + + +# --- 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 d86a44c..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,12 @@ 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") + app = web.Application() app.add_routes([ web.get('/discover', discover), @@ -120,6 +126,8 @@ 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('/refused', refused_endpoint), + web.get('/forbidden-discovery', forbidden_discovery), ]) return app @@ -341,3 +349,27 @@ 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] + + +@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/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/pyproject.toml b/plugins/communication_protocols/mcp/pyproject.toml index 87461b7..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" }, ] @@ -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", 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..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 @@ -1,5 +1,13 @@ +import asyncio +import contextvars +import copy +import functools +import os +import re import sys -from typing import Any, Dict, Optional, AsyncGenerator, TYPE_CHECKING, Tuple +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 @@ -23,6 +31,154 @@ 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 + + +# 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. + + 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" + +_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: + # 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 + + class McpCommunicationProtocol(CommunicationProtocol): """REQUIRED MCP transport implementation that connects to MCP servers via stdio or HTTP. @@ -33,7 +189,28 @@ class McpCommunicationProtocol(CommunicationProtocol): def __init__(self): self._oauth_tokens: Dict[str, Dict[str, Any]] = {} - self._mcp_client: Optional[MCPClient] = None + # 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]]" = {} + # 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[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 + # 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.""" @@ -47,38 +224,209 @@ 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 = MCPClient.from_dict(config) + @staticmethod + def _config_key(manual_call_template: 'McpCallTemplate') -> 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: + """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. + + 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 + # 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: + 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 + 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 _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 server_config.get("auth") + 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.""" - 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}") - session = await self._mcp_client.create_session(server_name, auto_initialize=True) - return session + pass - 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) + # 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. + # 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( + 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 + 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): @@ -111,7 +459,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 @@ -139,7 +487,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 @@ -159,7 +507,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 @@ -173,6 +521,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. @@ -246,6 +595,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. @@ -379,18 +729,31 @@ 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)}") - # 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 (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 # Process content if available (fallback) if hasattr(result, 'content'): @@ -464,6 +827,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): @@ -472,27 +836,82 @@ 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() + # 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") async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: - """Handles OAuth2 client credentials flow, trying both body and auth header methods.""" - client_id = auth_details.client_id - - # Return cached token if available - if client_id in self._oauth_tokens: - return self._oauth_tokens[client_id]["access_token"] + """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") + cache_key = self._oauth_cache_key(auth_details) + + # Return cached token if available. + 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(cache_key) + if task is None: + task = asyncio.ensure_future(self._fetch_oauth2_token(auth_details)) + self._oauth_inflight[cache_key] = task + 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))["access_token"] + + def _on_oauth_fetch_done(self, cache_key: str, task: "asyncio.Task") -> None: + """Settle handler for a shared token fetch. + + 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 async with aiohttp.ClientSession() as session: # Method 1: Send credentials in the request body try: @@ -503,11 +922,10 @@ 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 - return token_response["access_token"] + 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.") @@ -519,11 +937,52 @@ 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 - return token_response["access_token"] + 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. + """ + # 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 _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 + 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 new file mode 100644 index 0000000..7bb8356 --- /dev/null +++ b/plugins/communication_protocols/mcp/tests/test_mcp_oauth_security.py @@ -0,0 +1,312 @@ +"""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 asyncio + +import aiohttp +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_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. + 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(auth) + + +@pytest.mark.asyncio +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. + 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" + + +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 +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_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() + + 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_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() + 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") + + 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_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() + 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") + + 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( + 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) + + +@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 == {} + + +@pytest.mark.asyncio +async def test_fetch_landing_after_close_does_not_repopulate_cache(): + # 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): + 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)) + # 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() # 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(): + # 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) + # 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 +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 == {} 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..e933800 --- /dev/null +++ b/plugins/communication_protocols/mcp/tests/test_mcp_session_concurrency.py @@ -0,0 +1,95 @@ +"""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): + # 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 client + + 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 diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index d127791..df75aed 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -244,3 +244,151 @@ 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"}] + + +# --- 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, mcp_manual) + + +@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, mcp_manual) + + +@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} + 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 + + +@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/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" }, ] 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..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 @@ -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 @@ -327,11 +327,8 @@ 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): - # 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()