From 3870a81b89da66864298ba750753c586b5ad1608 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Sun, 31 May 2026 20:07:27 +0530 Subject: [PATCH] feat(storage): aggregate cost-ledger spend by dimension MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Store.SummarizeLedger — group spend (calls/tokens/dollars) across runs by provider, model, day, run name, or a run-metadata key (e.g. metadata.team), with optional since/until time bounds. Sums from the cost_ledger (the auditable source), not the run rollup. The grouping expression is chosen from a fixed whitelist (it's structural, not a bindable parameter); the metadata JSON path is bound as a value and its key is validated, so nothing is built unsafely from caller input. This is the storage foundation for the usage-rollup CLI/API (#37) — no public surface yet. --- internal/storage/sqlite.go | 93 +++++++++++++++++++++++++++++++ internal/storage/sqlite_test.go | 97 +++++++++++++++++++++++++++++++++ internal/storage/store.go | 32 +++++++++++ 3 files changed, 222 insertions(+) diff --git a/internal/storage/sqlite.go b/internal/storage/sqlite.go index 6697c8b..1ee02af 100644 --- a/internal/storage/sqlite.go +++ b/internal/storage/sqlite.go @@ -282,6 +282,99 @@ func (s *SQLite) Totals(ctx context.Context, runID string) (LedgerTotals, error) return t, nil } +// SummarizeLedger aggregates the cost ledger across runs, grouped by opts.By. +// The grouping expression is chosen from a fixed whitelist (it's structural and +// can't be a bound parameter); the metadata key is validated and its JSON path is +// bound as a value, so nothing here is built from raw user input unsafely. +func (s *SQLite) SummarizeLedger(ctx context.Context, opts SummarizeOptions) (UsageSummary, error) { + var groupExpr, metaArg string + needRuns := false + switch { + case opts.By == "provider": + groupExpr = "l.provider" + case opts.By == "model": + groupExpr = "l.model" + case opts.By == "day": + groupExpr = "substr(l.created_at, 1, 10)" // RFC3339 date prefix (UTC) + case opts.By == "name": + groupExpr, needRuns = "r.name", true + case strings.HasPrefix(opts.By, "metadata."): + key := strings.TrimPrefix(opts.By, "metadata.") + if !validMetaKey(key) { + return UsageSummary{}, fmt.Errorf("storage: invalid metadata key %q", key) + } + groupExpr, metaArg, needRuns = "json_extract(r.metadata, ?)", "$."+key, true + default: + return UsageSummary{}, fmt.Errorf("storage: unsupported group dimension %q", opts.By) + } + + from := "cost_ledger l" + if needRuns { + from += " JOIN runs r ON r.id = l.run_id" + } + + var args []any + if metaArg != "" { // the json_extract path is the first placeholder, in SELECT + args = append(args, metaArg) + } + var conds []string + if opts.Since != nil { + conds = append(conds, "l.created_at >= ?") + args = append(args, fmtTime(*opts.Since)) + } + if opts.Until != nil { + conds = append(conds, "l.created_at < ?") + args = append(args, fmtTime(*opts.Until)) + } + where := "" + if len(conds) > 0 { + where = " WHERE " + strings.Join(conds, " AND ") + } + + q := fmt.Sprintf(` + SELECT COALESCE(%s, '(none)') AS k, + COUNT(*), COALESCE(SUM(l.prompt_tokens), 0), COALESCE(SUM(l.completion_tokens), 0), + COALESCE(SUM(l.dollars), 0) + FROM %s%s + GROUP BY k + ORDER BY SUM(l.dollars) DESC, k ASC`, groupExpr, from, where) + + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return UsageSummary{}, fmt.Errorf("storage: summarize ledger: %w", err) + } + defer rows.Close() + + out := UsageSummary{By: opts.By, Groups: []UsageGroup{}, Total: UsageGroup{Key: "total"}} + for rows.Next() { + var g UsageGroup + if err := rows.Scan(&g.Key, &g.Calls, &g.PromptTokens, &g.CompletionTokens, &g.Dollars); err != nil { + return UsageSummary{}, err + } + out.Groups = append(out.Groups, g) + out.Total.Calls += g.Calls + out.Total.PromptTokens += g.PromptTokens + out.Total.CompletionTokens += g.CompletionTokens + out.Total.Dollars += g.Dollars + } + return out, rows.Err() +} + +// validMetaKey allows only flat, identifier-like metadata keys (no JSON-path +// metacharacters), so the bound "$." path can't be abused. +func validMetaKey(k string) bool { + if k == "" { + return false + } + for _, r := range k { + ok := r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r >= '0' && r <= '9' || r == '_' || r == '-' + if !ok { + return false + } + } + return true +} + // --- tool calls --- func (s *SQLite) AppendToolCall(ctx context.Context, t ToolCallRecord) error { diff --git a/internal/storage/sqlite_test.go b/internal/storage/sqlite_test.go index 0c3ac0f..23c3f24 100644 --- a/internal/storage/sqlite_test.go +++ b/internal/storage/sqlite_test.go @@ -150,6 +150,103 @@ func TestLedgerAndTotals(t *testing.T) { } } +func TestSummarizeLedger(t *testing.T) { + s := openTemp(t) + ctx := context.Background() + day1 := time.Date(2026, 5, 20, 10, 0, 0, 0, time.UTC) + day2 := time.Date(2026, 5, 21, 10, 0, 0, 0, time.UTC) + + mkRun := func(id, name, team string, created time.Time) { + if err := s.UpsertRun(ctx, RunRecord{ + ID: id, Name: name, Status: "running", + Metadata: map[string]string{"team": team}, + CreatedAt: created, UpdatedAt: created, + }); err != nil { + t.Fatalf("seed run %s: %v", id, err) + } + } + mkRun("run-a", "alice", "marketing", day1) + mkRun("run-b", "bob", "marketing", day1) + mkRun("run-c", "carol", "eng", day2) + + for _, e := range []LedgerEntry{ + {RunID: "run-a", StepIndex: 1, Provider: "anthropic", Model: "claude-haiku", PromptTokens: 100, CompletionTokens: 50, Dollars: 0.01, Priced: true, CreatedAt: day1}, + {RunID: "run-b", StepIndex: 1, Provider: "anthropic", Model: "claude-haiku", PromptTokens: 200, CompletionTokens: 80, Dollars: 0.012, Priced: true, CreatedAt: day1}, + {RunID: "run-b", StepIndex: 2, Provider: "anthropic", Model: "claude-haiku", PromptTokens: 50, CompletionTokens: 20, Dollars: 0.008, Priced: true, CreatedAt: day1.Add(time.Second)}, + {RunID: "run-c", StepIndex: 1, Provider: "openai", Model: "gpt-5", PromptTokens: 300, CompletionTokens: 100, Dollars: 0.05, Priced: true, CreatedAt: day2}, + } { + if err := s.AppendLedger(ctx, e); err != nil { + t.Fatal(err) + } + } + + asMap := func(sum UsageSummary) map[string]UsageGroup { + m := map[string]UsageGroup{} + for _, g := range sum.Groups { + m[g.Key] = g + } + return m + } + approx := func(got, want float64) bool { return got > want-1e-9 && got < want+1e-9 } + + // by metadata.team — spend rolled up across runs by a user-supplied tag. + sum, err := s.SummarizeLedger(ctx, SummarizeOptions{By: "metadata.team"}) + if err != nil { + t.Fatalf("by team: %v", err) + } + m := asMap(sum) + if g := m["marketing"]; g.Calls != 3 || g.PromptTokens != 350 || g.CompletionTokens != 150 || !approx(g.Dollars, 0.03) { + t.Fatalf("by team marketing = %+v", g) + } + if g := m["eng"]; g.Calls != 1 || !approx(g.Dollars, 0.05) { + t.Fatalf("by team eng = %+v", g) + } + if sum.Total.Calls != 4 || !approx(sum.Total.Dollars, 0.08) { + t.Fatalf("by team total = %+v", sum.Total) + } + if len(sum.Groups) != 2 || sum.Groups[0].Key != "eng" { // ordered by $ desc + t.Fatalf("group order = %+v", sum.Groups) + } + + // by provider + sum, _ = s.SummarizeLedger(ctx, SummarizeOptions{By: "provider"}) + m = asMap(sum) + if g := m["anthropic"]; g.Calls != 3 || !approx(g.Dollars, 0.03) { + t.Fatalf("by provider anthropic = %+v", g) + } + if g := m["openai"]; g.Calls != 1 || !approx(g.Dollars, 0.05) { + t.Fatalf("by provider openai = %+v", g) + } + + // by name → one group per run + sum, _ = s.SummarizeLedger(ctx, SummarizeOptions{By: "name"}) + if len(asMap(sum)) != 3 { + t.Fatalf("by name groups = %+v", sum.Groups) + } + + // by day, bounded to before day2 → only day1's three calls + until := day2 + sum, err = s.SummarizeLedger(ctx, SummarizeOptions{By: "day", Until: &until}) + if err != nil { + t.Fatalf("by day: %v", err) + } + m = asMap(sum) + if len(m) != 1 { + t.Fatalf("day groups (bounded) = %+v", sum.Groups) + } + if g := m["2026-05-20"]; g.Calls != 3 || !approx(g.Dollars, 0.03) { + t.Fatalf("day 2026-05-20 = %+v", g) + } + + // guardrails: unsupported dimension and unsafe metadata key are rejected. + if _, err := s.SummarizeLedger(ctx, SummarizeOptions{By: "bogus"}); err == nil { + t.Fatal("expected error for unsupported dimension") + } + if _, err := s.SummarizeLedger(ctx, SummarizeOptions{By: "metadata.bad key!"}); err == nil { + t.Fatal("expected error for invalid metadata key") + } +} + func TestToolCall(t *testing.T) { s := openTemp(t) ctx := context.Background() diff --git a/internal/storage/store.go b/internal/storage/store.go index 370514c..6627dc0 100644 --- a/internal/storage/store.go +++ b/internal/storage/store.go @@ -120,6 +120,35 @@ type LedgerTotals struct { Dollars float64 } +// UsageGroup is one bucket of aggregated spend — e.g. one team, one provider, or +// one day. Tokens/dollars are summed from the cost ledger (the auditable source). +type UsageGroup struct { + Key string `json:"key"` + Calls int64 `json:"calls"` + PromptTokens int64 `json:"promptTokens"` + CompletionTokens int64 `json:"completionTokens"` + Dollars float64 `json:"dollars"` +} + +// UsageSummary is cost-ledger spend grouped by one dimension, plus the grand total +// across all groups in range. +type UsageSummary struct { + By string `json:"by"` + Groups []UsageGroup `json:"groups"` + Total UsageGroup `json:"total"` +} + +// SummarizeOptions selects how SummarizeLedger aggregates the cost ledger. +type SummarizeOptions struct { + // By is the grouping dimension: "provider", "model", "day", "name", or + // "metadata." (e.g. "metadata.team"). Required. + By string + // Since/Until optionally bound the call time (created_at): Since is inclusive, + // Until exclusive. Nil means unbounded on that end. + Since *time.Time + Until *time.Time +} + // Store is the durable backend. Implementations must be safe for concurrent use. type Store interface { // UpsertRun inserts or replaces a run row by ID. @@ -143,6 +172,9 @@ type Store interface { LedgerForRun(ctx context.Context, runID string) ([]LedgerEntry, error) // Totals aggregates a run's ledger. Totals(ctx context.Context, runID string) (LedgerTotals, error) + // SummarizeLedger aggregates spend across runs, grouped by opts.By + // (provider/model/day/name/metadata.). The unit is the ledger row. + SummarizeLedger(ctx context.Context, opts SummarizeOptions) (UsageSummary, error) // AppendToolCall records a tool invocation. AppendToolCall(ctx context.Context, t ToolCallRecord) error