|
3 | 3 | import aiohttp |
4 | 4 | import json |
5 | 5 | import asyncio |
| 6 | +import codecs |
6 | 7 | import re |
7 | 8 | from urllib.parse import quote |
8 | 9 | import base64 |
|
28 | 29 |
|
29 | 30 | logger = logging.getLogger(__name__) |
30 | 31 |
|
| 32 | + |
| 33 | +class SseProtocolError(RuntimeError): |
| 34 | + """The server violated the SSE wire format. Not a connection loss, so never retried.""" |
| 35 | + |
| 36 | + |
31 | 37 | class SseCommunicationProtocol(CommunicationProtocol): |
32 | 38 | """REQUIRED |
33 | 39 | SSE communication protocol implementation for UTCP client. |
34 | 40 | |
35 | 41 | Handles Server-Sent Events based tool providers with streaming capabilities. |
36 | 42 | """ |
37 | 43 |
|
| 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 | + |
38 | 60 | def __init__(self, logger: Optional[Callable[[str], None]] = None): |
39 | 61 | self._oauth_tokens: Dict[str, Dict[str, Any]] = {} |
40 | 62 |
|
@@ -224,97 +246,187 @@ async def call_tool_streaming(self, caller, tool_name: str, tool_args: Dict[str, |
224 | 246 | token = await self._handle_oauth2(tool_call_template.auth) |
225 | 247 | request_headers["Authorization"] = f"Bearer {token}" |
226 | 248 |
|
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() |
264 | 343 |
|
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 | + """ |
267 | 354 | 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 |
308 | 355 |
|
| 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': |
309 | 377 | 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 |
318 | 430 |
|
319 | 431 | async def _handle_oauth2(self, auth_details: OAuth2Auth) -> str: |
320 | 432 | """Handle OAuth2 client credentials flow, trying both body and |
|
0 commit comments