Skip to content

Commit b824c07

Browse files
committed
Merge origin/dev into fix/mcp-quiet-child-stderr
Both #100 and this branch appended tests to test_mcp_transport.py; keep both.
2 parents 31a1dfb + 228d0a3 commit b824c07

10 files changed

Lines changed: 506 additions & 104 deletions

‎plugins/communication_protocols/cli/src/utcp_cli/cli_communication_protocol.py‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -972,9 +972,19 @@ async def call_tool(self, caller, tool_name: str, tool_args: Dict[str, Any], too
972972

973973
async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, Any], tool_call_template: CallTemplate) -> AsyncGenerator[Any, None]:
974974
"""REQUIRED
975-
Streaming calls are not supported for the CLI protocol.
975+
Execute a tool call through the CLI transport streamingly.
976976
977-
Raises:
978-
NotImplementedError: Always, as this functionality is not supported.
977+
The CLI protocol does not natively support streaming, so the command is
978+
executed to completion and the full result is yielded as a single chunk.
979+
980+
Args:
981+
caller: The UTCP client that is calling this method.
982+
tool_name: Name of the tool to call.
983+
tool_args: Dictionary of arguments to pass to the tool.
984+
tool_call_template: Call template of the tool to call.
985+
986+
Yields:
987+
The complete tool result as a single item.
979988
"""
980-
raise NotImplementedError("Streaming is not supported by the CLI communication protocol.")
989+
result = await self.call_tool(caller, tool_name, tool_args, tool_call_template)
990+
yield result

‎plugins/communication_protocols/cli/tests/test_cli_communication_protocol.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,22 @@ async def test_call_tool_json_output(transport: CliCommunicationProtocol, mock_c
276276
assert "Echo:" in result["result"] and "Hello" in result["result"]
277277

278278

279+
@pytest.mark.asyncio
280+
async def test_call_tool_streaming_yields_single_chunk(transport: CliCommunicationProtocol, mock_cli_script, python_executable):
281+
"""Streaming mode should emit the full result as one chunk instead of failing."""
282+
call_template = CliCallTemplate(
283+
commands=[
284+
{"command": f"{python_executable} {mock_cli_script} --message UTCP_ARG_message_UTCP_END"}
285+
]
286+
)
287+
288+
chunks = [chunk async for chunk in transport.call_tool_streaming(None, "echo", {"message": "Hello World"}, call_template)]
289+
290+
assert len(chunks) == 1
291+
assert isinstance(chunks[0], dict)
292+
assert "Echo:" in chunks[0]["result"] and "Hello" in chunks[0]["result"]
293+
294+
279295
@pytest.mark.asyncio
280296
async def test_call_tool_math_operation(transport: CliCommunicationProtocol, mock_cli_script, python_executable):
281297
"""Test calling a math tool with numeric arguments."""

‎plugins/communication_protocols/http/src/utcp_http/sse_communication_protocol.py‎

Lines changed: 199 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import aiohttp
44
import json
55
import asyncio
6+
import codecs
67
import re
78
from urllib.parse import quote
89
import base64
@@ -28,13 +29,34 @@
2829

2930
logger = logging.getLogger(__name__)
3031

32+
33+
class SseProtocolError(RuntimeError):
34+
"""The server violated the SSE wire format. Not a connection loss, so never retried."""
35+
36+
3137
class SseCommunicationProtocol(CommunicationProtocol):
3238
"""REQUIRED
3339
SSE communication protocol implementation for UTCP client.
3440
3541
Handles Server-Sent Events based tool providers with streaming capabilities.
3642
"""
3743

44+
# Upper bound on reconnection attempts for a single tool call when the
45+
# established stream drops and the call template has ``reconnect`` enabled.
46+
# Keeps a tool call bounded even if the server keeps dropping the connection.
47+
MAX_RECONNECT_ATTEMPTS: int = 5
48+
# Cap on the delay before a reconnect, whatever ``retry_timeout`` or a
49+
# server-sent ``retry:`` field asks for. Together with MAX_RECONNECT_ATTEMPTS
50+
# this bounds the total time a call can spend waiting to reconnect.
51+
MAX_RECONNECT_DELAY_MS: int = 60_000
52+
# Time allowed for the SSE handshake, i.e. until response headers arrive.
53+
# Reading the body is unbounded: an SSE stream may legitimately stay quiet.
54+
HANDSHAKE_TIMEOUT_SECONDS: float = 30.0
55+
# Largest partial event the parser buffers before declaring the stream
56+
# malformed. Guards against a server that streams data without ever sending
57+
# the blank-line event delimiter.
58+
MAX_EVENT_BUFFER_CHARS: int = 16 * 1024 * 1024
59+
3860
def __init__(self, logger: Optional[Callable[[str], None]] = None):
3961
self._oauth_tokens: Dict[str, Dict[str, Any]] = {}
4062

@@ -224,97 +246,187 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str,
224246
token = await self._handle_oauth2(tool_call_template.auth)
225247
request_headers["Authorization"] = f"Bearer {token}"
226248

227-
session = aiohttp.ClientSession()
228-
# Always close the session, success or failure. The previous
229-
# version only closed on the except path, leaking the session
230-
# on the (typical) success path.
231-
try:
232-
method = "POST" if body_content is not None else "GET"
233-
data = body_content if "application/json" not in request_headers.get("Content-Type", "") else None
234-
json_data = body_content if "application/json" in request_headers.get("Content-Type", "") else None
235-
236-
# SSE handshake must not follow redirects: the streaming
237-
# response has to stay open for the lifetime of the tool
238-
# call, which is incompatible with the per-hop validator's
239-
# release semantics, and SSE redirects are pathological in
240-
# practice. Reject 3xx outright so an attacker-controlled
241-
# endpoint cannot redirect the handshake into an internal
242-
# service (GHSA-9qhg-99ww-9mqc).
243-
response = await session.request(
244-
method, url, params=query_params, headers=request_headers,
245-
auth=auth, cookies=cookies, json=json_data, data=data,
246-
timeout=None, allow_redirects=False,
247-
)
248-
if 300 <= response.status < 400:
249-
response.release()
250-
raise RuntimeError(
251-
f"SSE endpoint at {url!r} returned a {response.status} "
252-
f"redirect. Redirects are not followed during SSE "
253-
f"handshakes; update the call template to point at "
254-
f"the final URL directly."
255-
)
256-
response.raise_for_status()
257-
async for event in self._process_sse_stream(response, tool_call_template.event_type):
258-
yield event
259-
except Exception as e:
260-
logger.error(f"Error establishing SSE connection to '{tool_call_template.name}': {e}")
261-
raise
262-
finally:
263-
await session.close()
249+
method = "POST" if body_content is not None else "GET"
250+
content_type = request_headers.get("Content-Type", "")
251+
data = body_content if "application/json" not in content_type else None
252+
json_data = body_content if "application/json" in content_type else None
253+
254+
reconnect = bool(tool_call_template.reconnect)
255+
retry_delay_ms = tool_call_template.retry_timeout
256+
last_event_id: Optional[str] = None
257+
reconnect_attempts = 0
258+
provider_name = tool_call_template.name
259+
260+
while True:
261+
attempt_headers = dict(request_headers)
262+
if last_event_id is not None:
263+
# Let the server resume from where we left off (SSE spec).
264+
attempt_headers["Last-Event-ID"] = last_event_id
265+
266+
session = aiohttp.ClientSession()
267+
try:
268+
try:
269+
# SSE handshake must not follow redirects: the streaming
270+
# response has to stay open for the lifetime of the tool
271+
# call, which is incompatible with the per-hop validator's
272+
# release semantics, and SSE redirects are pathological in
273+
# practice. Reject 3xx outright so an attacker-controlled
274+
# endpoint cannot redirect the handshake into an internal
275+
# service (GHSA-9qhg-99ww-9mqc).
276+
# Bound the handshake only (until response headers arrive);
277+
# the body read stays unbounded because a stream may be quiet.
278+
response = await asyncio.wait_for(
279+
session.request(
280+
method, url, params=query_params, headers=attempt_headers,
281+
auth=auth, cookies=cookies, json=json_data, data=data,
282+
timeout=None, allow_redirects=False,
283+
),
284+
timeout=self.HANDSHAKE_TIMEOUT_SECONDS,
285+
)
286+
if 300 <= response.status < 400:
287+
response.release()
288+
raise RuntimeError(
289+
f"SSE endpoint at {url!r} returned a {response.status} "
290+
f"redirect. Redirects are not followed during SSE "
291+
f"handshakes; update the call template to point at "
292+
f"the final URL directly."
293+
)
294+
response.raise_for_status()
295+
except Exception as e:
296+
if reconnect_attempts == 0:
297+
# The initial handshake failing (refused, timed out, non-2xx) is a
298+
# definitive answer about the endpoint: fail fast, no retry.
299+
logger.error(f"Error establishing SSE connection to '{provider_name}': {e}")
300+
raise
301+
# A reconnect handshake failing is part of the outage we are riding
302+
# out (the server may still be restarting): count it and try again.
303+
reconnect_attempts += 1
304+
if reconnect_attempts > self.MAX_RECONNECT_ATTEMPTS:
305+
logger.error(f"SSE reconnect to '{provider_name}' failed and attempts are exhausted: {e}")
306+
raise
307+
delay_ms = min(retry_delay_ms, self.MAX_RECONNECT_DELAY_MS)
308+
logger.warning(
309+
f"SSE reconnect to '{provider_name}' failed ({e}); retrying in {delay_ms} ms "
310+
f"(attempt {reconnect_attempts}/{self.MAX_RECONNECT_ATTEMPTS})"
311+
)
312+
await asyncio.sleep(delay_ms / 1000)
313+
continue
314+
315+
try:
316+
async for event in self._iter_sse_events(response):
317+
if event.get("id") is not None:
318+
last_event_id = event["id"]
319+
if event.get("retry") is not None:
320+
retry_delay_ms = event["retry"]
321+
if "data" not in event:
322+
continue
323+
if tool_call_template.event_type and event.get("event") != tool_call_template.event_type:
324+
continue
325+
yield self._parse_event_data(event["data"])
326+
# The server ended the stream cleanly: the tool call is complete.
327+
return
328+
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
329+
reconnect_attempts += 1
330+
if not reconnect or reconnect_attempts > self.MAX_RECONNECT_ATTEMPTS:
331+
logger.error(f"SSE connection to '{provider_name}' lost and not reconnecting: {e}")
332+
raise
333+
logger.warning(
334+
f"SSE connection to '{provider_name}' lost ({e}); reconnecting in "
335+
f"{min(retry_delay_ms, self.MAX_RECONNECT_DELAY_MS)} ms "
336+
f"(attempt {reconnect_attempts}/{self.MAX_RECONNECT_ATTEMPTS})"
337+
)
338+
finally:
339+
# Always release the connection, whether the stream completed, failed,
340+
# or the consumer stopped iterating early.
341+
if not session.closed:
342+
await session.close()
264343

265-
async def _process_sse_stream(self, response: aiohttp.ClientResponse, event_type=None):
266-
"""Process the SSE stream and yield events."""
344+
await asyncio.sleep(min(retry_delay_ms, self.MAX_RECONNECT_DELAY_MS) / 1000)
345+
346+
async def _iter_sse_events(self, response: aiohttp.ClientResponse) -> AsyncIterator[Dict[str, Any]]:
347+
"""Parse the SSE wire format and yield one dict per event block.
348+
349+
Each dict may contain ``event``, ``id``, ``retry`` (int) and ``data`` (str, with
350+
multi-line data joined by newlines). Blocks that only carry ``id``/``retry``
351+
are yielded too (without ``data``) so the caller can track reconnection
352+
state; comment-only blocks are skipped.
353+
"""
267354
buffer = ""
268-
try:
269-
async for chunk in response.content.iter_any():
270-
buffer += chunk.decode('utf-8')
271-
while '\n\n' in buffer:
272-
event_string, buffer = buffer.split('\n\n', 1)
273-
274-
# Ignore empty event strings
275-
if not event_string.strip():
276-
continue
277-
278-
# Process the event string
279-
lines = event_string.split('\n')
280-
current_event = {}
281-
data_lines = []
282-
for line in lines:
283-
if line.startswith(':'):
284-
continue # It's a comment
285-
286-
if ':' in line:
287-
field, value = line.split(':', 1)
288-
value = value.lstrip()
289-
if field == 'event':
290-
current_event['event'] = value
291-
elif field == 'data':
292-
data_lines.append(value)
293-
elif field == 'id':
294-
current_event['id'] = value
295-
elif field == 'retry':
296-
try:
297-
current_event['retry'] = int(value)
298-
except ValueError:
299-
pass
300-
301-
if not data_lines:
302-
continue
303-
304-
current_event['data'] = '\n'.join(data_lines)
305-
306-
if event_type and current_event.get('event') != event_type:
307-
continue
308355

356+
def flush(event_string: str):
357+
if not event_string.strip():
358+
return None
359+
current_event: Dict[str, Any] = {}
360+
data_lines: List[str] = []
361+
for line in event_string.split('\n'):
362+
if line.startswith(':'):
363+
continue # comment / keep-alive
364+
if ':' in line:
365+
field, value = line.split(':', 1)
366+
if value.startswith(' '):
367+
value = value[1:]
368+
else:
369+
field, value = line, ''
370+
if field == 'event':
371+
current_event['event'] = value
372+
elif field == 'data':
373+
data_lines.append(value)
374+
elif field == 'id':
375+
current_event['id'] = value
376+
elif field == 'retry':
309377
try:
310-
yield json.loads(current_event['data'])
311-
except json.JSONDecodeError:
312-
yield current_event['data']
313-
except Exception as e:
314-
logger.error(f"Error processing SSE stream: {e}")
315-
raise
316-
finally:
317-
pass # Session is managed and closed by deregister_tool_provider
378+
current_event['retry'] = int(value)
379+
except ValueError:
380+
pass
381+
if data_lines:
382+
current_event['data'] = '\n'.join(data_lines)
383+
return current_event or None
384+
385+
# Incremental decoding: a multi-byte UTF-8 character may straddle two chunks.
386+
decoder = codecs.getincrementaldecoder("utf-8")()
387+
# A "\r" that ended the previous chunk is held back until the next chunk
388+
# shows whether a "\n" follows; otherwise a CRLF split across two reads
389+
# would become two LFs and dispatch an event early.
390+
pending_cr = False
391+
392+
def normalise(text: str) -> str:
393+
nonlocal pending_cr
394+
if pending_cr:
395+
text = "\r" + text
396+
pending_cr = False
397+
if text.endswith("\r"):
398+
text = text[:-1]
399+
pending_cr = True
400+
# Normalise CRLF / CR line endings so the event delimiter is always "\n\n".
401+
return text.replace("\r\n", "\n").replace("\r", "\n")
402+
403+
async for chunk in response.content.iter_any():
404+
buffer += normalise(decoder.decode(chunk))
405+
while "\n\n" in buffer:
406+
event_string, buffer = buffer.split("\n\n", 1)
407+
event = flush(event_string)
408+
if event is not None:
409+
yield event
410+
if len(buffer) > self.MAX_EVENT_BUFFER_CHARS:
411+
raise SseProtocolError(
412+
f"SSE event exceeded {self.MAX_EVENT_BUFFER_CHARS} characters without a blank-line delimiter"
413+
)
414+
415+
# Flush a trailing event that was not terminated by a blank line.
416+
buffer += normalise(decoder.decode(b"", final=True))
417+
if pending_cr:
418+
buffer += "\n"
419+
event = flush(buffer)
420+
if event is not None:
421+
yield event
422+
423+
@staticmethod
424+
def _parse_event_data(data: str) -> Any:
425+
"""Return the JSON-decoded payload when possible, otherwise the raw string."""
426+
try:
427+
return json.loads(data)
428+
except json.JSONDecodeError:
429+
return data
318430

319431
async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str:
320432
"""Handle OAuth2 client credentials flow, trying both body and

0 commit comments

Comments
 (0)