spendwall 0.0.6__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.
- spendwall-0.0.6/PKG-INFO +33 -0
- spendwall-0.0.6/README.md +11 -0
- spendwall-0.0.6/pyproject.toml +25 -0
- spendwall-0.0.6/setup.cfg +4 -0
- spendwall-0.0.6/spendwall/__init__.py +11 -0
- spendwall-0.0.6/spendwall/anthropic.py +88 -0
- spendwall-0.0.6/spendwall/client.py +118 -0
- spendwall-0.0.6/spendwall/errors.py +37 -0
- spendwall-0.0.6/spendwall/modes.py +26 -0
- spendwall-0.0.6/spendwall/openai.py +98 -0
- spendwall-0.0.6/spendwall.egg-info/PKG-INFO +33 -0
- spendwall-0.0.6/spendwall.egg-info/SOURCES.txt +18 -0
- spendwall-0.0.6/spendwall.egg-info/dependency_links.txt +1 -0
- spendwall-0.0.6/spendwall.egg-info/requires.txt +17 -0
- spendwall-0.0.6/spendwall.egg-info/top_level.txt +1 -0
- spendwall-0.0.6/tests/test_anthropic_wrapper.py +159 -0
- spendwall-0.0.6/tests/test_client.py +253 -0
- spendwall-0.0.6/tests/test_errors.py +57 -0
- spendwall-0.0.6/tests/test_modes.py +28 -0
- spendwall-0.0.6/tests/test_openai_wrapper.py +178 -0
spendwall-0.0.6/PKG-INFO
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: spendwall
|
|
3
|
+
Version: 0.0.6
|
|
4
|
+
Summary: Spendwall SDK — local-first agent spend firewall
|
|
5
|
+
Author: Spendwall
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Requires-Python: >=3.10
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Requires-Dist: requests-unixsocket>=0.3
|
|
10
|
+
Requires-Dist: requests>=2.28
|
|
11
|
+
Provides-Extra: openai
|
|
12
|
+
Requires-Dist: openai>=1.0; extra == "openai"
|
|
13
|
+
Provides-Extra: anthropic
|
|
14
|
+
Requires-Dist: anthropic>=0.21; extra == "anthropic"
|
|
15
|
+
Provides-Extra: all
|
|
16
|
+
Requires-Dist: openai>=1.0; extra == "all"
|
|
17
|
+
Requires-Dist: anthropic>=0.21; extra == "all"
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: pytest>=7; extra == "dev"
|
|
20
|
+
Requires-Dist: pytest-asyncio>=0.21; extra == "dev"
|
|
21
|
+
Requires-Dist: responses>=0.23; extra == "dev"
|
|
22
|
+
|
|
23
|
+
# spendwall (Python)
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
import openai
|
|
27
|
+
import spendwall.openai as sw
|
|
28
|
+
|
|
29
|
+
client = openai.OpenAI(api_key="sk-...")
|
|
30
|
+
client = sw.wrap(client, mode="enforce")
|
|
31
|
+
|
|
32
|
+
# Now every client.chat.completions.create(...) is observed by spendwalld.
|
|
33
|
+
```
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "spendwall"
|
|
7
|
+
version = "0.0.6"
|
|
8
|
+
description = "Spendwall SDK — local-first agent spend firewall"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = { text = "Apache-2.0" }
|
|
12
|
+
authors = [{ name = "Spendwall" }]
|
|
13
|
+
dependencies = [
|
|
14
|
+
"requests-unixsocket>=0.3",
|
|
15
|
+
"requests>=2.28",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[project.optional-dependencies]
|
|
19
|
+
openai = ["openai>=1.0"]
|
|
20
|
+
anthropic = ["anthropic>=0.21"]
|
|
21
|
+
all = ["openai>=1.0", "anthropic>=0.21"]
|
|
22
|
+
dev = ["pytest>=7", "pytest-asyncio>=0.21", "responses>=0.23"]
|
|
23
|
+
|
|
24
|
+
[tool.setuptools.packages.find]
|
|
25
|
+
include = ["spendwall*"]
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
"""Spendwall — local-first agent spending firewall."""
|
|
2
|
+
__version__ = "0.0.2"
|
|
3
|
+
|
|
4
|
+
from spendwall.client import SpendwallClient # noqa: F401
|
|
5
|
+
from spendwall.errors import ( # noqa: F401
|
|
6
|
+
ApprovalRequiredError,
|
|
7
|
+
BlockedError,
|
|
8
|
+
SpendwallError,
|
|
9
|
+
SpendwallVersionMismatch,
|
|
10
|
+
)
|
|
11
|
+
from spendwall.modes import Mode # noqa: F401
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
"""Wrapper for ``anthropic.Anthropic`` clients.
|
|
2
|
+
|
|
3
|
+
Patches ``client.messages.create`` in-place so every upstream call is
|
|
4
|
+
observed by spendwalld. In ``enforce`` mode, the wrapper consults policy
|
|
5
|
+
*before* dispatching to the upstream and raises
|
|
6
|
+
``BlockedError``/``ApprovalRequiredError`` when the daemon refuses.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import json as _json
|
|
12
|
+
from typing import Any, Optional
|
|
13
|
+
|
|
14
|
+
from spendwall.client import SpendwallClient
|
|
15
|
+
from spendwall.modes import Mode, resolve as resolve_mode
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _estimate_tokens(messages_or_input: Any) -> int:
|
|
19
|
+
try:
|
|
20
|
+
return max(1, len(_json.dumps(messages_or_input)) // 4)
|
|
21
|
+
except (TypeError, ValueError):
|
|
22
|
+
return 1
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def wrap(
|
|
26
|
+
anthropic_client: Any,
|
|
27
|
+
*,
|
|
28
|
+
sw_client: Optional[SpendwallClient] = None,
|
|
29
|
+
agent_id: str = "anonymous",
|
|
30
|
+
mode: Optional[Mode] = None,
|
|
31
|
+
tags: Optional[dict[str, str]] = None,
|
|
32
|
+
) -> Any:
|
|
33
|
+
"""Patch ``anthropic_client.messages.create`` to call into spendwall
|
|
34
|
+
before/after the upstream request.
|
|
35
|
+
|
|
36
|
+
Mutates the supplied client in-place AND returns it for ergonomic chaining:
|
|
37
|
+
|
|
38
|
+
client = wrap(anthropic.Anthropic(...), agent_id="x")
|
|
39
|
+
"""
|
|
40
|
+
sw = sw_client or SpendwallClient()
|
|
41
|
+
configured_mode = mode
|
|
42
|
+
|
|
43
|
+
original = anthropic_client.messages.create
|
|
44
|
+
|
|
45
|
+
def patched(*args: Any, **kwargs: Any) -> Any:
|
|
46
|
+
call_mode = resolve_mode(configured=configured_mode)
|
|
47
|
+
model = kwargs.get("model", "unknown")
|
|
48
|
+
est_input_tokens = _estimate_tokens(kwargs.get("messages"))
|
|
49
|
+
|
|
50
|
+
if call_mode == "enforce":
|
|
51
|
+
sw.policy_check_strict(
|
|
52
|
+
agent_id=agent_id, provider="anthropic", model=model,
|
|
53
|
+
estimated_cost_cents=0, tags=tags or {},
|
|
54
|
+
)
|
|
55
|
+
elif call_mode == "advisory":
|
|
56
|
+
sw.policy_check(
|
|
57
|
+
agent_id=agent_id, provider="anthropic", model=model,
|
|
58
|
+
estimated_cost_cents=0, tags=tags or {},
|
|
59
|
+
)
|
|
60
|
+
# observe-mode skips policy_check entirely.
|
|
61
|
+
try:
|
|
62
|
+
result = original(*args, **kwargs)
|
|
63
|
+
except Exception:
|
|
64
|
+
sw.observe({
|
|
65
|
+
"agent_id": agent_id, "provider": "anthropic", "model": model,
|
|
66
|
+
"source": "sdk_python", "input_tokens": est_input_tokens,
|
|
67
|
+
"output_tokens": None, "estimated_cost_cents": 0,
|
|
68
|
+
"actual_cost_cents": None, "decision_str": "allowed",
|
|
69
|
+
"binding_caps": [], "tags": tags or {},
|
|
70
|
+
"outcome": "upstream_error",
|
|
71
|
+
})
|
|
72
|
+
raise
|
|
73
|
+
|
|
74
|
+
usage = getattr(result, "usage", None)
|
|
75
|
+
input_tokens = getattr(usage, "input_tokens", None) if usage else None
|
|
76
|
+
output_tokens = getattr(usage, "output_tokens", None) if usage else None
|
|
77
|
+
sw.observe({
|
|
78
|
+
"agent_id": agent_id, "provider": "anthropic", "model": model,
|
|
79
|
+
"source": "sdk_python", "input_tokens": input_tokens,
|
|
80
|
+
"output_tokens": output_tokens, "estimated_cost_cents": 0,
|
|
81
|
+
"actual_cost_cents": None, "decision_str": "allowed",
|
|
82
|
+
"binding_caps": [], "tags": tags or {},
|
|
83
|
+
"outcome": "success",
|
|
84
|
+
})
|
|
85
|
+
return result
|
|
86
|
+
|
|
87
|
+
anthropic_client.messages.create = patched
|
|
88
|
+
return anthropic_client
|
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
"""UDS HTTP client for spendwalld."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
from typing import Any, Optional
|
|
7
|
+
from urllib.parse import quote
|
|
8
|
+
|
|
9
|
+
import requests
|
|
10
|
+
import requests_unixsocket
|
|
11
|
+
|
|
12
|
+
from spendwall.errors import ApprovalRequiredError, BlockedError
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _default_socket_path() -> str:
|
|
16
|
+
return os.environ.get(
|
|
17
|
+
"SPENDWALL_SOCKET",
|
|
18
|
+
os.path.expanduser("~/.spendwall/spendwall.sock"),
|
|
19
|
+
)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def _default_bearer() -> Optional[str]:
|
|
23
|
+
return os.environ.get("SPENDWALL_TOKEN")
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class SpendwallClient:
|
|
27
|
+
"""Talks to the spendwall daemon over its UDS HTTP API.
|
|
28
|
+
|
|
29
|
+
Fail-open: when ``fail_open=True`` and the socket is unreachable, ``observe()``
|
|
30
|
+
is silently dropped and ``policy_check()`` returns a synthetic "allowed"
|
|
31
|
+
decision. Default is ``fail_open=True`` because a spending firewall MUST NOT
|
|
32
|
+
bring down user agents when the daemon is down.
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
def __init__(
|
|
36
|
+
self,
|
|
37
|
+
socket_path: Optional[str] = None,
|
|
38
|
+
bearer: Optional[str] = None,
|
|
39
|
+
timeout_seconds: float = 2.0,
|
|
40
|
+
fail_open: bool = True,
|
|
41
|
+
):
|
|
42
|
+
self.socket_path = socket_path or _default_socket_path()
|
|
43
|
+
self.bearer = bearer or _default_bearer()
|
|
44
|
+
self.timeout_seconds = timeout_seconds
|
|
45
|
+
self.fail_open = fail_open
|
|
46
|
+
self._session = requests_unixsocket.Session()
|
|
47
|
+
self._base = f"http+unix://{quote(self.socket_path, safe='')}"
|
|
48
|
+
|
|
49
|
+
def _headers(self) -> dict[str, str]:
|
|
50
|
+
h = {"content-type": "application/json"}
|
|
51
|
+
if self.bearer:
|
|
52
|
+
h["authorization"] = f"Bearer {self.bearer}"
|
|
53
|
+
return h
|
|
54
|
+
|
|
55
|
+
def policy_check(
|
|
56
|
+
self,
|
|
57
|
+
*,
|
|
58
|
+
agent_id: str,
|
|
59
|
+
provider: str,
|
|
60
|
+
model: str,
|
|
61
|
+
estimated_cost_cents: int,
|
|
62
|
+
tags: Optional[dict[str, str]] = None,
|
|
63
|
+
) -> dict[str, Any]:
|
|
64
|
+
body = {
|
|
65
|
+
"agent_id": agent_id,
|
|
66
|
+
"provider": provider,
|
|
67
|
+
"model": model,
|
|
68
|
+
"estimated_cost_cents": estimated_cost_cents,
|
|
69
|
+
"tags": tags or {},
|
|
70
|
+
}
|
|
71
|
+
try:
|
|
72
|
+
r = self._session.post(
|
|
73
|
+
f"{self._base}/v1/policy/check",
|
|
74
|
+
json=body,
|
|
75
|
+
headers=self._headers(),
|
|
76
|
+
timeout=self.timeout_seconds,
|
|
77
|
+
)
|
|
78
|
+
r.raise_for_status()
|
|
79
|
+
return r.json()
|
|
80
|
+
except (requests.RequestException, ValueError):
|
|
81
|
+
if self.fail_open:
|
|
82
|
+
return {"decision": "allowed", "alert_caps": []}
|
|
83
|
+
raise
|
|
84
|
+
|
|
85
|
+
def policy_check_strict(self, **kwargs: Any) -> dict[str, Any]:
|
|
86
|
+
"""Like ``policy_check`` but raises ``BlockedError`` /
|
|
87
|
+
``ApprovalRequiredError`` on non-allowed decisions."""
|
|
88
|
+
d = self.policy_check(**kwargs)
|
|
89
|
+
if d.get("decision") == "blocked":
|
|
90
|
+
raise BlockedError(
|
|
91
|
+
cap_id=d.get("cap_id", "?"),
|
|
92
|
+
projected_cents=d.get("projected_cents", 0),
|
|
93
|
+
limit_cents=d.get("limit_cents", 0),
|
|
94
|
+
)
|
|
95
|
+
if d.get("decision") == "approval_required":
|
|
96
|
+
raise ApprovalRequiredError(approval_caps=d.get("approval_caps", []))
|
|
97
|
+
return d
|
|
98
|
+
|
|
99
|
+
def observe(self, event: dict[str, Any]) -> None:
|
|
100
|
+
try:
|
|
101
|
+
self._session.post(
|
|
102
|
+
f"{self._base}/v1/spend/observe",
|
|
103
|
+
json=event,
|
|
104
|
+
headers=self._headers(),
|
|
105
|
+
timeout=self.timeout_seconds,
|
|
106
|
+
)
|
|
107
|
+
except requests.RequestException:
|
|
108
|
+
if not self.fail_open:
|
|
109
|
+
raise
|
|
110
|
+
# Fail-open: drop on the floor.
|
|
111
|
+
|
|
112
|
+
def about(self) -> dict[str, Any]:
|
|
113
|
+
r = self._session.get(
|
|
114
|
+
f"{self._base}/v1/about",
|
|
115
|
+
timeout=self.timeout_seconds,
|
|
116
|
+
)
|
|
117
|
+
r.raise_for_status()
|
|
118
|
+
return r.json()
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
"""Spendwall SDK exceptions."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class SpendwallError(Exception):
|
|
7
|
+
"""Base for all spendwall-raised exceptions."""
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class BlockedError(SpendwallError):
|
|
11
|
+
"""Raised in `enforce` mode when policy returned Blocked."""
|
|
12
|
+
|
|
13
|
+
def __init__(self, cap_id: str, projected_cents: int, limit_cents: int):
|
|
14
|
+
self.cap_id = cap_id
|
|
15
|
+
self.projected_cents = projected_cents
|
|
16
|
+
self.limit_cents = limit_cents
|
|
17
|
+
super().__init__(
|
|
18
|
+
f"Spendwall blocked by cap '{cap_id}': "
|
|
19
|
+
f"projected {projected_cents}c > limit {limit_cents}c"
|
|
20
|
+
)
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
class ApprovalRequiredError(SpendwallError):
|
|
24
|
+
"""Raised in `enforce` mode when policy returned ApprovalRequired and the
|
|
25
|
+
approval was denied or timed out."""
|
|
26
|
+
|
|
27
|
+
def __init__(self, approval_caps: list[str], reason: str = "denied"):
|
|
28
|
+
self.approval_caps = approval_caps
|
|
29
|
+
self.reason = reason
|
|
30
|
+
super().__init__(
|
|
31
|
+
f"Spendwall approval {reason} for caps={approval_caps}"
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class SpendwallVersionMismatch(SpendwallError):
|
|
36
|
+
"""Raised once when the daemon's API version does not match the SDK's.
|
|
37
|
+
SDK falls back to fail-open `observe` mode for the rest of the process."""
|
|
@@ -0,0 +1,26 @@
|
|
|
1
|
+
"""Mode resolution — observe, advisory, enforce."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
from typing import Literal, Optional
|
|
7
|
+
|
|
8
|
+
Mode = Literal["observe", "advisory", "enforce"]
|
|
9
|
+
_VALID = ("observe", "advisory", "enforce")
|
|
10
|
+
_ENV_VAR = "SPENDWALL_MODE"
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def resolve(per_call: Optional[str] = None, configured: Optional[str] = None) -> Mode:
|
|
14
|
+
"""Precedence (most specific wins):
|
|
15
|
+
1. `per_call` kwarg passed by user code at the call site
|
|
16
|
+
2. `configured` value from `wrap(..., mode=...)`
|
|
17
|
+
3. `SPENDWALL_MODE` env var
|
|
18
|
+
4. Default: `observe`
|
|
19
|
+
"""
|
|
20
|
+
for v in (per_call, configured, os.environ.get(_ENV_VAR)):
|
|
21
|
+
if v is None:
|
|
22
|
+
continue
|
|
23
|
+
if v not in _VALID:
|
|
24
|
+
raise ValueError(f"invalid spendwall mode: {v!r}; expected one of {_VALID}")
|
|
25
|
+
return v # type: ignore[return-value]
|
|
26
|
+
return "observe"
|
|
@@ -0,0 +1,98 @@
|
|
|
1
|
+
"""Wrapper for ``openai.OpenAI`` clients.
|
|
2
|
+
|
|
3
|
+
Patches ``client.chat.completions.create`` and ``client.responses.create``
|
|
4
|
+
in-place so every upstream call is observed by spendwalld. In ``enforce``
|
|
5
|
+
mode, the wrapper consults policy *before* dispatching to the upstream and
|
|
6
|
+
raises ``BlockedError``/``ApprovalRequiredError`` when the daemon refuses.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import json as _json
|
|
12
|
+
from typing import Any, Optional
|
|
13
|
+
|
|
14
|
+
from spendwall.client import SpendwallClient
|
|
15
|
+
from spendwall.modes import Mode, resolve as resolve_mode
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _estimate_tokens(messages_or_input: Any) -> int:
|
|
19
|
+
"""Cheap byte-based estimate; daemon reconciles cost post-flight."""
|
|
20
|
+
try:
|
|
21
|
+
return max(1, len(_json.dumps(messages_or_input)) // 4)
|
|
22
|
+
except (TypeError, ValueError):
|
|
23
|
+
return 1
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def wrap(
|
|
27
|
+
openai_client: Any,
|
|
28
|
+
*,
|
|
29
|
+
sw_client: Optional[SpendwallClient] = None,
|
|
30
|
+
agent_id: str = "anonymous",
|
|
31
|
+
mode: Optional[Mode] = None,
|
|
32
|
+
tags: Optional[dict[str, str]] = None,
|
|
33
|
+
) -> Any:
|
|
34
|
+
"""Patch ``openai_client.chat.completions.create`` (and ``.responses.create``)
|
|
35
|
+
to call into spendwall before/after the upstream request.
|
|
36
|
+
|
|
37
|
+
Mutates the supplied client in-place AND returns it for ergonomic chaining:
|
|
38
|
+
|
|
39
|
+
client = wrap(openai.OpenAI(...), agent_id="x")
|
|
40
|
+
"""
|
|
41
|
+
sw = sw_client or SpendwallClient()
|
|
42
|
+
configured_mode = mode
|
|
43
|
+
|
|
44
|
+
def _patch(parent: Any, name: str) -> None:
|
|
45
|
+
original = getattr(parent, name)
|
|
46
|
+
|
|
47
|
+
def patched(*args: Any, **kwargs: Any) -> Any:
|
|
48
|
+
call_mode = resolve_mode(configured=configured_mode)
|
|
49
|
+
model = kwargs.get("model", "unknown")
|
|
50
|
+
est_input_tokens = _estimate_tokens(
|
|
51
|
+
kwargs.get("messages") or kwargs.get("input")
|
|
52
|
+
)
|
|
53
|
+
# Daemon does pricing lookup; SDK sends estimated_cost_cents=0
|
|
54
|
+
# to defer cost estimation to the daemon side.
|
|
55
|
+
if call_mode == "enforce":
|
|
56
|
+
sw.policy_check_strict(
|
|
57
|
+
agent_id=agent_id, provider="openai", model=model,
|
|
58
|
+
estimated_cost_cents=0, tags=tags or {},
|
|
59
|
+
)
|
|
60
|
+
elif call_mode == "advisory":
|
|
61
|
+
sw.policy_check(
|
|
62
|
+
agent_id=agent_id, provider="openai", model=model,
|
|
63
|
+
estimated_cost_cents=0, tags=tags or {},
|
|
64
|
+
)
|
|
65
|
+
# observe-mode skips policy_check entirely (zero added latency).
|
|
66
|
+
try:
|
|
67
|
+
result = original(*args, **kwargs)
|
|
68
|
+
except Exception:
|
|
69
|
+
sw.observe({
|
|
70
|
+
"agent_id": agent_id, "provider": "openai", "model": model,
|
|
71
|
+
"source": "sdk_python", "input_tokens": est_input_tokens,
|
|
72
|
+
"output_tokens": None, "estimated_cost_cents": 0,
|
|
73
|
+
"actual_cost_cents": None, "decision_str": "allowed",
|
|
74
|
+
"binding_caps": [], "tags": tags or {},
|
|
75
|
+
"outcome": "upstream_error",
|
|
76
|
+
})
|
|
77
|
+
raise
|
|
78
|
+
|
|
79
|
+
usage = getattr(result, "usage", None)
|
|
80
|
+
input_tokens = getattr(usage, "prompt_tokens", None) if usage else None
|
|
81
|
+
output_tokens = getattr(usage, "completion_tokens", None) if usage else None
|
|
82
|
+
sw.observe({
|
|
83
|
+
"agent_id": agent_id, "provider": "openai", "model": model,
|
|
84
|
+
"source": "sdk_python", "input_tokens": input_tokens,
|
|
85
|
+
"output_tokens": output_tokens, "estimated_cost_cents": 0,
|
|
86
|
+
"actual_cost_cents": None, "decision_str": "allowed",
|
|
87
|
+
"binding_caps": [], "tags": tags or {},
|
|
88
|
+
"outcome": "success",
|
|
89
|
+
})
|
|
90
|
+
return result
|
|
91
|
+
|
|
92
|
+
setattr(parent, name, patched)
|
|
93
|
+
|
|
94
|
+
if hasattr(openai_client, "chat") and hasattr(openai_client.chat, "completions"):
|
|
95
|
+
_patch(openai_client.chat.completions, "create")
|
|
96
|
+
if hasattr(openai_client, "responses"):
|
|
97
|
+
_patch(openai_client.responses, "create")
|
|
98
|
+
return openai_client
|
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: spendwall
|
|
3
|
+
Version: 0.0.6
|
|
4
|
+
Summary: Spendwall SDK — local-first agent spend firewall
|
|
5
|
+
Author: Spendwall
|
|
6
|
+
License: Apache-2.0
|
|
7
|
+
Requires-Python: >=3.10
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Requires-Dist: requests-unixsocket>=0.3
|
|
10
|
+
Requires-Dist: requests>=2.28
|
|
11
|
+
Provides-Extra: openai
|
|
12
|
+
Requires-Dist: openai>=1.0; extra == "openai"
|
|
13
|
+
Provides-Extra: anthropic
|
|
14
|
+
Requires-Dist: anthropic>=0.21; extra == "anthropic"
|
|
15
|
+
Provides-Extra: all
|
|
16
|
+
Requires-Dist: openai>=1.0; extra == "all"
|
|
17
|
+
Requires-Dist: anthropic>=0.21; extra == "all"
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: pytest>=7; extra == "dev"
|
|
20
|
+
Requires-Dist: pytest-asyncio>=0.21; extra == "dev"
|
|
21
|
+
Requires-Dist: responses>=0.23; extra == "dev"
|
|
22
|
+
|
|
23
|
+
# spendwall (Python)
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
import openai
|
|
27
|
+
import spendwall.openai as sw
|
|
28
|
+
|
|
29
|
+
client = openai.OpenAI(api_key="sk-...")
|
|
30
|
+
client = sw.wrap(client, mode="enforce")
|
|
31
|
+
|
|
32
|
+
# Now every client.chat.completions.create(...) is observed by spendwalld.
|
|
33
|
+
```
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
spendwall/__init__.py
|
|
4
|
+
spendwall/anthropic.py
|
|
5
|
+
spendwall/client.py
|
|
6
|
+
spendwall/errors.py
|
|
7
|
+
spendwall/modes.py
|
|
8
|
+
spendwall/openai.py
|
|
9
|
+
spendwall.egg-info/PKG-INFO
|
|
10
|
+
spendwall.egg-info/SOURCES.txt
|
|
11
|
+
spendwall.egg-info/dependency_links.txt
|
|
12
|
+
spendwall.egg-info/requires.txt
|
|
13
|
+
spendwall.egg-info/top_level.txt
|
|
14
|
+
tests/test_anthropic_wrapper.py
|
|
15
|
+
tests/test_client.py
|
|
16
|
+
tests/test_errors.py
|
|
17
|
+
tests/test_modes.py
|
|
18
|
+
tests/test_openai_wrapper.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
spendwall
|
|
@@ -0,0 +1,159 @@
|
|
|
1
|
+
"""Tests for spendwall.anthropic.wrap()."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from unittest.mock import MagicMock
|
|
6
|
+
|
|
7
|
+
import pytest
|
|
8
|
+
|
|
9
|
+
from spendwall.errors import BlockedError
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
@pytest.fixture
|
|
13
|
+
def fake_anthropic_client():
|
|
14
|
+
client = MagicMock(name="anthropic.Anthropic")
|
|
15
|
+
msg = MagicMock(name="Message")
|
|
16
|
+
msg.id = "msg_test"
|
|
17
|
+
msg.model = "claude-3-5-sonnet-20241022"
|
|
18
|
+
msg.usage = MagicMock(input_tokens=10, output_tokens=5)
|
|
19
|
+
client.messages.create.return_value = msg
|
|
20
|
+
return client
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def test_wrap_observe_passes_through_and_observes(
|
|
24
|
+
fake_anthropic_client, monkeypatch
|
|
25
|
+
):
|
|
26
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
27
|
+
from spendwall.anthropic import wrap
|
|
28
|
+
|
|
29
|
+
sw = MagicMock()
|
|
30
|
+
wrapped = wrap(fake_anthropic_client, sw_client=sw, agent_id="a")
|
|
31
|
+
out = wrapped.messages.create(
|
|
32
|
+
model="claude-3-5-sonnet-20241022",
|
|
33
|
+
messages=[{"role": "user", "content": "hi"}],
|
|
34
|
+
)
|
|
35
|
+
assert out.id == "msg_test"
|
|
36
|
+
sw.policy_check_strict.assert_not_called()
|
|
37
|
+
sw.policy_check.assert_not_called()
|
|
38
|
+
sw.observe.assert_called_once()
|
|
39
|
+
payload = sw.observe.call_args.args[0]
|
|
40
|
+
assert payload["provider"] == "anthropic"
|
|
41
|
+
assert payload["source"] == "sdk_python"
|
|
42
|
+
assert payload["input_tokens"] == 10
|
|
43
|
+
assert payload["output_tokens"] == 5
|
|
44
|
+
assert payload["outcome"] == "success"
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def test_wrap_enforce_calls_policy_check_strict_then_observe(
|
|
48
|
+
fake_anthropic_client, monkeypatch
|
|
49
|
+
):
|
|
50
|
+
monkeypatch.setenv("SPENDWALL_MODE", "enforce")
|
|
51
|
+
from spendwall.anthropic import wrap
|
|
52
|
+
|
|
53
|
+
sw = MagicMock()
|
|
54
|
+
sw.policy_check_strict.return_value = {"decision": "allowed"}
|
|
55
|
+
wrapped = wrap(fake_anthropic_client, sw_client=sw, agent_id="a")
|
|
56
|
+
wrapped.messages.create(
|
|
57
|
+
model="claude-3-5-sonnet-20241022",
|
|
58
|
+
messages=[{"role": "user", "content": "hi"}],
|
|
59
|
+
)
|
|
60
|
+
sw.policy_check_strict.assert_called_once()
|
|
61
|
+
pcs_kwargs = sw.policy_check_strict.call_args.kwargs
|
|
62
|
+
assert pcs_kwargs["provider"] == "anthropic"
|
|
63
|
+
assert pcs_kwargs["model"] == "claude-3-5-sonnet-20241022"
|
|
64
|
+
sw.observe.assert_called_once()
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def test_wrap_advisory_calls_non_strict_policy_check(
|
|
68
|
+
fake_anthropic_client, monkeypatch
|
|
69
|
+
):
|
|
70
|
+
monkeypatch.setenv("SPENDWALL_MODE", "advisory")
|
|
71
|
+
from spendwall.anthropic import wrap
|
|
72
|
+
|
|
73
|
+
sw = MagicMock()
|
|
74
|
+
sw.policy_check.return_value = {"decision": "allowed"}
|
|
75
|
+
wrapped = wrap(fake_anthropic_client, sw_client=sw, agent_id="a")
|
|
76
|
+
wrapped.messages.create(
|
|
77
|
+
model="claude-3-5-sonnet-20241022",
|
|
78
|
+
messages=[{"role": "user", "content": "hi"}],
|
|
79
|
+
)
|
|
80
|
+
sw.policy_check.assert_called_once()
|
|
81
|
+
sw.policy_check_strict.assert_not_called()
|
|
82
|
+
sw.observe.assert_called_once()
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def test_wrap_enforce_blocked_raises_BlockedError_and_skips_upstream(
|
|
86
|
+
fake_anthropic_client, monkeypatch
|
|
87
|
+
):
|
|
88
|
+
monkeypatch.setenv("SPENDWALL_MODE", "enforce")
|
|
89
|
+
from spendwall.anthropic import wrap
|
|
90
|
+
|
|
91
|
+
original_create = fake_anthropic_client.messages.create
|
|
92
|
+
sw = MagicMock()
|
|
93
|
+
sw.policy_check_strict.side_effect = BlockedError("cap-y", 99, 1)
|
|
94
|
+
wrapped = wrap(fake_anthropic_client, sw_client=sw, agent_id="a")
|
|
95
|
+
with pytest.raises(BlockedError):
|
|
96
|
+
wrapped.messages.create(
|
|
97
|
+
model="claude-3-5-sonnet-20241022",
|
|
98
|
+
messages=[{"role": "user", "content": "hi"}],
|
|
99
|
+
)
|
|
100
|
+
original_create.assert_not_called()
|
|
101
|
+
sw.observe.assert_not_called()
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def test_wrap_observes_upstream_error(fake_anthropic_client, monkeypatch):
|
|
105
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
106
|
+
from spendwall.anthropic import wrap
|
|
107
|
+
|
|
108
|
+
fake_anthropic_client.messages.create.side_effect = RuntimeError("boom")
|
|
109
|
+
sw = MagicMock()
|
|
110
|
+
wrapped = wrap(fake_anthropic_client, sw_client=sw, agent_id="a")
|
|
111
|
+
with pytest.raises(RuntimeError):
|
|
112
|
+
wrapped.messages.create(
|
|
113
|
+
model="claude-3-5-sonnet-20241022",
|
|
114
|
+
messages=[{"role": "user", "content": "hi"}],
|
|
115
|
+
)
|
|
116
|
+
sw.observe.assert_called_once()
|
|
117
|
+
payload = sw.observe.call_args.args[0]
|
|
118
|
+
assert payload["outcome"] == "upstream_error"
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def test_wrap_per_call_mode_overrides_env(fake_anthropic_client, monkeypatch):
|
|
122
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
123
|
+
from spendwall.anthropic import wrap
|
|
124
|
+
|
|
125
|
+
sw = MagicMock()
|
|
126
|
+
sw.policy_check_strict.return_value = {"decision": "allowed"}
|
|
127
|
+
wrapped = wrap(
|
|
128
|
+
fake_anthropic_client, sw_client=sw, agent_id="a", mode="enforce"
|
|
129
|
+
)
|
|
130
|
+
wrapped.messages.create(
|
|
131
|
+
model="claude-3-5-sonnet-20241022",
|
|
132
|
+
messages=[{"role": "user", "content": "hi"}],
|
|
133
|
+
)
|
|
134
|
+
sw.policy_check_strict.assert_called_once()
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
def test_wrap_returns_same_client_instance(fake_anthropic_client):
|
|
138
|
+
from spendwall.anthropic import wrap
|
|
139
|
+
|
|
140
|
+
sw = MagicMock()
|
|
141
|
+
out = wrap(fake_anthropic_client, sw_client=sw, agent_id="a")
|
|
142
|
+
assert out is fake_anthropic_client
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def test_wrap_propagates_tags_into_observe(fake_anthropic_client, monkeypatch):
|
|
146
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
147
|
+
from spendwall.anthropic import wrap
|
|
148
|
+
|
|
149
|
+
sw = MagicMock()
|
|
150
|
+
wrapped = wrap(
|
|
151
|
+
fake_anthropic_client, sw_client=sw, agent_id="a",
|
|
152
|
+
tags={"env": "prod"},
|
|
153
|
+
)
|
|
154
|
+
wrapped.messages.create(
|
|
155
|
+
model="claude-3-5-sonnet-20241022",
|
|
156
|
+
messages=[{"role": "user", "content": "hi"}],
|
|
157
|
+
)
|
|
158
|
+
payload = sw.observe.call_args.args[0]
|
|
159
|
+
assert payload["tags"] == {"env": "prod"}
|
|
@@ -0,0 +1,253 @@
|
|
|
1
|
+
"""Tests for SpendwallClient — uses a stub UDS HTTP server.
|
|
2
|
+
|
|
3
|
+
The socket path lives under ``/tmp`` rather than pytest's ``tmp_path`` because
|
|
4
|
+
macOS limits AF_UNIX paths to ~104 characters and ``tmp_path`` exceeds that.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import os
|
|
10
|
+
import shutil
|
|
11
|
+
import socket
|
|
12
|
+
import tempfile
|
|
13
|
+
import threading
|
|
14
|
+
from contextlib import contextmanager
|
|
15
|
+
from typing import Iterator
|
|
16
|
+
|
|
17
|
+
import pytest
|
|
18
|
+
|
|
19
|
+
from spendwall.client import SpendwallClient
|
|
20
|
+
from spendwall.errors import ApprovalRequiredError, BlockedError
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@pytest.fixture
|
|
24
|
+
def short_tmp_path() -> Iterator[str]:
|
|
25
|
+
"""Return a short directory path safe for AF_UNIX sockets on macOS."""
|
|
26
|
+
d = tempfile.mkdtemp(prefix="sw_", dir="/tmp")
|
|
27
|
+
try:
|
|
28
|
+
yield d
|
|
29
|
+
finally:
|
|
30
|
+
shutil.rmtree(d, ignore_errors=True)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@contextmanager
|
|
34
|
+
def fake_uds_server(socket_path: str, responses: list[tuple[int, bytes]]) -> Iterator[list[bytes]]:
|
|
35
|
+
"""Tiny HTTP/1.1 server bound to a UDS path. ``responses`` is a list of
|
|
36
|
+
(status_code, body) tuples — popped one per accepted connection. Returns
|
|
37
|
+
a mutable list of raw request payloads received, for assertions."""
|
|
38
|
+
if os.path.exists(socket_path):
|
|
39
|
+
os.unlink(socket_path)
|
|
40
|
+
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
41
|
+
sock.bind(socket_path)
|
|
42
|
+
sock.listen(8)
|
|
43
|
+
received: list[bytes] = []
|
|
44
|
+
stop_event = threading.Event()
|
|
45
|
+
|
|
46
|
+
def serve() -> None:
|
|
47
|
+
sock.settimeout(0.1)
|
|
48
|
+
while not stop_event.is_set() and responses:
|
|
49
|
+
try:
|
|
50
|
+
conn, _ = sock.accept()
|
|
51
|
+
except socket.timeout:
|
|
52
|
+
continue
|
|
53
|
+
except OSError:
|
|
54
|
+
return
|
|
55
|
+
data = b""
|
|
56
|
+
conn.settimeout(1.0)
|
|
57
|
+
try:
|
|
58
|
+
while b"\r\n\r\n" not in data:
|
|
59
|
+
chunk = conn.recv(4096)
|
|
60
|
+
if not chunk:
|
|
61
|
+
break
|
|
62
|
+
data += chunk
|
|
63
|
+
hdr_end = data.find(b"\r\n\r\n") + 4
|
|
64
|
+
hdrs = data[:hdr_end].decode(errors="ignore")
|
|
65
|
+
body_len = 0
|
|
66
|
+
for line in hdrs.split("\r\n"):
|
|
67
|
+
if line.lower().startswith("content-length:"):
|
|
68
|
+
body_len = int(line.split(":", 1)[1].strip())
|
|
69
|
+
while len(data) - hdr_end < body_len:
|
|
70
|
+
chunk = conn.recv(4096)
|
|
71
|
+
if not chunk:
|
|
72
|
+
break
|
|
73
|
+
data += chunk
|
|
74
|
+
received.append(data)
|
|
75
|
+
if not responses:
|
|
76
|
+
conn.close()
|
|
77
|
+
continue
|
|
78
|
+
status, body = responses.pop(0)
|
|
79
|
+
reason = "OK" if status < 400 else "ERR"
|
|
80
|
+
if status == 204:
|
|
81
|
+
conn.send(
|
|
82
|
+
f"HTTP/1.1 {status} No Content\r\n"
|
|
83
|
+
f"Content-Length: 0\r\n\r\n".encode()
|
|
84
|
+
)
|
|
85
|
+
else:
|
|
86
|
+
conn.send(
|
|
87
|
+
f"HTTP/1.1 {status} {reason}\r\n"
|
|
88
|
+
f"Content-Length: {len(body)}\r\n"
|
|
89
|
+
f"Content-Type: application/json\r\n\r\n".encode()
|
|
90
|
+
+ body
|
|
91
|
+
)
|
|
92
|
+
finally:
|
|
93
|
+
conn.close()
|
|
94
|
+
|
|
95
|
+
t = threading.Thread(target=serve, daemon=True)
|
|
96
|
+
t.start()
|
|
97
|
+
try:
|
|
98
|
+
yield received
|
|
99
|
+
finally:
|
|
100
|
+
stop_event.set()
|
|
101
|
+
try:
|
|
102
|
+
sock.close()
|
|
103
|
+
except OSError:
|
|
104
|
+
pass
|
|
105
|
+
t.join(timeout=1.0)
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def test_policy_check_returns_decision_dict(short_tmp_path):
|
|
109
|
+
sp = os.path.join(short_tmp_path, "t.sock")
|
|
110
|
+
body = b'{"decision":"allowed","alert_caps":[]}'
|
|
111
|
+
with fake_uds_server(sp, [(200, body)]):
|
|
112
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test")
|
|
113
|
+
d = c.policy_check(
|
|
114
|
+
agent_id="a1", provider="openai", model="gpt-4o",
|
|
115
|
+
estimated_cost_cents=1,
|
|
116
|
+
)
|
|
117
|
+
assert d["decision"] == "allowed"
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def test_policy_check_strict_raises_blocked(short_tmp_path):
|
|
121
|
+
sp = os.path.join(short_tmp_path, "t.sock")
|
|
122
|
+
body = (
|
|
123
|
+
b'{"decision":"blocked","cap_id":"daily","projected_cents":100,'
|
|
124
|
+
b'"limit_cents":50,"alert_caps":[]}'
|
|
125
|
+
)
|
|
126
|
+
with fake_uds_server(sp, [(200, body)]):
|
|
127
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test")
|
|
128
|
+
with pytest.raises(BlockedError) as exc_info:
|
|
129
|
+
c.policy_check_strict(
|
|
130
|
+
agent_id="a1", provider="openai", model="gpt-4o",
|
|
131
|
+
estimated_cost_cents=1,
|
|
132
|
+
)
|
|
133
|
+
assert exc_info.value.cap_id == "daily"
|
|
134
|
+
assert exc_info.value.projected_cents == 100
|
|
135
|
+
assert exc_info.value.limit_cents == 50
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def test_policy_check_strict_raises_approval_required(short_tmp_path):
|
|
139
|
+
sp = os.path.join(short_tmp_path, "t.sock")
|
|
140
|
+
body = (
|
|
141
|
+
b'{"decision":"approval_required","approval_caps":["cap-x"],'
|
|
142
|
+
b'"alert_caps":[]}'
|
|
143
|
+
)
|
|
144
|
+
with fake_uds_server(sp, [(200, body)]):
|
|
145
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test")
|
|
146
|
+
with pytest.raises(ApprovalRequiredError) as exc_info:
|
|
147
|
+
c.policy_check_strict(
|
|
148
|
+
agent_id="a1", provider="openai", model="gpt-4o",
|
|
149
|
+
estimated_cost_cents=1,
|
|
150
|
+
)
|
|
151
|
+
assert exc_info.value.approval_caps == ["cap-x"]
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def test_observe_204_does_not_raise(short_tmp_path):
|
|
155
|
+
sp = os.path.join(short_tmp_path, "t.sock")
|
|
156
|
+
with fake_uds_server(sp, [(204, b"")]):
|
|
157
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test")
|
|
158
|
+
c.observe({
|
|
159
|
+
"agent_id": "a1", "provider": "openai", "model": "gpt-4o",
|
|
160
|
+
"source": "sdk_python", "estimated_cost_cents": 1,
|
|
161
|
+
"decision_str": "allowed", "outcome": "success",
|
|
162
|
+
})
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
def test_observe_swallows_connection_error_when_fail_open(short_tmp_path):
|
|
166
|
+
sp = os.path.join(short_tmp_path, "no.sock")
|
|
167
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test", fail_open=True)
|
|
168
|
+
# Should NOT raise.
|
|
169
|
+
c.observe({
|
|
170
|
+
"agent_id": "a1", "provider": "openai", "model": "gpt-4o",
|
|
171
|
+
"source": "sdk_python", "estimated_cost_cents": 1,
|
|
172
|
+
"decision_str": "allowed", "outcome": "success",
|
|
173
|
+
})
|
|
174
|
+
|
|
175
|
+
|
|
176
|
+
def test_observe_raises_when_fail_open_disabled(short_tmp_path):
|
|
177
|
+
sp = os.path.join(short_tmp_path, "no.sock")
|
|
178
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test", fail_open=False)
|
|
179
|
+
import requests as _requests
|
|
180
|
+
with pytest.raises(_requests.RequestException):
|
|
181
|
+
c.observe({"agent_id": "a1"})
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
def test_policy_check_fail_open_returns_allowed_when_daemon_down(short_tmp_path):
|
|
185
|
+
sp = os.path.join(short_tmp_path, "no.sock")
|
|
186
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test", fail_open=True)
|
|
187
|
+
d = c.policy_check(
|
|
188
|
+
agent_id="a1", provider="openai", model="gpt-4o",
|
|
189
|
+
estimated_cost_cents=1,
|
|
190
|
+
)
|
|
191
|
+
assert d["decision"] == "allowed"
|
|
192
|
+
assert d["alert_caps"] == []
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def test_policy_check_raises_when_fail_open_disabled(short_tmp_path):
|
|
196
|
+
sp = os.path.join(short_tmp_path, "no.sock")
|
|
197
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test", fail_open=False)
|
|
198
|
+
import requests as _requests
|
|
199
|
+
with pytest.raises(_requests.RequestException):
|
|
200
|
+
c.policy_check(
|
|
201
|
+
agent_id="a1", provider="openai", model="gpt-4o",
|
|
202
|
+
estimated_cost_cents=1,
|
|
203
|
+
)
|
|
204
|
+
|
|
205
|
+
|
|
206
|
+
def test_bearer_token_sent_on_every_request(short_tmp_path):
|
|
207
|
+
sp = os.path.join(short_tmp_path, "t.sock")
|
|
208
|
+
with fake_uds_server(sp, [
|
|
209
|
+
(200, b'{"decision":"allowed","alert_caps":[]}'),
|
|
210
|
+
(204, b""),
|
|
211
|
+
]) as received:
|
|
212
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_my_token_xyz")
|
|
213
|
+
c.policy_check(
|
|
214
|
+
agent_id="a1", provider="openai", model="gpt-4o",
|
|
215
|
+
estimated_cost_cents=1,
|
|
216
|
+
)
|
|
217
|
+
c.observe({"agent_id": "a1", "outcome": "success"})
|
|
218
|
+
|
|
219
|
+
# Both requests should have the bearer header.
|
|
220
|
+
assert len(received) == 2
|
|
221
|
+
for req in received:
|
|
222
|
+
lower = req.lower()
|
|
223
|
+
assert b"authorization: bearer swk_my_token_xyz" in lower
|
|
224
|
+
|
|
225
|
+
|
|
226
|
+
def test_about_returns_daemon_metadata(short_tmp_path):
|
|
227
|
+
sp = os.path.join(short_tmp_path, "t.sock")
|
|
228
|
+
body = b'{"daemon_version":"0.1.0","api_versions_supported":["v1"]}'
|
|
229
|
+
with fake_uds_server(sp, [(200, body)]):
|
|
230
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_test")
|
|
231
|
+
meta = c.about()
|
|
232
|
+
assert meta["daemon_version"] == "0.1.0"
|
|
233
|
+
assert meta["api_versions_supported"] == ["v1"]
|
|
234
|
+
|
|
235
|
+
|
|
236
|
+
def test_policy_check_sends_request_body(short_tmp_path):
|
|
237
|
+
sp = os.path.join(short_tmp_path, "t.sock")
|
|
238
|
+
body = b'{"decision":"allowed","alert_caps":[]}'
|
|
239
|
+
with fake_uds_server(sp, [(200, body)]) as received:
|
|
240
|
+
c = SpendwallClient(socket_path=sp, bearer="swk_t")
|
|
241
|
+
c.policy_check(
|
|
242
|
+
agent_id="agent-x", provider="anthropic",
|
|
243
|
+
model="claude-3-5-sonnet-20241022",
|
|
244
|
+
estimated_cost_cents=42,
|
|
245
|
+
tags={"env": "test"},
|
|
246
|
+
)
|
|
247
|
+
assert len(received) == 1
|
|
248
|
+
payload = received[0]
|
|
249
|
+
assert b"agent-x" in payload
|
|
250
|
+
assert b"anthropic" in payload
|
|
251
|
+
assert b"claude-3-5-sonnet-20241022" in payload
|
|
252
|
+
assert b"42" in payload
|
|
253
|
+
assert b"env" in payload
|
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
"""Tests for spendwall.errors — class hierarchy and message construction."""
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
|
|
5
|
+
from spendwall.errors import (
|
|
6
|
+
ApprovalRequiredError,
|
|
7
|
+
BlockedError,
|
|
8
|
+
SpendwallError,
|
|
9
|
+
SpendwallVersionMismatch,
|
|
10
|
+
)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def test_blocked_error_inherits_from_spendwall_error():
|
|
14
|
+
assert issubclass(BlockedError, SpendwallError)
|
|
15
|
+
assert issubclass(SpendwallError, Exception)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def test_approval_required_error_inherits_from_spendwall_error():
|
|
19
|
+
assert issubclass(ApprovalRequiredError, SpendwallError)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def test_version_mismatch_inherits_from_spendwall_error():
|
|
23
|
+
assert issubclass(SpendwallVersionMismatch, SpendwallError)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def test_blocked_error_stores_attrs_and_formats_message():
|
|
27
|
+
err = BlockedError(cap_id="daily-cap", projected_cents=200, limit_cents=100)
|
|
28
|
+
assert err.cap_id == "daily-cap"
|
|
29
|
+
assert err.projected_cents == 200
|
|
30
|
+
assert err.limit_cents == 100
|
|
31
|
+
msg = str(err)
|
|
32
|
+
assert "daily-cap" in msg
|
|
33
|
+
assert "200" in msg
|
|
34
|
+
assert "100" in msg
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_approval_required_error_stores_caps_and_reason():
|
|
38
|
+
err = ApprovalRequiredError(approval_caps=["c1", "c2"])
|
|
39
|
+
assert err.approval_caps == ["c1", "c2"]
|
|
40
|
+
assert err.reason == "denied"
|
|
41
|
+
assert "c1" in str(err)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def test_approval_required_error_custom_reason():
|
|
45
|
+
err = ApprovalRequiredError(approval_caps=["c1"], reason="timeout")
|
|
46
|
+
assert err.reason == "timeout"
|
|
47
|
+
assert "timeout" in str(err)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def test_blocked_error_can_be_caught_as_spendwall_error():
|
|
51
|
+
with pytest.raises(SpendwallError):
|
|
52
|
+
raise BlockedError(cap_id="x", projected_cents=1, limit_cents=0)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def test_approval_required_error_can_be_caught_as_spendwall_error():
|
|
56
|
+
with pytest.raises(SpendwallError):
|
|
57
|
+
raise ApprovalRequiredError(approval_caps=["x"])
|
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
import pytest
|
|
2
|
+
|
|
3
|
+
import spendwall.modes as m
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def test_resolve_default_is_observe(monkeypatch):
|
|
7
|
+
monkeypatch.delenv("SPENDWALL_MODE", raising=False)
|
|
8
|
+
assert m.resolve() == "observe"
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def test_resolve_env_var_wins_over_default(monkeypatch):
|
|
12
|
+
monkeypatch.setenv("SPENDWALL_MODE", "enforce")
|
|
13
|
+
assert m.resolve() == "enforce"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def test_resolve_configured_wins_over_env(monkeypatch):
|
|
17
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
18
|
+
assert m.resolve(configured="advisory") == "advisory"
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def test_resolve_per_call_wins_over_configured(monkeypatch):
|
|
22
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
23
|
+
assert m.resolve(per_call="enforce", configured="advisory") == "enforce"
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def test_resolve_rejects_invalid_mode():
|
|
27
|
+
with pytest.raises(ValueError):
|
|
28
|
+
m.resolve(per_call="block")
|
|
@@ -0,0 +1,178 @@
|
|
|
1
|
+
"""Tests for spendwall.openai.wrap()."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from unittest.mock import MagicMock
|
|
6
|
+
|
|
7
|
+
import pytest
|
|
8
|
+
|
|
9
|
+
from spendwall.errors import BlockedError
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
@pytest.fixture
|
|
13
|
+
def fake_openai_client():
|
|
14
|
+
client = MagicMock(name="openai.OpenAI")
|
|
15
|
+
completion = MagicMock(name="ChatCompletion")
|
|
16
|
+
completion.model = "gpt-4o"
|
|
17
|
+
completion.usage = MagicMock(prompt_tokens=10, completion_tokens=5)
|
|
18
|
+
completion.id = "chatcmpl_test"
|
|
19
|
+
client.chat.completions.create.return_value = completion
|
|
20
|
+
|
|
21
|
+
response = MagicMock(name="Response")
|
|
22
|
+
response.id = "resp_test"
|
|
23
|
+
response.usage = MagicMock(prompt_tokens=7, completion_tokens=3)
|
|
24
|
+
client.responses.create.return_value = response
|
|
25
|
+
return client
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def test_wrap_observe_mode_skips_policy_check_and_passes_through(
|
|
29
|
+
fake_openai_client, monkeypatch
|
|
30
|
+
):
|
|
31
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
32
|
+
from spendwall.openai import wrap
|
|
33
|
+
|
|
34
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
35
|
+
wrapped = wrap(fake_openai_client, sw_client=sw_client, agent_id="a-test")
|
|
36
|
+
out = wrapped.chat.completions.create(
|
|
37
|
+
model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
|
|
38
|
+
)
|
|
39
|
+
assert out.id == "chatcmpl_test"
|
|
40
|
+
sw_client.policy_check_strict.assert_not_called()
|
|
41
|
+
sw_client.policy_check.assert_not_called()
|
|
42
|
+
sw_client.observe.assert_called_once()
|
|
43
|
+
# Observe should carry post-flight token usage.
|
|
44
|
+
payload = sw_client.observe.call_args.args[0]
|
|
45
|
+
assert payload["agent_id"] == "a-test"
|
|
46
|
+
assert payload["provider"] == "openai"
|
|
47
|
+
assert payload["source"] == "sdk_python"
|
|
48
|
+
assert payload["input_tokens"] == 10
|
|
49
|
+
assert payload["output_tokens"] == 5
|
|
50
|
+
assert payload["outcome"] == "success"
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def test_wrap_enforce_mode_calls_policy_check_strict_then_observe(
|
|
54
|
+
fake_openai_client, monkeypatch
|
|
55
|
+
):
|
|
56
|
+
monkeypatch.setenv("SPENDWALL_MODE", "enforce")
|
|
57
|
+
from spendwall.openai import wrap
|
|
58
|
+
|
|
59
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
60
|
+
sw_client.policy_check_strict.return_value = {
|
|
61
|
+
"decision": "allowed", "alert_caps": [],
|
|
62
|
+
}
|
|
63
|
+
wrapped = wrap(fake_openai_client, sw_client=sw_client, agent_id="a-test")
|
|
64
|
+
out = wrapped.chat.completions.create(
|
|
65
|
+
model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
|
|
66
|
+
)
|
|
67
|
+
assert out.id == "chatcmpl_test"
|
|
68
|
+
sw_client.policy_check_strict.assert_called_once()
|
|
69
|
+
pcs_kwargs = sw_client.policy_check_strict.call_args.kwargs
|
|
70
|
+
assert pcs_kwargs["agent_id"] == "a-test"
|
|
71
|
+
assert pcs_kwargs["provider"] == "openai"
|
|
72
|
+
assert pcs_kwargs["model"] == "gpt-4o"
|
|
73
|
+
sw_client.observe.assert_called_once()
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def test_wrap_advisory_mode_calls_policy_check_non_strict(
|
|
77
|
+
fake_openai_client, monkeypatch
|
|
78
|
+
):
|
|
79
|
+
monkeypatch.setenv("SPENDWALL_MODE", "advisory")
|
|
80
|
+
from spendwall.openai import wrap
|
|
81
|
+
|
|
82
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
83
|
+
sw_client.policy_check.return_value = {"decision": "allowed", "alert_caps": []}
|
|
84
|
+
wrapped = wrap(fake_openai_client, sw_client=sw_client, agent_id="a-test")
|
|
85
|
+
wrapped.chat.completions.create(
|
|
86
|
+
model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
|
|
87
|
+
)
|
|
88
|
+
sw_client.policy_check.assert_called_once()
|
|
89
|
+
sw_client.policy_check_strict.assert_not_called()
|
|
90
|
+
sw_client.observe.assert_called_once()
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
def test_wrap_enforce_blocked_raises_BlockedError_and_skips_upstream(
|
|
94
|
+
fake_openai_client, monkeypatch
|
|
95
|
+
):
|
|
96
|
+
monkeypatch.setenv("SPENDWALL_MODE", "enforce")
|
|
97
|
+
from spendwall.openai import wrap
|
|
98
|
+
|
|
99
|
+
# Snapshot the original create-method mock before wrap() replaces it.
|
|
100
|
+
original_create = fake_openai_client.chat.completions.create
|
|
101
|
+
|
|
102
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
103
|
+
sw_client.policy_check_strict.side_effect = BlockedError("cap-x", 200, 100)
|
|
104
|
+
wrapped = wrap(fake_openai_client, sw_client=sw_client, agent_id="a-test")
|
|
105
|
+
with pytest.raises(BlockedError):
|
|
106
|
+
wrapped.chat.completions.create(
|
|
107
|
+
model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
|
|
108
|
+
)
|
|
109
|
+
original_create.assert_not_called()
|
|
110
|
+
sw_client.observe.assert_not_called()
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def test_wrap_observes_upstream_error(fake_openai_client, monkeypatch):
|
|
114
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
115
|
+
from spendwall.openai import wrap
|
|
116
|
+
|
|
117
|
+
fake_openai_client.chat.completions.create.side_effect = RuntimeError("boom")
|
|
118
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
119
|
+
wrapped = wrap(fake_openai_client, sw_client=sw_client, agent_id="a-test")
|
|
120
|
+
with pytest.raises(RuntimeError):
|
|
121
|
+
wrapped.chat.completions.create(
|
|
122
|
+
model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
|
|
123
|
+
)
|
|
124
|
+
sw_client.observe.assert_called_once()
|
|
125
|
+
payload = sw_client.observe.call_args.args[0]
|
|
126
|
+
assert payload["outcome"] == "upstream_error"
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def test_wrap_per_call_mode_overrides_configured(fake_openai_client, monkeypatch):
|
|
130
|
+
"""Configured mode should override env default."""
|
|
131
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
132
|
+
from spendwall.openai import wrap
|
|
133
|
+
|
|
134
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
135
|
+
sw_client.policy_check_strict.return_value = {"decision": "allowed"}
|
|
136
|
+
# configured=enforce overrides env=observe.
|
|
137
|
+
wrapped = wrap(
|
|
138
|
+
fake_openai_client, sw_client=sw_client, agent_id="a", mode="enforce",
|
|
139
|
+
)
|
|
140
|
+
wrapped.chat.completions.create(
|
|
141
|
+
model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
|
|
142
|
+
)
|
|
143
|
+
sw_client.policy_check_strict.assert_called_once()
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
def test_wrap_patches_responses_create_too(fake_openai_client, monkeypatch):
|
|
147
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
148
|
+
from spendwall.openai import wrap
|
|
149
|
+
|
|
150
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
151
|
+
wrapped = wrap(fake_openai_client, sw_client=sw_client, agent_id="a")
|
|
152
|
+
out = wrapped.responses.create(model="gpt-4o", input="hi")
|
|
153
|
+
assert out.id == "resp_test"
|
|
154
|
+
sw_client.observe.assert_called_once()
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def test_wrap_returns_same_client_instance(fake_openai_client):
|
|
158
|
+
from spendwall.openai import wrap
|
|
159
|
+
|
|
160
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
161
|
+
out = wrap(fake_openai_client, sw_client=sw_client, agent_id="a")
|
|
162
|
+
assert out is fake_openai_client
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
def test_wrap_propagates_tags_into_observe(fake_openai_client, monkeypatch):
|
|
166
|
+
monkeypatch.setenv("SPENDWALL_MODE", "observe")
|
|
167
|
+
from spendwall.openai import wrap
|
|
168
|
+
|
|
169
|
+
sw_client = MagicMock(name="SpendwallClient")
|
|
170
|
+
wrapped = wrap(
|
|
171
|
+
fake_openai_client, sw_client=sw_client, agent_id="a",
|
|
172
|
+
tags={"env": "test", "team": "ai"},
|
|
173
|
+
)
|
|
174
|
+
wrapped.chat.completions.create(
|
|
175
|
+
model="gpt-4o", messages=[{"role": "user", "content": "hi"}]
|
|
176
|
+
)
|
|
177
|
+
payload = sw_client.observe.call_args.args[0]
|
|
178
|
+
assert payload["tags"] == {"env": "test", "team": "ai"}
|