Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions robosystems_client/clients/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 7 additions & 1 deletion robosystems_client/clients/operator_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
78 changes: 64 additions & 14 deletions robosystems_client/clients/sse_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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"""
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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"""
Expand All @@ -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:
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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)
Expand All @@ -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()

Expand All @@ -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)"""
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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)"""
Expand All @@ -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:
Expand Down Expand Up @@ -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))
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand All @@ -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()

Expand Down
37 changes: 35 additions & 2 deletions tests/test_facade_wire_shapes.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
import json
import re

from unittest.mock import patch

import pytest

from robosystems_client.clients import OperatorRunError
Expand Down Expand Up @@ -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(),
Expand Down
Loading
Loading