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.
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/.gitignore +3 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/PKG-INFO +1 -1
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/README.md +37 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/__init__.py +8 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/_errors.py +46 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/_types.py +40 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/client.py +295 -19
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/pyproject.toml +1 -1
- oneshot_python-0.20.2/tests/test_budgets.py +184 -0
- oneshot_python-0.20.2/tests/test_poll_backoff.py +107 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_request_id.py +1 -4
- oneshot_python-0.20.2/tests/test_wait_false.py +88 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/oneshot/x402.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/__init__.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_balance.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_charge_amount.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_compute.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_domains.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_email_payload.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_emergency_error.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_max_cost_header.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_payment_rejection.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_phones_pending.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_tag_receipt_value.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/tests/test_x402.py +0 -0
- {oneshot_python-0.19.0 → oneshot_python-0.20.2}/uv.lock +0 -0
|
@@ -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.
|
|
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
|
-
|
|
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
|
-
#
|
|
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
|
-
|
|
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
|
-
|
|
881
|
-
the default ``agent@`` — provisions a mailbox
|
|
882
|
-
one-time ``mailbox_provisioning_fee`` to the quote;
|
|
883
|
-
that address is
|
|
884
|
-
|
|
885
|
-
|
|
886
|
-
|
|
887
|
-
|
|
888
|
-
|
|
889
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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.
|
|
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
|
-
#
|
|
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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|