diff --git a/robosystems_client/clients/README.md b/robosystems_client/clients/README.md index 97fceff..4e0e21b 100644 --- a/robosystems_client/clients/README.md +++ b/robosystems_client/clients/README.md @@ -114,9 +114,9 @@ result = clients.monitor_operation(operation_id, on_progress) print(result.status.value, result.result, result.error) ``` -`status` is an `OperationStatus`. `FAILED` can also mean the stream connection failed while the operation was still running, with the transport error in `error`. `get_operation_status(operation_id)` returns the server's view as a dict, and `cancel_operation(operation_id)` returns a bool. To set a limit, call `clients.operations.monitor_operation(operation_id, MonitorOptions(timeout=600))`, which raises `TimeoutError`. +`status` is an `OperationStatus`. A dropped stream reconnects, so a brief outage doesn't change the result. `FAILED` can still mean the reconnects ran out while the operation was running, with the transport error in `error`. `get_operation_status(operation_id)` returns the server's view as a dict, and `cancel_operation(operation_id)` returns a bool. To set a limit, call `clients.operations.monitor_operation(operation_id, MonitorOptions(timeout=600))`, which raises `TimeoutError`. -For custom event handling, use `create_sse_client()`, which copies the configured headers. An `SSEConfig` you build yourself needs credentials in `headers`, because the stream endpoint returns 401 without them. `connect` reports problems through the `error` event rather than raising. It retries with backoff only when opening the connection raises; a non-200 response or a drop mid-stream emits `error` and returns. +For custom event handling, use `create_sse_client()`, which copies the configured headers. An `SSEConfig` you build yourself needs credentials in `headers`, because the stream endpoint returns 401 without them. `connect` reports problems through events rather than raising. A dropped stream, a clean close before the operation finishes, or a `429` or `5xx` response reconnects with backoff at the original `from_sequence`, so a running operation's events during the gap are lost and a finished one replays. The attempts reset only when a real event arrives, and `max_retries_exceeded` fires when they run out. Any other non-200 response emits `error` and returns. `stream_end` ends the connection without a reconnect. ```python from robosystems_client.clients import EventType diff --git a/robosystems_client/clients/operator_client.py b/robosystems_client/clients/operator_client.py index 1fa552f..f6e8f43 100644 --- a/robosystems_client/clients/operator_client.py +++ b/robosystems_client/clients/operator_client.py @@ -246,7 +246,13 @@ def _rest_client(self) -> Client: def _sse_config(self) -> SSEConfig: """Stream config for one connect; headers carry the credential current now.""" - return SSEConfig(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + # No stream retries: `/status` polling already covers a lost stream, and + # retrying first would only delay it. + return SSEConfig( + base_url=self.base_url, + headers=resolve_auth_headers(self.config), + max_retries=0, + ) def execute_query( self, diff --git a/robosystems_client/clients/sse_client.py b/robosystems_client/clients/sse_client.py index 4e1eb10..9ddbc0c 100644 --- a/robosystems_client/clients/sse_client.py +++ b/robosystems_client/clients/sse_client.py @@ -74,6 +74,18 @@ class EventType(Enum): QUEUE_UPDATE = "queue_update" +# Sent on every open and while idle, so neither shows the stream got anywhere. +_NO_PROGRESS_EVENTS = {"connected", "keepalive"} + + +class _StreamInterrupted(Exception): + """The stream ended or failed before the server said it was finished.""" + + +def _is_retryable_status(status: int) -> bool: + return status == 429 or status >= 500 + + def event_error_message(err: Any) -> str: """Text of an ``error`` payload, whichever shape the stream emitted. @@ -104,6 +116,7 @@ def __init__(self, config: SSEConfig) -> None: self.closed = False self.listeners: Dict[str, Set[Callable]] = {} self._response = None + self._stream_ended = False def connect(self, operation_id: str, from_sequence: int = 0) -> None: """Connect to SSE stream for the given operation""" @@ -126,6 +139,10 @@ def connect(self, operation_id: str, from_sequence: int = 0) -> None: if self._response.status_code != 200: status = self._response.status_code body = self._response.read().decode("utf-8", errors="replace")[:500] + if _is_retryable_status(status): + raise _StreamInterrupted( + f"SSE connection failed: HTTP {status} {body}".strip() + ) self._context_manager.__exit__(None, None, None) self._context_manager = None self._response = None @@ -136,12 +153,15 @@ def connect(self, operation_id: str, from_sequence: int = 0) -> None: ) return - self.reconnect_attempts = 0 self.emit("connected", None) # Start processing events self._process_events() + # A deploy can close the stream cleanly while the operation runs on. + if not self.closed and not self._stream_ended: + raise _StreamInterrupted("SSE stream ended before the operation finished") + except Exception as error: if not self.closed: self._handle_error(error, operation_id, from_sequence) @@ -200,9 +220,9 @@ def _process_events(self) -> None: if event_buffer["data"] or event_buffer["event"]: self._dispatch_event(event_buffer) - except Exception as error: + except Exception: if not self.closed: - self.emit("error", error) + raise def _dispatch_event(self, event_buffer: Dict[str, Any]) -> None: """Dispatch a complete SSE event""" @@ -214,6 +234,11 @@ def _dispatch_event(self, event_buffer: Dict[str, Any]) -> None: event_type = event_buffer["event"] or "message" + if event_type == "stream_end": + self._stream_ended = True + if event_type not in _NO_PROGRESS_EVENTS: + self.reconnect_attempts = 0 + # Parse JSON data if possible parsed_data = data_str try: @@ -252,6 +277,8 @@ def _handle_error( if self.closed: return + self._release_transport() + if self.reconnect_attempts < self.config.max_retries: self.reconnect_attempts += 1 delay_ms = self.config.retry_delay * (2 ** (self.reconnect_attempts - 1)) @@ -308,10 +335,8 @@ def emit(self, event: str, data: Any) -> None: # Log error but don't stop other listeners print(f"Error in event listener for {event}: {e}") - def close(self): - """Close the SSE connection""" - self.closed = True - + def _release_transport(self) -> None: + """Drop the current connection, leaving the client able to reconnect.""" if hasattr(self, "_context_manager") and self._context_manager: try: self._context_manager.__exit__(None, None, None) @@ -324,6 +349,10 @@ def close(self): self.client.close() self.client = None + def close(self): + """Close the SSE connection""" + self.closed = True + self._release_transport() self.emit("closed", None) self.listeners.clear() @@ -348,6 +377,7 @@ def __init__(self, config: SSEConfig) -> None: self.closed = False self.listeners: Dict[str, Set[Callable]] = {} self._response = None + self._stream_ended = False async def connect(self, operation_id: str, from_sequence: int = 0) -> None: """Connect to SSE stream for the given operation (async)""" @@ -370,6 +400,10 @@ async def connect(self, operation_id: str, from_sequence: int = 0) -> None: if self._response.status_code != 200: status = self._response.status_code body = (await self._response.aread()).decode("utf-8", errors="replace")[:500] + if _is_retryable_status(status): + raise _StreamInterrupted( + f"SSE connection failed: HTTP {status} {body}".strip() + ) await self._context_manager.__aexit__(None, None, None) self._context_manager = None self._response = None @@ -380,12 +414,15 @@ async def connect(self, operation_id: str, from_sequence: int = 0) -> None: ) return - self.reconnect_attempts = 0 self.emit("connected", None) # Start processing events await self._process_events() + # A deploy can close the stream cleanly while the operation runs on. + if not self.closed and not self._stream_ended: + raise _StreamInterrupted("SSE stream ended before the operation finished") + except Exception as error: if not self.closed: await self._handle_error(error, operation_id, from_sequence) @@ -444,9 +481,9 @@ async def _process_events(self) -> None: if event_buffer["data"] or event_buffer["event"]: self._dispatch_event(event_buffer) - except Exception as error: + except Exception: if not self.closed: - self.emit("error", error) + raise def _dispatch_event(self, event_buffer: Dict[str, Any]): """Dispatch a complete SSE event (same as sync version)""" @@ -458,6 +495,11 @@ def _dispatch_event(self, event_buffer: Dict[str, Any]): event_type = event_buffer["event"] or "message" + if event_type == "stream_end": + self._stream_ended = True + if event_type not in _NO_PROGRESS_EVENTS: + self.reconnect_attempts = 0 + # Parse JSON data if possible parsed_data = data_str try: @@ -496,6 +538,8 @@ async def _handle_error( if self.closed: return + await self._release_transport() + if self.reconnect_attempts < self.config.max_retries: self.reconnect_attempts += 1 delay_ms = self.config.retry_delay * (2 ** (self.reconnect_attempts - 1)) @@ -512,6 +556,10 @@ async def _handle_error( await asyncio.sleep(delay_seconds) + # A close() during the backoff ends the reconnect. + if self.closed: + return + resume_from = 0 if self.last_event_id: try: @@ -546,10 +594,8 @@ def emit(self, event: str, data: Any) -> None: except Exception as e: print(f"Error in event listener for {event}: {e}") - async def close(self): - """Close the SSE connection (async)""" - self.closed = True - + async def _release_transport(self) -> None: + """Drop the current connection, leaving the client able to reconnect.""" if hasattr(self, "_context_manager") and self._context_manager: try: await self._context_manager.__aexit__(None, None, None) @@ -562,6 +608,10 @@ async def close(self): await self.client.aclose() self.client = None + async def close(self): + """Close the SSE connection (async)""" + self.closed = True + await self._release_transport() self.emit("closed", None) self.listeners.clear() diff --git a/tests/test_facade_wire_shapes.py b/tests/test_facade_wire_shapes.py index 6b72be0..226a79f 100644 --- a/tests/test_facade_wire_shapes.py +++ b/tests/test_facade_wire_shapes.py @@ -9,6 +9,8 @@ import json import re +from unittest.mock import patch + import pytest from robosystems_client.clients import OperatorRunError @@ -442,17 +444,48 @@ def test_cancelled_event_settles_the_monitor(self, httpx_mock): assert result.status == OperationStatus.CANCELLED - def test_stream_ending_without_a_verdict_raises(self, httpx_mock): + def test_stream_end_without_a_verdict_raises(self, httpx_mock): httpx_mock.add_response( url=_stream_url(), content=_sse( - ("operation_progress", {"message": "Working", "progress_percent": 5}) + ("operation_awaiting_input", {"prompt": "Confirm?"}), + ("stream_end", {"status": "awaiting_input"}), ), ) with pytest.raises(Exception, match="ended before a terminal event"): OperationClient(CONFIG).monitor_operation(OP) + def test_a_stream_that_keeps_closing_gives_up_after_its_retries(self, httpx_mock): + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("connected", {})), + is_reusable=True, + ) + + with patch("robosystems_client.clients.sse_client.time.sleep"): + result = OperationClient(CONFIG).monitor_operation(OP) + + assert result.status == OperationStatus.FAILED + assert "ended before the operation finished" in result.error + assert len(httpx_mock.get_requests()) == 6 + + def test_a_closed_stream_reconnects_and_reaches_the_verdict(self, httpx_mock): + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_progress", {"message": "Working"})), + ) + httpx_mock.add_response( + url=_stream_url(), + content=_sse(("operation_completed", {"result": {"ok": True}})), + ) + + with patch("robosystems_client.clients.sse_client.time.sleep"): + result = OperationClient(CONFIG).monitor_operation(OP) + + assert result.status == OperationStatus.COMPLETED + assert result.result == {"ok": True} + def test_progress_reads_progress_percent(self, httpx_mock): httpx_mock.add_response( url=_stream_url(), diff --git a/tests/test_sse_client.py b/tests/test_sse_client.py index b97074a..2b6db38 100644 --- a/tests/test_sse_client.py +++ b/tests/test_sse_client.py @@ -5,9 +5,14 @@ tests already exist in extensions/tests/test_unit.py. """ +import asyncio +import json + +import httpx import pytest from unittest.mock import Mock, patch, MagicMock from robosystems_client.clients.sse_client import ( + AsyncSSEClient, SSEClient, SSEConfig, ) @@ -536,3 +541,167 @@ def test_emit_to_nonexistent_event_is_noop(self, sse_config): """Test emitting to an event with no listeners doesn't raise.""" client = SSEClient(sse_config) client.emit("nonexistent", "data") # Should not raise + + +def _sse(*events): + """Encode (name, payload) pairs as a raw SSE body.""" + return b"".join( + f"event: {name}\ndata: {json.dumps(data)}\n\n".encode() for name, data in events + ) + + +class _DroppingStream(httpx.SyncByteStream, httpx.AsyncByteStream): + """Yields a body, then fails the way a dropped connection does.""" + + def __init__(self, body): + self.body = body + + def __iter__(self): + yield self.body + raise httpx.RemoteProtocolError("peer closed connection") + + async def __aiter__(self): + yield self.body + raise httpx.RemoteProtocolError("peer closed connection") + + +COMPLETED = _sse(("connected", {}), ("operation_completed", {"result": {"ok": True}})) + + +def _responses(*scripted): + """A handler replaying scripted responses, recording each request.""" + remaining = list(scripted) + seen = [] + + def handler(request): + seen.append(request) + return remaining.pop(0) if len(remaining) > 1 else remaining[0] + + return handler, seen + + +@pytest.fixture +def fast_config(): + return SSEConfig(base_url="https://api.test", max_retries=3, retry_delay=1) + + +def _run_sync(config, handler): + real_client = httpx.Client + with patch( + "robosystems_client.clients.sse_client.httpx.Client", + side_effect=lambda **kw: real_client(transport=httpx.MockTransport(handler), **kw), + ): + client = SSEClient(config) + events = {"completed": [], "errors": [], "exhausted": [], "reconnecting": []} + client.on("operation_completed", events["completed"].append) + client.on("error", events["errors"].append) + client.on("max_retries_exceeded", events["exhausted"].append) + client.on("reconnecting", events["reconnecting"].append) + client.connect("op_1") + return events + + +class TestStreamInterruption: + """A broken stream says nothing about the operation behind it.""" + + def test_a_dropped_stream_reconnects_and_reaches_the_result(self, fast_config): + handler, seen = _responses( + httpx.Response(200, stream=_DroppingStream(_sse(("operation_progress", {})))), + httpx.Response(200, content=COMPLETED), + ) + + events = _run_sync(fast_config, handler) + + assert len(seen) == 2 + assert events["completed"] == [{"result": {"ok": True}}] + assert events["errors"] == [] + + def test_a_clean_close_before_the_end_reconnects(self, fast_config): + handler, seen = _responses( + httpx.Response(200, content=_sse(("operation_progress", {}))), + httpx.Response(200, content=COMPLETED), + ) + + events = _run_sync(fast_config, handler) + + assert len(seen) == 2 + assert len(events["completed"]) == 1 + + def test_a_503_while_reconnecting_is_retried(self, fast_config): + handler, seen = _responses( + httpx.Response(200, stream=_DroppingStream(_sse(("operation_progress", {})))), + httpx.Response(503, content=b"unavailable"), + httpx.Response(200, content=COMPLETED), + ) + + events = _run_sync(fast_config, handler) + + assert len(seen) == 3 + assert len(events["completed"]) == 1 + assert events["errors"] == [] + + def test_a_refused_credential_is_not_retried(self, fast_config): + handler, seen = _responses(httpx.Response(401, content=b"denied")) + + events = _run_sync(fast_config, handler) + + assert len(seen) == 1 + assert "401" in str(events["errors"][0]) + assert events["reconnecting"] == [] + + def test_stream_end_stops_without_reconnecting(self, fast_config): + body = _sse(("operation_awaiting_input", {}), ("stream_end", {"status": "paused"})) + handler, seen = _responses(httpx.Response(200, content=body)) + + events = _run_sync(fast_config, handler) + + assert len(seen) == 1 + assert events["reconnecting"] == [] + + def test_reopening_without_progress_does_not_reset_the_budget(self, fast_config): + handler, seen = _responses( + httpx.Response(200, stream=_DroppingStream(_sse(("connected", {})))) + ) + + events = _run_sync(fast_config, handler) + + assert len(seen) == fast_config.max_retries + 1 + assert len(events["exhausted"]) == 1 + + def test_progress_resets_the_budget(self, fast_config): + dropping = [ + httpx.Response(200, stream=_DroppingStream(_sse(("operation_progress", {})))) + for _ in range(fast_config.max_retries + 2) + ] + handler, seen = _responses(*dropping, httpx.Response(200, content=COMPLETED)) + + events = _run_sync(fast_config, handler) + + assert len(seen) == fast_config.max_retries + 3 + assert events["exhausted"] == [] + assert len(events["completed"]) == 1 + + def test_the_async_client_reconnects_too(self, fast_config): + handler, seen = _responses( + httpx.Response(200, stream=_DroppingStream(_sse(("operation_progress", {})))), + httpx.Response(503, content=b"unavailable"), + httpx.Response(200, content=COMPLETED), + ) + real_client = httpx.AsyncClient + completed = [] + + async def run(): + with patch( + "robosystems_client.clients.sse_client.httpx.AsyncClient", + side_effect=lambda **kw: real_client( + transport=httpx.MockTransport(handler), **kw + ), + ): + client = AsyncSSEClient(fast_config) + client.on("operation_completed", completed.append) + await client.connect("op_1") + + asyncio.run(run()) + + assert len(seen) == 3 + assert completed == [{"result": {"ok": True}}]