From 7e97327cbcad4b2e50fb62c4619d02a49c9f7639 Mon Sep 17 00:00:00 2001 From: "Joseph T. French" Date: Sat, 29 Aug 2026 14:04:01 -0500 Subject: [PATCH] fix(clients): honour token_provider on the SSE-backed clients and never hang on a dead stream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The facade accepted a `token_provider`, but only the GraphQL facades used it: OperatorClient, OperationClient and QueryClient captured `token` / `headers` at construction for every REST call and every SSE connect. The backend revokes the previous JWT on each session refresh, so a rotating credential left those three clients dead after the first rotation — and a stream that then failed to open emitted the transport Exception, which the dict-only `on_error` handlers choked on inside `emit`, leaving the wait loop spinning forever. Python twin of typescript-client #206. - token_utils.resolve_auth_headers builds request/stream headers per call: static headers unchanged without a provider, the provider's current credential (routed by shape) replacing any stale auth header with one. apply_auth_header is the single routing rule; auth_integration aliases it. - Operator/Query/Operation clients build a fresh Client and SSEConfig from it per call and per connect. - OperatorClient splits run errors from transport errors and, when the stream gives no verdict, follows the run over /v1/operations/{id}/status (poll_interval; definitive 4xx ends it, transient failures retried). error_details passes through on every result path. - Operator/Query handlers accept Exception payloads, register error/max_retries_exceeded, take the cancellation payload, and raise when the stream ends without a result instead of spinning or returning None. --- robosystems_client/clients/README.md | 27 ++ .../clients/auth_integration.py | 18 +- .../clients/operation_client.py | 23 +- robosystems_client/clients/operator_client.py | 284 ++++++++++++----- robosystems_client/clients/query_client.py | 105 ++++--- robosystems_client/clients/sse_client.py | 14 + robosystems_client/clients/token_utils.py | 55 +++- tests/test_auth_header_resolution.py | 150 +++++++++ tests/test_operator_client_ops.py | 293 ++++++++++++++++++ tests/test_query_client_sse.py | 126 ++++++++ 10 files changed, 956 insertions(+), 139 deletions(-) create mode 100644 tests/test_auth_header_resolution.py create mode 100644 tests/test_operator_client_ops.py create mode 100644 tests/test_query_client_sse.py diff --git a/robosystems_client/clients/README.md b/robosystems_client/clients/README.md index 3f4f560..2d3d97e 100644 --- a/robosystems_client/clients/README.md +++ b/robosystems_client/clients/README.md @@ -115,6 +115,33 @@ from robosystems_client.clients import TokenClients extensions = TokenClients(token="your-jwt-token", base_url="https://api.robosystems.ai") ``` +### Rotating Credentials (`token_provider`) + +Short-lived JWTs rotate, and the backend revokes the previous token on every +session refresh — a credential captured when the facade was built stops +working the moment the session rotates. Pass a zero-arg callable instead and +every client resolves the credential fresh: the GraphQL facades (`ledger` / +`investor` / `library`) on each request, and the SSE-backed clients +(`operator` / `operations` / `query`) on each REST call _and_ each stream +connect. It wins over any static `token` or auth header in `headers`, and is +routed by shape (`rfs…` keys as `X-API-Key`, anything else as a Bearer JWT). + +```python +from robosystems_client.clients import RoboSystemsClients, RoboSystemsClientConfig + +extensions = RoboSystemsClients( + RoboSystemsClientConfig( + base_url="https://api.robosystems.ai", + token_provider=lambda: load_current_jwt(), # or `lambda: manager.token` + ) +) +``` + +`OperatorClient` also follows a queued run over `/v1/operations/{id}/status` +whenever its stream gives no verdict — it could not open, its reconnects ran +out, or it ended before a terminal event — so a run that is already executing +is never lost. `OperatorOptions.poll_interval` (seconds) tunes the interval. + ### Environment-Specific Configurations ```python diff --git a/robosystems_client/clients/auth_integration.py b/robosystems_client/clients/auth_integration.py index a4c78b4..1fffed7 100644 --- a/robosystems_client/clients/auth_integration.py +++ b/robosystems_client/clients/auth_integration.py @@ -6,26 +6,16 @@ from typing import Dict, Any from ..client import Client, AuthenticatedClient from .facade import RoboSystemsClients, RoboSystemsClientConfig +from .token_utils import apply_auth_header def _apply_auth_header(headers: Dict[str, str], credential: str) -> None: """Set the correct auth header for a credential, routed by shape. - The backend accepts two credential formats, and they go in DIFFERENT - headers — not interchangeable (see ``graphql/client.py``): - - - Long-lived API keys (``rfs…`` prefix) → ``X-API-Key``. Validated - against the api_keys table. - - Short-lived JWTs → ``Authorization: Bearer …``. Validated by the - JWT middleware. - - Sending a JWT as ``X-API-Key`` (or an API key as Bearer) both fail - with 401 "Invalid API key" — so exactly one header is set, never both. + Thin alias of :func:`token_utils.apply_auth_header`, which the per-call + header resolver in the SSE-backed clients shares — one routing rule. """ - if credential.startswith("rfs"): - headers["X-API-Key"] = credential - else: - headers["Authorization"] = f"Bearer {credential}" + apply_auth_header(headers, credential) def _build_sdk_client(base_url: str, credential: str, headers: Dict[str, str]): diff --git a/robosystems_client/clients/operation_client.py b/robosystems_client/clients/operation_client.py index fac8541..da05656 100644 --- a/robosystems_client/clients/operation_client.py +++ b/robosystems_client/clients/operation_client.py @@ -10,6 +10,7 @@ from enum import Enum from .sse_client import SSEClient, AsyncSSEClient, SSEConfig, EventType +from .token_utils import resolve_auth_headers logger = logging.getLogger(__name__) @@ -105,8 +106,11 @@ def monitor_operation( error = None # Set up SSE connection with event replay from the beginning - # This handles the race condition where the operation may have already completed - sse_config = SSEConfig(base_url=self.base_url, headers=self.headers) + # This handles the race condition where the operation may have already completed. + # Headers are resolved per connect so a rotated JWT reaches the stream. + sse_config = SSEConfig( + base_url=self.base_url, headers=resolve_auth_headers(self.config) + ) sse_client = SSEClient(sse_config) def on_operation_started(data): @@ -216,8 +220,8 @@ def get_operation_status(self, operation_id: str) -> Dict[str, Any]: ) from ..client import Client - # Use regular Client with headers instead of AuthenticatedClient - client = Client(base_url=self.base_url, headers=self.headers) + # Plain Client with the headers current now (`token_provider` wins). + client = Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) try: # Auth travels in self.headers (X-API-Key / Authorization). The generated # function takes no `token` kwarg — passing one raised TypeError, which @@ -246,8 +250,8 @@ def cancel_operation(self, operation_id: str) -> bool: from ..api.operations.cancel_operation import sync_detailed as cancel_operation from ..client import Client - # Use regular Client with headers instead of AuthenticatedClient - client = Client(base_url=self.base_url, headers=self.headers) + # Plain Client with the headers current now (`token_provider` wins). + client = Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) try: # See get_operation_status: no `token` kwarg on the generated function. response = cancel_operation(operation_id=operation_id, client=client) @@ -308,8 +312,11 @@ async def monitor_operation( completed = False error = None - # Set up SSE connection - sse_config = SSEConfig(base_url=self.base_url, headers=self.headers) + # Set up SSE connection; headers resolved per connect so a rotated JWT + # reaches the stream. + sse_config = SSEConfig( + base_url=self.base_url, headers=resolve_auth_headers(self.config) + ) sse_client = AsyncSSEClient(sse_config) def on_operation_started(data): diff --git a/robosystems_client/clients/operator_client.py b/robosystems_client/clients/operator_client.py index 06216e0..3bc35d5 100644 --- a/robosystems_client/clients/operator_client.py +++ b/robosystems_client/clients/operator_client.py @@ -3,17 +3,32 @@ Provides intelligent operator execution with automatic strategy selection. """ +import time from dataclasses import dataclass -from typing import Dict, Any, Optional, Callable +from typing import Dict, Any, Optional, Callable, cast from datetime import datetime from ..api.operator.auto_select_operator import sync_detailed as auto_select_operator from ..api.operator.execute_specific_operator import ( sync_detailed as execute_specific_operator, ) +from ..api.operations.get_operation_status import ( + sync_detailed as get_operation_status, +) +from ..client import Client from ..models.operator_request import OperatorRequest from ..models.operator_message import OperatorMessage -from .sse_client import SSEClient, SSEConfig, EventType +from ..types import UNSET +from .sse_client import SSEClient, SSEConfig, EventType, event_error_message +from .token_utils import resolve_auth_headers, resolve_config_token + +# Seconds between `/status` polls while following a run whose stream gave +# no verdict. Mirrors the TypeScript client's `pollIntervalMs` default. +DEFAULT_POLL_INTERVAL_SECONDS = 2.0 + +# Consecutive `/status` failures tolerated before the fallback gives up: one +# transient network error must not lose a run that is still going. +MAX_CONSECUTIVE_POLL_FAILURES = 3 @dataclass @@ -35,6 +50,10 @@ class OperatorOptions: mode: Optional[str] = "auto" # 'auto', 'sync', 'async' max_wait: Optional[int] = None on_progress: Optional[Callable[[str, Optional[int]], None]] = None + # Seconds between `/v1/operations/{id}/status` polls while the client + # follows a queued run its stream gave no verdict for. Only the fallback + # path uses it; defaults to DEFAULT_POLL_INTERVAL_SECONDS. + poll_interval: Optional[float] = None @dataclass @@ -49,6 +68,40 @@ class OperatorResult: confidence_score: Optional[float] = None execution_time: Optional[float] = None timestamp: Optional[str] = None + # Present when the API reports a failed run inside an otherwise successful + # response (credit pre-flight, operator timeouts, cancelled runs); `content` + # then carries the explanation rather than an answer. + error_details: Optional[Dict[str, Any]] = None + + +def _operator_result(data: Dict[str, Any]) -> OperatorResult: + """Shape an operator payload into the public result. + + The `operator_completed` event, the generic completion event's `result`, + and the `/status` `result` all carry the same fields — one mapper keeps + them consistent. + """ + return OperatorResult( + content=data.get("content") or "", + operator_used=data.get("operator_used") or "unknown", + mode_used=data.get("mode_used") or "standard", + metadata=data.get("metadata"), + tokens_used=data.get("tokens_used"), + confidence_score=data.get("confidence_score"), + execution_time=data.get("execution_time"), + timestamp=data.get("timestamp") or datetime.now().isoformat(), + error_details=data.get("error_details"), + ) + + +def _attr_or_none(data: Any, name: str) -> Any: + """An attrs-model field as a plain value: UNSET → None, models → dicts.""" + value = getattr(data, name, None) + if value is UNSET: + return None + if hasattr(value, "to_dict"): + return value.to_dict() + return value @dataclass @@ -61,6 +114,24 @@ class QueuedOperatorResponse: sse_endpoint: Optional[str] = None +class _PollAbort(Exception): + """Internal: a `/status` verdict that must not be retried.""" + + +def _response_detail(response: Any) -> str: + """Human-readable detail of a non-200 generated-client response.""" + parsed = getattr(response, "parsed", None) + detail = getattr(parsed, "detail", None) + if isinstance(detail, str) and detail: + return detail + content = getattr(response, "content", b"") + if isinstance(content, bytes): + text = content.decode("utf-8", errors="replace") + else: + text = str(content or "") + return text[:200] or "empty response" + + class QueuedOperatorError(Exception): """Exception thrown when operator execution is queued and maxWait is 0""" @@ -79,6 +150,21 @@ def __init__(self, config: Dict[str, Any]): self.token = config.get("token") self.sse_client: Optional[SSEClient] = None + def _rest_client(self) -> Client: + """A REST client for one call, carrying the credential current now. + + Resolved per call — `token_provider` wins over the static token — so a + rotated JWT is picked up without rebuilding the facade, the same way the + GraphQL facades do it. + """ + if not resolve_config_token(self.config): + raise Exception("No API key provided. Set X-API-Key in headers.") + return Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + + 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)) + def execute_query( self, graph_id: str, @@ -102,19 +188,8 @@ def execute_query( force_extended_analysis=request.force_extended_analysis, ) - # Execute through the generated client - from ..client import AuthenticatedClient - - if not self.token: - raise Exception("No API key provided. Set X-API-Key in headers.") - - client = AuthenticatedClient( - base_url=self.base_url, - token=self.token, - prefix="", - auth_header_name="X-API-Key", - headers=self.headers, - ) + # Execute through the generated client, with the credential current now + client = self._rest_client() try: response = auto_select_operator( @@ -153,11 +228,10 @@ def execute_query( confidence_score=data.get("confidence_score"), execution_time=data.get("execution_time"), timestamp=data.get("timestamp", datetime.now().isoformat()), + error_details=data.get("error_details"), ) else: # attrs object - access attributes directly - from ..types import UNSET - return OperatorResult( content=data.content if data.content is not UNSET else "", operator_used=data.operator_used @@ -179,6 +253,7 @@ def execute_query( timestamp=data.timestamp if hasattr(data, "timestamp") and data.timestamp is not UNSET else datetime.now().isoformat(), + error_details=_attr_or_none(data, "error_details"), ) # Check if this is a queued response (async background task execution) @@ -197,8 +272,6 @@ def execute_query( else: is_queued = hasattr(data, "operation_id") if is_queued: - from ..types import UNSET - queued_response = QueuedOperatorResponse( status=data.status if hasattr(data, "status") else "queued", operation_id=data.operation_id, @@ -260,19 +333,8 @@ def execute_operator( force_extended_analysis=request.force_extended_analysis, ) - # Execute through the generated client - from ..client import AuthenticatedClient - - if not self.token: - raise Exception("No API key provided. Set X-API-Key in headers.") - - client = AuthenticatedClient( - base_url=self.base_url, - token=self.token, - prefix="", - auth_header_name="X-API-Key", - headers=self.headers, - ) + # Execute through the generated client, with the credential current now + client = self._rest_client() try: response = execute_specific_operator( @@ -311,11 +373,10 @@ def execute_operator( confidence_score=data.get("confidence_score"), execution_time=data.get("execution_time"), timestamp=data.get("timestamp", datetime.now().isoformat()), + error_details=data.get("error_details"), ) else: # attrs object - from ..types import UNSET - return OperatorResult( content=data.content if data.content is not UNSET else "", operator_used=data.operator_used @@ -337,6 +398,7 @@ def execute_operator( timestamp=data.timestamp if hasattr(data, "timestamp") and data.timestamp is not UNSET else datetime.now().isoformat(), + error_details=_attr_or_none(data, "error_details"), ) # Check if this is a queued response @@ -355,8 +417,6 @@ def execute_operator( else: is_queued = hasattr(data, "operation_id") if is_queued: - from ..types import UNSET - queued_response = QueuedOperatorResponse( status=data.status if hasattr(data, "status") else "queued", operation_id=data.operation_id, @@ -396,14 +456,15 @@ def execute_operator( def _wait_for_operator_completion( self, operation_id: str, options: OperatorOptions ) -> OperatorResult: - """Wait for operator completion and return final result""" - result = None - error = None + """Follow a queued run to its result: over the stream, else over `/status`.""" + result: Optional[OperatorResult] = None + error: Optional[Exception] = None completed = False + transport_error: Optional[Exception] = None - # Set up SSE connection - sse_config = SSEConfig(base_url=self.base_url, headers=self.headers) - sse_client = SSEClient(sse_config) + # Headers are resolved per connect so a rotated JWT reaches the stream. + sse_client = SSEClient(self._sse_config()) + self.sse_client = sse_client def on_progress(data): if options.on_progress: @@ -421,45 +482,35 @@ def on_operator_initialized(data): def on_operator_completed(data): nonlocal result, completed - result = OperatorResult( - content=data.get("content", ""), - operator_used=data.get("operator_used", "unknown"), - mode_used=data.get("mode_used", "standard"), - metadata=data.get("metadata"), - tokens_used=data.get("tokens_used"), - confidence_score=data.get("confidence_score"), - execution_time=data.get("execution_time"), - timestamp=data.get("timestamp", datetime.now().isoformat()), - ) + result = _operator_result(data) completed = True def on_completed(data): nonlocal result, completed if not result: # Fallback to generic completion event - operator_result = data.get("result", data) - result = OperatorResult( - content=operator_result.get("content", ""), - operator_used=operator_result.get("operator_used", "unknown"), - mode_used=operator_result.get("mode_used", "standard"), - metadata=operator_result.get("metadata"), - tokens_used=operator_result.get("tokens_used"), - confidence_score=operator_result.get("confidence_score"), - execution_time=operator_result.get("execution_time"), - timestamp=operator_result.get("timestamp", datetime.now().isoformat()), - ) + result = _operator_result(data.get("result") or data) completed = True def on_error(err): + # The run itself failed — a verdict, not a transport problem. nonlocal error, completed - error = Exception(err.get("message", err.get("error", "Unknown error"))) + error = Exception(event_error_message(err)) completed = True - def on_cancelled(): + def on_cancelled(_data=None): nonlocal error, completed error = Exception("Operator execution cancelled") completed = True + def on_transport_error(err): + # The stream could not open (401/403/404/429) or its reconnects ran + # out. That is not a verdict on the run; `/status` gives one below. + nonlocal transport_error + transport_error = ( + err if isinstance(err, Exception) else Exception(event_error_message(err)) + ) + # Register event handlers sse_client.on(EventType.OPERATION_PROGRESS.value, on_progress) sse_client.on("operator_started", on_operator_started) @@ -468,23 +519,100 @@ def on_cancelled(): sse_client.on("operator_completed", on_operator_completed) sse_client.on(EventType.OPERATION_COMPLETED.value, on_completed) sse_client.on(EventType.OPERATION_ERROR.value, on_error) - sse_client.on("error", on_error) sse_client.on(EventType.OPERATION_CANCELLED.value, on_cancelled) + sse_client.on("error", on_transport_error) + sse_client.on("max_retries_exceeded", on_transport_error) - # Connect and wait - sse_client.connect(operation_id) - - # Wait for completion - import time + # connect() is blocking: it returns once the stream has ended, or right + # away when the stream never opened. + try: + sse_client.connect(operation_id) + finally: + sse_client.close() + if self.sse_client is sse_client: + self.sse_client = None + + if completed and error is not None: + raise error + if result is not None: + return result + + # No verdict from the stream: it never opened, its reconnects ran out, or + # it ended before a terminal event. The run is already queued and + # finishes regardless, so follow it over `/status` instead of losing it. + return self._poll_for_completion(operation_id, options, transport_error) + + def _poll_for_completion( + self, + operation_id: str, + options: OperatorOptions, + stream_error: Optional[Exception], + ) -> OperatorResult: + """Follow a queued run over `/v1/operations/{id}/status` until it settles. + + Used when the stream gave no verdict; `stream_error` is folded into the + failure message if polling cannot reach one either. + """ + interval = ( + options.poll_interval + if options.poll_interval is not None + else DEFAULT_POLL_INTERVAL_SECONDS + ) + stream_detail = ( + str(stream_error) if stream_error else "stream ended before a terminal event" + ) + consecutive_failures = 0 - while not completed: - if error: - sse_client.close() - raise error - time.sleep(0.1) + if options.on_progress: + options.on_progress("Live progress unavailable — waiting for the result", None) - sse_client.close() - return result + while True: + try: + response = get_operation_status( + operation_id=operation_id, client=self._rest_client() + ) + code = int(response.status_code) + parsed = response.parsed + if code != 200 or parsed is None: + detail = _response_detail(response) + # A definitive 4xx (expired, not ours, unauthenticated) ends the + # wait; anything else is treated as transient and retried below. + if 400 <= code < 500 and code != 429: + raise _PollAbort( + f"Operator stream failed ({stream_detail}); " + f"status check failed ({code}: {detail})" + ) + raise RuntimeError(f"{code}: {detail}") + status: Dict[str, Any] = ( + parsed.to_dict() + if hasattr(parsed, "to_dict") + else cast(Dict[str, Any], parsed) + ) + consecutive_failures = 0 + except _PollAbort as abort: + raise Exception(str(abort)) from None + except Exception as poll_error: + consecutive_failures += 1 + if consecutive_failures >= MAX_CONSECUTIVE_POLL_FAILURES: + raise Exception( + f"Operator stream failed ({stream_detail}); " + f"status polling failed ({poll_error})" + ) from poll_error + time.sleep(interval) + continue + + state = status.get("status") + if state == "completed": + return _operator_result(status.get("result") or {}) + if state == "failed": + raise Exception( + status.get("error") or status.get("message") or "Operator run failed" + ) + if state == "cancelled": + raise Exception("Operator execution cancelled") + if options.on_progress and status.get("message"): + options.on_progress(status["message"], None) + time.sleep(interval) def query( self, graph_id: str, message: str, context: Dict[str, Any] = None diff --git a/robosystems_client/clients/query_client.py b/robosystems_client/clients/query_client.py index 9b47c6f..6765a99 100644 --- a/robosystems_client/clients/query_client.py +++ b/robosystems_client/clients/query_client.py @@ -19,7 +19,15 @@ from ..api.query.execute_cypher import sync_detailed as execute_cypher_query from ..models.cypher_statement_request import CypherStatementRequest -from .sse_client import SSEClient, AsyncSSEClient, SSEConfig, EventType +from ..client import Client +from .sse_client import ( + SSEClient, + AsyncSSEClient, + SSEConfig, + EventType, + event_error_message, +) +from .token_utils import resolve_auth_headers, resolve_config_token @dataclass @@ -85,6 +93,16 @@ def __init__(self, config: Dict[str, Any]): self.token = config.get("token") self.sse_client: Optional[SSEClient] = None + def _rest_client(self) -> Client: + """A REST client for one call, carrying the credential current now.""" + if not resolve_config_token(self.config): + raise Exception("No API key provided. Set X-API-Key in headers.") + return Client(base_url=self.base_url, headers=resolve_auth_headers(self.config)) + + 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)) + def execute_query( self, graph_id: str, request: QueryRequest, options: QueryOptions = None ) -> Union[QueryResult, Iterator[Any]]: @@ -97,20 +115,9 @@ def execute_query( query=request.query, parameters=request.parameters or {} ) - # Execute the query through the generated client - from ..client import AuthenticatedClient - - # Create authenticated client with X-API-Key - if not self.token: - raise Exception("No API key provided. Set X-API-Key in headers.") - - client = AuthenticatedClient( - base_url=self.base_url, - token=self.token, - prefix="", - auth_header_name="X-API-Key", - headers=self.headers, - ) + # Execute the query through the generated client, with the credential + # current now (`token_provider` wins over the static token). + client = self._rest_client() try: kwargs = { @@ -346,9 +353,9 @@ def _stream_query_results( completed = False error = None - # Set up SSE connection - sse_config = SSEConfig(base_url=self.base_url, headers=self.headers) - self.sse_client = SSEClient(sse_config) + # Set up SSE connection; headers resolved per connect so a rotated JWT + # reaches the stream. + self.sse_client = SSEClient(self._sse_config()) # Set up event handlers def on_data_chunk(data): @@ -375,8 +382,10 @@ def on_completed(data): completed = True def on_error(err): + # Terminal events carry dicts; a stream that could not open (or whose + # reconnects ran out) emits the transport Exception itself. nonlocal error, completed - error = Exception(err.get("message", err.get("error", "Unknown error"))) + error = err if isinstance(err, Exception) else Exception(event_error_message(err)) completed = True # Register event handlers @@ -385,10 +394,25 @@ def on_error(err): self.sse_client.on(EventType.OPERATION_PROGRESS.value, on_progress) self.sse_client.on(EventType.OPERATION_COMPLETED.value, on_completed) self.sse_client.on(EventType.OPERATION_ERROR.value, on_error) + # Transport failures (bad status, retries exhausted) must end the wait too. + self.sse_client.on("error", on_error) + self.sse_client.on("max_retries_exceeded", on_error) - # Connect and start streaming + # Connect and start streaming. connect() is blocking: the stream has + # ended (or never opened) by the time it returns. self.sse_client.connect(operation_id) + # No terminal event means no verdict — say so rather than spin below. + if not completed and error is None: + error = Exception( + f"Query stream for operation {operation_id} ended before a result" + ) + completed = True + if error is not None: + self.sse_client.close() + self.sse_client = None + raise error + # Yield buffered results while not completed or buffer: if error: @@ -421,9 +445,8 @@ def _wait_for_query_completion( # Set up SSE connection. Headers carry the auth SSEClient.connect merges in; # omitting them made this stream anonymous, so every queued query 401'd. - # (_stream_query_results already passes them.) - sse_config = SSEConfig(base_url=self.base_url, headers=self.headers) - sse_client = SSEClient(sse_config) + # Resolved per connect so a rotated JWT reaches the stream. + sse_client = SSEClient(self._sse_config()) def on_queue_update(data): if options.on_queue_update: @@ -449,11 +472,13 @@ def on_completed(data): completed = True def on_error(err): + # Terminal events carry dicts; a stream that could not open (or whose + # reconnects ran out) emits the transport Exception itself. nonlocal error, completed - error = Exception(err.get("message", err.get("error", "Unknown error"))) + error = err if isinstance(err, Exception) else Exception(event_error_message(err)) completed = True - def on_cancelled(): + def on_cancelled(_data=None): nonlocal error, completed error = Exception("Query cancelled") completed = True @@ -464,20 +489,24 @@ def on_cancelled(): sse_client.on(EventType.OPERATION_COMPLETED.value, on_completed) sse_client.on(EventType.OPERATION_ERROR.value, on_error) sse_client.on(EventType.OPERATION_CANCELLED.value, on_cancelled) + # Transport failures (bad status, retries exhausted) must end the wait too. + sse_client.on("error", on_error) + sse_client.on("max_retries_exceeded", on_error) - # Connect and wait - sse_client.connect(operation_id) - - # Wait for completion - import time - - while not completed: - if error: - sse_client.close() - raise error - time.sleep(0.1) - - sse_client.close() + # Connect and wait. connect() is blocking: the stream has ended (or + # never opened) by the time it returns, so the verdict is in by now. + try: + sse_client.connect(operation_id) + finally: + sse_client.close() + + if error is not None: + raise error + if result is None: + # No terminal event means no verdict — say so rather than return None. + raise Exception( + f"Query stream for operation {operation_id} ended before a result" + ) return result def query( diff --git a/robosystems_client/clients/sse_client.py b/robosystems_client/clients/sse_client.py index 15e1666..2cf8c8b 100644 --- a/robosystems_client/clients/sse_client.py +++ b/robosystems_client/clients/sse_client.py @@ -74,6 +74,20 @@ class EventType(Enum): QUEUE_UPDATE = "queue_update" +def event_error_message(err: Any) -> str: + """Text of an ``error`` payload, whichever shape the stream emitted. + + The stream's own terminal events carry dicts (``{"message": …}`` / + ``{"error": …}``); a stream that could not open, or whose reconnects ran + out, emits the transport ``Exception`` itself. Handlers that assumed a + dict raised inside ``emit`` — which only logs — and left their wait loop + spinning with no verdict. + """ + if isinstance(err, dict): + return str(err.get("message", err.get("error", "Unknown error"))) + return str(err) + + class SSEClient: """SSE client for RoboSystems API with automatic reconnection""" diff --git a/robosystems_client/clients/token_utils.py b/robosystems_client/clients/token_utils.py index aa6f21e..443bc45 100644 --- a/robosystems_client/clients/token_utils.py +++ b/robosystems_client/clients/token_utils.py @@ -437,9 +437,62 @@ def resolve_config_token(config: Dict[str, Any]) -> Optional[str]: so JWT refreshes are picked up without rebuilding the facade. Pair with :class:`TokenManager` when you need refresh scheduling: keep - a manager instance and pass ``token_provider=manager.get_token``. + a manager instance and pass ``token_provider=lambda: manager.token`` + (the property refreshes on read when the token is near expiry). """ provider = config.get("token_provider") if provider is not None: return provider() return config.get("token") + + +_AUTH_HEADER_NAMES = ("x-api-key", "authorization") + + +def apply_auth_header(headers: Dict[str, str], credential: str) -> None: + """Set the correct auth header for a credential, routed by shape. + + The backend accepts two credential formats, and they go in DIFFERENT + headers — not interchangeable (see ``graphql/client.py``): + + - Long-lived API keys (``rfs…`` prefix) → ``X-API-Key``. Validated + against the api_keys table. + - Short-lived JWTs → ``Authorization: Bearer …``. Validated by the + JWT middleware. + + Sending a JWT as ``X-API-Key`` (or an API key as Bearer) both fail + with 401 "Invalid API key" — so exactly one header is set, never both. + """ + if credential.startswith("rfs"): + headers["X-API-Key"] = credential + else: + headers["Authorization"] = f"Bearer {credential}" + + +def resolve_auth_headers(config: Dict[str, Any]) -> Dict[str, str]: + """Headers for one request or stream, carrying the credential current *now*. + + Without a ``token_provider`` this is the configured static ``headers`` + unchanged (plus the static ``token``, routed by shape, when those headers + carry no credential of their own). With a provider, any ``X-API-Key`` / + ``Authorization`` the static headers hold is replaced by the provider's + current credential — a header captured at facade construction must not + outlive a rotation. The SSE-backed clients build their stream headers + from this on every connect; the REST paths build a fresh client from it + per call, the same way the GraphQL facades do. + """ + static = dict(config.get("headers") or {}) + provider = config.get("token_provider") + if provider is None: + if any(k.lower() in _AUTH_HEADER_NAMES for k in static): + return static + token = config.get("token") + if token: + apply_auth_header(static, token) + return static + + headers = {k: v for k, v in static.items() if k.lower() not in _AUTH_HEADER_NAMES} + credential = provider() + if credential: + apply_auth_header(headers, credential) + return headers diff --git a/tests/test_auth_header_resolution.py b/tests/test_auth_header_resolution.py new file mode 100644 index 0000000..2fa0abd --- /dev/null +++ b/tests/test_auth_header_resolution.py @@ -0,0 +1,150 @@ +"""Unit tests for per-call credential resolution shared by the SSE-backed clients. + +The stream endpoint authenticates whatever credential the request carries and +the backend revokes the previous JWT on every session refresh, so headers +captured at facade construction go dead the moment the session rotates. +`resolve_auth_headers` is what every stream connect and REST call now builds +its headers from. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from robosystems_client.clients.auth_integration import _apply_auth_header +from robosystems_client.clients.operation_client import OperationClient +from robosystems_client.clients.sse_client import SSEClient, event_error_message +from robosystems_client.clients.token_utils import ( + apply_auth_header, + resolve_auth_headers, +) + + +@pytest.mark.unit +class TestResolveAuthHeaders: + def test_static_headers_untouched_without_provider(self): + config = { + "headers": {"X-API-Key": "legacy-key", "X-Trace": "1"}, + "token": "legacy-key", + } + + assert resolve_auth_headers(config) == {"X-API-Key": "legacy-key", "X-Trace": "1"} + + def test_static_token_routed_when_headers_carry_no_credential(self): + assert resolve_auth_headers({"headers": {}, "token": "rfs_key"}) == { + "X-API-Key": "rfs_key" + } + assert resolve_auth_headers({"token": "eyJ.jwt"}) == { + "Authorization": "Bearer eyJ.jwt" + } + + def test_no_credential_at_all(self): + assert resolve_auth_headers({"headers": {"X-Trace": "1"}}) == {"X-Trace": "1"} + + def test_provider_replaces_stale_static_credential(self): + config = { + "headers": {"Authorization": "Bearer stale", "X-Trace": "1"}, + "token": "stale", + "token_provider": lambda: "fresh", + } + + assert resolve_auth_headers(config) == { + "Authorization": "Bearer fresh", + "X-Trace": "1", + } + + def test_provider_consulted_on_every_call(self): + tokens = iter(["jwt-1", "jwt-2"]) + config = {"token_provider": lambda: next(tokens)} + + assert resolve_auth_headers(config)["Authorization"] == "Bearer jwt-1" + assert resolve_auth_headers(config)["Authorization"] == "Bearer jwt-2" + + def test_provider_api_key_routes_to_x_api_key(self): + config = { + "headers": {"Authorization": "Bearer old"}, + "token_provider": lambda: "rfs_new", + } + + assert resolve_auth_headers(config) == {"X-API-Key": "rfs_new"} + + def test_provider_returning_none_sends_no_credential(self): + config = {"headers": {"X-API-Key": "stale"}, "token_provider": lambda: None} + + assert resolve_auth_headers(config) == {} + + def test_does_not_mutate_configured_headers(self): + static = {"Authorization": "Bearer stale"} + resolve_auth_headers({"headers": static, "token_provider": lambda: "fresh"}) + + assert static == {"Authorization": "Bearer stale"} + + +@pytest.mark.unit +class TestApplyAuthHeader: + def test_routes_by_credential_shape(self): + headers = {} + apply_auth_header(headers, "rfs_key") + assert headers == {"X-API-Key": "rfs_key"} + + headers = {} + apply_auth_header(headers, "eyJ.jwt") + assert headers == {"Authorization": "Bearer eyJ.jwt"} + + def test_auth_integration_alias_shares_the_rule(self): + headers = {} + _apply_auth_header(headers, "rfs_key") + assert headers == {"X-API-Key": "rfs_key"} + + +@pytest.mark.unit +class TestEventErrorMessage: + def test_dict_payloads(self): + assert event_error_message({"message": "boom"}) == "boom" + assert event_error_message({"error": "bad"}) == "bad" + assert event_error_message({}) == "Unknown error" + + def test_exception_payloads(self): + assert event_error_message(RuntimeError("HTTP 401")) == "HTTP 401" + + +@pytest.mark.unit +class TestOperationClientHeaders: + @patch("time.sleep") + @patch("robosystems_client.clients.operation_client.SSEClient") + def test_stream_headers_resolved_from_provider_at_connect( + self, MockSSE, mock_sleep, mock_config + ): + current = {"jwt": "jwt-old"} + config = { + **mock_config, + "headers": {"Authorization": "Bearer jwt-old"}, + "token": "jwt-old", + "token_provider": lambda: current["jwt"], + } + client = OperationClient(config) + current["jwt"] = "jwt-rotated" + + fake = MagicMock(spec=SSEClient) + listeners = {} + fake.on.side_effect = lambda event, handler: listeners.__setitem__(event, handler) + fake.connect.side_effect = lambda op_id: listeners["operation_completed"]( + {"result": {"ok": True}} + ) + MockSSE.return_value = fake + + client.monitor_operation("op-1") + + assert MockSSE.call_args[0][0].headers == {"Authorization": "Bearer jwt-rotated"} + + # `Client` is imported inside the method, so patch it at its source module. + @patch("robosystems_client.client.Client") + def test_status_call_uses_provider_credential(self, MockClient, mock_config): + config = {**mock_config, "token_provider": lambda: "rfs_fresh"} + with patch( + "robosystems_client.api.operations.get_operation_status.sync_detailed" + ) as mock_get: + mock_get.return_value.parsed = None + OperationClient(config).get_operation_status("op-1") + + assert MockClient.call_args.kwargs["headers"] == {"X-API-Key": "rfs_fresh"} diff --git a/tests/test_operator_client_ops.py b/tests/test_operator_client_ops.py new file mode 100644 index 0000000..9bdc218 --- /dev/null +++ b/tests/test_operator_client_ops.py @@ -0,0 +1,293 @@ +"""Unit tests for OperatorClient queued-run handling. + +Covers: following a queued run over the SSE stream, per-connect credential +resolution (token_provider), and the `/status` polling fallback that takes +over when the stream gives no verdict — it never opened (a revoked JWT +answers 401), its reconnects ran out, or it ended before a terminal event. + +Dataclass and sync-response tests live in tests/test_operator_client.py. +""" + +from http import HTTPStatus +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from robosystems_client.clients.operator_client import ( + OperatorClient, + OperatorOptions, + OperatorQueryRequest, +) +from robosystems_client.clients.sse_client import SSEClient + + +COMPLETED_RESULT = { + "content": "Burn is ~$1,500/month.", + "operator_used": "analyst", + "mode_used": "standard", + "metadata": {"sources": ["ledger"]}, + "tokens_used": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + "execution_time": 21.2, +} + + +def _queued(operation_id: str = "op-1") -> Mock: + resp = Mock() + resp.parsed = {"operation_id": operation_id, "status": "queued", "message": "Queued"} + return resp + + +def _status(code: int, payload=None, detail: str = "") -> Mock: + """A generated-client `Response` for `/status`: `parsed.to_dict()` is the body.""" + resp = Mock() + resp.status_code = HTTPStatus(code) + resp.content = detail.encode() + parsed = Mock() + parsed.to_dict = Mock(return_value=payload or {}) + parsed.detail = detail + resp.parsed = parsed + return resp + + +def _fake_sse(script): + """A mocked SSEClient whose connect() fires `script(listeners)`.""" + fake = MagicMock(spec=SSEClient) + listeners = {} + fake.on.side_effect = lambda event, handler: listeners.__setitem__(event, handler) + fake.connect.side_effect = lambda op_id: script(listeners) + return fake + + +def _run(client, graph_id, options=None): + return client.execute_query( + graph_id, OperatorQueryRequest(message="burn rate?"), options + ) + + +@pytest.mark.unit +@patch("time.sleep") +@patch("robosystems_client.clients.operator_client.get_operation_status") +@patch("robosystems_client.clients.operator_client.SSEClient") +@patch("robosystems_client.clients.operator_client.auto_select_operator") +class TestOperatorQueuedRuns: + def test_stream_completion_returns_result( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + MockSSE.return_value = _fake_sse( + lambda ls: ( + ls["operation_progress"]({"message": "Working", "percentage": 40}), + ls["operation_completed"]({"message": "done", "result": COMPLETED_RESULT}), + ) + ) + progress = [] + + result = _run( + OperatorClient(mock_config), + graph_id, + OperatorOptions(on_progress=lambda m, p: progress.append((m, p))), + ) + + assert result.content == COMPLETED_RESULT["content"] + assert result.operator_used == "analyst" + assert result.execution_time == 21.2 + assert result.error_details is None + assert progress == [("Working", 40)] + # The stream delivered the verdict — no polling. + mock_status.assert_not_called() + + def test_stream_headers_resolved_from_provider_at_connect( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + current = {"jwt": "jwt-captured"} + config = { + **mock_config, + "headers": {"Authorization": "Bearer jwt-captured"}, + "token": "jwt-captured", + "token_provider": lambda: current["jwt"], + } + client = OperatorClient(config) + # The session rotates after construction, before the run is submitted. + current["jwt"] = "jwt-rotated" + mock_auto.return_value = _queued() + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_completed"]({"result": COMPLETED_RESULT}) + ) + + _run(client, graph_id) + + sse_config = MockSSE.call_args[0][0] + assert sse_config.headers == {"Authorization": "Bearer jwt-rotated"} + rest_client = mock_auto.call_args.kwargs["client"] + assert rest_client._headers["Authorization"] == "Bearer jwt-rotated" + + def test_stream_that_cannot_open_falls_back_to_status_polling( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + # A 401 on the stream URL surfaces as the transport Exception itself. + MockSSE.return_value = _fake_sse( + lambda ls: ls["error"](RuntimeError("SSE connection failed: HTTP 401")) + ) + mock_status.side_effect = [ + _status( + 200, {"status": "running", "message": "Operation is currently executing"} + ), + _status(200, {"status": "completed", "result": COMPLETED_RESULT}), + ] + progress = [] + + result = _run( + OperatorClient(mock_config), + graph_id, + OperatorOptions(on_progress=lambda m, p: progress.append(m), poll_interval=0), + ) + + assert result.content == COMPLETED_RESULT["content"] + assert result.operator_used == "analyst" + assert mock_status.call_count == 2 + assert mock_status.call_args.kwargs["operation_id"] == "op-1" + assert progress == [ + "Live progress unavailable — waiting for the result", + "Operation is currently executing", + ] + + def test_stream_ending_without_verdict_falls_back_to_polling( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + # Progress arrived, then the stream ended with no terminal event. + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_progress"]({"message": "Working"}) + ) + mock_status.return_value = _status( + 200, {"status": "completed", "result": COMPLETED_RESULT} + ) + + result = _run( + OperatorClient(mock_config), graph_id, OperatorOptions(poll_interval=0) + ) + + assert result.content == COMPLETED_RESULT["content"] + mock_status.assert_called_once() + + def test_run_error_event_raises_without_polling( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_error"]({"message": "model timeout"}) + ) + + with pytest.raises(Exception, match="model timeout"): + _run(OperatorClient(mock_config), graph_id) + mock_status.assert_not_called() + + def test_cancelled_event_raises_without_hanging( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + # The stream dispatches cancellation with a payload, like every event. + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_cancelled"]({"message": "Cancelled by user"}) + ) + + with pytest.raises(Exception, match="cancelled"): + _run(OperatorClient(mock_config), graph_id) + mock_status.assert_not_called() + + def test_polling_surfaces_failed_run( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + MockSSE.return_value = _fake_sse(lambda ls: ls["error"](RuntimeError("HTTP 401"))) + mock_status.return_value = _status( + 200, {"status": "failed", "error": "Operator run failed: model timeout"} + ) + + with pytest.raises(Exception, match="model timeout"): + _run(OperatorClient(mock_config), graph_id, OperatorOptions(poll_interval=0)) + + def test_polling_stops_on_definitive_4xx( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + MockSSE.return_value = _fake_sse(lambda ls: ls["error"](RuntimeError("HTTP 401"))) + mock_status.return_value = _status( + 404, detail="Operation not found. It may have expired or been cancelled." + ) + + with pytest.raises(Exception, match=r"404: Operation not found"): + _run(OperatorClient(mock_config), graph_id, OperatorOptions(poll_interval=0)) + # A definitive answer is not retried. + mock_status.assert_called_once() + + def test_polling_rides_out_transient_failure( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + MockSSE.return_value = _fake_sse(lambda ls: ls["error"](RuntimeError("HTTP 401"))) + mock_status.side_effect = [ + ConnectionError("network down"), + _status(503, detail="upstream unavailable"), + _status(200, {"status": "completed", "result": COMPLETED_RESULT}), + ] + + result = _run( + OperatorClient(mock_config), graph_id, OperatorOptions(poll_interval=0) + ) + + assert result.content == COMPLETED_RESULT["content"] + assert mock_status.call_count == 3 + + def test_polling_gives_up_after_repeated_failures( + self, mock_auto, MockSSE, mock_status, mock_sleep, mock_config, graph_id + ): + mock_auto.return_value = _queued() + MockSSE.return_value = _fake_sse(lambda ls: ls["error"](RuntimeError("HTTP 401"))) + mock_status.side_effect = ConnectionError("network down") + + with pytest.raises(Exception, match=r"status polling failed \(network down\)"): + _run(OperatorClient(mock_config), graph_id, OperatorOptions(poll_interval=0)) + assert mock_status.call_count == 3 + + +@pytest.mark.unit +class TestOperatorSyncResponse: + @patch("robosystems_client.clients.operator_client.auto_select_operator") + def test_error_details_pass_through(self, mock_auto, mock_config, graph_id): + resp = Mock() + resp.parsed = { + "content": "Not enough credits to perform AI analysis", + "operator_used": "analyst", + "mode_used": "standard", + "error_details": { + "code": "INSUFFICIENT_CREDITS", + "message": "Not enough credits", + }, + } + mock_auto.return_value = resp + + result = _run(OperatorClient(mock_config), graph_id) + + assert result.error_details == { + "code": "INSUFFICIENT_CREDITS", + "message": "Not enough credits", + } + + @patch("robosystems_client.clients.operator_client.auto_select_operator") + def test_rest_client_uses_provider_credential(self, mock_auto, mock_config, graph_id): + config = { + **mock_config, + "headers": {"X-API-Key": "stale"}, + "token": "stale", + "token_provider": lambda: "rfs_fresh", + } + resp = Mock() + resp.parsed = {"content": "ok", "operator_used": "analyst", "mode_used": "quick"} + mock_auto.return_value = resp + + _run(OperatorClient(config), graph_id) + + rest_client = mock_auto.call_args.kwargs["client"] + assert rest_client._headers == {"X-API-Key": "rfs_fresh"} diff --git a/tests/test_query_client_sse.py b/tests/test_query_client_sse.py new file mode 100644 index 0000000..fe81278 --- /dev/null +++ b/tests/test_query_client_sse.py @@ -0,0 +1,126 @@ +"""Unit tests for QueryClient's SSE-backed waits. + +The stream's terminal events carry dicts, but a stream that cannot open (a +revoked JWT answers 401) or whose reconnects run out emits the transport +Exception itself. Handlers that assumed a dict raised inside `emit` — which +only logs — and left the wait loop spinning forever with no verdict. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from robosystems_client.clients.query_client import QueryClient, QueryOptions +from robosystems_client.clients.sse_client import SSEClient + + +def _fake_sse(script): + fake = MagicMock(spec=SSEClient) + listeners = {} + fake.on.side_effect = lambda event, handler: listeners.__setitem__(event, handler) + fake.connect.side_effect = lambda op_id: script(listeners) + return fake + + +@pytest.mark.unit +@patch("time.sleep") +@patch("robosystems_client.clients.query_client.SSEClient") +class TestWaitForQueryCompletion: + def test_completed_event_returns_result(self, MockSSE, mock_sleep, mock_config): + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_completed"]( + {"result": {"data": [{"n": 1}], "columns": ["n"], "row_count": 1}} + ) + ) + + result = QueryClient(mock_config)._wait_for_query_completion("op-1", QueryOptions()) + + assert result.data == [{"n": 1}] + assert result.columns == ["n"] + + def test_transport_error_object_raises_instead_of_hanging( + self, MockSSE, mock_sleep, mock_config + ): + MockSSE.return_value = _fake_sse( + lambda ls: ls["error"](RuntimeError("SSE connection failed: HTTP 401")) + ) + + with pytest.raises(RuntimeError, match="HTTP 401"): + QueryClient(mock_config)._wait_for_query_completion("op-1", QueryOptions()) + + def test_retries_exhausted_raises(self, MockSSE, mock_sleep, mock_config): + MockSSE.return_value = _fake_sse( + lambda ls: ls["max_retries_exceeded"](ConnectionError("connection reset")) + ) + + with pytest.raises(ConnectionError, match="connection reset"): + QueryClient(mock_config)._wait_for_query_completion("op-1", QueryOptions()) + + def test_stream_ending_without_verdict_raises(self, MockSSE, mock_sleep, mock_config): + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_progress"]({"message": "Working"}) + ) + + with pytest.raises(Exception, match="ended before a result"): + QueryClient(mock_config)._wait_for_query_completion("op-1", QueryOptions()) + + def test_cancelled_event_raises(self, MockSSE, mock_sleep, mock_config): + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_cancelled"]({"message": "Cancelled by user"}) + ) + + with pytest.raises(Exception, match="Query cancelled"): + QueryClient(mock_config)._wait_for_query_completion("op-1", QueryOptions()) + + def test_stream_headers_resolved_from_provider_at_connect( + self, MockSSE, mock_sleep, mock_config + ): + current = {"jwt": "jwt-old"} + config = { + **mock_config, + "headers": {"Authorization": "Bearer jwt-old"}, + "token": "jwt-old", + "token_provider": lambda: current["jwt"], + } + client = QueryClient(config) + current["jwt"] = "jwt-rotated" + MockSSE.return_value = _fake_sse( + lambda ls: ls["operation_completed"]({"result": {"data": []}}) + ) + + client._wait_for_query_completion("op-1", QueryOptions()) + + assert MockSSE.call_args[0][0].headers == {"Authorization": "Bearer jwt-rotated"} + + +@pytest.mark.unit +@patch("time.sleep") +@patch("robosystems_client.clients.query_client.SSEClient") +class TestStreamQueryResults: + def test_yields_buffered_rows(self, MockSSE, mock_sleep, mock_config): + MockSSE.return_value = _fake_sse( + lambda ls: ( + ls["data_chunk"]({"rows": [{"n": 1}, {"n": 2}]}), + ls["operation_completed"]({"result": {}}), + ) + ) + + rows = list(QueryClient(mock_config)._stream_query_results("op-1", QueryOptions())) + + assert rows == [{"n": 1}, {"n": 2}] + + def test_transport_error_object_raises_instead_of_hanging( + self, MockSSE, mock_sleep, mock_config + ): + MockSSE.return_value = _fake_sse( + lambda ls: ls["error"](RuntimeError("SSE connection failed: HTTP 401")) + ) + + with pytest.raises(RuntimeError, match="HTTP 401"): + list(QueryClient(mock_config)._stream_query_results("op-1", QueryOptions())) + + def test_stream_ending_without_verdict_raises(self, MockSSE, mock_sleep, mock_config): + MockSSE.return_value = _fake_sse(lambda ls: None) + + with pytest.raises(Exception, match="ended before a result"): + list(QueryClient(mock_config)._stream_query_results("op-1", QueryOptions()))