agent-tally 0.5.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.
- agent_tally-0.5.0/LICENSE +21 -0
- agent_tally-0.5.0/PKG-INFO +153 -0
- agent_tally-0.5.0/README.md +129 -0
- agent_tally-0.5.0/agent_tally/__init__.py +32 -0
- agent_tally-0.5.0/agent_tally/budget.py +193 -0
- agent_tally-0.5.0/agent_tally/cli.py +569 -0
- agent_tally-0.5.0/agent_tally/config.py +136 -0
- agent_tally-0.5.0/agent_tally/dashboard.py +317 -0
- agent_tally-0.5.0/agent_tally/detector.py +235 -0
- agent_tally-0.5.0/agent_tally/display.py +180 -0
- agent_tally-0.5.0/agent_tally/notifier.py +259 -0
- agent_tally-0.5.0/agent_tally/pricing.py +176 -0
- agent_tally-0.5.0/agent_tally/py.typed +0 -0
- agent_tally-0.5.0/agent_tally/storage.py +201 -0
- agent_tally-0.5.0/agent_tally/ticker.py +265 -0
- agent_tally-0.5.0/agent_tally/wrapper.py +237 -0
- agent_tally-0.5.0/agent_tally.egg-info/PKG-INFO +153 -0
- agent_tally-0.5.0/agent_tally.egg-info/SOURCES.txt +33 -0
- agent_tally-0.5.0/agent_tally.egg-info/dependency_links.txt +1 -0
- agent_tally-0.5.0/agent_tally.egg-info/entry_points.txt +2 -0
- agent_tally-0.5.0/agent_tally.egg-info/requires.txt +7 -0
- agent_tally-0.5.0/agent_tally.egg-info/top_level.txt +1 -0
- agent_tally-0.5.0/pyproject.toml +43 -0
- agent_tally-0.5.0/setup.cfg +4 -0
- agent_tally-0.5.0/setup.py +28 -0
- agent_tally-0.5.0/tests/test_budget.py +198 -0
- agent_tally-0.5.0/tests/test_config.py +175 -0
- agent_tally-0.5.0/tests/test_detector.py +215 -0
- agent_tally-0.5.0/tests/test_edge_cases.py +371 -0
- agent_tally-0.5.0/tests/test_extended.py +243 -0
- agent_tally-0.5.0/tests/test_integration.py +343 -0
- agent_tally-0.5.0/tests/test_notifier.py +250 -0
- agent_tally-0.5.0/tests/test_pricing.py +54 -0
- agent_tally-0.5.0/tests/test_storage.py +104 -0
- agent_tally-0.5.0/tests/test_ticker.py +192 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Hiren Thakore
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,153 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: agent-tally
|
|
3
|
+
Version: 0.5.0
|
|
4
|
+
Summary: Real-time cost tracking for AI coding agents with budget limits, kill switch, and TUI dashboard.
|
|
5
|
+
Author: Hiren Thakore
|
|
6
|
+
Author-email: Hiren Thakore <thakoreh@users.noreply.github.com>
|
|
7
|
+
License: MIT
|
|
8
|
+
Project-URL: Homepage, https://github.com/thakoreh/agent-tally
|
|
9
|
+
Project-URL: Repository, https://github.com/thakoreh/agent-tally
|
|
10
|
+
Project-URL: Issues, https://github.com/thakoreh/agent-tally/issues
|
|
11
|
+
Keywords: ai,agents,cost-tracking,cli,claude-code,codex,openclaw
|
|
12
|
+
Requires-Python: >=3.9
|
|
13
|
+
Description-Content-Type: text/markdown
|
|
14
|
+
License-File: LICENSE
|
|
15
|
+
Requires-Dist: click>=8.0
|
|
16
|
+
Requires-Dist: rich>=13.0
|
|
17
|
+
Requires-Dist: pyyaml>=6.0
|
|
18
|
+
Provides-Extra: dev
|
|
19
|
+
Requires-Dist: pytest>=7.0; extra == "dev"
|
|
20
|
+
Requires-Dist: pytest-cov>=4.0; extra == "dev"
|
|
21
|
+
Dynamic: author
|
|
22
|
+
Dynamic: license-file
|
|
23
|
+
Dynamic: requires-python
|
|
24
|
+
|
|
25
|
+
# agent-tally
|
|
26
|
+
|
|
27
|
+
[](https://pypi.org/project/agent-tally/)
|
|
28
|
+
[](https://pypi.org/project/agent-tally/)
|
|
29
|
+
[](LICENSE)
|
|
30
|
+
[](https://github.com/thakoreh/agent-tally)
|
|
31
|
+
|
|
32
|
+
**Real-time cost tracker for AI coding agents.** Wrap any CLI agent, see a live cost ticker while it runs, and enforce budgets with an automatic kill switch.
|
|
33
|
+
|
|
34
|
+
## Install
|
|
35
|
+
|
|
36
|
+
```bash
|
|
37
|
+
pip install agent-tally
|
|
38
|
+
```
|
|
39
|
+
|
|
40
|
+
## Quick Start
|
|
41
|
+
|
|
42
|
+
```bash
|
|
43
|
+
# Track a Claude Code session
|
|
44
|
+
agent-tally run claude "refactor auth module"
|
|
45
|
+
|
|
46
|
+
# Set budgets
|
|
47
|
+
agent-tally budget set --daily 5.00 --session 1.00
|
|
48
|
+
|
|
49
|
+
# Watch costs live
|
|
50
|
+
agent-tally dashboard
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
Output while your agent runs:
|
|
54
|
+
|
|
55
|
+
```
|
|
56
|
+
$ agent-tally run claude "fix the login bug"
|
|
57
|
+
⠋ Running claude... $0.42 │ 2.1k tokens │ 14s elapsed
|
|
58
|
+
```
|
|
59
|
+
|
|
60
|
+
## Features
|
|
61
|
+
|
|
62
|
+
- **Live cost ticker** — dollar amount updates in real-time as the agent runs, like a taxi meter
|
|
63
|
+
- **Budget enforcement** — set daily and per-session limits; auto-kills the process at 100%, warns at 80% and 95%
|
|
64
|
+
- **Shell-level interception** — wraps subprocess calls. No API hooks, no SDK changes, no config files
|
|
65
|
+
- **Cross-agent** — works with Claude Code, Codex CLI, Gemini CLI, OpenClaw, Cursor, and any CLI tool
|
|
66
|
+
- **TUI dashboard** — `agent-tally dashboard` for a live overview of all tracked sessions
|
|
67
|
+
- **Webhook alerts** — get Discord or Slack notifications when you hit budget thresholds
|
|
68
|
+
- **Export** — dump cost history as JSON or CSV for your own analytics
|
|
69
|
+
|
|
70
|
+
## CLI Reference
|
|
71
|
+
|
|
72
|
+
| Command | Description |
|
|
73
|
+
|---|---|
|
|
74
|
+
| `agent-tally run <command> [args...]` | Wrap and track an agent |
|
|
75
|
+
| `agent-tally budget set --daily 5.00 --session 1.00` | Set budget limits |
|
|
76
|
+
| `agent-tally budget show` | Show current budget and usage |
|
|
77
|
+
| `agent-tally dashboard` | Live TUI cost dashboard |
|
|
78
|
+
| `agent-tally summary [--by-agent\|--by-model] [--since 7d]` | Usage summary |
|
|
79
|
+
| `agent-tally export [--format json\|csv]` | Export cost data |
|
|
80
|
+
| `agent-tally agents` | List supported agents |
|
|
81
|
+
|
|
82
|
+
## How It Works
|
|
83
|
+
|
|
84
|
+
```
|
|
85
|
+
┌─────────────────────────────────────────────────────┐
|
|
86
|
+
│ agent-tally │
|
|
87
|
+
│ │
|
|
88
|
+
│ ┌──────────┐ ┌──────────────┐ ┌─────────────┐ │
|
|
89
|
+
│ │ CLI run │──▶│ Wrapper │──▶│ Subprocess │ │
|
|
90
|
+
│ │ command │ │ intercepts │ │ (agent CLI) │ │
|
|
91
|
+
│ └──────────┘ │ stdout/err │ └──────┬───────┘ │
|
|
92
|
+
│ └──────┬───────┘ │ │
|
|
93
|
+
│ │ │ │
|
|
94
|
+
│ ▼ │ │
|
|
95
|
+
│ ┌──────────────┐ │ │
|
|
96
|
+
│ │ Token parser │◀─────────┘ │
|
|
97
|
+
│ │ counts tokens│ (reads agent │
|
|
98
|
+
│ │ from output │ output stream) │
|
|
99
|
+
│ └──────┬───────┘ │
|
|
100
|
+
│ │ │
|
|
101
|
+
│ ┌──────────┴──────────┐ │
|
|
102
|
+
│ ▼ ▼ │
|
|
103
|
+
│ ┌──────────────┐ ┌────────────────┐ │
|
|
104
|
+
│ │ Live ticker │ │ Budget checker │ │
|
|
105
|
+
│ │ $X.XX ████ │ │ warn @ 80/95% │ │
|
|
106
|
+
│ └──────────────┘ │ kill @ 100% │ │
|
|
107
|
+
│ └───────┬────────┘ │
|
|
108
|
+
│ ▼ │
|
|
109
|
+
│ ┌────────────────────┐ │
|
|
110
|
+
│ │ Webhook alerts │ │
|
|
111
|
+
│ │ Discord / Slack │ │
|
|
112
|
+
│ └────────────────────┘ │
|
|
113
|
+
└─────────────────────────────────────────────────────┘
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
agent-tally wraps your agent command as a subprocess. It parses the agent's output stream to count tokens in real-time, calculates cost using model-specific pricing, and displays a live ticker. If spending hits your budget limit, it kills the process automatically.
|
|
117
|
+
|
|
118
|
+
## Supported Agents
|
|
119
|
+
|
|
120
|
+
| Agent | Status |
|
|
121
|
+
|---|---|
|
|
122
|
+
| Claude Code | ✅ Supported |
|
|
123
|
+
| Codex CLI | ✅ Supported |
|
|
124
|
+
| Gemini CLI | ✅ Supported |
|
|
125
|
+
| OpenClaw | ✅ Supported |
|
|
126
|
+
| Cursor | ✅ Supported |
|
|
127
|
+
| Any CLI agent | ✅ Generic support |
|
|
128
|
+
|
|
129
|
+
## Why agent-tally over tokscale?
|
|
130
|
+
|
|
131
|
+
| | agent-tally | tokscale |
|
|
132
|
+
|---|---|---|
|
|
133
|
+
| **Approach** | Active — wraps commands in real-time | Passive — reads session logs after the fact |
|
|
134
|
+
| **Live cost display** | ✅ Ticker updates as the agent runs | ❌ Post-hoc only |
|
|
135
|
+
| **Budget enforcement** | ✅ Kill switch at budget limit | ❌ No enforcement |
|
|
136
|
+
| **Kill switch** | ✅ Auto-terminates runaway agents | ❌ Not possible (after the fact) |
|
|
137
|
+
| **Setup** | Wrap your command, done | Parse exported logs |
|
|
138
|
+
|
|
139
|
+
tokscale is great for retroactive analysis. agent-tally is for when you want to **control spend while it's happening**.
|
|
140
|
+
|
|
141
|
+
## Contributing
|
|
142
|
+
|
|
143
|
+
1. Fork the repo
|
|
144
|
+
2. Create a branch: `git checkout -b feat/my-feature`
|
|
145
|
+
3. Make changes and add tests
|
|
146
|
+
4. Run `pytest`
|
|
147
|
+
5. Open a PR
|
|
148
|
+
|
|
149
|
+
Bug reports and feature requests welcome in [Issues](https://github.com/thakoreh/agent-tally/issues).
|
|
150
|
+
|
|
151
|
+
## License
|
|
152
|
+
|
|
153
|
+
[MIT](LICENSE)
|
|
@@ -0,0 +1,129 @@
|
|
|
1
|
+
# agent-tally
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/agent-tally/)
|
|
4
|
+
[](https://pypi.org/project/agent-tally/)
|
|
5
|
+
[](LICENSE)
|
|
6
|
+
[](https://github.com/thakoreh/agent-tally)
|
|
7
|
+
|
|
8
|
+
**Real-time cost tracker for AI coding agents.** Wrap any CLI agent, see a live cost ticker while it runs, and enforce budgets with an automatic kill switch.
|
|
9
|
+
|
|
10
|
+
## Install
|
|
11
|
+
|
|
12
|
+
```bash
|
|
13
|
+
pip install agent-tally
|
|
14
|
+
```
|
|
15
|
+
|
|
16
|
+
## Quick Start
|
|
17
|
+
|
|
18
|
+
```bash
|
|
19
|
+
# Track a Claude Code session
|
|
20
|
+
agent-tally run claude "refactor auth module"
|
|
21
|
+
|
|
22
|
+
# Set budgets
|
|
23
|
+
agent-tally budget set --daily 5.00 --session 1.00
|
|
24
|
+
|
|
25
|
+
# Watch costs live
|
|
26
|
+
agent-tally dashboard
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
Output while your agent runs:
|
|
30
|
+
|
|
31
|
+
```
|
|
32
|
+
$ agent-tally run claude "fix the login bug"
|
|
33
|
+
⠋ Running claude... $0.42 │ 2.1k tokens │ 14s elapsed
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
## Features
|
|
37
|
+
|
|
38
|
+
- **Live cost ticker** — dollar amount updates in real-time as the agent runs, like a taxi meter
|
|
39
|
+
- **Budget enforcement** — set daily and per-session limits; auto-kills the process at 100%, warns at 80% and 95%
|
|
40
|
+
- **Shell-level interception** — wraps subprocess calls. No API hooks, no SDK changes, no config files
|
|
41
|
+
- **Cross-agent** — works with Claude Code, Codex CLI, Gemini CLI, OpenClaw, Cursor, and any CLI tool
|
|
42
|
+
- **TUI dashboard** — `agent-tally dashboard` for a live overview of all tracked sessions
|
|
43
|
+
- **Webhook alerts** — get Discord or Slack notifications when you hit budget thresholds
|
|
44
|
+
- **Export** — dump cost history as JSON or CSV for your own analytics
|
|
45
|
+
|
|
46
|
+
## CLI Reference
|
|
47
|
+
|
|
48
|
+
| Command | Description |
|
|
49
|
+
|---|---|
|
|
50
|
+
| `agent-tally run <command> [args...]` | Wrap and track an agent |
|
|
51
|
+
| `agent-tally budget set --daily 5.00 --session 1.00` | Set budget limits |
|
|
52
|
+
| `agent-tally budget show` | Show current budget and usage |
|
|
53
|
+
| `agent-tally dashboard` | Live TUI cost dashboard |
|
|
54
|
+
| `agent-tally summary [--by-agent\|--by-model] [--since 7d]` | Usage summary |
|
|
55
|
+
| `agent-tally export [--format json\|csv]` | Export cost data |
|
|
56
|
+
| `agent-tally agents` | List supported agents |
|
|
57
|
+
|
|
58
|
+
## How It Works
|
|
59
|
+
|
|
60
|
+
```
|
|
61
|
+
┌─────────────────────────────────────────────────────┐
|
|
62
|
+
│ agent-tally │
|
|
63
|
+
│ │
|
|
64
|
+
│ ┌──────────┐ ┌──────────────┐ ┌─────────────┐ │
|
|
65
|
+
│ │ CLI run │──▶│ Wrapper │──▶│ Subprocess │ │
|
|
66
|
+
│ │ command │ │ intercepts │ │ (agent CLI) │ │
|
|
67
|
+
│ └──────────┘ │ stdout/err │ └──────┬───────┘ │
|
|
68
|
+
│ └──────┬───────┘ │ │
|
|
69
|
+
│ │ │ │
|
|
70
|
+
│ ▼ │ │
|
|
71
|
+
│ ┌──────────────┐ │ │
|
|
72
|
+
│ │ Token parser │◀─────────┘ │
|
|
73
|
+
│ │ counts tokens│ (reads agent │
|
|
74
|
+
│ │ from output │ output stream) │
|
|
75
|
+
│ └──────┬───────┘ │
|
|
76
|
+
│ │ │
|
|
77
|
+
│ ┌──────────┴──────────┐ │
|
|
78
|
+
│ ▼ ▼ │
|
|
79
|
+
│ ┌──────────────┐ ┌────────────────┐ │
|
|
80
|
+
│ │ Live ticker │ │ Budget checker │ │
|
|
81
|
+
│ │ $X.XX ████ │ │ warn @ 80/95% │ │
|
|
82
|
+
│ └──────────────┘ │ kill @ 100% │ │
|
|
83
|
+
│ └───────┬────────┘ │
|
|
84
|
+
│ ▼ │
|
|
85
|
+
│ ┌────────────────────┐ │
|
|
86
|
+
│ │ Webhook alerts │ │
|
|
87
|
+
│ │ Discord / Slack │ │
|
|
88
|
+
│ └────────────────────┘ │
|
|
89
|
+
└─────────────────────────────────────────────────────┘
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
agent-tally wraps your agent command as a subprocess. It parses the agent's output stream to count tokens in real-time, calculates cost using model-specific pricing, and displays a live ticker. If spending hits your budget limit, it kills the process automatically.
|
|
93
|
+
|
|
94
|
+
## Supported Agents
|
|
95
|
+
|
|
96
|
+
| Agent | Status |
|
|
97
|
+
|---|---|
|
|
98
|
+
| Claude Code | ✅ Supported |
|
|
99
|
+
| Codex CLI | ✅ Supported |
|
|
100
|
+
| Gemini CLI | ✅ Supported |
|
|
101
|
+
| OpenClaw | ✅ Supported |
|
|
102
|
+
| Cursor | ✅ Supported |
|
|
103
|
+
| Any CLI agent | ✅ Generic support |
|
|
104
|
+
|
|
105
|
+
## Why agent-tally over tokscale?
|
|
106
|
+
|
|
107
|
+
| | agent-tally | tokscale |
|
|
108
|
+
|---|---|---|
|
|
109
|
+
| **Approach** | Active — wraps commands in real-time | Passive — reads session logs after the fact |
|
|
110
|
+
| **Live cost display** | ✅ Ticker updates as the agent runs | ❌ Post-hoc only |
|
|
111
|
+
| **Budget enforcement** | ✅ Kill switch at budget limit | ❌ No enforcement |
|
|
112
|
+
| **Kill switch** | ✅ Auto-terminates runaway agents | ❌ Not possible (after the fact) |
|
|
113
|
+
| **Setup** | Wrap your command, done | Parse exported logs |
|
|
114
|
+
|
|
115
|
+
tokscale is great for retroactive analysis. agent-tally is for when you want to **control spend while it's happening**.
|
|
116
|
+
|
|
117
|
+
## Contributing
|
|
118
|
+
|
|
119
|
+
1. Fork the repo
|
|
120
|
+
2. Create a branch: `git checkout -b feat/my-feature`
|
|
121
|
+
3. Make changes and add tests
|
|
122
|
+
4. Run `pytest`
|
|
123
|
+
5. Open a PR
|
|
124
|
+
|
|
125
|
+
Bug reports and feature requests welcome in [Issues](https://github.com/thakoreh/agent-tally/issues).
|
|
126
|
+
|
|
127
|
+
## License
|
|
128
|
+
|
|
129
|
+
[MIT](LICENSE)
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
"""agent-tally — track costs across every AI coding agent in real-time."""
|
|
2
|
+
|
|
3
|
+
__version__ = "0.5.0"
|
|
4
|
+
|
|
5
|
+
from .budget import BudgetManager, BudgetConfig, BudgetStatus
|
|
6
|
+
from .config import AgentTallyConfig, load_config, save_config
|
|
7
|
+
from .ticker import CostTicker, IncrementalCostTracker
|
|
8
|
+
from .notifier import Notifier, Alert
|
|
9
|
+
from .dashboard import Dashboard
|
|
10
|
+
from .storage import Storage, Session
|
|
11
|
+
from .pricing import PricingConfig
|
|
12
|
+
from .detector import detect_agent, parse_tokens, parse_model
|
|
13
|
+
|
|
14
|
+
__all__ = [
|
|
15
|
+
"BudgetManager",
|
|
16
|
+
"BudgetConfig",
|
|
17
|
+
"BudgetStatus",
|
|
18
|
+
"AgentTallyConfig",
|
|
19
|
+
"load_config",
|
|
20
|
+
"save_config",
|
|
21
|
+
"CostTicker",
|
|
22
|
+
"IncrementalCostTracker",
|
|
23
|
+
"Notifier",
|
|
24
|
+
"Alert",
|
|
25
|
+
"Dashboard",
|
|
26
|
+
"Storage",
|
|
27
|
+
"Session",
|
|
28
|
+
"PricingConfig",
|
|
29
|
+
"detect_agent",
|
|
30
|
+
"parse_tokens",
|
|
31
|
+
"parse_model",
|
|
32
|
+
]
|
|
@@ -0,0 +1,193 @@
|
|
|
1
|
+
"""Budget tracking and kill switch for agent-tally."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
import json
|
|
5
|
+
import os
|
|
6
|
+
import signal
|
|
7
|
+
from dataclasses import dataclass, field, asdict
|
|
8
|
+
from datetime import datetime, date
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
from typing import Optional, Callable
|
|
11
|
+
|
|
12
|
+
import yaml
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
DEFAULT_CONFIG_DIR = Path.home() / ".agent-tally"
|
|
16
|
+
DEFAULT_BUDGET_FILE = DEFAULT_CONFIG_DIR / "budget.yaml"
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
@dataclass
|
|
20
|
+
class BudgetConfig:
|
|
21
|
+
"""Budget limits configuration."""
|
|
22
|
+
daily_limit: Optional[float] = None # USD
|
|
23
|
+
session_limit: Optional[float] = None # USD
|
|
24
|
+
warn_at_80: bool = True
|
|
25
|
+
warn_at_95: bool = True
|
|
26
|
+
kill_at_100: bool = True
|
|
27
|
+
webhook_url: Optional[str] = None # For notifications
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@dataclass
|
|
31
|
+
class BudgetStatus:
|
|
32
|
+
"""Current budget status."""
|
|
33
|
+
session_cost: float = 0.0
|
|
34
|
+
daily_cost: float = 0.0
|
|
35
|
+
daily_limit: Optional[float] = None
|
|
36
|
+
session_limit: Optional[float] = None
|
|
37
|
+
session_pct: float = 0.0
|
|
38
|
+
daily_pct: float = 0.0
|
|
39
|
+
session_exceeded: bool = False
|
|
40
|
+
daily_exceeded: bool = False
|
|
41
|
+
session_warning: Optional[str] = None # "80" or "95"
|
|
42
|
+
daily_warning: Optional[str] = None
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class BudgetManager:
|
|
46
|
+
"""Manages budget limits and tracking."""
|
|
47
|
+
|
|
48
|
+
def __init__(self, config_path: Optional[Path] = None):
|
|
49
|
+
self.config_path = config_path or DEFAULT_BUDGET_FILE
|
|
50
|
+
self.config = self._load_config()
|
|
51
|
+
self._session_costs: dict[str, float] = {} # session_id -> cost
|
|
52
|
+
self._warned_sessions: set[str] = set() # sessions that hit 80%
|
|
53
|
+
_warned_95_sessions: set[str] = set() # sessions that hit 95%
|
|
54
|
+
|
|
55
|
+
def _load_config(self) -> BudgetConfig:
|
|
56
|
+
"""Load budget configuration from file."""
|
|
57
|
+
if not self.config_path.exists():
|
|
58
|
+
return BudgetConfig()
|
|
59
|
+
|
|
60
|
+
with open(self.config_path) as f:
|
|
61
|
+
data = yaml.safe_load(f) or {}
|
|
62
|
+
|
|
63
|
+
return BudgetConfig(
|
|
64
|
+
daily_limit=data.get("daily_limit"),
|
|
65
|
+
session_limit=data.get("session_limit"),
|
|
66
|
+
warn_at_80=data.get("warn_at_80", True),
|
|
67
|
+
warn_at_95=data.get("warn_at_95", True),
|
|
68
|
+
kill_at_100=data.get("kill_at_100", True),
|
|
69
|
+
webhook_url=data.get("webhook_url"),
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
def save_config(self) -> None:
|
|
73
|
+
"""Save current config to file."""
|
|
74
|
+
self.config_path.parent.mkdir(parents=True, exist_ok=True)
|
|
75
|
+
with open(self.config_path, "w") as f:
|
|
76
|
+
yaml.dump({
|
|
77
|
+
"daily_limit": self.config.daily_limit,
|
|
78
|
+
"session_limit": self.config.session_limit,
|
|
79
|
+
"warn_at_80": self.config.warn_at_80,
|
|
80
|
+
"warn_at_95": self.config.warn_at_95,
|
|
81
|
+
"kill_at_100": self.config.kill_at_100,
|
|
82
|
+
"webhook_url": self.config.webhook_url,
|
|
83
|
+
}, f, default_flow_style=False)
|
|
84
|
+
|
|
85
|
+
def set_limits(self, daily: Optional[float] = None, session: Optional[float] = None) -> None:
|
|
86
|
+
"""Set budget limits.
|
|
87
|
+
|
|
88
|
+
Raises:
|
|
89
|
+
ValueError: If daily or session limit is negative.
|
|
90
|
+
"""
|
|
91
|
+
if daily is not None:
|
|
92
|
+
if daily < 0:
|
|
93
|
+
raise ValueError(f"daily_limit cannot be negative, got {daily}")
|
|
94
|
+
self.config.daily_limit = daily
|
|
95
|
+
if session is not None:
|
|
96
|
+
if session < 0:
|
|
97
|
+
raise ValueError(f"session_limit cannot be negative, got {session}")
|
|
98
|
+
self.config.session_limit = session
|
|
99
|
+
self.save_config()
|
|
100
|
+
|
|
101
|
+
def check(self, session_id: str, current_cost: float, daily_total: float) -> BudgetStatus:
|
|
102
|
+
"""Check budget status and return warnings/exceeded flags."""
|
|
103
|
+
config = self.config
|
|
104
|
+
|
|
105
|
+
# Calculate percentages
|
|
106
|
+
session_pct = 0.0
|
|
107
|
+
daily_pct = 0.0
|
|
108
|
+
|
|
109
|
+
if config.session_limit and config.session_limit > 0:
|
|
110
|
+
session_pct = (current_cost / config.session_limit) * 100
|
|
111
|
+
|
|
112
|
+
if config.daily_limit and config.daily_limit > 0:
|
|
113
|
+
daily_pct = (daily_total / config.daily_limit) * 100
|
|
114
|
+
|
|
115
|
+
# Check warnings
|
|
116
|
+
session_warning = None
|
|
117
|
+
daily_warning = None
|
|
118
|
+
|
|
119
|
+
if config.warn_at_80:
|
|
120
|
+
if session_pct >= 80 and session_id not in self._warned_sessions:
|
|
121
|
+
session_warning = "80"
|
|
122
|
+
self._warned_sessions.add(session_id)
|
|
123
|
+
if daily_pct >= 80:
|
|
124
|
+
daily_warning = "80"
|
|
125
|
+
|
|
126
|
+
if config.warn_at_95:
|
|
127
|
+
if session_pct >= 95:
|
|
128
|
+
session_warning = "95"
|
|
129
|
+
if daily_pct >= 95:
|
|
130
|
+
daily_warning = "95"
|
|
131
|
+
|
|
132
|
+
# Check exceeded
|
|
133
|
+
session_exceeded = config.kill_at_100 and config.session_limit and current_cost >= config.session_limit
|
|
134
|
+
daily_exceeded = config.kill_at_100 and config.daily_limit and daily_total >= config.daily_limit
|
|
135
|
+
|
|
136
|
+
return BudgetStatus(
|
|
137
|
+
session_cost=current_cost,
|
|
138
|
+
daily_cost=daily_total,
|
|
139
|
+
daily_limit=config.daily_limit,
|
|
140
|
+
session_limit=config.session_limit,
|
|
141
|
+
session_pct=session_pct,
|
|
142
|
+
daily_pct=daily_pct,
|
|
143
|
+
session_exceeded=session_exceeded,
|
|
144
|
+
daily_exceeded=daily_exceeded,
|
|
145
|
+
session_warning=session_warning,
|
|
146
|
+
daily_warning=daily_warning,
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
def should_kill(self, status: BudgetStatus) -> bool:
|
|
150
|
+
"""Return True if process should be killed due to budget exceeded."""
|
|
151
|
+
return status.session_exceeded or status.daily_exceeded
|
|
152
|
+
|
|
153
|
+
def get_warning_level(self, status: BudgetStatus) -> str:
|
|
154
|
+
"""Get warning level: 'none', 'yellow' (80%), 'red' (95%), 'kill' (100%)."""
|
|
155
|
+
if status.session_exceeded or status.daily_exceeded:
|
|
156
|
+
return "kill"
|
|
157
|
+
if (status.session_warning == "95" or status.daily_warning == "95"):
|
|
158
|
+
return "red"
|
|
159
|
+
if (status.session_warning == "80" or status.daily_warning == "80"):
|
|
160
|
+
return "yellow"
|
|
161
|
+
return "none"
|
|
162
|
+
|
|
163
|
+
@staticmethod
|
|
164
|
+
def kill_process(pid: int) -> None:
|
|
165
|
+
"""Kill a process by PID."""
|
|
166
|
+
try:
|
|
167
|
+
os.kill(pid, signal.SIGTERM)
|
|
168
|
+
except ProcessLookupError:
|
|
169
|
+
pass
|
|
170
|
+
except PermissionError:
|
|
171
|
+
# Try SIGKILL if SIGTERM fails
|
|
172
|
+
try:
|
|
173
|
+
os.kill(pid, signal.SIGKILL)
|
|
174
|
+
except (ProcessLookupError, PermissionError):
|
|
175
|
+
pass
|
|
176
|
+
|
|
177
|
+
def get_status_text(self, status: BudgetStatus) -> str:
|
|
178
|
+
"""Get human-readable status text."""
|
|
179
|
+
lines = []
|
|
180
|
+
|
|
181
|
+
if status.session_limit:
|
|
182
|
+
pct = status.session_pct
|
|
183
|
+
cost = status.session_cost
|
|
184
|
+
limit = status.session_limit
|
|
185
|
+
lines.append(f"Session: ${cost:.4f}/${limit:.2f} ({pct:.1f}%)")
|
|
186
|
+
|
|
187
|
+
if status.daily_limit:
|
|
188
|
+
pct = status.daily_pct
|
|
189
|
+
cost = status.daily_cost
|
|
190
|
+
limit = status.daily_limit
|
|
191
|
+
lines.append(f"Daily: ${cost:.4f}/${limit:.2f} ({pct:.1f}%)")
|
|
192
|
+
|
|
193
|
+
return " | ".join(lines) if lines else "No budget limits set"
|