diff --git a/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py b/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py index bca07fc..422d70a 100644 --- a/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py +++ b/plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py @@ -972,9 +972,19 @@ async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], too async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: """REQUIRED - Streaming calls are not supported for the CLI protocol. + Execute a tool call through the CLI transport streamingly. - Raises: - NotImplementedError: Always, as this functionality is not supported. + The CLI protocol does not natively support streaming, so the command is + executed to completion and the full result is yielded as a single chunk. + + Args: + caller: The UTCP client that is calling this method. + tool_name: Name of the tool to call. + tool_args: Dictionary of arguments to pass to the tool. + tool_call_template: Call template of the tool to call. + + Yields: + The complete tool result as a single item. """ - raise NotImplementedError("Streaming is not supported by the CLI communication protocol.") + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result diff --git a/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py b/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py index f96ffa7..0d98ee8 100644 --- a/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py +++ b/plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py @@ -276,6 +276,22 @@ async def test_call_tool_json_output(transport: CliCommunicationProtocol, mock_c assert "Echo:" in result["result"] and "Hello" in result["result"] +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk(transport: CliCommunicationProtocol, mock_cli_script, python_executable): + """Streaming mode should emit the full result as one chunk instead of failing.""" + call_template = CliCallTemplate( + commands=[ + {"command": f"{python_executable} {mock_cli_script} --message UTCP_ARG_message_UTCP_END"} + ] + ) + + chunks = [chunk async for chunk in transport.call_tool_streaming(None, "echo", {"message": "Hello World"}, call_template)] + + assert len(chunks) == 1 + assert isinstance(chunks[0], dict) + assert "Echo:" in chunks[0]["result"] and "Hello" in chunks[0]["result"] + + @pytest.mark.asyncio async def test_call_tool_math_operation(transport: CliCommunicationProtocol, mock_cli_script, python_executable): """Test calling a math tool with numeric arguments.""" diff --git a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py index 83afac2..a4d284f 100644 --- a/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py +++ b/plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py @@ -3,6 +3,7 @@ import aiohttp import json import asyncio +import codecs import re from urllib.parse import quote import base64 @@ -28,6 +29,11 @@ logger = logging.getLogger(__name__) + +class SseProtocolError(RuntimeError): + """The server violated the SSE wire format. Not a connection loss, so never retried.""" + + class SseCommunicationProtocol(CommunicationProtocol): """REQUIRED SSE communication protocol implementation for UTCP client. @@ -35,6 +41,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]] = {} @@ -224,97 +246,187 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, token = await self._handle_oauth2(tool_call_template.auth) request_headers["Authorization"] = f"Bearer {token}" - session = aiohttp.ClientSession() - # Always close the session, success or failure. The previous - # version only closed on the except path, leaking the session - # on the (typical) success path. - try: - method = "POST" if body_content is not None else "GET" - data = body_content if "application/json" not in request_headers.get("Content-Type", "") else None - json_data = body_content if "application/json" in request_headers.get("Content-Type", "") else None - - # SSE handshake must not follow redirects: the streaming - # response has to stay open for the lifetime of the tool - # call, which is incompatible with the per-hop validator's - # release semantics, and SSE redirects are pathological in - # practice. Reject 3xx outright so an attacker-controlled - # endpoint cannot redirect the handshake into an internal - # service (GHSA-9qhg-99ww-9mqc). - response = await session.request( - method, url, params=query_params, headers=request_headers, - auth=auth, cookies=cookies, json=json_data, data=data, - timeout=None, allow_redirects=False, - ) - if 300 <= response.status < 400: - response.release() - raise RuntimeError( - f"SSE endpoint at {url!r} returned a {response.status} " - f"redirect. Redirects are not followed during SSE " - f"handshakes; update the call template to point at " - f"the final URL directly." - ) - response.raise_for_status() - async for event in self._process_sse_stream(response, tool_call_template.event_type): - yield event - except Exception as e: - logger.error(f"Error establishing SSE connection to '{tool_call_template.name}': {e}") - raise - finally: - await session.close() + method = "POST" if body_content is not None else "GET" + content_type = request_headers.get("Content-Type", "") + data = body_content if "application/json" not in content_type else None + json_data = body_content if "application/json" in content_type else None + + reconnect = bool(tool_call_template.reconnect) + retry_delay_ms = tool_call_template.retry_timeout + last_event_id: Optional[str] = None + reconnect_attempts = 0 + provider_name = tool_call_template.name + + while True: + attempt_headers = dict(request_headers) + if last_event_id is not None: + # Let the server resume from where we left off (SSE spec). + attempt_headers["Last-Event-ID"] = last_event_id + + session = aiohttp.ClientSession() + try: + try: + # SSE handshake must not follow redirects: the streaming + # response has to stay open for the lifetime of the tool + # call, which is incompatible with the per-hop validator's + # release semantics, and SSE redirects are pathological in + # practice. Reject 3xx outright so an attacker-controlled + # endpoint cannot redirect the handshake into an internal + # service (GHSA-9qhg-99ww-9mqc). + # 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." + ) + response.raise_for_status() + 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): + if event.get("id") is not None: + last_event_id = event["id"] + if event.get("retry") is not None: + retry_delay_ms = event["retry"] + if "data" not in event: + continue + if tool_call_template.event_type and event.get("event") != tool_call_template.event_type: + continue + yield self._parse_event_data(event["data"]) + # The server ended the stream cleanly: the tool call is complete. + return + except (aiohttp.ClientError, asyncio.TimeoutError) as e: + reconnect_attempts += 1 + if not reconnect or reconnect_attempts > self.MAX_RECONNECT_ATTEMPTS: + logger.error(f"SSE connection to '{provider_name}' lost and not reconnecting: {e}") + raise + logger.warning( + f"SSE connection to '{provider_name}' lost ({e}); reconnecting in " + 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() - async def _process_sse_stream(self, response: aiohttp.ClientResponse, event_type=None): - """Process the SSE stream and yield events.""" + 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 = "" - try: - async for chunk in response.content.iter_any(): - buffer += chunk.decode('utf-8') - while '\n\n' in buffer: - event_string, buffer = buffer.split('\n\n', 1) - - # Ignore empty event strings - if not event_string.strip(): - continue - - # Process the event string - lines = event_string.split('\n') - current_event = {} - data_lines = [] - for line in lines: - if line.startswith(':'): - continue # It's a comment - - if ':' in line: - field, value = line.split(':', 1) - value = value.lstrip() - if field == 'event': - current_event['event'] = value - elif field == 'data': - data_lines.append(value) - elif field == 'id': - current_event['id'] = value - elif field == 'retry': - try: - current_event['retry'] = int(value) - except ValueError: - pass - - if not data_lines: - continue - - current_event['data'] = '\n'.join(data_lines) - - if event_type and current_event.get('event') != event_type: - continue + def flush(event_string: str): + if not event_string.strip(): + return None + current_event: Dict[str, Any] = {} + data_lines: List[str] = [] + for line in event_string.split('\n'): + if line.startswith(':'): + continue # comment / keep-alive + if ':' in line: + field, value = line.split(':', 1) + if value.startswith(' '): + value = value[1:] + else: + field, value = line, '' + if field == 'event': + current_event['event'] = value + elif field == 'data': + data_lines.append(value) + elif field == 'id': + current_event['id'] = value + elif field == 'retry': try: - yield json.loads(current_event['data']) - except json.JSONDecodeError: - yield current_event['data'] - except Exception as e: - logger.error(f"Error processing SSE stream: {e}") - raise - finally: - pass # Session is managed and closed by deregister_tool_provider + current_event['retry'] = int(value) + except ValueError: + pass + if data_lines: + current_event['data'] = '\n'.join(data_lines) + return current_event or None + + # 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" + ) + + # Flush a trailing event that was not terminated by a blank line. + buffer += normalise(decoder.decode(b"", final=True)) + if pending_cr: + buffer += "\n" + event = flush(buffer) + if event is not None: + yield event + + @staticmethod + def _parse_event_data(data: str) -> Any: + """Return the JSON-decoded payload when possible, otherwise the raw string.""" + try: + return json.loads(data) + except json.JSONDecodeError: + return data async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: """Handle OAuth2 client credentials flow, trying both body and diff --git a/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py b/plugins/communication_protocols/http/tests/test_sse_communication_protocol.py index 76cb41e..48f73c6 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,87 @@ async def token_header_auth_handler(request): async def error_handler(request): return web.Response(status=500, text="Internal Server Error") + +async def crlf_split_events_handler(request): + """One multi-line CRLF event whose CRLF is split across two writes.""" + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + await response.write(b"data: line1\r") + await asyncio.sleep(0.05) + await response.write(b"\ndata: line2\r\n\r\n") + return response + + +async def no_delimiter_events_handler(request): + """Streams data lines without ever sending the blank-line event delimiter.""" + 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')) + await asyncio.sleep(0.01) + request.transport.close() + return response + for event in SAMPLE_SSE_EVENTS[1:]: + await response.write(event.encode('utf-8')) + return response + + +async def slow_handshake_handler(request): + """Accepts the connection but does not send response headers for a long time.""" + await asyncio.sleep(5) + return web.Response(status=204) + + +async def huge_retry_events_handler(request): + """First connection asks for a very long retry delay, then drops.""" + state = request.app["huge_retry"] + state["connections"] += 1 + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + if state["connections"] == 1: + await response.write(b'id: 1\nretry: 100000\ndata: {"seq": 1}\n\n') + await asyncio.sleep(0.01) + request.transport.close() + return response + await response.write(b'id: 2\ndata: {"seq": 2}\n\n') + return response + +async def flaky_events_handler(request): + """Serves the first event then drops the TCP connection on the first connection + (or on every connection when ``always_drop`` is set). A reconnecting client is + served the remaining events and a clean end of stream.""" + state = request.app["flaky"] + state["connections"] += 1 + state["last_event_ids"].append(request.headers.get("Last-Event-ID")) + + response = web.StreamResponse(status=200, headers={'Content-Type': 'text/event-stream'}) + await response.prepare(request) + + if state["always_drop"] or state["connections"] == 1: + await response.write(SAMPLE_SSE_EVENTS[0].encode('utf-8')) + await asyncio.sleep(0.01) + request.transport.close() + return response + + for event in SAMPLE_SSE_EVENTS[1:]: + await response.write(event.encode('utf-8')) + return response + # --- Pytest Fixtures --- @pytest_asyncio.fixture @@ -121,6 +202,16 @@ def app(): app.router.add_post("/token", token_handler) app.router.add_post("/token_header_auth", token_header_auth_handler) app.router.add_get("/error", error_handler) + app.router.add_get("/flaky_events", flaky_events_handler) + app["flaky"] = {"connections": 0, "last_event_ids": [], "always_drop": False} + 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 +468,115 @@ 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 diff --git a/plugins/communication_protocols/mcp/pyproject.toml b/plugins/communication_protocols/mcp/pyproject.toml index 87461b7..ec66664 100644 --- a/plugins/communication_protocols/mcp/pyproject.toml +++ b/plugins/communication_protocols/mcp/pyproject.toml @@ -13,7 +13,7 @@ readme = "README.md" requires-python = ">=3.11" dependencies = [ "pydantic>=2.0", - "mcp>=1.12", + "mcp>=1.12,<2", "utcp>=1.1", "mcp-use>=1.3", "langchain>=0.3.27,<0.4.0", diff --git a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py index 7204b43..2c6c5d1 100644 --- a/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py +++ b/plugins/communication_protocols/mcp/src/utcp_mcp/mcp_communication_protocol.py @@ -379,7 +379,8 @@ async def _get_resource_server(self, resource_name: str, tool_call_template: Mcp async def call_tool_streaming(self, caller: 'UtcpClient', tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: """REQUIRED Streaming calls are not supported for MCP protocol, so we just call the tool and return the result as one item.""" - yield self.call_tool(caller, tool_name, tool_args, tool_call_template) + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result def _process_tool_result(self, result, tool_name: str) -> Any: self._log_info(f"Processing tool result for '{tool_name}', type: {type(result)}") diff --git a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py index d127791..ce9872d 100644 --- a/plugins/communication_protocols/mcp/tests/test_mcp_transport.py +++ b/plugins/communication_protocols/mcp/tests/test_mcp_transport.py @@ -244,3 +244,10 @@ async def test_resource_tool_without_registration(transport: McpCommunicationPro # Should still work and return content assert isinstance(result, dict) assert "contents" in result + + +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk(transport: McpCommunicationProtocol, mcp_manual: McpCallTemplate): + """Streaming mode should emit the awaited result as one chunk, not a coroutine.""" + chunks = [chunk async for chunk in transport.call_tool_streaming(None, f"{SERVER_NAME}.echo", {"message": "test"}, mcp_manual)] + assert chunks == [{"reply": "you said: test"}] diff --git a/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py b/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py index b2f08c3..4d9ff0a 100644 --- a/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py +++ b/plugins/communication_protocols/socket/src/utcp_socket/tcp_communication_protocol.py @@ -8,7 +8,7 @@ import socket import struct import sys -from typing import Dict, Any, List, Optional, Callable, Union +from typing import Dict, Any, List, Optional, Callable, Union, AsyncGenerator from utcp.interfaces.communication_protocol import CommunicationProtocol from utcp_socket.tcp_call_template import TCPProvider, TCPProviderSerializer @@ -404,10 +404,11 @@ async def deregister_manual(self, caller, manual_call_template: CallTemplate) -> raise ValueError("TCPTransport can only be used with TCPProvider") self._log_info(f"Deregistering TCP provider '{manual_call_template.name}' (no-op)") - async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate): - async def _generator(): - yield await self.call_tool(caller, tool_name, tool_args, tool_call_template) - return _generator() + async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: + """REQUIRED + Streaming variant: the TCP protocol does not natively stream, so the full result is yielded as a single chunk.""" + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> Any: """Call a TCP tool.""" diff --git a/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py b/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py index 89ae3e3..fa1f98e 100644 --- a/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py +++ b/plugins/communication_protocols/socket/src/utcp_socket/udp_communication_protocol.py @@ -7,7 +7,7 @@ import json import socket import traceback -from typing import Dict, Any, List, Optional, Callable, Union +from typing import Dict, Any, List, Optional, Callable, Union, AsyncGenerator from utcp.interfaces.communication_protocol import CommunicationProtocol from utcp_socket.udp_call_template import UDPProvider, UDPProviderSerializer @@ -331,7 +331,8 @@ async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], too # While this works, it's inconsistent with the other implementation in tcp_communication_protocol.py (lines 384-387) which properly uses async def with an inner generator. # For consistency and clarity, this should also use async def directly: # - # async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate): - # yield await self.call_tool(caller, tool_name, tool_args, tool_call_template) - async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate): - yield await self.call_tool(caller, tool_name, tool_args, tool_call_template) + async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]: + """REQUIRED + Streaming variant: the UDP protocol does not natively stream, so the full result is yielded as a single chunk.""" + result = await self.call_tool(caller, tool_name, tool_args, tool_call_template) + yield result diff --git a/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py b/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py index d359fd9..a82d14f 100644 --- a/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py +++ b/plugins/communication_protocols/socket/tests/test_tcp_communication_protocol.py @@ -177,4 +177,30 @@ async def test_register_manual_fallbacks_to_manual_template_tcp(): assert tool.tool_call_template.name == provider.name finally: server.close() - await server.wait_closed() \ No newline at end of file + await server.wait_closed() + + +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk_tcp(): + """Streaming mode should be an async generator that yields the full result once.""" + server, port, set_response = await start_tcp_server() + set_response({"echo": "hello"}) + + try: + provider = TCPProvider( + name="tcp-provider", + host="127.0.0.1", + port=port, + request_data_format="json", + response_byte_format="utf-8", + framing_strategy="stream", + timeout=2000 + ) + transport_client = TCPTransport() + expected = await transport_client.call_tool(None, "tcp-provider.tcp_tool", {"x": 1}, provider) + chunks = [chunk async for chunk in transport_client.call_tool_streaming(None, "tcp-provider.tcp_tool", {"x": 1}, provider)] + + assert chunks == [expected] + finally: + server.close() + await server.wait_closed() diff --git a/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py b/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py index d6a770c..26fd402 100644 --- a/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py +++ b/plugins/communication_protocols/socket/tests/test_udp_communication_protocol.py @@ -173,4 +173,29 @@ async def test_register_manual_fallbacks_to_manual_template_udp(): assert tool.tool_call_template.port == provider.port assert tool.tool_call_template.name == provider.name finally: - transport.close() \ No newline at end of file + transport.close() + + +@pytest.mark.asyncio +async def test_call_tool_streaming_yields_single_chunk_udp(): + """Streaming mode should be an async generator that yields the full result once.""" + transport, port, set_response = await start_udp_server() + set_response({"echo": "hello"}) + + try: + provider = UDPProvider( + name="udp-provider", + host="127.0.0.1", + port=port, + number_of_response_datagrams=1, + request_data_format="json", + response_byte_format="utf-8", + timeout=2000 + ) + transport_client = UDPTransport() + expected = await transport_client.call_tool(None, "udp-provider.udp_tool", {"x": 1}, provider) + chunks = [chunk async for chunk in transport_client.call_tool_streaming(None, "udp-provider.udp_tool", {"x": 1}, provider)] + + assert chunks == [expected] + finally: + transport.close()