oneshot-python 0.19.0__tar.gz → 0.20.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (26) hide show
  1. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/.gitignore +3 -0
  2. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/PKG-INFO +1 -1
  3. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/README.md +37 -0
  4. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/__init__.py +8 -0
  5. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/_errors.py +46 -0
  6. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/_types.py +40 -0
  7. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/client.py +295 -19
  8. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/pyproject.toml +1 -1
  9. oneshot_python-0.20.2/tests/test_budgets.py +184 -0
  10. oneshot_python-0.20.2/tests/test_poll_backoff.py +107 -0
  11. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_request_id.py +1 -4
  12. oneshot_python-0.20.2/tests/test_wait_false.py +88 -0
  13. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/x402.py +0 -0
  14. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/__init__.py +0 -0
  15. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_balance.py +0 -0
  16. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_charge_amount.py +0 -0
  17. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_compute.py +0 -0
  18. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_domains.py +0 -0
  19. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_email_payload.py +0 -0
  20. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_emergency_error.py +0 -0
  21. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_max_cost_header.py +0 -0
  22. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_payment_rejection.py +0 -0
  23. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_phones_pending.py +0 -0
  24. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_tag_receipt_value.py +0 -0
  25. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_x402.py +0 -0
  26. {oneshot_python-0.19.0 → oneshot_python-0.20.2}/uv.lock +0 -0
@@ -90,3 +90,6 @@ output/
90
90
 
91
91
  # Claude Code agent metadata (lockfile + local settings)
92
92
  .claude/
93
+
94
+ # generated by scripts/deploy-blog-generator.sh
95
+ scripts/humanizer.vendored.ts
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: oneshot-python
3
- Version: 0.19.0
3
+ Version: 0.20.2
4
4
  Summary: Core Python SDK for the OneShot API — HTTP client with x402 payment handling
5
5
  License-Expression: MIT
6
6
  Requires-Python: >=3.10
@@ -66,6 +66,43 @@ Read endpoints (inbox, SMS inbox, notifications, balance, browser profiles) retu
66
66
 
67
67
  This is automatic — no extra code beyond `private_key`. **Requires `oneshot-python >= 0.17.0`.** Older versions keep working until the API enables enforcement, after which they must send a valid `x-agent-proof`.
68
68
 
69
+ ## Spend Budgets
70
+
71
+ Cap what the agent can spend so a runaway loop can't drain the wallet overnight:
72
+
73
+ ```python
74
+ client = OneShotClient(
75
+ "0x...",
76
+ budgets={
77
+ "daily": 50, # max USDC per UTC day
78
+ "per_transaction": 5, # max USDC for any single call
79
+ "alert_at": 0.8, # warn at 80% of daily
80
+ "pause_at": 1.0, # stop paying at 100%
81
+ },
82
+ alert_email="you@example.com",
83
+ )
84
+ ```
85
+
86
+ Unlike `max_cost` (a per-call client-side check), budgets are **enforced server-side** against your receipt ledger. The config is synced once before your first paid call, so the daily figure is a true per-agent total — shared across every process and restart using the same wallet.
87
+
88
+ At `alert_at` you get a `budget_warning` notification (in-app via `/v1/tools/notifications`, plus email if set). At `pause_at` paid calls raise `BudgetExceededError` until the window resets at the next UTC midnight:
89
+
90
+ ```python
91
+ from oneshot import BudgetExceededError
92
+
93
+ try:
94
+ client.research(topic="...")
95
+ except BudgetExceededError as e:
96
+ print(f"{e.reason} budget hit — ${e.spent} spent, resumes {e.resets_at}")
97
+
98
+ client.get_budgets()
99
+ # {"daily_usdc": 50, "spent_today_usdc": "41.20", "pct_used": 0.824, "resets_at": "..."}
100
+ ```
101
+
102
+ Omit `budgets` to leave whatever is stored server-side untouched — agents without a budget are unlimited. **Requires `oneshot-python >= 0.20.0`.**
103
+
104
+ The sync **fails closed**: if the budget can't be confirmed with the server, the paid call raises `BudgetSyncError` and is not made (retried on the next call). An invalid config (`daily=-1`, `alert_at=2`) raises `ValidationError` at construction.
105
+
69
106
  ## Configuration
70
107
 
71
108
  ```python
@@ -1,6 +1,8 @@
1
1
  """oneshot-python — Core Python SDK for the OneShot API."""
2
2
 
3
3
  from oneshot._errors import (
4
+ BudgetExceededError,
5
+ BudgetSyncError,
4
6
  ContentBlockedError,
5
7
  EmergencyNumberError,
6
8
  JobError,
@@ -11,6 +13,8 @@ from oneshot._errors import (
11
13
  ValidationError,
12
14
  )
13
15
  from oneshot._types import (
16
+ AgentBudgetConfig,
17
+ AgentBudgetStatus,
14
18
  ComputeBudgetStatus,
15
19
  ComputeCancelResult,
16
20
  ComputeFundResult,
@@ -34,12 +38,16 @@ __all__ = [
34
38
  "OneShotError",
35
39
  "ToolError",
36
40
  "PaymentError",
41
+ "BudgetExceededError",
42
+ "BudgetSyncError",
37
43
  "JobError",
38
44
  "JobTimeoutError",
39
45
  "ValidationError",
40
46
  "ContentBlockedError",
41
47
  "EmergencyNumberError",
42
48
  "sign_payment_authorization",
49
+ "AgentBudgetConfig",
50
+ "AgentBudgetStatus",
43
51
  "ComputeSchedule",
44
52
  "ComputeGoalResult",
45
53
  "ComputeGoalStatus",
@@ -81,3 +81,49 @@ class EmergencyNumberError(OneShotError):
81
81
  def __init__(self, message: str, blocked_number: str) -> None:
82
82
  super().__init__(message)
83
83
  self.blocked_number = blocked_number
84
+
85
+
86
+ class BudgetExceededError(OneShotError):
87
+ """The agent's own spend budget stopped this call (HTTP 403 ``budget_exceeded``).
88
+
89
+ Distinct from :class:`PaymentError`: nothing was signed and nothing was
90
+ charged. A retry cannot succeed until the daily window resets
91
+ (``resets_at``) or the budget is raised — which is why the server answers
92
+ 403 rather than 402. ``reason`` is ``"daily"`` when cumulative UTC-day
93
+ spend would cross the budget, or ``"per_transaction"`` when this single
94
+ call is larger than the per-call cap.
95
+
96
+ Mirrors ``BudgetExceededError`` in ``libs/agent-sdk/src/errors.ts``.
97
+ """
98
+
99
+ def __init__(
100
+ self,
101
+ message: str,
102
+ reason: str,
103
+ cap: float | None = None,
104
+ spent: float | None = None,
105
+ charge: float | None = None,
106
+ resets_at: str | None = None,
107
+ ) -> None:
108
+ super().__init__(message)
109
+ self.reason = reason
110
+ self.cap = cap
111
+ self.spent = spent
112
+ self.charge = charge
113
+ self.resets_at = resets_at
114
+
115
+
116
+ class BudgetSyncError(OneShotError):
117
+ """The configured spend budget could not be confirmed with the server, so
118
+ the paid call was NOT made.
119
+
120
+ Budgets are enforced server-side; if the sync fails (network error, 5xx,
121
+ rate limit, rejected config) proceeding would silently drop the guardrail
122
+ the developer configured. The sync is retried on the next paid call.
123
+ ``status`` is the HTTP status when the server answered, else None.
124
+ """
125
+
126
+ def __init__(self, message: str, status: int | None = None, body: str | None = None) -> None:
127
+ super().__init__(message)
128
+ self.status = status
129
+ self.body = body
@@ -181,6 +181,12 @@ class DomainPoolEntry(TypedDict, total=False):
181
181
 
182
182
  domain: str
183
183
  default_from: str
184
+ # 'mailbox' = a real mailbox is provisioned per address, so every address not
185
+ # already active in addresses[] bills the one-time mailbox_provisioning_fee on
186
+ # its first send — the canonical default_from included. 'relay' = header-only,
187
+ # no per-address fee. A property of the DOMAIN, fixed at setup: omitting
188
+ # mailbox_mode on a quote does not make the send relay.
189
+ mailbox_mode: str # relay | mailbox
184
190
  addresses: list[DomainAddressEntry]
185
191
  pool_status: str # active | paused | removed
186
192
  warmup_state: str # warming | warmed | degraded
@@ -227,3 +233,37 @@ __all__ = [
227
233
  "DomainPoolListResult",
228
234
  "DomainPoolStatusResult",
229
235
  ]
236
+
237
+
238
+ # ── Agent spend budgets ──────────────────────────────────────────────────
239
+
240
+
241
+ class AgentBudgetConfig(TypedDict, total=False):
242
+ """Spend budget passed to ``OneShotClient(budgets=...)``.
243
+
244
+ Mirrors ``AgentBudgetConfig`` in TS (snake_case here; camelCase keys are
245
+ accepted too and normalised). Synced to the server once, before the first
246
+ paid call, and enforced there against the receipt ledger.
247
+ """
248
+
249
+ daily: float
250
+ """Max USDC per UTC day."""
251
+ per_transaction: float
252
+ """Max USDC for any single call."""
253
+ alert_at: float
254
+ """Fraction of ``daily`` that triggers a ``budget_warning`` (default 0.8)."""
255
+ pause_at: float
256
+ """Fraction of ``daily`` at which paid calls stop (default 1.0)."""
257
+
258
+
259
+ class AgentBudgetStatus(TypedDict):
260
+ """Live utilization from ``GET /v1/agents/me/budgets``. Mirrors ``AgentBudgetStatus`` in TS."""
261
+
262
+ daily_usdc: Optional[float]
263
+ per_transaction_usdc: Optional[float]
264
+ alert_at: Optional[float]
265
+ pause_at: Optional[float]
266
+ spent_today_usdc: str
267
+ remaining_usdc: Optional[str]
268
+ pct_used: Optional[float]
269
+ resets_at: str
@@ -8,6 +8,7 @@ GET/POST for free endpoints.
8
8
  from __future__ import annotations
9
9
 
10
10
  import asyncio
11
+ import math
11
12
  import json
12
13
  import time
13
14
  from decimal import Decimal
@@ -18,6 +19,8 @@ import httpx
18
19
  from eth_account import Account
19
20
 
20
21
  from oneshot._errors import (
22
+ BudgetExceededError,
23
+ BudgetSyncError,
21
24
  ContentBlockedError,
22
25
  EmergencyNumberError,
23
26
  JobError,
@@ -28,6 +31,8 @@ from oneshot._errors import (
28
31
  ValidationError,
29
32
  )
30
33
  from oneshot._types import (
34
+ AgentBudgetConfig,
35
+ AgentBudgetStatus,
31
36
  DomainPoolListResult,
32
37
  DomainPoolStatusResult,
33
38
  )
@@ -46,7 +51,7 @@ try:
46
51
 
47
52
  SDK_VERSION = _pkg_version("oneshot-python")
48
53
  except Exception: # pragma: no cover - editable/source runs without dist metadata
49
- SDK_VERSION = "0.19.0"
54
+ SDK_VERSION = "0.20.2"
50
55
 
51
56
  # ---------------------------------------------------------------------------
52
57
  # Environment configuration
@@ -56,7 +61,10 @@ _BASE_URL = "https://win.oneshotagent.com"
56
61
  _CHAIN_ID = 8453
57
62
  _USDC_ADDRESS = "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913"
58
63
 
59
- _POLL_INTERVAL = 2.0 # seconds
64
+ # Poll cadence for GET /v1/requests/{id}: fast first checks (most jobs finish in
65
+ # a few seconds), settling at 2s. Index = number of polls already made.
66
+ _POLL_BACKOFF = (0.3, 0.6, 1.0, 2.0) # seconds
67
+ _POLL_RETRY_BASE = 2.0 # seconds, transient-error retry backoff base
60
68
  _MAX_POLL_RETRIES = 3
61
69
 
62
70
 
@@ -107,6 +115,39 @@ def _build_email_payload(
107
115
  return payload
108
116
 
109
117
 
118
+ def _parse_budget_rejection(resp: Any) -> Optional[BudgetExceededError]:
119
+ """Map a 403 ``budget_exceeded`` body onto :class:`BudgetExceededError`.
120
+
121
+ Lets callers catch "my own budget stopped this" separately from an auth
122
+ failure or a payment rejection. Returns None for any other response so the
123
+ caller falls back to ``ToolError``.
124
+ """
125
+ if resp.status_code != 403:
126
+ return None
127
+ try:
128
+ data = resp.json()
129
+ except Exception: # noqa: BLE001 - non-JSON body
130
+ return None
131
+ if not isinstance(data, dict) or data.get("error") != "budget_exceeded":
132
+ return None
133
+ b = data.get("budget") or {}
134
+
135
+ def _num(v: Any) -> float | None:
136
+ try:
137
+ return float(v) if v is not None else None
138
+ except (TypeError, ValueError):
139
+ return None
140
+
141
+ return BudgetExceededError(
142
+ data.get("message") or "Agent spend budget exceeded",
143
+ "per_transaction" if b.get("reason") == "per_transaction" else "daily",
144
+ cap=_num(b.get("cap")),
145
+ spent=_num(b.get("spent")),
146
+ charge=_num(b.get("charge")),
147
+ resets_at=b.get("resets_at"),
148
+ )
149
+
150
+
110
151
  def _parse_payment_rejection(resp: Any) -> Optional[PaymentError]:
111
152
  """Build a ``PaymentError`` from a 402 that names its own cause.
112
153
 
@@ -183,6 +224,44 @@ def _resolve_charge_amount(accepted: Any, payment_request: dict[str, Any]) -> st
183
224
  return from_header
184
225
 
185
226
 
227
+ _BUDGET_KEY_ALIASES = {
228
+ "perTransaction": "per_transaction",
229
+ "alertAt": "alert_at",
230
+ "pauseAt": "pause_at",
231
+ }
232
+
233
+
234
+ def _normalize_budgets(budgets: Optional[AgentBudgetConfig]) -> Optional[dict[str, float]]:
235
+ """Accept the TS camelCase keys too; keep only known, numeric fields."""
236
+ if not budgets:
237
+ return None
238
+ out: dict[str, float] = {}
239
+ for k, v in budgets.items():
240
+ key = _BUDGET_KEY_ALIASES.get(k, k)
241
+ # Validate here, at construction, what the server would reject — a
242
+ # typo'd cap must never silently become "no cap" on the first call.
243
+ # That includes a typo'd KEY: `{"daliy": 50}` used to normalise to
244
+ # "no budget" and skip the sync entirely.
245
+ if key not in ("daily", "per_transaction", "alert_at", "pause_at"):
246
+ raise ValidationError(f"budgets.{k} is not a recognized field", f"budgets.{k}")
247
+ if v is None:
248
+ continue
249
+ if isinstance(v, bool): # float(True) == 1.0 would pass as a real cap
250
+ raise ValidationError(f"budgets.{key} must be a number", f"budgets.{key}")
251
+ try:
252
+ value = float(v)
253
+ except (TypeError, ValueError):
254
+ raise ValidationError(f"budgets.{key} must be a number", f"budgets.{key}") from None
255
+ if not math.isfinite(value):
256
+ raise ValidationError(f"budgets.{key} must be finite", f"budgets.{key}")
257
+ if key in ("daily", "per_transaction") and value <= 0:
258
+ raise ValidationError(f"budgets.{key} must be positive", f"budgets.{key}")
259
+ if key in ("alert_at", "pause_at") and not 0 < value <= 1:
260
+ raise ValidationError(f"budgets.{key} must be in (0, 1]", f"budgets.{key}")
261
+ out[key] = value
262
+ return out or None
263
+
264
+
186
265
  class OneShotClient:
187
266
  """Synchronous + async HTTP client for the OneShot API.
188
267
 
@@ -196,13 +275,22 @@ class OneShotClient:
196
275
  *,
197
276
  base_url: Optional[str] = None,
198
277
  debug: bool = False,
278
+ budgets: Optional[AgentBudgetConfig] = None,
279
+ alert_email: Optional[str] = None,
199
280
  ) -> None:
200
281
  self.base_url = base_url or _BASE_URL
201
282
  self.chain_id: int = _CHAIN_ID
202
283
  self.usdc_address: str = _USDC_ADDRESS
203
284
  self.debug = debug
204
285
 
205
- # Derive wallet address from private key
286
+ # Spend budget: synced to the server once, before the first paid call,
287
+ # and enforced there against the receipt ledger. Omit to leave whatever
288
+ # is stored server-side untouched (no budget = unlimited).
289
+ self._budgets = _normalize_budgets(budgets)
290
+ self._alert_email = alert_email
291
+ self._budget_synced = False
292
+ self._budget_sync_lock: Optional[asyncio.Lock] = None
293
+
206
294
  self._private_key = private_key
207
295
  acct = Account.from_key(private_key)
208
296
  self.address: str = acct.address
@@ -239,10 +327,93 @@ class OneShotClient:
239
327
  self._log(f"Failed to sign read proof (continuing without): {e}")
240
328
  return headers
241
329
 
330
+ def _write_headers(self) -> dict[str, str]:
331
+ """Like ``_read_headers`` but with a ``write``-scoped proof, for the
332
+ signed mutation routes (budgets). The server enforces the proof on
333
+ those routes outright — there is no log-only fallback — so a signing
334
+ failure here is surfaced rather than swallowed."""
335
+ headers = self._headers()
336
+ headers["x-agent-proof"] = sign_read_proof(self._private_key, self.address, scope="write")
337
+ return headers
338
+
242
339
  def _log(self, msg: str) -> None:
243
340
  if self.debug:
244
341
  print(f"[OneShot] {msg}")
245
342
 
343
+ # ------------------------------------------------------------------
344
+ # Spend budget
345
+ # ------------------------------------------------------------------
346
+
347
+ async def aensure_budgets_synced(self) -> None:
348
+ """Push ``budgets`` / ``alert_email`` to the server once (async).
349
+
350
+ Called at the top of every paid call so the server-side gate knows the
351
+ budget on the very first request — including for a brand-new wallet,
352
+ which the budget route registers. Guarded by a lock so concurrent
353
+ first calls issue one PUT.
354
+
355
+ FAILS CLOSED: if the server can't confirm the budget (network error,
356
+ 5xx, 429, rejected config) this raises :class:`BudgetSyncError` and the
357
+ paid call is not made — proceeding would silently drop the guardrail
358
+ the developer configured. ``_budget_synced`` is set only on success,
359
+ so the next paid call retries.
360
+ """
361
+ if self._budget_synced or (self._budgets is None and self._alert_email is None):
362
+ return
363
+ if self._budget_sync_lock is None:
364
+ self._budget_sync_lock = asyncio.Lock()
365
+ async with self._budget_sync_lock:
366
+ if self._budget_synced:
367
+ return
368
+ body: dict[str, Any] = dict(self._budgets or {})
369
+ if self._alert_email is not None:
370
+ body["alert_email"] = self._alert_email
371
+ try:
372
+ url = f"{self.base_url}/v1/agents/me/budgets"
373
+ async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
374
+ resp = await client.put(url, headers=self._write_headers(), json=body)
375
+ except Exception as e: # noqa: BLE001
376
+ self._log(f"Budget sync failed ({e}) — paid call refused until the budget is confirmed")
377
+ raise BudgetSyncError(f"Could not sync spend budget (network): {e}") from e
378
+ if not resp.is_success:
379
+ self._log(f"Budget sync rejected ({resp.status_code}) — paid call refused")
380
+ raise BudgetSyncError(
381
+ f"Could not sync spend budget ({resp.status_code}): {resp.text}",
382
+ status=resp.status_code,
383
+ body=resp.text,
384
+ )
385
+ self._budget_synced = True
386
+ self._log("Budget synced")
387
+
388
+ def _assert_within_budget(self, total: Any) -> None:
389
+ """Local fast-fail on the per-transaction cap — no round-trip for an
390
+ oversized call. The server enforces the same cap (and the daily one,
391
+ which needs the ledger) regardless of client; this mirrors the
392
+ ``max_cost`` SDK-checks/server-enforces split."""
393
+ cap = (self._budgets or {}).get("per_transaction")
394
+ if not cap or cap <= 0:
395
+ return
396
+ try:
397
+ amount = float(total)
398
+ except (TypeError, ValueError):
399
+ return
400
+ if amount > cap:
401
+ raise BudgetExceededError(
402
+ f"Quote ${total} exceeds this agent's per-transaction budget of ${cap}",
403
+ "per_transaction",
404
+ cap=cap,
405
+ charge=amount,
406
+ )
407
+
408
+ def get_budgets(self) -> AgentBudgetStatus:
409
+ """Current spend budget and today's utilization. Blocking."""
410
+ return asyncio.get_event_loop().run_until_complete(self.aget_budgets())
411
+
412
+ async def aget_budgets(self) -> AgentBudgetStatus:
413
+ """Current spend budget and today's utilization. Async."""
414
+ body = await self.acall_free_get("/v1/agents/me/budgets")
415
+ return body.get("data", body) if isinstance(body, dict) else body
416
+
246
417
  # ------------------------------------------------------------------
247
418
  # Paid tool flow (POST -> 402 -> sign -> POST w/ payment -> poll)
248
419
  # ------------------------------------------------------------------
@@ -254,11 +425,16 @@ class OneShotClient:
254
425
  *,
255
426
  max_cost: Optional[float] = None,
256
427
  timeout_sec: int = 120,
428
+ wait: bool = True,
257
429
  wait_for_phones: bool = False,
258
430
  phone_timeout_sec: int = 360,
259
431
  ) -> Any:
260
432
  """Execute a paid tool call (blocking). Handles the full x402 flow.
261
433
 
434
+ Set ``wait=False`` to return as soon as the job is queued — you get
435
+ ``{"request_id": ..., "status": ...}`` and can resolve it later with
436
+ :meth:`wait_for_result`.
437
+
262
438
  Set ``wait_for_phones=True`` to keep polling AFTER status=completed
263
439
  until Apollo's async phone-reveal webhook delivers phone numbers
264
440
  (or until ``phone_timeout_sec`` expires, default 6 min — Apollo says
@@ -270,11 +446,50 @@ class OneShotClient:
270
446
  payload,
271
447
  max_cost=max_cost,
272
448
  timeout_sec=timeout_sec,
449
+ wait=wait,
273
450
  wait_for_phones=wait_for_phones,
274
451
  phone_timeout_sec=phone_timeout_sec,
275
452
  )
276
453
  )
277
454
 
455
+ def wait_for_result(
456
+ self,
457
+ request_id: str,
458
+ *,
459
+ timeout_sec: int = 120,
460
+ wait_for_phones: bool = False,
461
+ phone_timeout_sec: int = 360,
462
+ ) -> Any:
463
+ """Wait for a previously dispatched job (e.g. from ``wait=False``) and return its result."""
464
+ return asyncio.get_event_loop().run_until_complete(
465
+ self.await_result(
466
+ request_id,
467
+ timeout_sec=timeout_sec,
468
+ wait_for_phones=wait_for_phones,
469
+ phone_timeout_sec=phone_timeout_sec,
470
+ )
471
+ )
472
+
473
+ async def await_result(
474
+ self,
475
+ request_id: str,
476
+ *,
477
+ timeout_sec: int = 120,
478
+ wait_for_phones: bool = False,
479
+ phone_timeout_sec: int = 360,
480
+ ) -> Any:
481
+ """Async form of :meth:`wait_for_result`: poll ``GET /v1/requests/{id}`` until terminal."""
482
+ if not request_id:
483
+ raise ValidationError("request_id is required", "request_id")
484
+ async with httpx.AsyncClient(timeout=30.0) as client:
485
+ return await self._poll_job(
486
+ client,
487
+ request_id,
488
+ timeout_sec,
489
+ wait_for_phones=wait_for_phones,
490
+ phone_timeout_sec=phone_timeout_sec,
491
+ )
492
+
278
493
  async def acall_tool(
279
494
  self,
280
495
  endpoint: str,
@@ -282,11 +497,15 @@ class OneShotClient:
282
497
  *,
283
498
  max_cost: Optional[float] = None,
284
499
  timeout_sec: int = 120,
500
+ wait: bool = True,
285
501
  wait_for_phones: bool = False,
286
502
  phone_timeout_sec: int = 360,
287
503
  ) -> Any:
288
- """Execute a paid tool call (async). Handles the full x402 flow."""
289
- # Validate memo
504
+ """Execute a paid tool call (async). Handles the full x402 flow.
505
+
506
+ With ``wait=False`` a queued job is returned as
507
+ ``{"request_id": ..., "status": ...}`` instead of being polled.
508
+ """
290
509
  memo = payload.get("memo")
291
510
  if memo is not None:
292
511
  if not isinstance(memo, str) or not memo.strip():
@@ -297,7 +516,6 @@ class OneShotClient:
297
516
  elif "/inbox" not in endpoint and "/notifications" not in endpoint and "/balance" not in endpoint:
298
517
  self._log("No memo provided — consider adding a reason for audit trail")
299
518
 
300
- # Validate decisionContext
301
519
  dc = payload.get("decisionContext") or payload.get("decision_context")
302
520
  if dc is not None:
303
521
  # Normalize to camelCase key for API
@@ -312,10 +530,20 @@ class OneShotClient:
312
530
 
313
531
  url = f"{self.base_url}{endpoint}"
314
532
 
533
+ # One-time push of the budget config before the first paid call, so
534
+ # the server-side gate knows about it on this very request.
535
+ await self.aensure_budgets_synced()
536
+
315
537
  async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client:
316
538
  # Step 1 — Initial POST (expect 402 for paid tools)
317
539
  resp = await client.post(url, headers=self._headers(max_cost=max_cost), json=payload)
318
540
 
541
+ # The spend-budget gate rejects at quote time (403) so nothing is
542
+ # signed; surface it as the typed error.
543
+ budget_rejection = _parse_budget_rejection(resp)
544
+ if budget_rejection is not None:
545
+ raise budget_rejection
546
+
319
547
  # Handle validation / content-blocked / emergency-number errors
320
548
  if resp.status_code == 400:
321
549
  data = resp.json()
@@ -343,6 +571,8 @@ class OneShotClient:
343
571
  "pending",
344
572
  "processing",
345
573
  ):
574
+ if not wait:
575
+ return {"request_id": result["request_id"], "status": result["status"]}
346
576
  return await self._poll_job(
347
577
  client,
348
578
  result["request_id"],
@@ -364,7 +594,6 @@ class OneShotClient:
364
594
  context = quote_data.get("context", {})
365
595
  quote_id = context.get("quote_id")
366
596
 
367
- # Check max_cost
368
597
  total = context.get("total") or context.get("pricing", {}).get("total")
369
598
  if max_cost is not None and total is not None:
370
599
  if float(total) > max_cost:
@@ -376,6 +605,7 @@ class OneShotClient:
376
605
 
377
606
  # Step 3 — Sign x402 payment (zero-cost auth if credits cover full cost)
378
607
  charge = _resolve_charge_amount(accepted, payment_request)
608
+ self._assert_within_budget(charge)
379
609
  amount = float(charge)
380
610
  if amount == 0:
381
611
  self._log("Credits cover full cost — sending zero-cost authorization")
@@ -424,6 +654,9 @@ class OneShotClient:
424
654
  rejection = _parse_payment_rejection(resp2)
425
655
  if rejection is not None:
426
656
  raise rejection
657
+ budget_rejection = _parse_budget_rejection(resp2)
658
+ if budget_rejection is not None:
659
+ raise budget_rejection
427
660
  raise ToolError(
428
661
  "Tool request failed after payment",
429
662
  resp2.status_code,
@@ -438,6 +671,8 @@ class OneShotClient:
438
671
  "processing",
439
672
  ):
440
673
  self._log(f"Job queued: {result['request_id']}")
674
+ if not wait:
675
+ return {"request_id": result["request_id"], "status": result["status"]}
441
676
  return await self._poll_job(
442
677
  client,
443
678
  result["request_id"],
@@ -524,6 +759,29 @@ class OneShotClient:
524
759
  return {"success": True}
525
760
  return resp.json()
526
761
 
762
+ def call_free_put(
763
+ self,
764
+ endpoint: str,
765
+ payload: Optional[dict[str, Any]] = None,
766
+ ) -> Any:
767
+ """PUT to a signed mutation endpoint (blocking)."""
768
+ return asyncio.get_event_loop().run_until_complete(
769
+ self.acall_free_put(endpoint, payload)
770
+ )
771
+
772
+ async def acall_free_put(
773
+ self,
774
+ endpoint: str,
775
+ payload: Optional[dict[str, Any]] = None,
776
+ ) -> Any:
777
+ """PUT to a signed mutation endpoint (async). Sends a write-scoped proof."""
778
+ url = f"{self.base_url}{endpoint}"
779
+ async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
780
+ resp = await client.put(url, headers=self._write_headers(), json=payload or {})
781
+ if not resp.is_success:
782
+ raise ToolError(f"PUT {endpoint} failed", resp.status_code, resp.text)
783
+ return resp.json()
784
+
527
785
  def call_free_delete(
528
786
  self,
529
787
  endpoint: str,
@@ -877,16 +1135,25 @@ class OneShotClient:
877
1135
  reply within a thread — ``to``/``subject`` are then derived from the
878
1136
  inbound message unless given.
879
1137
 
880
- The first send from any ``from_mailbox`` on an owned domain including
881
- the default ``agent@`` — provisions a mailbox per address and adds a
882
- one-time ``mailbox_provisioning_fee`` to the quote; it's free only once
883
- that address is already provisioned (``active``).
884
-
885
- ``mailbox_mode`` ('relay' | 'mailbox', default 'relay') selects how a
886
- domain sends: 'relay' is header-only send (no per-address mailbox, no
887
- mailbox fee); 'mailbox' provisions a real dedicated mailbox per address
888
- (better deliverability + per-address warmup). Only applies when a domain
889
- is first set up use it on a new ``from_domain``.
1138
+ On a domain whose ``mailbox_mode`` is ``'mailbox'``, the first send from
1139
+ any address — including the default ``agent@`` — provisions a mailbox
1140
+ for it and adds a one-time ``mailbox_provisioning_fee`` to the quote;
1141
+ it's free once that address is ``active``. Domains in ``'relay'`` mode
1142
+ never charge it. ``list_domains()`` reports the mode per domain.
1143
+
1144
+ ``mailbox_mode`` ('relay' | 'mailbox') selects how a domain sends:
1145
+ 'relay' is header-only send (no per-address mailbox, no mailbox fee);
1146
+ 'mailbox' provisions a real dedicated mailbox per address (better
1147
+ deliverability + per-address warmup). It is honored **only while a domain
1148
+ is unprovisioned** — use it on a new ``from_domain``.
1149
+
1150
+ So omitting it is not a way to avoid the fee: an already-provisioned
1151
+ domain always uses the mode it was created with, and ``from_mailbox`` has
1152
+ no bearing on it. (A domain you own that is still unprovisioned does
1153
+ honor it — asking for 'mailbox' there opts it in.) Check
1154
+ ``list_domains()`` for which of your domains report
1155
+ ``mailbox_mode == 'mailbox'``, and read ``mailbox_provisioning_fee`` on
1156
+ the quote for the exact amount before paying.
890
1157
  """
891
1158
  return self.call_tool(
892
1159
  "/v1/tools/email/send",
@@ -1218,6 +1485,7 @@ class OneShotClient:
1218
1485
  ) -> Any:
1219
1486
  start = time.monotonic()
1220
1487
  retries = 0
1488
+ polls = 0
1221
1489
 
1222
1490
  while (time.monotonic() - start) < timeout_sec:
1223
1491
  try:
@@ -1262,7 +1530,12 @@ class OneShotClient:
1262
1530
  )
1263
1531
 
1264
1532
  retries = 0
1265
- await asyncio.sleep(_POLL_INTERVAL)
1533
+ interval = _POLL_BACKOFF[min(polls, len(_POLL_BACKOFF) - 1)]
1534
+ polls += 1
1535
+ remaining = timeout_sec - (time.monotonic() - start)
1536
+ if remaining <= 0:
1537
+ break
1538
+ await asyncio.sleep(min(interval, remaining))
1266
1539
 
1267
1540
  except (OneShotError, JobError):
1268
1541
  raise
@@ -1270,7 +1543,10 @@ class OneShotClient:
1270
1543
  retries += 1
1271
1544
  if retries > _MAX_POLL_RETRIES:
1272
1545
  raise
1273
- await asyncio.sleep(_POLL_INTERVAL * (2 ** (retries - 1)))
1546
+ remaining = timeout_sec - (time.monotonic() - start)
1547
+ if remaining <= 0:
1548
+ break
1549
+ await asyncio.sleep(min(_POLL_RETRY_BASE * (2 ** (retries - 1)), remaining))
1274
1550
 
1275
1551
  elapsed_ms = int((time.monotonic() - start) * 1000)
1276
1552
  raise JobTimeoutError(request_id, elapsed_ms)
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "oneshot-python"
3
- version = "0.19.0"
3
+ version = "0.20.2"
4
4
  description = "Core Python SDK for the OneShot API — HTTP client with x402 payment handling"
5
5
  readme = {text = "Core Python SDK for the OneShot API", content-type = "text/plain"}
6
6
  license = "MIT"
@@ -0,0 +1,184 @@
1
+ """Tests for the spend-budget surface of OneShotClient.
2
+
3
+ Mirrors tests/unit/sdk-budget-config.test.ts on the TS side: config is synced
4
+ once via a write-scoped PUT, a 403 budget_exceeded maps to BudgetExceededError,
5
+ the per-transaction cap fails fast locally, and get_budgets returns `data`.
6
+ The server-side arithmetic is covered by tests/unit/spend-budget-guard.test.ts.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import json
12
+ from unittest.mock import AsyncMock, MagicMock, patch
13
+
14
+ import pytest
15
+
16
+ from oneshot import BudgetExceededError, BudgetSyncError, ToolError, ValidationError
17
+ from oneshot.client import OneShotClient, _parse_budget_rejection
18
+
19
+ TEST_PRIVATE_KEY = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"
20
+
21
+
22
+ class _Resp:
23
+ def __init__(self, status_code: int, body=None, text: str = ""):
24
+ self.status_code = status_code
25
+ self._body = body
26
+ self.text = text or (json.dumps(body) if body is not None else "")
27
+ self.is_success = 200 <= status_code < 300
28
+
29
+ def json(self):
30
+ if self._body is None:
31
+ raise ValueError("no json")
32
+ return self._body
33
+
34
+
35
+ BUDGET_403 = {
36
+ "error": "budget_exceeded",
37
+ "message": "over",
38
+ "budget": {"reason": "daily", "cap": "50.000000", "spent": "49.700000", "charge": "0.500000", "resets_at": "2026-08-23T00:00:00.000Z"},
39
+ }
40
+
41
+
42
+ # ── config normalisation ────────────────────────────────────────────────
43
+
44
+ def test_accepts_snake_and_camel_keys():
45
+ c = OneShotClient(TEST_PRIVATE_KEY, budgets={"daily": 50, "perTransaction": 5, "alertAt": 0.5})
46
+ assert c._budgets == {"daily": 50.0, "per_transaction": 5.0, "alert_at": 0.5}
47
+
48
+
49
+ def test_no_budget_means_no_sync():
50
+ c = OneShotClient(TEST_PRIVATE_KEY)
51
+ assert c._budgets is None and c._alert_email is None
52
+
53
+
54
+ # ── 403 mapping ─────────────────────────────────────────────────────────
55
+
56
+ def test_parse_budget_rejection_maps_fields():
57
+ err = _parse_budget_rejection(_Resp(403, BUDGET_403))
58
+ assert isinstance(err, BudgetExceededError)
59
+ assert err.reason == "daily"
60
+ assert err.cap == 50.0 and err.spent == 49.7 and err.charge == 0.5
61
+ assert err.resets_at == "2026-08-23T00:00:00.000Z"
62
+
63
+
64
+ def test_parse_budget_rejection_ignores_other_403s():
65
+ # An auth 403 must not be reported as a budget problem.
66
+ assert _parse_budget_rejection(_Resp(403, {"error": "forbidden"})) is None
67
+ assert _parse_budget_rejection(_Resp(403, None, text="nope")) is None
68
+ assert _parse_budget_rejection(_Resp(402, BUDGET_403)) is None
69
+
70
+
71
+ # ── sync ────────────────────────────────────────────────────────────────
72
+
73
+ def _client_with_put(status: int = 200, exc: Exception | None = None):
74
+ c = OneShotClient(TEST_PRIVATE_KEY, budgets={"daily": 50, "per_transaction": 5}, alert_email="a@b.co")
75
+ put = AsyncMock(side_effect=exc) if exc else AsyncMock(return_value=_Resp(status, {"success": True}))
76
+ return c, put
77
+
78
+
79
+ def _patched_httpx(put_mock):
80
+ client = MagicMock()
81
+ client.put = put_mock
82
+ cm = MagicMock()
83
+ cm.__aenter__ = AsyncMock(return_value=client)
84
+ cm.__aexit__ = AsyncMock(return_value=False)
85
+ return patch("oneshot.client.httpx.AsyncClient", return_value=cm)
86
+
87
+
88
+ @pytest.mark.asyncio
89
+ async def test_sync_sends_one_write_scoped_put_with_snake_case_body():
90
+ c, put = _client_with_put()
91
+ with _patched_httpx(put):
92
+ await c.aensure_budgets_synced()
93
+ await c.aensure_budgets_synced() # second call is a no-op
94
+ assert put.await_count == 1
95
+ kwargs = put.await_args.kwargs
96
+ assert kwargs["json"] == {"daily": 50.0, "per_transaction": 5.0, "alert_email": "a@b.co"}
97
+ assert put.await_args.args[0].endswith("/v1/agents/me/budgets")
98
+ # write-scoped proof, not the read one
99
+ proof = json.loads(__import__("base64").b64decode(kwargs["headers"]["x-agent-proof"]))
100
+ assert proof["scope"] == "write"
101
+
102
+
103
+ @pytest.mark.asyncio
104
+ async def test_sync_fails_closed_on_5xx_network_and_4xx_and_retries_next_call():
105
+ # Proceeding without the configured guardrail would be the exact failure
106
+ # the budget exists to prevent — every failure refuses the paid call and
107
+ # leaves the client unsynced so the next call retries.
108
+ for status, exc in ((503, None), (429, None), (400, None), (None, ConnectionError("boom"))):
109
+ c, put = _client_with_put(status=status or 200, exc=exc)
110
+ with _patched_httpx(put):
111
+ with pytest.raises(BudgetSyncError) as excinfo:
112
+ await c.aensure_budgets_synced()
113
+ assert c._budget_synced is False
114
+ assert excinfo.value.status == status
115
+
116
+ # ...and a later success clears the condition.
117
+ c, put = _client_with_put(status=200)
118
+ with _patched_httpx(put):
119
+ await c.aensure_budgets_synced()
120
+ assert c._budget_synced is True
121
+
122
+
123
+ @pytest.mark.asyncio
124
+ async def test_paid_call_is_refused_when_sync_fails():
125
+ c, put = _client_with_put(status=503)
126
+ with _patched_httpx(put):
127
+ with pytest.raises(BudgetSyncError):
128
+ await c.acall_tool("/v1/tools/research", {"topic": "x"})
129
+ # the PUT was attempted, no tool POST followed
130
+ assert put.await_count == 1
131
+
132
+
133
+ # ── config validation ───────────────────────────────────────────────────
134
+
135
+ @pytest.mark.parametrize(
136
+ "budgets, field",
137
+ [
138
+ ({"daily": -1}, "budgets.daily"),
139
+ ({"daily": 0}, "budgets.daily"),
140
+ ({"per_transaction": float("inf")}, "budgets.per_transaction"),
141
+ ({"alert_at": 2}, "budgets.alert_at"),
142
+ ({"pauseAt": 0}, "budgets.pause_at"),
143
+ ({"daily": "5oops"}, "budgets.daily"),
144
+ ({"daily": True}, "budgets.daily"), # float(True) == 1.0
145
+ ({"daliy": 50}, "budgets.daliy"), # typo'd key must not mean "no cap"
146
+ ({"daily": 50, "per_tx": 5}, "budgets.per_tx"),
147
+ ],
148
+ )
149
+ def test_invalid_budget_config_is_rejected_at_construction(budgets, field):
150
+ with pytest.raises(ValidationError) as excinfo:
151
+ OneShotClient(TEST_PRIVATE_KEY, budgets=budgets)
152
+ assert excinfo.value.field == field
153
+
154
+
155
+ # ── local per-transaction pre-flight ────────────────────────────────────
156
+
157
+ def test_per_transaction_preflight_raises_before_signing():
158
+ c = OneShotClient(TEST_PRIVATE_KEY, budgets={"per_transaction": 1})
159
+ with pytest.raises(BudgetExceededError) as excinfo:
160
+ c._assert_within_budget("1.500000")
161
+ assert excinfo.value.reason == "per_transaction"
162
+ assert excinfo.value.cap == 1.0
163
+ c._assert_within_budget("1.000000") # exactly at the cap is allowed
164
+ OneShotClient(TEST_PRIVATE_KEY)._assert_within_budget("999") # no cap → no-op
165
+
166
+
167
+ # ── get_budgets ─────────────────────────────────────────────────────────
168
+
169
+ @pytest.mark.asyncio
170
+ async def test_get_budgets_unwraps_data():
171
+ c = OneShotClient(TEST_PRIVATE_KEY)
172
+ c.acall_free_get = AsyncMock(return_value={"success": True, "data": {"daily_usdc": 50, "spent_today_usdc": "1.0"}}) # type: ignore[method-assign]
173
+ out = await c.aget_budgets()
174
+ assert out == {"daily_usdc": 50, "spent_today_usdc": "1.0"}
175
+ c.acall_free_get.assert_awaited_once_with("/v1/agents/me/budgets")
176
+
177
+
178
+ @pytest.mark.asyncio
179
+ async def test_put_helper_raises_tool_error_on_failure():
180
+ c = OneShotClient(TEST_PRIVATE_KEY)
181
+ put = AsyncMock(return_value=_Resp(401, {"error": "proof_required"}))
182
+ with _patched_httpx(put):
183
+ with pytest.raises(ToolError):
184
+ await c.acall_free_put("/v1/agents/me/budgets", {"daily": 1})
@@ -0,0 +1,107 @@
1
+ """Tests for the job poll cadence in OneShotClient._poll_job.
2
+
3
+ The client polls GET /v1/requests/{id} immediately, then backs off
4
+ 0.3s → 0.6s → 1.0s → 2.0s (and stays at 2.0s). Most jobs finish within a few
5
+ seconds, so the first checks are fast; the previous flat 2s interval added up
6
+ to 2s of dead time to every call. Transient errors keep their own 2s-based
7
+ retry backoff, and no sleep may overshoot the caller's deadline.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import asyncio
13
+ from unittest.mock import AsyncMock, MagicMock
14
+
15
+ import pytest
16
+
17
+ from oneshot import client as client_module
18
+ from oneshot._errors import JobTimeoutError
19
+ from oneshot.client import OneShotClient
20
+
21
+ TEST_PRIVATE_KEY = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"
22
+ REQ = "req_backoff_1"
23
+
24
+
25
+ def _resp(payload: dict) -> MagicMock:
26
+ resp = MagicMock()
27
+ resp.is_success = True
28
+ resp.status_code = 200
29
+ resp.json.return_value = payload
30
+ return resp
31
+
32
+
33
+ def _processing() -> MagicMock:
34
+ return _resp({"status": "processing", "request_id": REQ})
35
+
36
+
37
+ def _completed(result: dict | None = None) -> MagicMock:
38
+ return _resp({"status": "completed", "request_id": REQ, "result": result or {"ok": True}})
39
+
40
+
41
+ @pytest.fixture
42
+ def fake_time(monkeypatch):
43
+ """Virtual clock: asyncio.sleep records the delay and advances time.monotonic."""
44
+ clock = {"t": 0.0}
45
+ sleeps: list[float] = []
46
+ real_sleep = asyncio.sleep
47
+
48
+ async def fake_sleep(delay: float) -> None:
49
+ sleeps.append(round(delay, 6))
50
+ clock["t"] += delay
51
+ await real_sleep(0)
52
+
53
+ monkeypatch.setattr(client_module.asyncio, "sleep", fake_sleep)
54
+ monkeypatch.setattr(client_module.time, "monotonic", lambda: clock["t"])
55
+ return sleeps
56
+
57
+
58
+ async def _run_poll(http_responses, timeout_sec: int = 120):
59
+ client = OneShotClient(TEST_PRIVATE_KEY)
60
+ http = MagicMock()
61
+ http.get = AsyncMock(side_effect=http_responses)
62
+ return await client._poll_job(http, REQ, timeout_sec), http
63
+
64
+
65
+ @pytest.mark.asyncio
66
+ class TestPollBackoff:
67
+ async def test_first_poll_is_immediate_and_backoff_ramps(self, fake_time):
68
+ result, http = await _run_poll([_processing()] * 4 + [_completed({"answer": 42})])
69
+
70
+ assert result == {"answer": 42, "request_id": REQ}
71
+ assert fake_time == [0.3, 0.6, 1.0, 2.0]
72
+ assert http.get.await_count == 5
73
+
74
+ async def test_backoff_caps_at_two_seconds(self, fake_time):
75
+ await _run_poll([_processing()] * 6 + [_completed()])
76
+
77
+ assert fake_time == [0.3, 0.6, 1.0, 2.0, 2.0, 2.0]
78
+
79
+ async def test_completed_on_first_poll_never_sleeps(self, fake_time):
80
+ result, http = await _run_poll([_completed()])
81
+
82
+ assert result["request_id"] == REQ
83
+ assert fake_time == []
84
+ assert http.get.await_count == 1
85
+
86
+ async def test_transient_error_uses_retry_backoff(self, fake_time):
87
+ result, http = await _run_poll([RuntimeError("connection reset"), _completed()])
88
+
89
+ assert result["request_id"] == REQ
90
+ assert fake_time == [2.0]
91
+ assert http.get.await_count == 2
92
+
93
+ async def test_sleep_never_overshoots_deadline(self, fake_time):
94
+ with pytest.raises(JobTimeoutError):
95
+ await _run_poll([_processing()] * 50, timeout_sec=1)
96
+
97
+ # 0.3 + 0.6 = 0.9 elapsed; the next interval (1.0) is clamped to the
98
+ # remaining 0.1 so the wait ends at the deadline, not 0.9s past it.
99
+ assert fake_time == [0.3, 0.6, 0.1]
100
+
101
+ async def test_retry_sleep_never_overshoots_deadline(self, fake_time):
102
+ with pytest.raises(JobTimeoutError):
103
+ await _run_poll([_processing(), RuntimeError("reset")] + [_processing()] * 50, timeout_sec=1)
104
+
105
+ # 0.3 elapsed after the first poll; the 2.0s retry backoff is clamped
106
+ # to the remaining 0.7 so the timeout fires on time, not ~1.3s late.
107
+ assert fake_time == [0.3, 0.7]
@@ -187,16 +187,13 @@ class TestMemoValidation:
187
187
  def test_valid_memo_passes_through(self):
188
188
  client = make_client()
189
189
  payload = {"query": "test", "memo": "valid reason"}
190
- # Run validation by calling acall_tool internals
191
- # We test by inspecting payload mutation before the HTTP call
192
- # Simulating the validation block directly
190
+ # Simulating acall_tool's validation block directly
193
191
  memo = payload.get("memo")
194
192
  assert memo == "valid reason"
195
193
 
196
194
  def test_empty_memo_is_dropped(self):
197
195
  client = make_client()
198
196
  payload = {"query": "test", "memo": ""}
199
- # Simulate validation
200
197
  memo = payload.get("memo")
201
198
  if memo is not None and (not isinstance(memo, str) or not memo.strip()):
202
199
  payload.pop("memo", None)
@@ -0,0 +1,88 @@
1
+ """Tests for ``wait=False`` and ``wait_for_result`` / ``await_result``.
2
+
3
+ Parity with the TS SDK: a paid call can return as soon as the job is queued,
4
+ and the ``request_id`` can be resolved later against ``GET /v1/requests/{id}``.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ from unittest.mock import AsyncMock, MagicMock, patch
10
+
11
+ import pytest
12
+
13
+ from oneshot._errors import ValidationError
14
+ from oneshot.client import OneShotClient
15
+
16
+ TEST_PRIVATE_KEY = "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"
17
+ REQ = "req_wait_1"
18
+
19
+
20
+ def _resp(status_code: int, payload: dict) -> MagicMock:
21
+ resp = MagicMock()
22
+ resp.status_code = status_code
23
+ resp.is_success = 200 <= status_code < 300
24
+ resp.json.return_value = payload
25
+ resp.text = ""
26
+ resp.raise_for_status = MagicMock()
27
+ return resp
28
+
29
+
30
+ @pytest.mark.asyncio
31
+ class TestWaitFalse:
32
+ async def test_free_route_queued_job_returns_immediately(self):
33
+ client = OneShotClient(TEST_PRIVATE_KEY)
34
+ client._budget_synced = True
35
+ http = MagicMock()
36
+ http.post = AsyncMock(return_value=_resp(202, {"request_id": REQ, "status": "processing"}))
37
+ http.__aenter__ = AsyncMock(return_value=http)
38
+ http.__aexit__ = AsyncMock(return_value=False)
39
+ with patch("oneshot.client.httpx.AsyncClient", return_value=http), \
40
+ patch.object(client, "_poll_job", new=AsyncMock()) as poll:
41
+ out = await client.acall_tool("/v1/tools/web-read", {"url": "https://example.com"}, wait=False)
42
+
43
+ assert out == {"request_id": REQ, "status": "processing"}
44
+ poll.assert_not_awaited()
45
+
46
+ async def test_default_still_waits(self):
47
+ client = OneShotClient(TEST_PRIVATE_KEY)
48
+ client._budget_synced = True
49
+ http = MagicMock()
50
+ http.post = AsyncMock(return_value=_resp(202, {"request_id": REQ, "status": "processing"}))
51
+ http.__aenter__ = AsyncMock(return_value=http)
52
+ http.__aexit__ = AsyncMock(return_value=False)
53
+ with patch("oneshot.client.httpx.AsyncClient", return_value=http), \
54
+ patch.object(client, "_poll_job", new=AsyncMock(return_value={"ok": True})) as poll:
55
+ out = await client.acall_tool("/v1/tools/web-read", {"url": "https://example.com"})
56
+
57
+ assert out == {"ok": True}
58
+ poll.assert_awaited_once()
59
+
60
+ async def test_await_result_polls_the_request(self):
61
+ client = OneShotClient(TEST_PRIVATE_KEY)
62
+ http = MagicMock()
63
+ http.get = AsyncMock(side_effect=[
64
+ _resp(200, {"status": "processing", "request_id": REQ}),
65
+ _resp(200, {"status": "completed", "request_id": REQ, "result": {"answer": 42}}),
66
+ ])
67
+ http.__aenter__ = AsyncMock(return_value=http)
68
+ http.__aexit__ = AsyncMock(return_value=False)
69
+ with patch("oneshot.client.httpx.AsyncClient", return_value=http), \
70
+ patch("oneshot.client.asyncio.sleep", new=AsyncMock()):
71
+ out = await client.await_result(REQ, timeout_sec=30)
72
+
73
+ assert out == {"answer": 42, "request_id": REQ}
74
+ assert http.get.await_count == 2
75
+ assert http.get.await_args_list[0].args[0].endswith(f"/v1/requests/{REQ}")
76
+
77
+ async def test_await_result_requires_request_id(self):
78
+ client = OneShotClient(TEST_PRIVATE_KEY)
79
+ with pytest.raises(ValidationError):
80
+ await client.await_result("")
81
+
82
+
83
+ class TestSyncWrapper:
84
+ def test_wait_for_result_delegates(self):
85
+ client = OneShotClient(TEST_PRIVATE_KEY)
86
+ with patch.object(client, "await_result", new=AsyncMock(return_value={"done": True})) as inner:
87
+ assert client.wait_for_result(REQ, timeout_sec=5) == {"done": True}
88
+ inner.assert_awaited_once_with(REQ, timeout_sec=5, wait_for_phones=False, phone_timeout_sec=360)
File without changes