From 4640d1faf39290337d7881f9c2a0d0bda7358fa0 Mon Sep 17 00:00:00 2001 From: smaramwbc <145447586+smaramwbc@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:26:37 +0100 Subject: [PATCH 1/2] feat: per-call timeout override and CancelledError propagation Add optional `timeout: float | None = None` to every public method on both StatewaveClient and AsyncStatewaveClient. When provided it overrides the constructor-level timeout for that individual HTTP call only; when omitted the client's configured default applies unchanged. compile_memories_wait() uses `request_timeout` to avoid collision with its existing `timeout` (total polling duration) parameter. asyncio.CancelledError now propagates immediately from the async retry loop instead of being caught and retried. KeyboardInterrupt and SystemExit receive the same treatment in the sync loop. 10 new unit tests cover: timeout forwarding to httpx, no kwarg when None, public-method passthrough, interrupt non-retry, and compile_memories_wait request_timeout plumbing (sync + async). --- statewave/client.py | 490 +++++++++++++++++++++++++++------ tests/test_per_call_timeout.py | 186 +++++++++++++ 2 files changed, 595 insertions(+), 81 deletions(-) create mode 100644 tests/test_per_call_timeout.py diff --git a/statewave/client.py b/statewave/client.py index 8534f2b..8536739 100644 --- a/statewave/client.py +++ b/statewave/client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import random import time from dataclasses import dataclass @@ -194,6 +195,7 @@ def create_episode( provenance: dict[str, Any] | None = None, session_id: str | None = None, idempotency_key: str | None = None, + timeout: float | None = None, ) -> Episode: """Record a raw interaction episode. @@ -205,6 +207,8 @@ def create_episode( Pass ``idempotency_key`` to make re-ingest a no-op: a later episode with the same key (re-running a backfill, retrying a failed request) returns the existing episode instead of inserting a duplicate. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = { "subject_id": subject_id, @@ -218,48 +222,82 @@ def create_episode( body["session_id"] = session_id if idempotency_key is not None: body["idempotency_key"] = idempotency_key - return self._request("POST", "/v1/episodes", json=body, model=Episode) + return self._request("POST", "/v1/episodes", json=body, model=Episode, timeout=timeout) def create_episodes_batch( self, episodes: list[dict[str, Any]], + *, + timeout: float | None = None, ) -> BatchCreateResult: - """Record multiple episodes in a single request (max 100).""" + """Record multiple episodes in a single request (max 100). + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return self._request( "POST", "/v1/episodes/batch", json={"episodes": episodes}, model=BatchCreateResult, + timeout=timeout, ) # -- Memories ---------------------------------------------------------- - def compile_memories(self, subject_id: str) -> CompileResult: - """Compile memories from unprocessed episodes. Idempotent.""" + def compile_memories( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> CompileResult: + """Compile memories from unprocessed episodes. Idempotent. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return self._request( "POST", "/v1/memories/compile", json={"subject_id": subject_id}, model=CompileResult, + timeout=timeout, ) - def compile_memories_async(self, subject_id: str) -> CompileJob: + def compile_memories_async( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> CompileJob: """Submit async compilation. Returns immediately with a job_id for polling. Use `get_compile_status()` to poll for completion. + + Pass ``timeout`` to override the client-level timeout for this call only. """ + kwargs: dict[str, Any] = {} + if timeout is not None: + kwargs["timeout"] = timeout resp = self._http.request( "POST", "/v1/memories/compile", json={"subject_id": subject_id, "async": True}, + **kwargs, ) if not resp.is_success: raise _parse_error(resp) return CompileJob.model_validate(resp.json()) - def get_compile_status(self, job_id: str) -> CompileJob: - """Poll the status of an async compile job.""" + def get_compile_status( + self, + job_id: str, + *, + timeout: float | None = None, + ) -> CompileJob: + """Poll the status of an async compile job. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return self._request( - "GET", f"/v1/memories/compile/{job_id}", model=CompileJob, + "GET", f"/v1/memories/compile/{job_id}", model=CompileJob, timeout=timeout, ) def compile_memories_wait( @@ -268,18 +306,23 @@ def compile_memories_wait( *, poll_interval: float = 0.5, timeout: float = 60.0, + request_timeout: float | None = None, ) -> CompileJob: """Submit async compilation and poll until completion or timeout. Convenience method that combines submit + polling. Raises TimeoutError if job doesn't complete within timeout. + + ``timeout`` is the total polling duration in seconds. + ``request_timeout`` overrides the client-level HTTP timeout for each + individual poll request. """ - job = self.compile_memories_async(subject_id) + job = self.compile_memories_async(subject_id, timeout=request_timeout) elapsed = 0.0 while elapsed < timeout: time.sleep(poll_interval) elapsed += poll_interval - job = self.get_compile_status(job.job_id) + job = self.get_compile_status(job.job_id, timeout=request_timeout) if job.status in ("completed", "failed"): return job raise TimeoutError(f"Compile job {job.job_id} did not complete within {timeout}s") @@ -292,8 +335,12 @@ def search_memories( query: str | None = None, semantic: bool = False, limit: int = 20, + timeout: float | None = None, ) -> SearchResult: - """Search memories by kind, text query, or semantic similarity.""" + """Search memories by kind, text query, or semantic similarity. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ params: dict[str, Any] = {"subject_id": subject_id, "limit": limit} if kind: params["kind"] = kind @@ -301,7 +348,9 @@ def search_memories( params["q"] = query if semantic: params["semantic"] = "true" - return self._request("GET", "/v1/memories/search", params=params, model=SearchResult) + return self._request( + "GET", "/v1/memories/search", params=params, model=SearchResult, timeout=timeout, + ) # -- Context ----------------------------------------------------------- @@ -318,6 +367,7 @@ def get_context( parent_receipt_id: str | None = None, caller_id: str | None = None, caller_type: str | None = None, + timeout: float | None = None, ) -> ContextBundle: """Assemble a ranked, token-bounded context bundle. @@ -329,6 +379,8 @@ def get_context( `always`), the returned bundle carries a `receipt_id` that can be fetched via `get_receipt()`. See `Receipt` for the body schema. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = {"subject_id": subject_id, "task": task} if max_tokens is not None: @@ -347,7 +399,9 @@ def get_context( body["caller_id"] = caller_id if caller_type is not None: body["caller_type"] = caller_type - return self._request("POST", "/v1/context", json=body, model=ContextBundle) + return self._request( + "POST", "/v1/context", json=body, model=ContextBundle, timeout=timeout, + ) # -- Memory labels (#50) ---------------------------------------------- @@ -355,6 +409,8 @@ def set_memory_labels( self, memory_id: str, labels: list[str], + *, + timeout: float | None = None, ) -> Memory: """Replace a memory's sensitivity_labels with the supplied list. @@ -363,12 +419,15 @@ def set_memory_labels( write boundary is the only place to do it safely. An empty list clears all labels (memory becomes untagged → policy default-allow). + + Pass ``timeout`` to override the client-level timeout for this call only. """ return self._request( "PATCH", f"/v1/memories/{memory_id}/labels", json={"sensitivity_labels": labels}, model=Memory, + timeout=timeout, ) def list_suggested_labels( @@ -379,6 +438,7 @@ def list_suggested_labels( label: str | None = None, limit: int = 50, offset: int = 0, + timeout: float | None = None, ) -> SuggestedLabelsList: """List memories carrying auto-derived suggested labels — the review surface for the promote workflow (admin/governance, v0.9 #158). @@ -386,6 +446,8 @@ def list_suggested_labels( Requires auto-labeling enabled server-side (``STATEWAVE_AUTO_LABELING_ENABLED=true``). Pair with :meth:`promote_suggested_labels` to commit a reviewed suggestion. + + Pass ``timeout`` to override the client-level timeout for this call only. """ params: dict[str, Any] = {"limit": limit, "offset": offset} if tenant_id is not None: @@ -399,6 +461,7 @@ def list_suggested_labels( "/admin/memories/with-suggested-labels", params=params, model=SuggestedLabelsList, + timeout=timeout, ) def promote_suggested_labels( @@ -407,6 +470,7 @@ def promote_suggested_labels( labels: list[str], *, tenant_id: str | None = None, + timeout: float | None = None, ) -> PromoteLabelsResult: """Promote a subset of a memory's suggested labels into authoritative ``sensitivity_labels`` (admin/governance, v0.9 #160). @@ -414,6 +478,8 @@ def promote_suggested_labels( Review-only: every label in ``labels`` must already be a current suggestion on the memory, else the server returns 422. Promoted labels are removed from the suggestions so they don't re-surface in review. + + Pass ``timeout`` to override the client-level timeout for this call only. """ params: dict[str, Any] = {} if tenant_id is not None: @@ -424,6 +490,7 @@ def promote_suggested_labels( json={"labels": labels}, params=params, model=PromoteLabelsResult, + timeout=timeout, ) def get_context_string( @@ -432,16 +499,30 @@ def get_context_string( task: str, *, max_tokens: int | None = None, + timeout: float | None = None, ) -> str: - """Return just the assembled context string, ready to inject into a prompt.""" - bundle = self.get_context(subject_id, task, max_tokens=max_tokens) + """Return just the assembled context string, ready to inject into a prompt. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ + bundle = self.get_context(subject_id, task, max_tokens=max_tokens, timeout=timeout) return bundle.assembled_context # -- Receipts ---------------------------------------------------------- - def get_receipt(self, receipt_id: str) -> Receipt: - """Fetch a single state-assembly receipt by ULID.""" - return self._request("GET", f"/v1/receipts/{receipt_id}", model=Receipt) + def get_receipt( + self, + receipt_id: str, + *, + timeout: float | None = None, + ) -> Receipt: + """Fetch a single state-assembly receipt by ULID. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ + return self._request( + "GET", f"/v1/receipts/{receipt_id}", model=Receipt, timeout=timeout, + ) def list_receipts( self, @@ -451,9 +532,13 @@ def list_receipts( until: str | None = None, cursor: str | None = None, limit: int = 50, + timeout: float | None = None, ) -> ReceiptList: """List receipts for a subject, newest first. Cursor-paginated — - pass the previous response's `next_cursor` to fetch the next page.""" + pass the previous response's `next_cursor` to fetch the next page. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ params: dict[str, Any] = {"subject_id": subject_id, "limit": limit} if since is not None: params["since"] = since @@ -461,9 +546,16 @@ def list_receipts( params["until"] = until if cursor is not None: params["cursor"] = cursor - return self._request("GET", "/v1/receipts", params=params, model=ReceiptList) + return self._request( + "GET", "/v1/receipts", params=params, model=ReceiptList, timeout=timeout, + ) - def verify_receipt(self, receipt_id: str) -> ReceiptVerifyResult: + def verify_receipt( + self, + receipt_id: str, + *, + timeout: float | None = None, + ) -> ReceiptVerifyResult: """Verify the HMAC signature on a stored receipt (v0.9+ #157). Calls ``GET /v1/receipts/{receipt_id}/verify``. Returns a @@ -486,14 +578,22 @@ def verify_receipt(self, receipt_id: str) -> ReceiptVerifyResult: Raises :class:`StatewaveAPIError` on 404 (receipt not found or belongs to a different tenant — indistinguishable on the wire) and other non-2xx responses. + + Pass ``timeout`` to override the client-level timeout for this call only. """ return self._request( "GET", f"/v1/receipts/{receipt_id}/verify", model=ReceiptVerifyResult, + timeout=timeout, ) - def replay_receipt(self, receipt_id: str) -> ReceiptReplayResult: + def replay_receipt( + self, + receipt_id: str, + *, + timeout: float | None = None, + ) -> ReceiptReplayResult: """Re-run the original retrieval against current memories using the original policy bundle captured in the receipt's ``policy_snapshot`` (v0.9+ #159). @@ -522,23 +622,33 @@ def replay_receipt(self, receipt_id: str) -> ReceiptReplayResult: Raises :class:`StatewaveAPIError` on 404 and other non-2xx responses. + + Pass ``timeout`` to override the client-level timeout for this call only. """ return self._request( "POST", f"/v1/receipts/{receipt_id}/replay", model=ReceiptReplayResult, + timeout=timeout, ) # -- Support: health, SLA, handoff, resolutions ------------------------ - def get_health(self, subject_id: str) -> Health: + def get_health( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> Health: """Compute the customer health score (0-100) for a subject. Returns the score, the state bucket (``healthy`` | ``watch`` | ``at_risk``), and the explainable factors that drove it. + + Pass ``timeout`` to override the client-level timeout for this call only. """ return self._request( - "GET", f"/v1/subjects/{subject_id}/health", model=Health, + "GET", f"/v1/subjects/{subject_id}/health", model=Health, timeout=timeout, ) def get_sla( @@ -547,6 +657,7 @@ def get_sla( *, first_response_threshold_minutes: float | None = None, resolution_threshold_hours: float | None = None, + timeout: float | None = None, ) -> SLASummary: """Compute SLA metrics for a subject. @@ -554,6 +665,8 @@ def get_sla( subject's sessions and flags breaches against the supplied thresholds. Both thresholds fall back to the server defaults (5 minutes / 24 hours) when omitted. + + Pass ``timeout`` to override the client-level timeout for this call only. """ params: dict[str, Any] = {} if first_response_threshold_minutes is not None: @@ -562,6 +675,7 @@ def get_sla( params["resolution_threshold_hours"] = resolution_threshold_hours return self._request( "GET", f"/v1/subjects/{subject_id}/sla", params=params, model=SLASummary, + timeout=timeout, ) def create_handoff( @@ -577,6 +691,7 @@ def create_handoff( parent_receipt_id: str | None = None, caller_id: str | None = None, caller_type: str | None = None, + timeout: float | None = None, ) -> Handoff: """Generate a handoff context pack for escalation or shift change. @@ -586,6 +701,8 @@ def create_handoff( ``caller_id`` and ``caller_type`` are consumed by the sensitivity-label policy layer (#50); when the tenant config sets ``require_caller_identity: true``, both are mandatory. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = {"subject_id": subject_id, "session_id": session_id} if reason is not None: @@ -604,7 +721,9 @@ def create_handoff( body["caller_id"] = caller_id if caller_type is not None: body["caller_type"] = caller_type - return self._request("POST", "/v1/handoff", json=body, model=Handoff) + return self._request( + "POST", "/v1/handoff", json=body, model=Handoff, timeout=timeout, + ) def create_resolution( self, @@ -614,11 +733,14 @@ def create_resolution( status: str = "open", resolution_summary: str | None = None, metadata: dict[str, Any] | None = None, + timeout: float | None = None, ) -> Resolution: """Create or update a resolution record for a support session. Upserts by ``subject_id`` + ``session_id``. ``status`` is one of ``open``, ``resolved``, or ``unresolved``. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = { "subject_id": subject_id, @@ -629,46 +751,81 @@ def create_resolution( body["resolution_summary"] = resolution_summary if metadata is not None: body["metadata"] = metadata - return self._request("POST", "/v1/resolutions", json=body, model=Resolution) + return self._request( + "POST", "/v1/resolutions", json=body, model=Resolution, timeout=timeout, + ) def list_resolutions( self, subject_id: str, *, status: str | None = None, + timeout: float | None = None, ) -> list[Resolution]: """List resolution records for a subject. Optionally filter to a single ``status`` (``open`` | ``resolved`` | ``unresolved``). + + Pass ``timeout`` to override the client-level timeout for this call only. """ params: dict[str, Any] = {"subject_id": subject_id} if status is not None: params["status"] = status return self._request( "GET", "/v1/resolutions", params=params, model=Resolution, is_list=True, + timeout=timeout, ) # -- Timeline ---------------------------------------------------------- - def get_timeline(self, subject_id: str) -> Timeline: - """Get chronological subject timeline.""" + def get_timeline( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> Timeline: + """Get chronological subject timeline. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return self._request( - "GET", "/v1/timeline", params={"subject_id": subject_id}, model=Timeline + "GET", "/v1/timeline", params={"subject_id": subject_id}, model=Timeline, + timeout=timeout, ) # -- Subjects ---------------------------------------------------------- - def delete_subject(self, subject_id: str) -> DeleteResult: - """Permanently delete all data for a subject.""" - return self._request("DELETE", f"/v1/subjects/{subject_id}", model=DeleteResult) + def delete_subject( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> DeleteResult: + """Permanently delete all data for a subject. - def list_subjects(self, *, limit: int = 50, offset: int = 0) -> ListSubjectsResult: - """List all known subjects with episode/memory counts.""" + Pass ``timeout`` to override the client-level timeout for this call only. + """ + return self._request( + "DELETE", f"/v1/subjects/{subject_id}", model=DeleteResult, timeout=timeout, + ) + + def list_subjects( + self, + *, + limit: int = 50, + offset: int = 0, + timeout: float | None = None, + ) -> ListSubjectsResult: + """List all known subjects with episode/memory counts. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return self._request( "GET", "/v1/subjects", params={"limit": limit, "offset": offset}, model=ListSubjectsResult, + timeout=timeout, ) # -- Lifecycle --------------------------------------------------------- @@ -693,11 +850,17 @@ def _request( json: Any = None, params: Any = None, is_list: bool = False, + timeout: float | None = None, ): last_exc: Exception | None = None + extra: dict[str, Any] = {} + if timeout is not None: + extra["timeout"] = timeout for attempt in range(self._retry.max_retries + 1): try: - resp = self._http.request(method, path, json=json, params=params) + resp = self._http.request(method, path, json=json, params=params, **extra) + except (KeyboardInterrupt, SystemExit): + raise except httpx.HTTPStatusError: raise except Exception as exc: @@ -766,6 +929,7 @@ async def create_episode( provenance: dict[str, Any] | None = None, session_id: str | None = None, idempotency_key: str | None = None, + timeout: float | None = None, ) -> Episode: """Record a raw interaction episode. @@ -777,6 +941,8 @@ async def create_episode( Pass ``idempotency_key`` to make re-ingest a no-op: a later episode with the same key (re-running a backfill, retrying a failed request) returns the existing episode instead of inserting a duplicate. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = { "subject_id": subject_id, @@ -790,45 +956,82 @@ async def create_episode( body["session_id"] = session_id if idempotency_key is not None: body["idempotency_key"] = idempotency_key - return await self._request("POST", "/v1/episodes", json=body, model=Episode) + return await self._request( + "POST", "/v1/episodes", json=body, model=Episode, timeout=timeout, + ) async def create_episodes_batch( self, episodes: list[dict[str, Any]], + *, + timeout: float | None = None, ) -> BatchCreateResult: - """Record multiple episodes in a single request (max 100).""" + """Record multiple episodes in a single request (max 100). + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( "POST", "/v1/episodes/batch", json={"episodes": episodes}, model=BatchCreateResult, + timeout=timeout, ) # -- Memories ---------------------------------------------------------- - async def compile_memories(self, subject_id: str) -> CompileResult: - """Compile memories from unprocessed episodes. Idempotent.""" + async def compile_memories( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> CompileResult: + """Compile memories from unprocessed episodes. Idempotent. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( "POST", "/v1/memories/compile", json={"subject_id": subject_id}, model=CompileResult, + timeout=timeout, ) - async def compile_memories_async(self, subject_id: str) -> CompileJob: - """Submit async compilation. Returns immediately with a job_id for polling.""" + async def compile_memories_async( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> CompileJob: + """Submit async compilation. Returns immediately with a job_id for polling. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ + kwargs: dict[str, Any] = {} + if timeout is not None: + kwargs["timeout"] = timeout resp = await self._http.request( "POST", "/v1/memories/compile", json={"subject_id": subject_id, "async": True}, + **kwargs, ) if not resp.is_success: raise _parse_error(resp) return CompileJob.model_validate(resp.json()) - async def get_compile_status(self, job_id: str) -> CompileJob: - """Poll the status of an async compile job.""" + async def get_compile_status( + self, + job_id: str, + *, + timeout: float | None = None, + ) -> CompileJob: + """Poll the status of an async compile job. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( - "GET", f"/v1/memories/compile/{job_id}", model=CompileJob, + "GET", f"/v1/memories/compile/{job_id}", model=CompileJob, timeout=timeout, ) async def compile_memories_wait( @@ -837,16 +1040,20 @@ async def compile_memories_wait( *, poll_interval: float = 0.5, timeout: float = 60.0, + request_timeout: float | None = None, ) -> CompileJob: - """Submit async compilation and poll until completion or timeout.""" - import asyncio as _asyncio + """Submit async compilation and poll until completion or timeout. - job = await self.compile_memories_async(subject_id) + ``timeout`` is the total polling duration in seconds. + ``request_timeout`` overrides the client-level HTTP timeout for each + individual poll request. + """ + job = await self.compile_memories_async(subject_id, timeout=request_timeout) elapsed = 0.0 while elapsed < timeout: - await _asyncio.sleep(poll_interval) + await asyncio.sleep(poll_interval) elapsed += poll_interval - job = await self.get_compile_status(job.job_id) + job = await self.get_compile_status(job.job_id, timeout=request_timeout) if job.status in ("completed", "failed"): return job raise TimeoutError(f"Compile job {job.job_id} did not complete within {timeout}s") @@ -859,8 +1066,12 @@ async def search_memories( query: str | None = None, semantic: bool = False, limit: int = 20, + timeout: float | None = None, ) -> SearchResult: - """Search memories by kind, text query, or semantic similarity.""" + """Search memories by kind, text query, or semantic similarity. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ params: dict[str, Any] = {"subject_id": subject_id, "limit": limit} if kind: params["kind"] = kind @@ -868,7 +1079,9 @@ async def search_memories( params["q"] = query if semantic: params["semantic"] = "true" - return await self._request("GET", "/v1/memories/search", params=params, model=SearchResult) + return await self._request( + "GET", "/v1/memories/search", params=params, model=SearchResult, timeout=timeout, + ) # -- Context ----------------------------------------------------------- @@ -885,6 +1098,7 @@ async def get_context( parent_receipt_id: str | None = None, caller_id: str | None = None, caller_type: str | None = None, + timeout: float | None = None, ) -> ContextBundle: """Assemble a ranked, token-bounded context bundle. @@ -895,6 +1109,8 @@ async def get_context( When `emit_receipt=True` (or the tenant's receipts config is `always`), the returned bundle carries a `receipt_id` that can be fetched via `get_receipt()`. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = {"subject_id": subject_id, "task": task} if max_tokens is not None: @@ -913,19 +1129,27 @@ async def get_context( body["caller_id"] = caller_id if caller_type is not None: body["caller_type"] = caller_type - return await self._request("POST", "/v1/context", json=body, model=ContextBundle) + return await self._request( + "POST", "/v1/context", json=body, model=ContextBundle, timeout=timeout, + ) async def set_memory_labels( self, memory_id: str, labels: list[str], + *, + timeout: float | None = None, ) -> Memory: - """Replace a memory's sensitivity_labels (#50).""" + """Replace a memory's sensitivity_labels (#50). + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( "PATCH", f"/v1/memories/{memory_id}/labels", json={"sensitivity_labels": labels}, model=Memory, + timeout=timeout, ) async def list_suggested_labels( @@ -936,8 +1160,12 @@ async def list_suggested_labels( label: str | None = None, limit: int = 50, offset: int = 0, + timeout: float | None = None, ) -> SuggestedLabelsList: - """List memories carrying auto-derived suggested labels (admin/governance, v0.9 #158).""" + """List memories carrying auto-derived suggested labels (admin/governance, v0.9 #158). + + Pass ``timeout`` to override the client-level timeout for this call only. + """ params: dict[str, Any] = {"limit": limit, "offset": offset} if tenant_id is not None: params["tenant_id"] = tenant_id @@ -950,6 +1178,7 @@ async def list_suggested_labels( "/admin/memories/with-suggested-labels", params=params, model=SuggestedLabelsList, + timeout=timeout, ) async def promote_suggested_labels( @@ -958,10 +1187,14 @@ async def promote_suggested_labels( labels: list[str], *, tenant_id: str | None = None, + timeout: float | None = None, ) -> PromoteLabelsResult: """Promote a subset of a memory's suggested labels into authoritative ``sensitivity_labels`` (admin/governance, v0.9 #160). Review-only — - every label must already be a current suggestion (else 422).""" + every label must already be a current suggestion (else 422). + + Pass ``timeout`` to override the client-level timeout for this call only. + """ params: dict[str, Any] = {} if tenant_id is not None: params["tenant_id"] = tenant_id @@ -971,6 +1204,7 @@ async def promote_suggested_labels( json={"labels": labels}, params=params, model=PromoteLabelsResult, + timeout=timeout, ) async def get_context_string( @@ -979,17 +1213,29 @@ async def get_context_string( task: str, *, max_tokens: int | None = None, + timeout: float | None = None, ) -> str: - """Return just the assembled context string, ready to inject into a prompt.""" - bundle = await self.get_context(subject_id, task, max_tokens=max_tokens) + """Return just the assembled context string, ready to inject into a prompt. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ + bundle = await self.get_context(subject_id, task, max_tokens=max_tokens, timeout=timeout) return bundle.assembled_context # -- Receipts ---------------------------------------------------------- - async def get_receipt(self, receipt_id: str) -> Receipt: - """Fetch a single state-assembly receipt by ULID.""" + async def get_receipt( + self, + receipt_id: str, + *, + timeout: float | None = None, + ) -> Receipt: + """Fetch a single state-assembly receipt by ULID. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( - "GET", f"/v1/receipts/{receipt_id}", model=Receipt + "GET", f"/v1/receipts/{receipt_id}", model=Receipt, timeout=timeout, ) async def list_receipts( @@ -1000,8 +1246,12 @@ async def list_receipts( until: str | None = None, cursor: str | None = None, limit: int = 50, + timeout: float | None = None, ) -> ReceiptList: - """List receipts for a subject, newest first. Cursor-paginated.""" + """List receipts for a subject, newest first. Cursor-paginated. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ params: dict[str, Any] = {"subject_id": subject_id, "limit": limit} if since is not None: params["since"] = since @@ -1009,38 +1259,65 @@ async def list_receipts( params["until"] = until if cursor is not None: params["cursor"] = cursor - return await self._request("GET", "/v1/receipts", params=params, model=ReceiptList) + return await self._request( + "GET", "/v1/receipts", params=params, model=ReceiptList, timeout=timeout, + ) - async def verify_receipt(self, receipt_id: str) -> ReceiptVerifyResult: + async def verify_receipt( + self, + receipt_id: str, + *, + timeout: float | None = None, + ) -> ReceiptVerifyResult: """Async — verify the HMAC signature on a stored receipt (v0.9+ #157). - See :meth:`StatewaveClient.verify_receipt` for the full contract.""" + See :meth:`StatewaveClient.verify_receipt` for the full contract. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( "GET", f"/v1/receipts/{receipt_id}/verify", model=ReceiptVerifyResult, + timeout=timeout, ) - async def replay_receipt(self, receipt_id: str) -> ReceiptReplayResult: + async def replay_receipt( + self, + receipt_id: str, + *, + timeout: float | None = None, + ) -> ReceiptReplayResult: """Async — replay a receipt against current memories with the receipt's original policy bundle (v0.9+ #159). See :meth:`StatewaveClient.replay_receipt` for the full contract, - including the :class:`StatewaveUnreplayableError` cases.""" + including the :class:`StatewaveUnreplayableError` cases. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( "POST", f"/v1/receipts/{receipt_id}/replay", model=ReceiptReplayResult, + timeout=timeout, ) # -- Support: health, SLA, handoff, resolutions ------------------------ - async def get_health(self, subject_id: str) -> Health: + async def get_health( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> Health: """Compute the customer health score (0-100) for a subject. Returns the score, the state bucket (``healthy`` | ``watch`` | ``at_risk``), and the explainable factors that drove it. + + Pass ``timeout`` to override the client-level timeout for this call only. """ return await self._request( - "GET", f"/v1/subjects/{subject_id}/health", model=Health, + "GET", f"/v1/subjects/{subject_id}/health", model=Health, timeout=timeout, ) async def get_sla( @@ -1049,6 +1326,7 @@ async def get_sla( *, first_response_threshold_minutes: float | None = None, resolution_threshold_hours: float | None = None, + timeout: float | None = None, ) -> SLASummary: """Compute SLA metrics for a subject. @@ -1056,6 +1334,8 @@ async def get_sla( subject's sessions and flags breaches against the supplied thresholds. Both thresholds fall back to the server defaults (5 minutes / 24 hours) when omitted. + + Pass ``timeout`` to override the client-level timeout for this call only. """ params: dict[str, Any] = {} if first_response_threshold_minutes is not None: @@ -1064,6 +1344,7 @@ async def get_sla( params["resolution_threshold_hours"] = resolution_threshold_hours return await self._request( "GET", f"/v1/subjects/{subject_id}/sla", params=params, model=SLASummary, + timeout=timeout, ) async def create_handoff( @@ -1079,6 +1360,7 @@ async def create_handoff( parent_receipt_id: str | None = None, caller_id: str | None = None, caller_type: str | None = None, + timeout: float | None = None, ) -> Handoff: """Generate a handoff context pack for escalation or shift change. @@ -1088,6 +1370,8 @@ async def create_handoff( ``caller_id`` and ``caller_type`` are consumed by the sensitivity-label policy layer (#50); when the tenant config sets ``require_caller_identity: true``, both are mandatory. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = {"subject_id": subject_id, "session_id": session_id} if reason is not None: @@ -1106,7 +1390,9 @@ async def create_handoff( body["caller_id"] = caller_id if caller_type is not None: body["caller_type"] = caller_type - return await self._request("POST", "/v1/handoff", json=body, model=Handoff) + return await self._request( + "POST", "/v1/handoff", json=body, model=Handoff, timeout=timeout, + ) async def create_resolution( self, @@ -1116,11 +1402,14 @@ async def create_resolution( status: str = "open", resolution_summary: str | None = None, metadata: dict[str, Any] | None = None, + timeout: float | None = None, ) -> Resolution: """Create or update a resolution record for a support session. Upserts by ``subject_id`` + ``session_id``. ``status`` is one of ``open``, ``resolved``, or ``unresolved``. + + Pass ``timeout`` to override the client-level timeout for this call only. """ body: dict[str, Any] = { "subject_id": subject_id, @@ -1131,46 +1420,81 @@ async def create_resolution( body["resolution_summary"] = resolution_summary if metadata is not None: body["metadata"] = metadata - return await self._request("POST", "/v1/resolutions", json=body, model=Resolution) + return await self._request( + "POST", "/v1/resolutions", json=body, model=Resolution, timeout=timeout, + ) async def list_resolutions( self, subject_id: str, *, status: str | None = None, + timeout: float | None = None, ) -> list[Resolution]: """List resolution records for a subject. Optionally filter to a single ``status`` (``open`` | ``resolved`` | ``unresolved``). + + Pass ``timeout`` to override the client-level timeout for this call only. """ params: dict[str, Any] = {"subject_id": subject_id} if status is not None: params["status"] = status return await self._request( "GET", "/v1/resolutions", params=params, model=Resolution, is_list=True, + timeout=timeout, ) # -- Timeline ---------------------------------------------------------- - async def get_timeline(self, subject_id: str) -> Timeline: - """Get chronological subject timeline.""" + async def get_timeline( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> Timeline: + """Get chronological subject timeline. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( - "GET", "/v1/timeline", params={"subject_id": subject_id}, model=Timeline + "GET", "/v1/timeline", params={"subject_id": subject_id}, model=Timeline, + timeout=timeout, ) # -- Subjects ---------------------------------------------------------- - async def delete_subject(self, subject_id: str) -> DeleteResult: - """Permanently delete all data for a subject.""" - return await self._request("DELETE", f"/v1/subjects/{subject_id}", model=DeleteResult) + async def delete_subject( + self, + subject_id: str, + *, + timeout: float | None = None, + ) -> DeleteResult: + """Permanently delete all data for a subject. - async def list_subjects(self, *, limit: int = 50, offset: int = 0) -> ListSubjectsResult: - """List all known subjects with episode/memory counts.""" + Pass ``timeout`` to override the client-level timeout for this call only. + """ + return await self._request( + "DELETE", f"/v1/subjects/{subject_id}", model=DeleteResult, timeout=timeout, + ) + + async def list_subjects( + self, + *, + limit: int = 50, + offset: int = 0, + timeout: float | None = None, + ) -> ListSubjectsResult: + """List all known subjects with episode/memory counts. + + Pass ``timeout`` to override the client-level timeout for this call only. + """ return await self._request( "GET", "/v1/subjects", params={"limit": limit, "offset": offset}, model=ListSubjectsResult, + timeout=timeout, ) # -- Lifecycle --------------------------------------------------------- @@ -1195,13 +1519,17 @@ async def _request( json: Any = None, params: Any = None, is_list: bool = False, + timeout: float | None = None, ): - import asyncio - last_exc: Exception | None = None + extra: dict[str, Any] = {} + if timeout is not None: + extra["timeout"] = timeout for attempt in range(self._retry.max_retries + 1): try: - resp = await self._http.request(method, path, json=json, params=params) + resp = await self._http.request(method, path, json=json, params=params, **extra) + except asyncio.CancelledError: + raise except Exception as exc: if attempt < self._retry.max_retries: last_exc = exc diff --git a/tests/test_per_call_timeout.py b/tests/test_per_call_timeout.py new file mode 100644 index 0000000..7c3544c --- /dev/null +++ b/tests/test_per_call_timeout.py @@ -0,0 +1,186 @@ +"""Tests for per-call timeout override and asyncio.CancelledError propagation.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from statewave import AsyncStatewaveClient, StatewaveClient +from statewave.exceptions import StatewaveTimeoutError + + +def _mock_response(status: int, *, json_body: dict | None = None): + resp = MagicMock(spec=httpx.Response) + resp.status_code = status + resp.is_success = 200 <= status < 300 + resp.headers = {} + resp.reason_phrase = "OK" if resp.is_success else "Error" + resp.json.return_value = json_body or { + "id": "ep-1", "subject_id": "s1", "source": "t", "type": "t", + "payload": {}, "metadata": {}, "provenance": {}, "created_at": "2026-01-01T00:00:00Z", + } + return resp + + +class TestSyncPerCallTimeout: + def test_timeout_passed_to_httpx(self): + """Per-call timeout is forwarded to httpx as a request kwarg.""" + client = StatewaveClient() + resp = _mock_response(200) + + with patch.object(client._http, "request", return_value=resp) as mock_req: + from statewave.models import Episode + with patch("statewave.models.Episode.model_validate", return_value=MagicMock()): + client._request("POST", "/v1/episodes", model=Episode, timeout=5.0) + + _args, kwargs = mock_req.call_args + assert kwargs.get("timeout") == 5.0 + + def test_no_timeout_kwarg_when_none(self): + """When timeout=None, httpx is not given an explicit timeout kwarg.""" + client = StatewaveClient() + resp = _mock_response(200) + + with patch.object(client._http, "request", return_value=resp) as mock_req: + from statewave.models import Episode + with patch("statewave.models.Episode.model_validate", return_value=MagicMock()): + client._request("POST", "/v1/episodes", model=Episode) + + _args, kwargs = mock_req.call_args + assert "timeout" not in kwargs + + def test_public_method_passes_timeout(self): + """get_context() accepts and forwards timeout to _request.""" + client = StatewaveClient() + resp = _mock_response(200, json_body={ + "assembled_context": "ctx", "items": [], "token_count": 0, + "subject_id": "s1", "task": "t", + }) + + with patch.object(client._http, "request", return_value=resp) as mock_req: + from statewave.models import ContextBundle + with patch("statewave.models.ContextBundle.model_validate", return_value=MagicMock()): + client.get_context("s1", "task", timeout=3.0) + + _args, kwargs = mock_req.call_args + assert kwargs.get("timeout") == 3.0 + + def test_keyboard_interrupt_not_retried(self): + """KeyboardInterrupt propagates immediately without retry.""" + from statewave import RetryConfig + client = StatewaveClient(retry=RetryConfig(max_retries=3, backoff_base=0.0, jitter=False)) + + with patch.object(client._http, "request", side_effect=KeyboardInterrupt): + from statewave.models import Episode + with pytest.raises(KeyboardInterrupt): + client._request("POST", "/v1/episodes", model=Episode) + + def test_compile_memories_wait_request_timeout(self): + """compile_memories_wait() passes request_timeout to each poll.""" + client = StatewaveClient() + job_pending = MagicMock() + job_pending.job_id = "job-1" + job_pending.status = "pending" + job_done = MagicMock() + job_done.job_id = "job-1" + job_done.status = "completed" + + with patch.object(client, "compile_memories_async", return_value=job_pending) as mock_submit: + with patch.object(client, "get_compile_status", return_value=job_done) as mock_poll: + with patch("time.sleep"): + client.compile_memories_wait("s1", request_timeout=2.0) + + mock_submit.assert_called_once_with("s1", timeout=2.0) + mock_poll.assert_called_once_with("job-1", timeout=2.0) + + +class TestAsyncPerCallTimeout: + @pytest.mark.asyncio + async def test_timeout_passed_to_httpx(self): + """Async per-call timeout is forwarded to httpx.""" + client = AsyncStatewaveClient() + resp = _mock_response(200) + + with patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp) as mock_req: + from statewave.models import Episode + with patch("statewave.models.Episode.model_validate", return_value=MagicMock()): + await client._request("POST", "/v1/episodes", model=Episode, timeout=7.0) + + _args, kwargs = mock_req.call_args + assert kwargs.get("timeout") == 7.0 + + @pytest.mark.asyncio + async def test_no_timeout_kwarg_when_none(self): + """When timeout=None, async httpx is not given an explicit timeout kwarg.""" + client = AsyncStatewaveClient() + resp = _mock_response(200) + + with patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp) as mock_req: + from statewave.models import Episode + with patch("statewave.models.Episode.model_validate", return_value=MagicMock()): + await client._request("POST", "/v1/episodes", model=Episode) + + _args, kwargs = mock_req.call_args + assert "timeout" not in kwargs + + @pytest.mark.asyncio + async def test_cancelled_error_not_retried(self): + """asyncio.CancelledError propagates immediately without retry.""" + from statewave import RetryConfig + client = AsyncStatewaveClient( + retry=RetryConfig(max_retries=3, backoff_base=0.0, jitter=False), + ) + + with patch.object( + client._http, "request", + new_callable=AsyncMock, side_effect=asyncio.CancelledError, + ): + from statewave.models import Episode + with pytest.raises(asyncio.CancelledError): + await client._request("POST", "/v1/episodes", model=Episode) + + @pytest.mark.asyncio + async def test_cancelled_error_not_swallowed_on_retry_sleep(self): + """CancelledError raised during retry sleep propagates correctly.""" + from statewave import RetryConfig + client = AsyncStatewaveClient( + retry=RetryConfig(max_retries=3, backoff_base=0.01, jitter=False), + ) + resp_503 = _mock_response(503, json_body={"error": {"code": "unavailable", "message": "down"}}) + + async def raise_on_second_sleep(*_a, **_kw): + raise asyncio.CancelledError + + with patch.object(client._http, "request", new_callable=AsyncMock, return_value=resp_503): + with patch("asyncio.sleep", side_effect=raise_on_second_sleep): + from statewave.models import Episode + with pytest.raises(asyncio.CancelledError): + await client._request("POST", "/v1/episodes", model=Episode) + + @pytest.mark.asyncio + async def test_compile_memories_wait_request_timeout(self): + """async compile_memories_wait() passes request_timeout to each poll.""" + client = AsyncStatewaveClient() + job_pending = MagicMock() + job_pending.job_id = "job-1" + job_pending.status = "pending" + job_done = MagicMock() + job_done.job_id = "job-1" + job_done.status = "completed" + + with patch.object( + client, "compile_memories_async", + new_callable=AsyncMock, return_value=job_pending, + ) as mock_submit: + with patch.object( + client, "get_compile_status", + new_callable=AsyncMock, return_value=job_done, + ) as mock_poll: + with patch("asyncio.sleep", new_callable=AsyncMock): + await client.compile_memories_wait("s1", request_timeout=2.0) + + mock_submit.assert_called_once_with("s1", timeout=2.0) + mock_poll.assert_called_once_with("job-1", timeout=2.0) From 54bc296024183e49dfa00d331fbf39853233b2df Mon Sep 17 00:00:00 2001 From: smaramwbc <145447586+smaramwbc@users.noreply.github.com> Date: Tue, 16 Jun 2026 16:40:02 +0100 Subject: [PATCH 2/2] fix: remove unused imports flagged by ruff --- tests/test_per_call_timeout.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_per_call_timeout.py b/tests/test_per_call_timeout.py index 7c3544c..14d22dd 100644 --- a/tests/test_per_call_timeout.py +++ b/tests/test_per_call_timeout.py @@ -9,7 +9,6 @@ import pytest from statewave import AsyncStatewaveClient, StatewaveClient -from statewave.exceptions import StatewaveTimeoutError def _mock_response(status: int, *, json_body: dict | None = None): @@ -61,7 +60,6 @@ def test_public_method_passes_timeout(self): }) with patch.object(client._http, "request", return_value=resp) as mock_req: - from statewave.models import ContextBundle with patch("statewave.models.ContextBundle.model_validate", return_value=MagicMock()): client.get_context("s1", "task", timeout=3.0)