Skip to content

Commit c1ced31

Browse files
author
Alex
committed
feat(usage): implement Phases B+C of LLM cost tracking #1075
Phase B - Persistence: - Add configurable pricing module (pricing.rs) with 22 model defaults, TOML config support, and longest-prefix glob matching - Add terraphim_usage dep to terraphim_multi_agent - Add TokenUsageTracker::drain_records() for record extraction - Add TerraphimAgent::flush_usage() to persist records to UsageStore - Wire production usage recording in agent.rs after each LLM call (pricing lookup + TokenUsageRecord creation) Phase C - Provider completion: - Implement Claude provider via ccusage (7-day usage aggregation) - Implement OpenCode Go provider via sqlite3 CLI (cost aggregation) - Create CcusageProvider adapter wrapping terraphim_ccusage - Add terraphim_ccusage as dependency of terraphim_usage - Register all 5 providers in CLI handle_usage() - Kimi remains stub (undocumented API) 37 tests pass, clippy clean, fmt clean. Refs #1075
1 parent 39d7079 commit c1ced31

14 files changed

Lines changed: 971 additions & 22 deletions

File tree

‎.docs/research-design-phases-bc-llm-cost-tracking.md‎

Lines changed: 268 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,268 @@
1+
# Research & Design: LLM Cost Tracking -- Phases B+C
2+
3+
## 1. Problem Restatement and Scope
4+
5+
### Problem
6+
Phase A delivered `LlmUsage` types, `chat_completion_with_usage()`, and the `GenAiClient` adapter. Usage metadata now flows from API responses back to callers. However:
7+
8+
- **Pricing is still hardcoded** in two disconnected systems (`terraphim_multi_agent::tracking::ModelPricing` with per-1K pricing for 3 models, `terraphim_types::llm_usage::ModelPricing` with per-million pricing and model patterns)
9+
- **`TokenUsageTracker` has zero persistence** -- all records lost on process exit. No `flush()` method exists.
10+
- **`token_tracker` is never written to in production code** -- `record_usage()` only called from tests
11+
- **3 of 5 providers are stubs** (Claude, OpenCode Go, Kimi return "Not yet implemented")
12+
- **CLI `handle_usage()` creates an empty `UsageRegistry`** -- no providers registered, so `usage show` returns nothing
13+
- **`terraphim_usage` does not depend on `terraphim_ccusage`** -- Claude provider can't access ccusage data
14+
- **`terraphim_multi_agent` does not depend on `terraphim_usage`** -- no bridge exists
15+
16+
### IN Scope (Phases B+C)
17+
- Configurable pricing module (TOML file + embedded defaults)
18+
- Wire `TokenUsageTracker` to persist records into `UsageStore`
19+
- Complete Claude provider (via ccusage)
20+
- Complete OpenCode Go provider (via SQLite)
21+
- Add ccusage as dependency of terraphim_usage and create ccusage adapter provider
22+
- Register all providers in CLI `handle_usage()`
23+
- Keep Kimi as stub (API unknown, low priority)
24+
25+
### OUT of Scope
26+
- Phase D CLI enhancements (`usage history --by model`, `usage alert --budget N`) -- separate PR
27+
- Web dashboard
28+
- Rate limiting based on costs
29+
- Kimi provider implementation (Moonshot API undocumented)
30+
31+
## 2. System Elements and Dependencies
32+
33+
### Current State After Phase A
34+
35+
| Element | Crate | Status | Notes |
36+
|---|---|---|---|
37+
| `LlmUsage` | `terraphim_types` | DONE | input/output tokens, model, provider, cost_usd, latency_ms |
38+
| `LlmResult` | `terraphim_types` | DONE | content + optional LlmUsage |
39+
| `ModelPricing` (types) | `terraphim_types` | DONE | per-million pricing, model_pattern |
40+
| `ModelPricing` (multi_agent) | `terraphim_multi_agent` | LEGACY | per-1K pricing, 3 hardcoded models, to be replaced |
41+
| `chat_completion_with_usage()` | `terraphim_service` | DONE | GenAiClient extracts usage |
42+
| `ExecutionRecord::from_llm_usage()` | `terraphim_usage` | DONE | Bridge from LlmUsage to persistent record |
43+
| `UsageStore::save_execution()` | `terraphim_usage` | DONE | Persists ExecutionRecord |
44+
| `UsageProvider` trait | `terraphim_usage` | DONE | fetch_usage() -> ProviderUsage |
45+
| MiniMax provider | `terraphim_usage` | DONE | Full API integration |
46+
| ZAI provider | `terraphim_usage` | DONE | Full API integration |
47+
| Claude provider | `terraphim_usage` | STUB | Returns error |
48+
| OpenCode Go provider | `terraphim_usage` | STUB | Returns error |
49+
| Kimi provider | `terraphim_usage` | STUB | Returns error |
50+
| `CcusageClient` | `terraphim_ccusage` | DONE | Shells out to `bun dlx ccusage`, returns DailyUsageReport |
51+
| `TokenUsageTracker` | `terraphim_multi_agent` | IN-MEMORY ONLY | No persistence |
52+
| `CostTracker` | `terraphim_multi_agent` | IN-MEMORY ONLY | Hardcoded pricing |
53+
| `handle_usage()` | `terraphim_cli` | BROKEN | Empty registry, no providers registered |
54+
55+
### Dependency Graph (Desired)
56+
57+
```
58+
terraphim_types (LlmUsage, ModelPricing)
59+
^ ^
60+
| |
61+
terraphim_service terraphim_usage ----> terraphim_ccusage (new dep)
62+
(return usage) (persist usage) (Claude token data)
63+
| |
64+
v v
65+
terraphim_multi_agent ---> terraphim_usage (new dep)
66+
(flush TokenUsageTracker) (persist records)
67+
```
68+
69+
### Key Constraint: terraphim_multi_agent adding terraphim_usage dependency
70+
Currently `terraphim_multi_agent` does NOT depend on `terraphim_usage`. Adding this dep is safe (no circular risk) since `terraphim_usage` depends only on `terraphim_types`, `terraphim_persistence`, `terraphim_settings`, and `opendal` -- none of which create cycles back to `terraphim_multi_agent`.
71+
72+
## 3. Constraints
73+
74+
| Constraint | Implication |
75+
|---|---|
76+
| No circular deps | terraphim_usage cannot depend on terraphim_multi_agent (safe: it doesn't) |
77+
| Fire-and-forget persistence | flush_to_store must not block LLM calls |
78+
| Sub-cent precision | ExecutionRecord uses cost_sub_cents (1/1,000,000 cent) |
79+
| Feature flags | providers gated behind `providers` feature in terraphim_usage |
80+
| ccusage requires bun/pnpm | Claude provider graceful degradation if runner missing |
81+
| SQLite dependency for OpenCode Go | Need rusqlite or shell out to sqlite3 CLI |
82+
| Existing test suite must pass | 128+ tests, clippy clean |
83+
84+
## 4. Risks and Unknowns
85+
86+
| Risk | Likelihood | Mitigation |
87+
|---|---|---|
88+
| Adding terraphim_usage dep to multi_agent causes compile issues | Low | Check dep graph first; terraphim_usage only pulls in terraphim_types/persistence |
89+
| ccusage not installed / bun missing | Medium | Return ProviderNotFound gracefully, same as OpenCode Go pattern |
90+
| OpenCode SQLite schema differs from assumed | Medium | Query sqlite_master first; handle missing columns gracefully |
91+
| Two ModelPricing types confuse callers | High | Unify: use terraphim_types::ModelPricing everywhere, deprecate multi_agent version |
92+
| flush_to_store loses data on crash | Low | Acceptable trade-off per design; records accumulate in memory until flush |
93+
94+
### Assumptions
95+
- [ASSUMPTION] OpenCode Go SQLite at `~/.local/share/opencode/opencode.db` has a `message` table with cost fields
96+
- [ASSUMPTION] ccusage@18.0.10 JSON output format is stable (DailyUsageReport already handles it)
97+
- [ASSUMPTION] Kimi/Moonshot usage API is not publicly documented -- keep as stub
98+
99+
## 5. Design: Step-by-Step Plan
100+
101+
### Step B1: Configurable Pricing Module
102+
103+
**File**: `crates/terraphim_usage/src/pricing.rs` (NEW)
104+
105+
Create `PricingTable` struct:
106+
- `entries: Vec<terraphim_types::ModelPricing>` -- reuses existing type
107+
- `load(path: &Path) -> Result<Self>` -- reads TOML, falls back to embedded defaults
108+
- `embedded_defaults() -> Self` -- comprehensive default pricing for 20+ models
109+
- `find_pricing(model: &str) -> Option<&ModelPricing>` -- glob/substring match on model_pattern
110+
- `calculate_cost(model: &str, input: u64, output: u64) -> Option<f64>` -- convenience wrapper
111+
112+
TOML format:
113+
```toml
114+
[[models]]
115+
pattern = "openai/gpt-4o*"
116+
input_per_m = 2.50
117+
output_per_m = 10.00
118+
119+
[[models]]
120+
pattern = "anthropic/claude-sonnet*"
121+
input_per_m = 3.00
122+
output_per_m = 15.00
123+
```
124+
125+
Config file location: `~/.config/terraphim/pricing.toml` (optional, falls back to embedded)
126+
127+
Also add `toml` dependency to terraphim_usage Cargo.toml.
128+
129+
### Step B2: TokenUsageTracker flush_to_store
130+
131+
**File**: `crates/terraphim_multi_agent/src/tracking.rs` (MODIFY)
132+
133+
Add to `TokenUsageTracker`:
134+
```rust
135+
pub fn drain_records(&mut self) -> Vec<TokenUsageRecord> {
136+
std::mem::take(&mut self.records)
137+
// Keep totals intact for in-memory queries
138+
}
139+
```
140+
141+
**File**: `crates/terraphim_multi_agent/src/agent.rs` (MODIFY)
142+
143+
Add a `flush_usage` method to `TerraphimAgent` that:
144+
1. Takes `self.token_tracker.write().drain_records()`
145+
2. Converts each `TokenUsageRecord` into `LlmUsage` (already has all fields)
146+
3. Creates `ExecutionRecord::from_llm_usage()` for each
147+
4. Calls `UsageStore::save_execution()` for each
148+
5. Fire-and-forget via `tokio::spawn`
149+
150+
Requires: Add `terraphim_usage` dependency to `terraphim_multi_agent/Cargo.toml`
151+
152+
### Step B3: Wire Production Usage Recording
153+
154+
**File**: `crates/terraphim_multi_agent/src/genai_llm_client.rs` (MODIFY)
155+
156+
After each `exec_chat()` call that returns usage metadata, record it into `token_tracker`:
157+
- Extract usage from genai response
158+
- Create `TokenUsageRecord`
159+
- Call `token_tracker.write().record_usage(record)`
160+
161+
This is the missing production write path -- currently only tests call `record_usage()`.
162+
163+
### Step C1: Add ccusage Dependency + Adapter Provider
164+
165+
**File**: `crates/terraphim_usage/Cargo.toml` (MODIFY)
166+
- Add `terraphim_ccusage = { path = "../terraphim_ccusage", optional = true }`
167+
- Add to `providers` feature: `"dep:terraphim_ccusage"`
168+
169+
**File**: `crates/terraphim_usage/src/providers/ccusage.rs` (NEW)
170+
- `CcusageProvider` struct wrapping `terraphim_ccusage::CcusageClient`
171+
- Implements `UsageProvider` trait
172+
- `fetch_usage()` queries ccusage for last 30 days
173+
- Builds `MetricLine::Progress` entries for daily/weekly/monthly spend
174+
- `id() -> "ccusage"`, `display_name() -> "Claude Code (ccusage)"`
175+
176+
**File**: `crates/terraphim_usage/src/providers/mod.rs` (MODIFY)
177+
- Add `#[cfg(feature = "providers")] pub mod ccusage;`
178+
179+
### Step C2: Complete Claude Provider
180+
181+
**File**: `crates/terraphim_usage/src/providers/claude.rs` (MODIFY)
182+
183+
Two-layer approach:
184+
1. Primary: Read OAuth token from `~/.claude/.credentials.json`, call Anthropic usage API
185+
2. Fallback: If OAuth fails or API unavailable, delegate to ccusage adapter
186+
187+
Since Anthropic's OAuth usage API (`/api/oauth/usage`) endpoint and response format are not publicly documented, implement via ccusage as primary method for now. The Claude provider will:
188+
- Create internal `CcusageClient::new(CcusageProvider::Claude)`
189+
- Query last 7 days of usage
190+
- Build `MetricLine::Progress` for daily and weekly windows
191+
- Add `terraphim_ccusage` as internal dependency
192+
193+
### Step C3: Complete OpenCode Go Provider
194+
195+
**File**: `crates/terraphim_usage/src/providers/opencode_go.rs` (MODIFY)
196+
197+
SQLite query approach (avoid rusqlite heavy dep):
198+
- Shell out to `sqlite3` CLI to query the database
199+
- `SELECT role, COUNT(*), SUM(cost) FROM message GROUP BY role`
200+
- Parse TSV output
201+
- If `sqlite3` not available, return ProviderNotFound
202+
- Build `MetricLine::Progress` entries
203+
204+
Alternative: Use `rusqlite` with `bundled-sqlite` feature. This is more robust but adds a heavy C dependency. Decision: use `sqlite3` CLI shell-out for now (matches ccusage pattern of shelling out).
205+
206+
### Step C4: Register All Providers in CLI
207+
208+
**File**: `crates/terraphim_cli/src/main.rs` (MODIFY)
209+
210+
In `handle_usage()`, register all available providers:
211+
```rust
212+
let mut registry = UsageRegistry::new();
213+
registry.register(Box::new(ClaudeProvider::new()));
214+
registry.register(Box::new(OpenCodeGoProvider::new()));
215+
registry.register(Box::new(MiniMaxProvider::new()));
216+
registry.register(Box::new(ZaiProvider::new()));
217+
registry.register(Box::new(CcusageProvider::new()));
218+
```
219+
220+
This fixes the "empty registry" bug where `usage show` returns nothing.
221+
222+
### Step C5: Keep Kimi as Stub
223+
224+
No changes -- Kimi/Moonshot usage API is undocumented. Leave as-is with clear TODO.
225+
226+
## 6. File Change Summary
227+
228+
| File | Action | Purpose |
229+
|---|---|---|
230+
| `terraphim_usage/src/pricing.rs` | CREATE | Configurable pricing from TOML + defaults |
231+
| `terraphim_usage/Cargo.toml` | MODIFY | Add `toml` dep, add `terraphim_ccusage` dep |
232+
| `terraphim_usage/src/lib.rs` | MODIFY | Add `pub mod pricing;` |
233+
| `terraphim_usage/src/providers/ccusage.rs` | CREATE | CcusageProvider wrapping CcusageClient |
234+
| `terraphim_usage/src/providers/claude.rs` | MODIFY | Implement via ccusage |
235+
| `terraphim_usage/src/providers/opencode_go.rs` | MODIFY | Implement via sqlite3 CLI |
236+
| `terraphim_usage/src/providers/mod.rs` | MODIFY | Add ccusage module |
237+
| `terraphim_multi_agent/Cargo.toml` | MODIFY | Add `terraphim_usage` dependency |
238+
| `terraphim_multi_agent/src/tracking.rs` | MODIFY | Add `drain_records()` method |
239+
| `terraphim_multi_agent/src/agent.rs` | MODIFY | Add `flush_usage()` method |
240+
| `terraphim_multi_agent/src/genai_llm_client.rs` | MODIFY | Record usage after each LLM call |
241+
| `terraphim_cli/src/main.rs` | MODIFY | Register all providers in handle_usage() |
242+
243+
## 7. Testing Strategy
244+
245+
| AC / Invariant | Test Type | Location |
246+
|---|---|---|
247+
| PricingTable loads from TOML | Unit | `terraphim_usage/src/pricing.rs` |
248+
| PricingTable embedded defaults cover 20+ models | Unit | `terraphim_usage/src/pricing.rs` |
249+
| PricingTable pattern matching (gpt-4o-mini matches openai/gpt-4o*) | Unit | `terraphim_usage/src/pricing.rs` |
250+
| TokenUsageTracker::drain_records() empties vec, keeps totals | Unit | `terraphim_multi_agent/src/tracking.rs` |
251+
| ExecutionRecord::from_llm_usage() preserves all fields | Unit | `terraphim_usage/src/store.rs` (existing) |
252+
| CcusageProvider returns ProviderNotFound if no runner | Unit | `terraphim_usage/src/providers/ccusage.rs` |
253+
| OpenCode Go provider returns ProviderNotFound if no db | Unit | `terraphim_usage/src/providers/opencode_go.rs` |
254+
| Full test suite passes | Integration | `cargo test --workspace` |
255+
| Clippy clean | Lint | `cargo clippy --workspace` |
256+
257+
## 8. Implementation Order
258+
259+
1. **B1**: Pricing module (no deps on other steps)
260+
2. **C1**: ccusage dependency + adapter (no deps on other steps)
261+
3. **C2**: Claude provider (depends on C1)
262+
4. **C3**: OpenCode Go provider (independent)
263+
5. **B2**: TokenUsageTracker drain + flush (depends on B1 for pricing)
264+
6. **B3**: Wire production recording in genai_llm_client (depends on B2)
265+
7. **C4**: Register providers in CLI (depends on C1, C2, C3)
266+
8. **Verify**: Full test suite + clippy
267+
268+
Steps 1-4 can be done in parallel. Steps 5-6 are sequential. Step 7 is last.

‎Cargo.lock‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎crates/terraphim_cli/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ terraphim_update = { path = "../terraphim_update", version = "1.0.0" }
2626
terraphim_hooks = { path = "../terraphim_hooks", version = "1.0.0" }
2727

2828
# Usage tracking (feature-gated)
29-
terraphim_usage = { path = "../terraphim_usage", optional = true, features = ["cli"] }
29+
terraphim_usage = { path = "../terraphim_usage", optional = true, features = ["cli", "providers"] }
3030

3131
# CLI framework
3232
clap = { version = "4.5", features = ["derive", "cargo", "env"] }

‎crates/terraphim_cli/src/main.rs‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -944,7 +944,25 @@ fn format_as_text(value: &serde_json::Value) -> Result<String> {
944944

945945
#[cfg(feature = "usage")]
946946
async fn handle_usage(action: terraphim_usage::cli::UsageAction) -> Result<serde_json::Value> {
947-
let registry = terraphim_usage::UsageRegistry::new();
947+
let mut registry = terraphim_usage::UsageRegistry::new();
948+
949+
#[cfg(feature = "usage")]
950+
{
951+
registry.register(Box::new(
952+
terraphim_usage::providers::claude::ClaudeProvider::new(),
953+
));
954+
registry.register(Box::new(
955+
terraphim_usage::providers::opencode_go::OpenCodeGoProvider::new(),
956+
));
957+
registry.register(Box::new(
958+
terraphim_usage::providers::minimax::MiniMaxProvider::new(),
959+
));
960+
registry.register(Box::new(terraphim_usage::providers::zai::ZaiProvider::new()));
961+
registry.register(Box::new(
962+
terraphim_usage::providers::ccusage::CcusageProvider::new(),
963+
));
964+
}
965+
948966
let output = terraphim_usage::cli::execute_usage_action(action, &registry)
949967
.await
950968
.map_err(|e| anyhow::anyhow!("{}", e))?;

‎crates/terraphim_multi_agent/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ terraphim_persistence = { path = "../terraphim_persistence" }
4747
terraphim_agent_evolution = { path = "../terraphim_agent_evolution" }
4848
terraphim_service = { path = "../terraphim_service" }
4949
terraphim_agent_registry = { path = "../terraphim_agent_registry" }
50+
terraphim_usage = { path = "../terraphim_usage" }
5051

5152
# Firecracker VM dependencies (optional, feature-gated)
5253
# Note: fcctl-repl has its own workspace dependencies that may conflict

0 commit comments

Comments
 (0)