fix(clients): retry the remaining REST clients, and repair the operations status/cancel calls

robosystems_client/clients/operation_client.py

@@ -5,7 +5,7 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Dict, Any, Optional, Callable, List
+from typing import Dict, Any, Optional, Callable, List, cast
 from datetime import datetime
 from enum import Enum
 
@@ -69,6 +69,24 @@ class MonitorOptions:
   poll_interval: Optional[int] = None
 
 
+def _parsed_dict(parsed: Any) -> Dict[str, Any] | None:
+  """Body of a generated operations response as a plain dict.
+
+  The operations endpoints are declared as free-form objects, so the
+  generator emits models whose only field is ``additional_properties``
+  — attribute access on them raises, whatever the payload contained.
+  ``to_dict()`` is the accessor that works, and is what
+  ``operator_client._poll_for_completion`` already uses.
+  """
+  if parsed is None:
+    return None
+  if hasattr(parsed, "to_dict"):
+    return cast(Dict[str, Any], parsed.to_dict())
+  if isinstance(parsed, dict):
+    return cast(Dict[str, Any], parsed)
+  return None
+
+
 class OperationClient:
   """Client for monitoring operations via SSE"""
 
@@ -218,23 +236,33 @@ def get_operation_status(self, operation_id: str) -> Dict[str, Any]:
     from ..api.operations.get_operation_status import (
       sync_detailed as get_operation_status,
     )
-    from ..client import Client
+    from .retry import retrying_client
 
     # Plain Client with the headers current now (`token_provider` wins).
-    client = Client(base_url=self.base_url, headers=resolve_auth_headers(self.config))
+    client = retrying_client(
+      base_url=self.base_url,
+      headers=resolve_auth_headers(self.config),
+      config=self.config,
+    )
     try:
       # Auth travels in self.headers (X-API-Key / Authorization). The generated
       # function takes no `token` kwarg — passing one raised TypeError, which
       # the handler below laundered into a fake {"status": "error"} result, so
       # this call never succeeded whenever a token was configured.
+      #
+      # Read through `to_dict()`, not attributes: the second incarnation of
+      # that same bug was `response.parsed.status` raising AttributeError on
+      # an additional-properties model, laundered into the same fake result
+      # for *every* response — including successful ones.
       response = get_operation_status(operation_id=operation_id, client=client)
-      if response.parsed:
+      payload = _parsed_dict(response.parsed)
+      if payload is not None:
         return {
           "operation_id": operation_id,
-          "status": response.parsed.status,
-          "progress": getattr(response.parsed, "progress", None),
-          "result": getattr(response.parsed, "result", None),
-          "error": getattr(response.parsed, "error", None),
+          "status": payload.get("status", "unknown"),
+          "progress": payload.get("progress"),
+          "result": payload.get("result"),
+          "error": payload.get("error"),
         }
     except Exception as e:
       # Logged rather than silently shaped into a result: swallowing here is
@@ -248,26 +276,34 @@ def cancel_operation(self, operation_id: str) -> bool:
     """Cancel an operation"""
     # This would use the generated SDK to call /v1/operations/{operation_id}/cancel
     from ..api.operations.cancel_operation import sync_detailed as cancel_operation
-    from ..client import Client
+    from .retry import retrying_client
 
     # Plain Client with the headers current now (`token_provider` wins).
-    client = Client(base_url=self.base_url, headers=resolve_auth_headers(self.config))
+    client = retrying_client(
+      base_url=self.base_url,
+      headers=resolve_auth_headers(self.config),
+      config=self.config,
+    )
     try:
-      # See get_operation_status: no `token` kwarg on the generated function.
+      # See get_operation_status: no `token` kwarg, and the body is read
+      # through `to_dict()` because attribute access always raises.
       response = cancel_operation(operation_id=operation_id, client=client)
-      if response.parsed:
-        return response.parsed.cancelled or False
+      payload = _parsed_dict(response.parsed)
+      cancelled = bool(payload.get("cancelled")) if payload is not None else False
     except Exception as e:
       logger.warning("Failed to cancel operation %s: %s", operation_id, e)
       return False
 
-    # Also close any active SSE connection with thread safety
+    # Close any active SSE connection with thread safety. This used to sit
+    # after an early `return` on the success path, so the one case that
+    # needs it — the cancel actually landed — was the one case that skipped
+    # it, leaking the stream.
     with self._lock:
       if operation_id in self.active_operations:
         self.active_operations[operation_id].close()
         del self.active_operations[operation_id]
 
-    return False
+    return cancelled
 
   def list_operations(self) -> List[Dict[str, Any]]:
     """List all operations (if supported by the API)"""

robosystems_client/clients/operator_client.py

@@ -16,6 +16,7 @@
   sync_detailed as get_operation_status,
 )
 from ..client import Client
+from .retry import retrying_client
 from ..models.operator_request import OperatorRequest
 from ..models.operator_message import OperatorMessage
 from ..types import UNSET
@@ -159,7 +160,11 @@ def _rest_client(self) -> Client:
     """
     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))
+    return retrying_client(
+      base_url=self.base_url,
+      headers=resolve_auth_headers(self.config),
+      config=self.config,
+    )
 
   def _sse_config(self) -> SSEConfig:
     """Stream config for one connect; headers carry the credential current now."""

robosystems_client/clients/query_client.py

@@ -20,6 +20,7 @@
 from ..api.query.execute_cypher import sync_detailed as execute_cypher_query
 from ..models.cypher_statement_request import CypherStatementRequest
 from ..client import Client
+from .retry import retrying_client
 from .sse_client import (
   SSEClient,
   AsyncSSEClient,
@@ -97,7 +98,11 @@ 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))
+    return retrying_client(
+      base_url=self.base_url,
+      headers=resolve_auth_headers(self.config),
+      config=self.config,
+    )
 
   def _sse_config(self) -> SSEConfig:
     """Stream config for one connect; headers carry the credential current now."""

robosystems_client/clients/retry.py

@@ -22,7 +22,7 @@
 
 import httpx
 
-from ..client import AuthenticatedClient
+from ..client import AuthenticatedClient, Client
 
 RETRY_STATUS_CODES = frozenset({429})
 
@@ -174,3 +174,28 @@ def retrying_authenticated_client(
       base_url=base_url, headers=request_headers, timeout=None, config=config
     )
   )
+
+
+def retrying_client(
+  *,
+  base_url: str,
+  headers: dict[str, str] | None = None,
+  config: dict[str, Any] | None = None,
+) -> Client:
+  """A plain :class:`Client` whose transport replays 429s.
+
+  The unauthenticated sibling of :func:`retrying_authenticated_client`,
+  for the facades that resolve their credential themselves and pass it
+  in ``headers`` (query / operator / operations) rather than letting
+  ``AuthenticatedClient`` stamp it. Those paths are rate-limited like
+  any other — Cypher queries, operator runs and operation-status polls
+  each draw on their own category budget — so they need the same replay.
+  """
+  return Client(
+    base_url=base_url,
+    headers=dict(headers or {}),
+  ).set_httpx_client(
+    build_httpx_client(
+      base_url=base_url, headers=headers or {}, timeout=None, config=config
+    )
+  )

tests/test_auth_header_resolution.py

@@ -13,6 +13,7 @@
 
 from robosystems_client.clients.auth_integration import _apply_auth_header
 from robosystems_client.clients.operation_client import OperationClient
+from robosystems_client.clients.retry import RetryingClient
 from robosystems_client.clients.sse_client import SSEClient, event_error_message
 from robosystems_client.clients.token_utils import (
   apply_auth_header,
@@ -137,14 +138,18 @@ def test_stream_headers_resolved_from_provider_at_connect(
 
     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):
+  def test_status_call_uses_provider_credential(self, mock_config):
+    # Asserts the credential on the client actually handed to the
+    # generated op, rather than on a patched constructor — the status
+    # call now builds its client through `retrying_client`, and a
+    # constructor patch would only pin today's construction path.
     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"}
+    passed = mock_get.call_args.kwargs["client"]
+    assert passed.get_httpx_client().headers["X-API-Key"] == "rfs_fresh"
+    assert isinstance(passed.get_httpx_client(), RetryingClient)

tests/test_operation_client_ops.py

@@ -16,6 +16,12 @@
   MonitorOptions,
 )
 from robosystems_client.clients.sse_client import SSEClient
+from robosystems_client.models.cancel_operation_response_canceloperation import (
+  CancelOperationResponseCanceloperation,
+)
+from robosystems_client.models.get_operation_status_response_getoperationstatus import (
+  GetOperationStatusResponseGetoperationstatus,
+)
 
 
 # ── Helpers ──────────────────────────────────────────────────────────
@@ -243,12 +249,14 @@ class TestGetOperationStatus:
   @patch("robosystems_client.api.operations.get_operation_status.sync_detailed")
   def test_get_status_success(self, mock_get, mock_config):
     """Test successful status retrieval."""
+    # The real model, not a bare Mock: the endpoint is a free-form object,
+    # so the generated model carries only `additional_properties` and
+    # attribute access on it raises. A Mock makes `.status` work and let
+    # this test pass against code that could never work in production.
     mock_resp = Mock()
-    mock_resp.parsed = Mock()
-    mock_resp.parsed.status = "running"
-    mock_resp.parsed.progress = 50
-    mock_resp.parsed.result = None
-    mock_resp.parsed.error = None
+    mock_resp.parsed = GetOperationStatusResponseGetoperationstatus.from_dict(
+      {"status": "running", "progress": 50, "result": None, "error": None}
+    )
     mock_get.return_value = mock_resp
 
     client = OperationClient(mock_config)
@@ -292,8 +300,9 @@ class TestCancelOperation:
   def test_cancel_success(self, mock_cancel, mock_config):
     """Test successful cancellation."""
     mock_resp = Mock()
-    mock_resp.parsed = Mock()
-    mock_resp.parsed.cancelled = True
+    mock_resp.parsed = CancelOperationResponseCanceloperation.from_dict(
+      {"cancelled": True}
+    )
     mock_cancel.return_value = mock_resp
 
     client = OperationClient(mock_config)