subagent-budget-team 0.1.0__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.
- subagent_budget_team-0.1.0/LICENSE +18 -0
- subagent_budget_team-0.1.0/PKG-INFO +97 -0
- subagent_budget_team-0.1.0/README.md +77 -0
- subagent_budget_team-0.1.0/pyproject.toml +31 -0
- subagent_budget_team-0.1.0/setup.cfg +4 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team/__init__.py +8 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team/cli.py +188 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team/core.py +214 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team.egg-info/PKG-INFO +97 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team.egg-info/SOURCES.txt +12 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team.egg-info/dependency_links.txt +1 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team.egg-info/entry_points.txt +2 -0
- subagent_budget_team-0.1.0/src/subagent_budget_team.egg-info/top_level.txt +1 -0
- subagent_budget_team-0.1.0/tests/test_smoke.py +126 -0
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 hao li
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
|
|
6
|
+
associated documentation files (the "Software"), to deal in the Software without restriction, including
|
|
7
|
+
without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
8
|
+
copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the
|
|
9
|
+
following conditions:
|
|
10
|
+
|
|
11
|
+
The above copyright notice and this permission notice shall be included in all copies or substantial
|
|
12
|
+
portions of the Software.
|
|
13
|
+
|
|
14
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
|
|
15
|
+
LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO
|
|
16
|
+
EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
17
|
+
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
|
|
18
|
+
USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: subagent-budget-team
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Team/CI budget control for subagent token spend: shared budget pools, per-agent quotas, overspend checks, and a kill switch.
|
|
5
|
+
Author-email: hao li <haoli5933@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/hahahahahahahahah6/subagent-budget-team
|
|
8
|
+
Project-URL: Issues, https://github.com/hahahahahahahahah6/subagent-budget-team/issues
|
|
9
|
+
Keywords: subagent,budget,token,cost-control,ci,agents
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Environment :: Console
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Topic :: Software Development :: Quality Assurance
|
|
16
|
+
Requires-Python: >=3.9
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
License-File: LICENSE
|
|
19
|
+
Dynamic: license-file
|
|
20
|
+
|
|
21
|
+
# subagent-budget-team
|
|
22
|
+
|
|
23
|
+
Team/CI budget control for subagent token spend. When several agents run
|
|
24
|
+
concurrently — a CI matrix, a fleet of coding agents, a team sharing one
|
|
25
|
+
token pool — costs can spiral before anyone notices. `subagent-budget-team`
|
|
26
|
+
gives you a **pre-spend control plane**: shared budget pools, per-agent
|
|
27
|
+
quotas, overspend detection, and a kill switch that freezes spend *before*
|
|
28
|
+
the next invoice.
|
|
29
|
+
|
|
30
|
+
Measured context: subagents can burn ~48% of an agent's total token budget
|
|
31
|
+
(see the `subagent-tax` write-up). This tool is what you deploy once you
|
|
32
|
+
know that number and want to do something about it at the team level.
|
|
33
|
+
|
|
34
|
+
## Install
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install subagent-budget-team
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
Requires Python 3.9+. Zero dependencies (stdlib only).
|
|
41
|
+
|
|
42
|
+
## Usage
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
# 1. Create a shared budget pool for a CI job or team
|
|
46
|
+
subagent-budget-team pool create --name ci-nightly --limit-usd 25
|
|
47
|
+
|
|
48
|
+
# 2. Give each agent its own quota inside the pool
|
|
49
|
+
subagent-budget-team quota set --pool ci-nightly --agent coder-1 --limit-usd 5
|
|
50
|
+
subagent-budget-team quota set --pool ci-nightly --agent reviewer-1 --limit-usd 3
|
|
51
|
+
|
|
52
|
+
# 3. Record spend (call this from your agent wrapper / CI step)
|
|
53
|
+
subagent-budget-team spend --pool ci-nightly --agent coder-1 --tokens 12000 --usd 1.20
|
|
54
|
+
|
|
55
|
+
# 4. Check everything; exit code 1 means over budget — gate your CI on it
|
|
56
|
+
subagent-budget-team check
|
|
57
|
+
# or machine-readable:
|
|
58
|
+
subagent-budget-team check --json
|
|
59
|
+
|
|
60
|
+
# 5. Kill switch: freeze a pool (or a single agent) immediately
|
|
61
|
+
subagent-budget-team kill --pool ci-nightly
|
|
62
|
+
subagent-budget-team spend --pool ci-nightly --agent coder-1 --usd 1 # rejected
|
|
63
|
+
subagent-budget-team release --pool ci-nightly # unfreeze
|
|
64
|
+
|
|
65
|
+
# List pools
|
|
66
|
+
subagent-budget-team pool list
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
State lives in `~/.subagent-budget-team/budgets.json` (override with
|
|
70
|
+
`--state` or the `SBT_HOME` env var, handy for per-job isolation in CI).
|
|
71
|
+
|
|
72
|
+
### CI example (GitHub Actions)
|
|
73
|
+
|
|
74
|
+
```yaml
|
|
75
|
+
- run: pip install subagent-budget-team
|
|
76
|
+
- run: subagent-budget-team pool create --name pr-${{ github.event.number }} --limit-usd 10
|
|
77
|
+
- run: ./run-agents.sh # agents call `spend` as they work
|
|
78
|
+
- run: subagent-budget-team check # fails the job on overspend
|
|
79
|
+
- if: failure()
|
|
80
|
+
run: subagent-budget-team kill --pool pr-${{ github.event.number }}
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
## Differentiation
|
|
84
|
+
|
|
85
|
+
- **vs `subagent-tax`**: that tool is a *personal audit* lens — it measures
|
|
86
|
+
how much of *your* token budget subagents burn after the fact.
|
|
87
|
+
`subagent-budget-team` is the *team/CI control* lens: a shared budget
|
|
88
|
+
pool with per-agent quotas and a circuit breaker, designed for many
|
|
89
|
+
agents spending from one pot.
|
|
90
|
+
- **vs cloud-vendor budget alerts**: those are *post-hoc alarms* — you get
|
|
91
|
+
paged after the money is spent. This is a *pre-spend circuit breaker*:
|
|
92
|
+
`kill` rejects new spend immediately and `check` can gate a pipeline
|
|
93
|
+
before the next batch of work runs.
|
|
94
|
+
|
|
95
|
+
## License
|
|
96
|
+
|
|
97
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
# subagent-budget-team
|
|
2
|
+
|
|
3
|
+
Team/CI budget control for subagent token spend. When several agents run
|
|
4
|
+
concurrently — a CI matrix, a fleet of coding agents, a team sharing one
|
|
5
|
+
token pool — costs can spiral before anyone notices. `subagent-budget-team`
|
|
6
|
+
gives you a **pre-spend control plane**: shared budget pools, per-agent
|
|
7
|
+
quotas, overspend detection, and a kill switch that freezes spend *before*
|
|
8
|
+
the next invoice.
|
|
9
|
+
|
|
10
|
+
Measured context: subagents can burn ~48% of an agent's total token budget
|
|
11
|
+
(see the `subagent-tax` write-up). This tool is what you deploy once you
|
|
12
|
+
know that number and want to do something about it at the team level.
|
|
13
|
+
|
|
14
|
+
## Install
|
|
15
|
+
|
|
16
|
+
```bash
|
|
17
|
+
pip install subagent-budget-team
|
|
18
|
+
```
|
|
19
|
+
|
|
20
|
+
Requires Python 3.9+. Zero dependencies (stdlib only).
|
|
21
|
+
|
|
22
|
+
## Usage
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
# 1. Create a shared budget pool for a CI job or team
|
|
26
|
+
subagent-budget-team pool create --name ci-nightly --limit-usd 25
|
|
27
|
+
|
|
28
|
+
# 2. Give each agent its own quota inside the pool
|
|
29
|
+
subagent-budget-team quota set --pool ci-nightly --agent coder-1 --limit-usd 5
|
|
30
|
+
subagent-budget-team quota set --pool ci-nightly --agent reviewer-1 --limit-usd 3
|
|
31
|
+
|
|
32
|
+
# 3. Record spend (call this from your agent wrapper / CI step)
|
|
33
|
+
subagent-budget-team spend --pool ci-nightly --agent coder-1 --tokens 12000 --usd 1.20
|
|
34
|
+
|
|
35
|
+
# 4. Check everything; exit code 1 means over budget — gate your CI on it
|
|
36
|
+
subagent-budget-team check
|
|
37
|
+
# or machine-readable:
|
|
38
|
+
subagent-budget-team check --json
|
|
39
|
+
|
|
40
|
+
# 5. Kill switch: freeze a pool (or a single agent) immediately
|
|
41
|
+
subagent-budget-team kill --pool ci-nightly
|
|
42
|
+
subagent-budget-team spend --pool ci-nightly --agent coder-1 --usd 1 # rejected
|
|
43
|
+
subagent-budget-team release --pool ci-nightly # unfreeze
|
|
44
|
+
|
|
45
|
+
# List pools
|
|
46
|
+
subagent-budget-team pool list
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
State lives in `~/.subagent-budget-team/budgets.json` (override with
|
|
50
|
+
`--state` or the `SBT_HOME` env var, handy for per-job isolation in CI).
|
|
51
|
+
|
|
52
|
+
### CI example (GitHub Actions)
|
|
53
|
+
|
|
54
|
+
```yaml
|
|
55
|
+
- run: pip install subagent-budget-team
|
|
56
|
+
- run: subagent-budget-team pool create --name pr-${{ github.event.number }} --limit-usd 10
|
|
57
|
+
- run: ./run-agents.sh # agents call `spend` as they work
|
|
58
|
+
- run: subagent-budget-team check # fails the job on overspend
|
|
59
|
+
- if: failure()
|
|
60
|
+
run: subagent-budget-team kill --pool pr-${{ github.event.number }}
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
## Differentiation
|
|
64
|
+
|
|
65
|
+
- **vs `subagent-tax`**: that tool is a *personal audit* lens — it measures
|
|
66
|
+
how much of *your* token budget subagents burn after the fact.
|
|
67
|
+
`subagent-budget-team` is the *team/CI control* lens: a shared budget
|
|
68
|
+
pool with per-agent quotas and a circuit breaker, designed for many
|
|
69
|
+
agents spending from one pot.
|
|
70
|
+
- **vs cloud-vendor budget alerts**: those are *post-hoc alarms* — you get
|
|
71
|
+
paged after the money is spent. This is a *pre-spend circuit breaker*:
|
|
72
|
+
`kill` rejects new spend immediately and `check` can gate a pipeline
|
|
73
|
+
before the next batch of work runs.
|
|
74
|
+
|
|
75
|
+
## License
|
|
76
|
+
|
|
77
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "subagent-budget-team"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Team/CI budget control for subagent token spend: shared budget pools, per-agent quotas, overspend checks, and a kill switch."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = { text = "MIT" }
|
|
11
|
+
requires-python = ">=3.9"
|
|
12
|
+
authors = [{ name = "hao li", email = "haoli5933@gmail.com" }]
|
|
13
|
+
keywords = ["subagent", "budget", "token", "cost-control", "ci", "agents"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 4 - Beta",
|
|
16
|
+
"Environment :: Console",
|
|
17
|
+
"Intended Audience :: Developers",
|
|
18
|
+
"License :: OSI Approved :: MIT License",
|
|
19
|
+
"Programming Language :: Python :: 3",
|
|
20
|
+
"Topic :: Software Development :: Quality Assurance",
|
|
21
|
+
]
|
|
22
|
+
|
|
23
|
+
[project.urls]
|
|
24
|
+
Homepage = "https://github.com/hahahahahahahahah6/subagent-budget-team"
|
|
25
|
+
Issues = "https://github.com/hahahahahahahahah6/subagent-budget-team/issues"
|
|
26
|
+
|
|
27
|
+
[project.scripts]
|
|
28
|
+
subagent-budget-team = "subagent_budget_team.cli:main"
|
|
29
|
+
|
|
30
|
+
[tool.setuptools.packages.find]
|
|
31
|
+
where = ["src"]
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
"""subagent-budget-team: team/CI budget control for subagent token spend.
|
|
2
|
+
|
|
3
|
+
A stdlib-only CLI that manages shared token-budget pools, per-agent
|
|
4
|
+
quotas, pre-spend circuit breaking (kill switch), and overspend checks.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
__version__ = "0.1.0"
|
|
8
|
+
__all__ = ["__version__"]
|
|
@@ -0,0 +1,188 @@
|
|
|
1
|
+
"""Command-line interface for subagent-budget-team."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import json
|
|
7
|
+
import sys
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
|
|
10
|
+
from . import __version__
|
|
11
|
+
from .core import BudgetError, Store, default_state_path
|
|
12
|
+
|
|
13
|
+
EXIT_OK = 0
|
|
14
|
+
EXIT_OVER_BUDGET = 1
|
|
15
|
+
EXIT_USAGE = 2
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _store(args: argparse.Namespace) -> Store:
|
|
19
|
+
path = Path(args.state) if getattr(args, "state", None) else default_state_path()
|
|
20
|
+
return Store(path).load()
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def cmd_pool_create(args: argparse.Namespace) -> int:
|
|
24
|
+
try:
|
|
25
|
+
_store(args).create_pool(args.name, args.limit_usd, args.limit_tokens)
|
|
26
|
+
except BudgetError as exc:
|
|
27
|
+
print(f"error: {exc}", file=sys.stderr)
|
|
28
|
+
return EXIT_USAGE
|
|
29
|
+
print(f"pool '{args.name}' created")
|
|
30
|
+
return EXIT_OK
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def cmd_pool_list(args: argparse.Namespace) -> int:
|
|
34
|
+
pools = _store(args).list_pools()
|
|
35
|
+
if args.json:
|
|
36
|
+
print(json.dumps(pools, indent=2))
|
|
37
|
+
return EXIT_OK
|
|
38
|
+
if not pools:
|
|
39
|
+
print("no pools yet; create one with: pool create --name NAME --limit-usd X")
|
|
40
|
+
return EXIT_OK
|
|
41
|
+
for p in pools:
|
|
42
|
+
status = "KILLED" if p["killed"] else "active"
|
|
43
|
+
lim = []
|
|
44
|
+
if p["limit_usd"] is not None:
|
|
45
|
+
lim.append(f"${p['limit_usd']:.2f}")
|
|
46
|
+
if p["limit_tokens"] is not None:
|
|
47
|
+
lim.append(f"{p['limit_tokens']} tokens")
|
|
48
|
+
print(f"{p['name']}: limit {' + '.join(lim)} | "
|
|
49
|
+
f"spent ${p['spent_usd']:.4f} / {p['spent_tokens']} tokens | "
|
|
50
|
+
f"{status} | agents: {len(p['agents'])}")
|
|
51
|
+
return EXIT_OK
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def cmd_quota_set(args: argparse.Namespace) -> int:
|
|
55
|
+
try:
|
|
56
|
+
_store(args).set_quota(args.pool, args.agent,
|
|
57
|
+
args.limit_usd, args.limit_tokens)
|
|
58
|
+
except BudgetError as exc:
|
|
59
|
+
print(f"error: {exc}", file=sys.stderr)
|
|
60
|
+
return EXIT_USAGE
|
|
61
|
+
print(f"quota set: agent '{args.agent}' in pool '{args.pool}'")
|
|
62
|
+
return EXIT_OK
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def cmd_spend(args: argparse.Namespace) -> int:
|
|
66
|
+
try:
|
|
67
|
+
_store(args).record_spend(args.pool, args.agent,
|
|
68
|
+
args.tokens, args.usd)
|
|
69
|
+
except BudgetError as exc:
|
|
70
|
+
print(f"error: {exc}", file=sys.stderr)
|
|
71
|
+
return EXIT_USAGE
|
|
72
|
+
print(f"recorded: agent '{args.agent}' +{args.tokens} tokens "
|
|
73
|
+
f"+${args.usd:.4f} in pool '{args.pool}'")
|
|
74
|
+
return EXIT_OK
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def cmd_check(args: argparse.Namespace) -> int:
|
|
78
|
+
violations = _store(args).check()
|
|
79
|
+
if args.json:
|
|
80
|
+
print(json.dumps([v.to_dict() for v in violations], indent=2))
|
|
81
|
+
elif violations:
|
|
82
|
+
print(f"OVER BUDGET: {len(violations)} violation(s)")
|
|
83
|
+
for v in violations:
|
|
84
|
+
print(f" - {v.detail}")
|
|
85
|
+
else:
|
|
86
|
+
print("OK: all pools and quotas within budget")
|
|
87
|
+
return EXIT_OVER_BUDGET if violations else EXIT_OK
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def cmd_kill(args: argparse.Namespace) -> int:
|
|
91
|
+
try:
|
|
92
|
+
_store(args).kill(args.pool, args.agent)
|
|
93
|
+
except BudgetError as exc:
|
|
94
|
+
print(f"error: {exc}", file=sys.stderr)
|
|
95
|
+
return EXIT_USAGE
|
|
96
|
+
target = f"agent '{args.agent}' in " if args.agent else ""
|
|
97
|
+
print(f"kill switch engaged: {target}pool '{args.pool}'")
|
|
98
|
+
return EXIT_OK
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def cmd_release(args: argparse.Namespace) -> int:
|
|
102
|
+
try:
|
|
103
|
+
_store(args).release(args.pool, args.agent)
|
|
104
|
+
except BudgetError as exc:
|
|
105
|
+
print(f"error: {exc}", file=sys.stderr)
|
|
106
|
+
return EXIT_USAGE
|
|
107
|
+
target = f"agent '{args.agent}' in " if args.agent else ""
|
|
108
|
+
print(f"released: {target}pool '{args.pool}'")
|
|
109
|
+
return EXIT_OK
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def _add_limit_opts(p: argparse.ArgumentParser) -> None:
|
|
113
|
+
p.add_argument("--limit-usd", type=float, default=None,
|
|
114
|
+
help="budget limit in USD (e.g. 25.50)")
|
|
115
|
+
p.add_argument("--limit-tokens", type=int, default=None,
|
|
116
|
+
help="budget limit in tokens")
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def build_parser() -> argparse.ArgumentParser:
|
|
120
|
+
ap = argparse.ArgumentParser(
|
|
121
|
+
prog="subagent-budget-team",
|
|
122
|
+
description="Team/CI budget control for subagent token spend: "
|
|
123
|
+
"shared budget pools, per-agent quotas, overspend checks, "
|
|
124
|
+
"and a kill switch.",
|
|
125
|
+
)
|
|
126
|
+
ap.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
|
|
127
|
+
ap.add_argument("--state", default=None,
|
|
128
|
+
help="path to state file (default: ~/.subagent-budget-team/budgets.json)")
|
|
129
|
+
sub = ap.add_subparsers(dest="command", required=True)
|
|
130
|
+
|
|
131
|
+
# pool create / pool list
|
|
132
|
+
pool_p = sub.add_parser("pool", help="manage budget pools")
|
|
133
|
+
pool_sub = pool_p.add_subparsers(dest="pool_command", required=True)
|
|
134
|
+
pc = pool_sub.add_parser("create", help="create a budget pool")
|
|
135
|
+
pc.add_argument("--name", required=True, help="pool name")
|
|
136
|
+
_add_limit_opts(pc)
|
|
137
|
+
pc.set_defaults(func=cmd_pool_create)
|
|
138
|
+
pl = pool_sub.add_parser("list", help="list budget pools")
|
|
139
|
+
pl.add_argument("--json", action="store_true", help="JSON output")
|
|
140
|
+
pl.set_defaults(func=cmd_pool_list)
|
|
141
|
+
|
|
142
|
+
# quota set
|
|
143
|
+
quota_p = sub.add_parser("quota", help="manage per-agent quotas")
|
|
144
|
+
quota_sub = quota_p.add_subparsers(dest="quota_command", required=True)
|
|
145
|
+
qs = quota_sub.add_parser("set", help="set an agent's quota inside a pool")
|
|
146
|
+
qs.add_argument("--pool", required=True, help="pool name")
|
|
147
|
+
qs.add_argument("--agent", required=True, help="agent id")
|
|
148
|
+
_add_limit_opts(qs)
|
|
149
|
+
qs.set_defaults(func=cmd_quota_set)
|
|
150
|
+
|
|
151
|
+
# spend
|
|
152
|
+
sp = sub.add_parser("spend", help="record token/USD spend")
|
|
153
|
+
sp.add_argument("--pool", required=True)
|
|
154
|
+
sp.add_argument("--agent", required=True)
|
|
155
|
+
sp.add_argument("--tokens", type=int, default=0)
|
|
156
|
+
sp.add_argument("--usd", type=float, default=0.0)
|
|
157
|
+
sp.set_defaults(func=cmd_spend)
|
|
158
|
+
|
|
159
|
+
# check
|
|
160
|
+
ch = sub.add_parser("check", help="detect overspend (exit 1 if over budget)")
|
|
161
|
+
ch.add_argument("--json", action="store_true", help="JSON output")
|
|
162
|
+
ch.set_defaults(func=cmd_check)
|
|
163
|
+
|
|
164
|
+
# kill / release
|
|
165
|
+
k = sub.add_parser("kill", help="engage the kill switch (freeze a pool or agent)")
|
|
166
|
+
k.add_argument("--pool", required=True)
|
|
167
|
+
k.add_argument("--agent", default=None, help="omit to kill the whole pool")
|
|
168
|
+
k.set_defaults(func=cmd_kill)
|
|
169
|
+
r = sub.add_parser("release", help="release a kill switch")
|
|
170
|
+
r.add_argument("--pool", required=True)
|
|
171
|
+
r.add_argument("--agent", default=None, help="omit to release the whole pool")
|
|
172
|
+
r.set_defaults(func=cmd_release)
|
|
173
|
+
|
|
174
|
+
return ap
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def main(argv: list[str] | None = None) -> int:
|
|
178
|
+
parser = build_parser()
|
|
179
|
+
args = parser.parse_args(argv)
|
|
180
|
+
try:
|
|
181
|
+
return int(args.func(args))
|
|
182
|
+
except BudgetError as exc: # pragma: no cover - defensive
|
|
183
|
+
print(f"error: {exc}", file=sys.stderr)
|
|
184
|
+
return EXIT_USAGE
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
if __name__ == "__main__": # pragma: no cover
|
|
188
|
+
sys.exit(main())
|
|
@@ -0,0 +1,214 @@
|
|
|
1
|
+
"""Core budget-pool logic. Pure stdlib: JSON state file, no dependencies."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import os
|
|
7
|
+
from dataclasses import dataclass, field
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def default_state_path() -> Path:
|
|
12
|
+
base = os.environ.get("SBT_HOME") or os.path.expanduser("~")
|
|
13
|
+
return Path(base) / ".subagent-budget-team" / "budgets.json"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass
|
|
17
|
+
class BudgetError(Exception):
|
|
18
|
+
message: str
|
|
19
|
+
|
|
20
|
+
def __str__(self) -> str: # pragma: no cover - trivial
|
|
21
|
+
return self.message
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@dataclass
|
|
25
|
+
class Violation:
|
|
26
|
+
kind: str # "pool_over" | "agent_over" | "killed"
|
|
27
|
+
pool: str
|
|
28
|
+
agent: str | None = None
|
|
29
|
+
detail: str = ""
|
|
30
|
+
|
|
31
|
+
def to_dict(self) -> dict:
|
|
32
|
+
return {"kind": self.kind, "pool": self.pool,
|
|
33
|
+
"agent": self.agent, "detail": self.detail}
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
@dataclass
|
|
37
|
+
class Store:
|
|
38
|
+
path: Path = field(default_factory=default_state_path)
|
|
39
|
+
data: dict = field(default_factory=lambda: {"pools": {}})
|
|
40
|
+
|
|
41
|
+
# ---- persistence -------------------------------------------------
|
|
42
|
+
def load(self) -> "Store":
|
|
43
|
+
if self.path.exists():
|
|
44
|
+
try:
|
|
45
|
+
self.data = json.loads(self.path.read_text(encoding="utf-8"))
|
|
46
|
+
except (json.JSONDecodeError, OSError) as exc:
|
|
47
|
+
raise BudgetError(f"cannot read state file {self.path}: {exc}")
|
|
48
|
+
self.data.setdefault("pools", {})
|
|
49
|
+
return self
|
|
50
|
+
|
|
51
|
+
def save(self) -> None:
|
|
52
|
+
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
53
|
+
tmp = self.path.with_suffix(".tmp")
|
|
54
|
+
tmp.write_text(json.dumps(self.data, indent=2, sort_keys=True),
|
|
55
|
+
encoding="utf-8")
|
|
56
|
+
os.replace(tmp, self.path)
|
|
57
|
+
|
|
58
|
+
# ---- helpers -----------------------------------------------------
|
|
59
|
+
def _pool(self, name: str) -> dict:
|
|
60
|
+
pool = self.data["pools"].get(name)
|
|
61
|
+
if pool is None:
|
|
62
|
+
raise BudgetError(f"pool '{name}' does not exist")
|
|
63
|
+
return pool
|
|
64
|
+
|
|
65
|
+
@staticmethod
|
|
66
|
+
def _check_positive(name: str, value: float | int | None) -> None:
|
|
67
|
+
if value is not None and value <= 0:
|
|
68
|
+
raise BudgetError(f"{name} must be positive, got {value}")
|
|
69
|
+
|
|
70
|
+
# ---- pool management ---------------------------------------------
|
|
71
|
+
def create_pool(self, name: str, limit_usd: float | None = None,
|
|
72
|
+
limit_tokens: int | None = None) -> None:
|
|
73
|
+
if not name or not name.strip():
|
|
74
|
+
raise BudgetError("pool name must not be empty")
|
|
75
|
+
if name in self.data["pools"]:
|
|
76
|
+
raise BudgetError(f"pool '{name}' already exists")
|
|
77
|
+
self._check_positive("limit_usd", limit_usd)
|
|
78
|
+
self._check_positive("limit_tokens", limit_tokens)
|
|
79
|
+
if limit_usd is None and limit_tokens is None:
|
|
80
|
+
raise BudgetError("at least one of --limit-usd / --limit-tokens "
|
|
81
|
+
"is required")
|
|
82
|
+
self.data["pools"][name] = {
|
|
83
|
+
"limit_usd": limit_usd,
|
|
84
|
+
"limit_tokens": limit_tokens,
|
|
85
|
+
"killed": False,
|
|
86
|
+
"spent_usd": 0.0,
|
|
87
|
+
"spent_tokens": 0,
|
|
88
|
+
"agents": {},
|
|
89
|
+
}
|
|
90
|
+
self.save()
|
|
91
|
+
|
|
92
|
+
def list_pools(self) -> list[dict]:
|
|
93
|
+
out = []
|
|
94
|
+
for name, pool in sorted(self.data["pools"].items()):
|
|
95
|
+
out.append({
|
|
96
|
+
"name": name,
|
|
97
|
+
"limit_usd": pool.get("limit_usd"),
|
|
98
|
+
"limit_tokens": pool.get("limit_tokens"),
|
|
99
|
+
"killed": pool.get("killed", False),
|
|
100
|
+
"spent_usd": pool.get("spent_usd", 0.0),
|
|
101
|
+
"spent_tokens": pool.get("spent_tokens", 0),
|
|
102
|
+
"agents": sorted(pool.get("agents", {})),
|
|
103
|
+
})
|
|
104
|
+
return out
|
|
105
|
+
|
|
106
|
+
# ---- quotas ------------------------------------------------------
|
|
107
|
+
def set_quota(self, pool_name: str, agent: str,
|
|
108
|
+
limit_usd: float | None = None,
|
|
109
|
+
limit_tokens: int | None = None) -> None:
|
|
110
|
+
pool = self._pool(pool_name)
|
|
111
|
+
if not agent or not agent.strip():
|
|
112
|
+
raise BudgetError("agent name must not be empty")
|
|
113
|
+
self._check_positive("limit_usd", limit_usd)
|
|
114
|
+
self._check_positive("limit_tokens", limit_tokens)
|
|
115
|
+
if limit_usd is None and limit_tokens is None:
|
|
116
|
+
raise BudgetError("at least one of --limit-usd / --limit-tokens "
|
|
117
|
+
"is required")
|
|
118
|
+
agents = pool.setdefault("agents", {})
|
|
119
|
+
entry = agents.get(agent, {"spent_usd": 0.0, "spent_tokens": 0})
|
|
120
|
+
entry["limit_usd"] = limit_usd
|
|
121
|
+
entry["limit_tokens"] = limit_tokens
|
|
122
|
+
agents[agent] = entry
|
|
123
|
+
self.save()
|
|
124
|
+
|
|
125
|
+
# ---- spend -------------------------------------------------------
|
|
126
|
+
def record_spend(self, pool_name: str, agent: str, tokens: int = 0,
|
|
127
|
+
usd: float = 0.0) -> None:
|
|
128
|
+
pool = self._pool(pool_name)
|
|
129
|
+
if pool.get("killed"):
|
|
130
|
+
raise BudgetError(f"pool '{pool_name}' is killed: spend rejected")
|
|
131
|
+
if tokens < 0 or usd < 0:
|
|
132
|
+
raise BudgetError("tokens and usd must be non-negative")
|
|
133
|
+
pool["spent_tokens"] = pool.get("spent_tokens", 0) + tokens
|
|
134
|
+
pool["spent_usd"] = pool.get("spent_usd", 0.0) + usd
|
|
135
|
+
agents = pool.setdefault("agents", {})
|
|
136
|
+
entry = agents.setdefault(agent,
|
|
137
|
+
{"spent_usd": 0.0, "spent_tokens": 0,
|
|
138
|
+
"limit_usd": None, "limit_tokens": None})
|
|
139
|
+
entry["spent_tokens"] = entry.get("spent_tokens", 0) + tokens
|
|
140
|
+
entry["spent_usd"] = entry.get("spent_usd", 0.0) + usd
|
|
141
|
+
if entry.get("killed"):
|
|
142
|
+
raise BudgetError(f"agent '{agent}' in pool '{pool_name}' is "
|
|
143
|
+
"killed: spend rejected")
|
|
144
|
+
self.save()
|
|
145
|
+
|
|
146
|
+
# ---- kill switch -------------------------------------------------
|
|
147
|
+
def kill(self, pool_name: str, agent: str | None = None) -> None:
|
|
148
|
+
pool = self._pool(pool_name)
|
|
149
|
+
if agent is None:
|
|
150
|
+
pool["killed"] = True
|
|
151
|
+
else:
|
|
152
|
+
entry = pool.get("agents", {}).get(agent)
|
|
153
|
+
if entry is None:
|
|
154
|
+
raise BudgetError(f"agent '{agent}' has no quota in pool "
|
|
155
|
+
f"'{pool_name}'")
|
|
156
|
+
entry["killed"] = True
|
|
157
|
+
self.save()
|
|
158
|
+
|
|
159
|
+
def release(self, pool_name: str, agent: str | None = None) -> None:
|
|
160
|
+
pool = self._pool(pool_name)
|
|
161
|
+
if agent is None:
|
|
162
|
+
pool["killed"] = False
|
|
163
|
+
else:
|
|
164
|
+
entry = pool.get("agents", {}).get(agent)
|
|
165
|
+
if entry is None:
|
|
166
|
+
raise BudgetError(f"agent '{agent}' has no quota in pool "
|
|
167
|
+
f"'{pool_name}'")
|
|
168
|
+
entry["killed"] = False
|
|
169
|
+
self.save()
|
|
170
|
+
|
|
171
|
+
# ---- check -------------------------------------------------------
|
|
172
|
+
def check(self) -> list[Violation]:
|
|
173
|
+
violations: list[Violation] = []
|
|
174
|
+
for name, pool in self.data["pools"].items():
|
|
175
|
+
if pool.get("killed"):
|
|
176
|
+
violations.append(Violation(
|
|
177
|
+
"killed", name,
|
|
178
|
+
detail=f"pool '{name}' is killed (circuit breaker open)"))
|
|
179
|
+
continue
|
|
180
|
+
limit_usd = pool.get("limit_usd")
|
|
181
|
+
spent_usd = pool.get("spent_usd", 0.0)
|
|
182
|
+
limit_tokens = pool.get("limit_tokens")
|
|
183
|
+
spent_tokens = pool.get("spent_tokens", 0)
|
|
184
|
+
if limit_usd is not None and spent_usd > limit_usd:
|
|
185
|
+
violations.append(Violation(
|
|
186
|
+
"pool_over", name,
|
|
187
|
+
detail=f"pool '{name}': spent ${spent_usd:.4f} > "
|
|
188
|
+
f"limit ${limit_usd:.4f}"))
|
|
189
|
+
if limit_tokens is not None and spent_tokens > limit_tokens:
|
|
190
|
+
violations.append(Violation(
|
|
191
|
+
"pool_over", name,
|
|
192
|
+
detail=f"pool '{name}': spent {spent_tokens} tokens > "
|
|
193
|
+
f"limit {limit_tokens}"))
|
|
194
|
+
for agent, entry in pool.get("agents", {}).items():
|
|
195
|
+
if entry.get("killed"):
|
|
196
|
+
violations.append(Violation(
|
|
197
|
+
"killed", name, agent,
|
|
198
|
+
detail=f"agent '{agent}' in pool '{name}' is killed"))
|
|
199
|
+
continue
|
|
200
|
+
a_usd = entry.get("limit_usd")
|
|
201
|
+
a_spent_usd = entry.get("spent_usd", 0.0)
|
|
202
|
+
a_tokens = entry.get("limit_tokens")
|
|
203
|
+
a_spent_tokens = entry.get("spent_tokens", 0)
|
|
204
|
+
if a_usd is not None and a_spent_usd > a_usd:
|
|
205
|
+
violations.append(Violation(
|
|
206
|
+
"agent_over", name, agent,
|
|
207
|
+
detail=f"agent '{agent}' in pool '{name}': spent "
|
|
208
|
+
f"${a_spent_usd:.4f} > limit ${a_usd:.4f}"))
|
|
209
|
+
if a_tokens is not None and a_spent_tokens > a_tokens:
|
|
210
|
+
violations.append(Violation(
|
|
211
|
+
"agent_over", name, agent,
|
|
212
|
+
detail=f"agent '{agent}' in pool '{name}': spent "
|
|
213
|
+
f"{a_spent_tokens} tokens > limit {a_tokens}"))
|
|
214
|
+
return violations
|
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: subagent-budget-team
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Team/CI budget control for subagent token spend: shared budget pools, per-agent quotas, overspend checks, and a kill switch.
|
|
5
|
+
Author-email: hao li <haoli5933@gmail.com>
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/hahahahahahahahah6/subagent-budget-team
|
|
8
|
+
Project-URL: Issues, https://github.com/hahahahahahahahah6/subagent-budget-team/issues
|
|
9
|
+
Keywords: subagent,budget,token,cost-control,ci,agents
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: Environment :: Console
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Topic :: Software Development :: Quality Assurance
|
|
16
|
+
Requires-Python: >=3.9
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
License-File: LICENSE
|
|
19
|
+
Dynamic: license-file
|
|
20
|
+
|
|
21
|
+
# subagent-budget-team
|
|
22
|
+
|
|
23
|
+
Team/CI budget control for subagent token spend. When several agents run
|
|
24
|
+
concurrently — a CI matrix, a fleet of coding agents, a team sharing one
|
|
25
|
+
token pool — costs can spiral before anyone notices. `subagent-budget-team`
|
|
26
|
+
gives you a **pre-spend control plane**: shared budget pools, per-agent
|
|
27
|
+
quotas, overspend detection, and a kill switch that freezes spend *before*
|
|
28
|
+
the next invoice.
|
|
29
|
+
|
|
30
|
+
Measured context: subagents can burn ~48% of an agent's total token budget
|
|
31
|
+
(see the `subagent-tax` write-up). This tool is what you deploy once you
|
|
32
|
+
know that number and want to do something about it at the team level.
|
|
33
|
+
|
|
34
|
+
## Install
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install subagent-budget-team
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
Requires Python 3.9+. Zero dependencies (stdlib only).
|
|
41
|
+
|
|
42
|
+
## Usage
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
# 1. Create a shared budget pool for a CI job or team
|
|
46
|
+
subagent-budget-team pool create --name ci-nightly --limit-usd 25
|
|
47
|
+
|
|
48
|
+
# 2. Give each agent its own quota inside the pool
|
|
49
|
+
subagent-budget-team quota set --pool ci-nightly --agent coder-1 --limit-usd 5
|
|
50
|
+
subagent-budget-team quota set --pool ci-nightly --agent reviewer-1 --limit-usd 3
|
|
51
|
+
|
|
52
|
+
# 3. Record spend (call this from your agent wrapper / CI step)
|
|
53
|
+
subagent-budget-team spend --pool ci-nightly --agent coder-1 --tokens 12000 --usd 1.20
|
|
54
|
+
|
|
55
|
+
# 4. Check everything; exit code 1 means over budget — gate your CI on it
|
|
56
|
+
subagent-budget-team check
|
|
57
|
+
# or machine-readable:
|
|
58
|
+
subagent-budget-team check --json
|
|
59
|
+
|
|
60
|
+
# 5. Kill switch: freeze a pool (or a single agent) immediately
|
|
61
|
+
subagent-budget-team kill --pool ci-nightly
|
|
62
|
+
subagent-budget-team spend --pool ci-nightly --agent coder-1 --usd 1 # rejected
|
|
63
|
+
subagent-budget-team release --pool ci-nightly # unfreeze
|
|
64
|
+
|
|
65
|
+
# List pools
|
|
66
|
+
subagent-budget-team pool list
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
State lives in `~/.subagent-budget-team/budgets.json` (override with
|
|
70
|
+
`--state` or the `SBT_HOME` env var, handy for per-job isolation in CI).
|
|
71
|
+
|
|
72
|
+
### CI example (GitHub Actions)
|
|
73
|
+
|
|
74
|
+
```yaml
|
|
75
|
+
- run: pip install subagent-budget-team
|
|
76
|
+
- run: subagent-budget-team pool create --name pr-${{ github.event.number }} --limit-usd 10
|
|
77
|
+
- run: ./run-agents.sh # agents call `spend` as they work
|
|
78
|
+
- run: subagent-budget-team check # fails the job on overspend
|
|
79
|
+
- if: failure()
|
|
80
|
+
run: subagent-budget-team kill --pool pr-${{ github.event.number }}
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
## Differentiation
|
|
84
|
+
|
|
85
|
+
- **vs `subagent-tax`**: that tool is a *personal audit* lens — it measures
|
|
86
|
+
how much of *your* token budget subagents burn after the fact.
|
|
87
|
+
`subagent-budget-team` is the *team/CI control* lens: a shared budget
|
|
88
|
+
pool with per-agent quotas and a circuit breaker, designed for many
|
|
89
|
+
agents spending from one pot.
|
|
90
|
+
- **vs cloud-vendor budget alerts**: those are *post-hoc alarms* — you get
|
|
91
|
+
paged after the money is spent. This is a *pre-spend circuit breaker*:
|
|
92
|
+
`kill` rejects new spend immediately and `check` can gate a pipeline
|
|
93
|
+
before the next batch of work runs.
|
|
94
|
+
|
|
95
|
+
## License
|
|
96
|
+
|
|
97
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
src/subagent_budget_team/__init__.py
|
|
5
|
+
src/subagent_budget_team/cli.py
|
|
6
|
+
src/subagent_budget_team/core.py
|
|
7
|
+
src/subagent_budget_team.egg-info/PKG-INFO
|
|
8
|
+
src/subagent_budget_team.egg-info/SOURCES.txt
|
|
9
|
+
src/subagent_budget_team.egg-info/dependency_links.txt
|
|
10
|
+
src/subagent_budget_team.egg-info/entry_points.txt
|
|
11
|
+
src/subagent_budget_team.egg-info/top_level.txt
|
|
12
|
+
tests/test_smoke.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
subagent_budget_team
|
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
"""Smoke tests for subagent-budget-team. stdlib only (unittest)."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import os
|
|
5
|
+
import subprocess
|
|
6
|
+
import sys
|
|
7
|
+
import tempfile
|
|
8
|
+
import unittest
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
|
|
11
|
+
CLI = [sys.executable, "-m", "subagent_budget_team.cli"]
|
|
12
|
+
SRC = Path(__file__).resolve().parent.parent / "src"
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _env(tmp: Path) -> dict:
|
|
16
|
+
env = dict(os.environ)
|
|
17
|
+
env["SBT_HOME"] = str(tmp)
|
|
18
|
+
env["PYTHONPATH"] = str(SRC) + os.pathsep + env.get("PYTHONPATH", "")
|
|
19
|
+
return env
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class SmokeTest(unittest.TestCase):
|
|
23
|
+
def setUp(self):
|
|
24
|
+
self.tmp = Path(tempfile.mkdtemp(prefix="sbt-test-"))
|
|
25
|
+
self.env = _env(self.tmp)
|
|
26
|
+
|
|
27
|
+
def run_cli(self, *args: str):
|
|
28
|
+
return subprocess.run(CLI + list(args), capture_output=True,
|
|
29
|
+
text=True, env=self.env)
|
|
30
|
+
|
|
31
|
+
def test_01_pool_create_and_list(self):
|
|
32
|
+
r = self.run_cli("pool", "create", "--name", "ci-nightly",
|
|
33
|
+
"--limit-usd", "25")
|
|
34
|
+
self.assertEqual(r.returncode, 0, r.stderr)
|
|
35
|
+
r = self.run_cli("pool", "list")
|
|
36
|
+
self.assertEqual(r.returncode, 0, r.stderr)
|
|
37
|
+
self.assertIn("ci-nightly", r.stdout)
|
|
38
|
+
|
|
39
|
+
def test_02_quota_set(self):
|
|
40
|
+
self.run_cli("pool", "create", "--name", "team", "--limit-usd", "100")
|
|
41
|
+
r = self.run_cli("quota", "set", "--pool", "team",
|
|
42
|
+
"--agent", "coder-1", "--limit-usd", "10")
|
|
43
|
+
self.assertEqual(r.returncode, 0, r.stderr)
|
|
44
|
+
r = self.run_cli("pool", "list", "--json")
|
|
45
|
+
pools = json.loads(r.stdout)
|
|
46
|
+
self.assertIn("coder-1", pools[0]["agents"])
|
|
47
|
+
|
|
48
|
+
def test_03_spend_within_budget_check_ok(self):
|
|
49
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-usd", "50")
|
|
50
|
+
r = self.run_cli("spend", "--pool", "p", "--agent", "a",
|
|
51
|
+
"--tokens", "1000", "--usd", "2.50")
|
|
52
|
+
self.assertEqual(r.returncode, 0, r.stderr)
|
|
53
|
+
r = self.run_cli("check")
|
|
54
|
+
self.assertEqual(r.returncode, 0, r.stdout + r.stderr)
|
|
55
|
+
self.assertIn("OK", r.stdout)
|
|
56
|
+
|
|
57
|
+
def test_04_pool_overspend_check_fails(self):
|
|
58
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-usd", "5")
|
|
59
|
+
self.run_cli("spend", "--pool", "p", "--agent", "a", "--usd", "6.01")
|
|
60
|
+
r = self.run_cli("check")
|
|
61
|
+
self.assertEqual(r.returncode, 1)
|
|
62
|
+
self.assertIn("OVER BUDGET", r.stdout)
|
|
63
|
+
|
|
64
|
+
def test_05_agent_quota_overspend_check_fails(self):
|
|
65
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-usd", "100")
|
|
66
|
+
self.run_cli("quota", "set", "--pool", "p", "--agent", "a",
|
|
67
|
+
"--limit-usd", "1")
|
|
68
|
+
self.run_cli("spend", "--pool", "p", "--agent", "a", "--usd", "1.50")
|
|
69
|
+
r = self.run_cli("check")
|
|
70
|
+
self.assertEqual(r.returncode, 1)
|
|
71
|
+
self.assertIn("agent_over", r.stdout + self.run_cli(
|
|
72
|
+
"check", "--json").stdout)
|
|
73
|
+
|
|
74
|
+
def test_06_kill_switch_breaks_check(self):
|
|
75
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-usd", "50")
|
|
76
|
+
self.run_cli("spend", "--pool", "p", "--agent", "a", "--usd", "1")
|
|
77
|
+
r = self.run_cli("check")
|
|
78
|
+
self.assertEqual(r.returncode, 0)
|
|
79
|
+
r = self.run_cli("kill", "--pool", "p")
|
|
80
|
+
self.assertEqual(r.returncode, 0, r.stderr)
|
|
81
|
+
r = self.run_cli("check")
|
|
82
|
+
self.assertEqual(r.returncode, 1)
|
|
83
|
+
self.assertIn("killed", r.stdout.lower())
|
|
84
|
+
|
|
85
|
+
def test_07_kill_rejects_new_spend(self):
|
|
86
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-usd", "50")
|
|
87
|
+
self.run_cli("kill", "--pool", "p")
|
|
88
|
+
r = self.run_cli("spend", "--pool", "p", "--agent", "a", "--usd", "1")
|
|
89
|
+
self.assertNotEqual(r.returncode, 0)
|
|
90
|
+
|
|
91
|
+
def test_08_release_restores_ok(self):
|
|
92
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-usd", "50")
|
|
93
|
+
self.run_cli("kill", "--pool", "p")
|
|
94
|
+
self.run_cli("release", "--pool", "p")
|
|
95
|
+
r = self.run_cli("check")
|
|
96
|
+
self.assertEqual(r.returncode, 0)
|
|
97
|
+
|
|
98
|
+
def test_09_duplicate_pool_rejected(self):
|
|
99
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-usd", "1")
|
|
100
|
+
r = self.run_cli("pool", "create", "--name", "p", "--limit-usd", "1")
|
|
101
|
+
self.assertNotEqual(r.returncode, 0)
|
|
102
|
+
self.assertIn("already exists", r.stderr)
|
|
103
|
+
|
|
104
|
+
def test_10_negative_limit_rejected(self):
|
|
105
|
+
r = self.run_cli("pool", "create", "--name", "p",
|
|
106
|
+
"--limit-usd", "-5")
|
|
107
|
+
self.assertNotEqual(r.returncode, 0)
|
|
108
|
+
self.assertIn("positive", r.stderr)
|
|
109
|
+
|
|
110
|
+
def test_11_check_json_output(self):
|
|
111
|
+
self.run_cli("pool", "create", "--name", "p", "--limit-tokens", "100")
|
|
112
|
+
self.run_cli("spend", "--pool", "p", "--agent", "a", "--tokens", "150")
|
|
113
|
+
r = self.run_cli("check", "--json")
|
|
114
|
+
self.assertEqual(r.returncode, 1)
|
|
115
|
+
data = json.loads(r.stdout)
|
|
116
|
+
self.assertEqual(data[0]["kind"], "pool_over")
|
|
117
|
+
self.assertEqual(data[0]["pool"], "p")
|
|
118
|
+
|
|
119
|
+
def test_12_version(self):
|
|
120
|
+
r = self.run_cli("--version")
|
|
121
|
+
self.assertEqual(r.returncode, 0)
|
|
122
|
+
self.assertIn("0.1.0", r.stdout)
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
if __name__ == "__main__":
|
|
126
|
+
unittest.main()
|