You've already forked agentic-coding-workflow
bugfix
This commit is contained in:
@@ -0,0 +1,396 @@
|
||||
// usage.go implements ACP session cost accounting:
|
||||
// - ExtractUsage: a tolerant parser pulling per-turn billable token usage out
|
||||
// of agent messages (session/prompt RESPONSE result.usage, snake/camel) and
|
||||
// opportunistic cost out of usage_update notifications.
|
||||
// - usageAccumulator: applies one parsed turn — prices it (ModelPriceLookup or
|
||||
// agent self-reported cost fallback), persists ledger + running totals via
|
||||
// the repo, and reports a budget breach when a cap is now exceeded.
|
||||
//
|
||||
// Billing source of truth (spec §7 risk #1): claude-agent-acp returns billable
|
||||
// usage in the session/prompt response ({stopReason, usage:{input_tokens,...}}).
|
||||
// The ACP usage_update notification carries a context-window GAUGE (used/size),
|
||||
// NOT a per-turn delta — summing it would massively overcount, so it is used
|
||||
// only for the optional cost{amount} fallback, never for token deltas.
|
||||
package acp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ParsedUsage is one extracted turn. HasUsage=false means no billable usage was
|
||||
// present (ordinary chunk / malformed payload) and the caller must skip it.
|
||||
type ParsedUsage struct {
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
ThinkingTokens int
|
||||
// CostUSD is the agent-self-reported cost (usage_update.cost.amount), used as
|
||||
// a pricing fallback when no model price is configured. nil = not reported.
|
||||
CostUSD *float64
|
||||
HasUsage bool
|
||||
}
|
||||
|
||||
// ModelPrice is the per-million-token price triple for a model.
|
||||
type ModelPrice struct {
|
||||
PromptPerM float64
|
||||
CompletionPerM float64
|
||||
ThinkingPerM float64
|
||||
Found bool
|
||||
}
|
||||
|
||||
// ModelPriceLookup resolves a model_id to its prices. Implemented in app.go by
|
||||
// an adapter over the chat endpoint/model service.
|
||||
type ModelPriceLookup interface {
|
||||
PriceForModel(ctx context.Context, modelID uuid.UUID) (ModelPrice, error)
|
||||
}
|
||||
|
||||
// UsageAccumulator applies parsed turns to one session: persists the ledger +
|
||||
// running totals and reports budget breaches.
|
||||
type UsageAccumulator interface {
|
||||
// Observe applies one parsed turn. sourceEventID (>0) de-dups via the ledger
|
||||
// unique index. Returns a non-empty breach reason when a cap is now exceeded.
|
||||
Observe(ctx context.Context, p ParsedUsage, sourceEventID int64) (breach BreachReason, breached bool)
|
||||
}
|
||||
|
||||
// usageSink is fed raw messages from relay.reader after persistence; it extracts
|
||||
// usage and forwards to the UsageAccumulator. Separated from the accumulator so
|
||||
// the relay does not depend on pricing/repo details.
|
||||
type usageSink interface {
|
||||
Observe(ctx context.Context, msg *Message, eventID int64)
|
||||
}
|
||||
|
||||
// accumulatorSink adapts a UsageAccumulator to the relay's usageSink: it parses
|
||||
// each agent->client message and, on a real billable turn, forwards to Observe.
|
||||
// A breach invokes onBreach exactly once on a detached goroutine (so the relay
|
||||
// reader never blocks on sup.Kill).
|
||||
type accumulatorSink struct {
|
||||
acc UsageAccumulator
|
||||
onBreach func(reason BreachReason)
|
||||
fired sync.Once
|
||||
log *slog.Logger
|
||||
}
|
||||
|
||||
// newAccumulatorSink builds the relay usage sink.
|
||||
func newAccumulatorSink(acc UsageAccumulator, onBreach func(reason BreachReason), log *slog.Logger) *accumulatorSink {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &accumulatorSink{acc: acc, onBreach: onBreach, log: log}
|
||||
}
|
||||
|
||||
func (s *accumulatorSink) Observe(ctx context.Context, msg *Message, eventID int64) {
|
||||
if s == nil || s.acc == nil {
|
||||
return
|
||||
}
|
||||
p := ExtractUsage(msg)
|
||||
if !p.HasUsage {
|
||||
return
|
||||
}
|
||||
breach, ok := s.acc.Observe(ctx, p, eventID)
|
||||
if !ok || s.onBreach == nil {
|
||||
return
|
||||
}
|
||||
// Fire the kill exactly once, detached so the reader keeps draining.
|
||||
s.fired.Do(func() {
|
||||
reason := breach
|
||||
go s.onBreach(reason)
|
||||
})
|
||||
}
|
||||
|
||||
// ===== ExtractUsage =====
|
||||
|
||||
// promptResponseUsage matches both snake_case (claude-agent-acp stable) and
|
||||
// camelCase (unstable v2 Usage) usage objects on a session/prompt response.
|
||||
type promptResponseResult struct {
|
||||
StopReason string `json:"stopReason"`
|
||||
Usage *struct {
|
||||
// snake_case (stable)
|
||||
InputTokens *int `json:"input_tokens"`
|
||||
OutputTokens *int `json:"output_tokens"`
|
||||
// camelCase (unstable v2)
|
||||
InputTokensCamel *int `json:"inputTokens"`
|
||||
OutputTokensCamel *int `json:"outputTokens"`
|
||||
ThoughtTokens *int `json:"thoughtTokens"`
|
||||
CachedReadTokens *int `json:"cachedReadTokens"`
|
||||
CachedWriteTokens *int `json:"cachedWriteTokens"`
|
||||
// thinking_tokens snake fallback
|
||||
ThinkingTokens *int `json:"thinking_tokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
// usageUpdateParams matches a session/update notification carrying usage_update.
|
||||
type usageUpdateParams struct {
|
||||
Update *struct {
|
||||
SessionUpdate string `json:"sessionUpdate"`
|
||||
Cost *struct {
|
||||
Amount float64 `json:"amount"`
|
||||
Currency string `json:"currency"`
|
||||
} `json:"cost"`
|
||||
} `json:"update"`
|
||||
}
|
||||
|
||||
// ExtractUsage is the tolerant entry point. It inspects an agent->client message
|
||||
// for billable usage. Never panics on garbage — returns HasUsage=false instead.
|
||||
func ExtractUsage(msg *Message) ParsedUsage {
|
||||
if msg == nil {
|
||||
return ParsedUsage{}
|
||||
}
|
||||
// 1) session/prompt RESPONSE result.usage — the reliable billable source.
|
||||
if len(msg.Result) > 0 {
|
||||
if p, ok := parsePromptResponseUsage(msg.Result); ok {
|
||||
return p
|
||||
}
|
||||
}
|
||||
// 2) usage_update notification — opportunistic cost only (used/size is a
|
||||
// context-window gauge, NOT a billable token delta; do not sum it).
|
||||
if msg.Method == "session/update" && len(msg.Params) > 0 {
|
||||
if p, ok := parseUsageUpdateCost(msg.Params); ok {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return ParsedUsage{}
|
||||
}
|
||||
|
||||
func parsePromptResponseUsage(result json.RawMessage) (ParsedUsage, bool) {
|
||||
var r promptResponseResult
|
||||
if err := json.Unmarshal(result, &r); err != nil || r.Usage == nil {
|
||||
return ParsedUsage{}, false
|
||||
}
|
||||
u := r.Usage
|
||||
prompt := firstNonNil(u.InputTokens, u.InputTokensCamel)
|
||||
completion := firstNonNil(u.OutputTokens, u.OutputTokensCamel)
|
||||
thinking := firstNonNil(u.ThoughtTokens, u.ThinkingTokens)
|
||||
// cachedReadTokens/cachedWriteTokens are reported but not separately priced;
|
||||
// fold cached writes into prompt-equivalent only if no input token present.
|
||||
if prompt == 0 && completion == 0 && thinking == 0 {
|
||||
// No billable token signal in the usage object.
|
||||
return ParsedUsage{}, false
|
||||
}
|
||||
return ParsedUsage{
|
||||
PromptTokens: prompt,
|
||||
CompletionTokens: completion,
|
||||
ThinkingTokens: thinking,
|
||||
HasUsage: true,
|
||||
}, true
|
||||
}
|
||||
|
||||
func parseUsageUpdateCost(params json.RawMessage) (ParsedUsage, bool) {
|
||||
var p usageUpdateParams
|
||||
if err := json.Unmarshal(params, &p); err != nil || p.Update == nil {
|
||||
return ParsedUsage{}, false
|
||||
}
|
||||
if p.Update.SessionUpdate != "usage_update" || p.Update.Cost == nil {
|
||||
return ParsedUsage{}, false
|
||||
}
|
||||
amt := p.Update.Cost.Amount
|
||||
if amt <= 0 {
|
||||
return ParsedUsage{}, false
|
||||
}
|
||||
return ParsedUsage{CostUSD: &amt, HasUsage: true}, true
|
||||
}
|
||||
|
||||
func firstNonNil(a, b *int) int {
|
||||
if a != nil {
|
||||
return *a
|
||||
}
|
||||
if b != nil {
|
||||
return *b
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// ===== usageAccumulator (per-session) =====
|
||||
|
||||
// usageRepo is the narrow repo surface the accumulator needs.
|
||||
type usageRepo interface {
|
||||
InsertSessionUsage(ctx context.Context, r *SessionUsageRecord) (bool, error)
|
||||
AddSessionUsageTotals(ctx context.Context, sid uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error)
|
||||
SumProjectCostUSD(ctx context.Context, projectID uuid.UUID, since time.Time) (float64, error)
|
||||
}
|
||||
|
||||
// pgUsageAccumulator is the concrete per-session accumulator. Thread-safe: the
|
||||
// relay reader is single-goroutine today, but Observe is guarded so future
|
||||
// concurrent callers (or tests) stay correct.
|
||||
type pgUsageAccumulator struct {
|
||||
repo usageRepo
|
||||
prices ModelPriceLookup
|
||||
log *slog.Logger
|
||||
now func() time.Time
|
||||
mu sync.Mutex
|
||||
dedup map[int64]struct{}
|
||||
|
||||
sessionID uuid.UUID
|
||||
userID uuid.UUID
|
||||
projectID uuid.UUID
|
||||
agentKindID uuid.UUID
|
||||
modelID *uuid.UUID
|
||||
startedAt time.Time
|
||||
caps BudgetCaps
|
||||
// projectBudgetUSD (>0) enforces a cumulative per-project ACP spend soft cap
|
||||
// over [projectSince, now]; 0 disables it.
|
||||
projectBudgetUSD float64
|
||||
projectSince time.Time
|
||||
}
|
||||
|
||||
// AccumulatorParams seeds a per-session accumulator.
|
||||
type AccumulatorParams struct {
|
||||
Repo usageRepo
|
||||
Prices ModelPriceLookup
|
||||
Log *slog.Logger
|
||||
SessionID uuid.UUID
|
||||
UserID uuid.UUID
|
||||
ProjectID uuid.UUID
|
||||
AgentKindID uuid.UUID
|
||||
ModelID *uuid.UUID
|
||||
StartedAt time.Time
|
||||
Caps BudgetCaps
|
||||
// ProjectBudgetUSD, when >0, enforces a per-project soft cap on cumulative
|
||||
// ACP spend since ProjectSince.
|
||||
ProjectBudgetUSD float64
|
||||
ProjectSince time.Time
|
||||
Now func() time.Time
|
||||
}
|
||||
|
||||
func newUsageAccumulator(p AccumulatorParams) *pgUsageAccumulator {
|
||||
log := p.Log
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
now := p.Now
|
||||
if now == nil {
|
||||
now = time.Now
|
||||
}
|
||||
since := p.ProjectSince
|
||||
if since.IsZero() {
|
||||
since = now().Add(-30 * 24 * time.Hour)
|
||||
}
|
||||
return &pgUsageAccumulator{
|
||||
repo: p.Repo,
|
||||
prices: p.Prices,
|
||||
log: log,
|
||||
now: now,
|
||||
dedup: map[int64]struct{}{},
|
||||
sessionID: p.SessionID,
|
||||
userID: p.UserID,
|
||||
projectID: p.ProjectID,
|
||||
agentKindID: p.AgentKindID,
|
||||
modelID: p.ModelID,
|
||||
startedAt: p.StartedAt,
|
||||
caps: p.Caps,
|
||||
projectBudgetUSD: p.ProjectBudgetUSD,
|
||||
projectSince: since,
|
||||
}
|
||||
}
|
||||
|
||||
func (a *pgUsageAccumulator) Observe(ctx context.Context, p ParsedUsage, sourceEventID int64) (BreachReason, bool) {
|
||||
if !p.HasUsage {
|
||||
return "", false
|
||||
}
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
|
||||
// In-memory de-dup guard (DB unique index is the durable de-dup).
|
||||
if sourceEventID > 0 {
|
||||
if _, seen := a.dedup[sourceEventID]; seen {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
cost := a.price(ctx, p)
|
||||
|
||||
// Persist ledger row first (idempotent per source_event_id), then accumulate
|
||||
// session totals only when the ledger row was actually inserted (no re-add on
|
||||
// a duplicate event).
|
||||
var srcID *int64
|
||||
if sourceEventID > 0 {
|
||||
v := sourceEventID
|
||||
srcID = &v
|
||||
}
|
||||
inserted, err := a.repo.InsertSessionUsage(ctx, &SessionUsageRecord{
|
||||
SessionID: a.sessionID,
|
||||
UserID: a.userID,
|
||||
ProjectID: a.projectID,
|
||||
AgentKindID: a.agentKindID,
|
||||
ModelID: a.modelID,
|
||||
PromptTokens: int64(p.PromptTokens),
|
||||
CompletionTokens: int64(p.CompletionTokens),
|
||||
ThinkingTokens: int64(p.ThinkingTokens),
|
||||
CostUSD: cost,
|
||||
SourceEventID: srcID,
|
||||
})
|
||||
if err != nil {
|
||||
a.log.Error("acp.usage.insert_ledger", "session_id", a.sessionID, "err", err.Error())
|
||||
return "", false
|
||||
}
|
||||
if !inserted {
|
||||
// Duplicate event id — already accounted.
|
||||
if sourceEventID > 0 {
|
||||
a.dedup[sourceEventID] = struct{}{}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
if sourceEventID > 0 {
|
||||
a.dedup[sourceEventID] = struct{}{}
|
||||
}
|
||||
|
||||
totPrompt, totCompletion, totThinking, totCost, err := a.repo.AddSessionUsageTotals(
|
||||
ctx, a.sessionID,
|
||||
int64(p.PromptTokens), int64(p.CompletionTokens), int64(p.ThinkingTokens), cost)
|
||||
if err != nil {
|
||||
a.log.Error("acp.usage.add_totals", "session_id", a.sessionID, "err", err.Error())
|
||||
return "", false
|
||||
}
|
||||
|
||||
return a.checkBreach(ctx, totPrompt, totCompletion, totThinking, totCost)
|
||||
}
|
||||
|
||||
// price computes the turn cost: prefer the configured model price; otherwise
|
||||
// fall back to the agent-self-reported cost; otherwise 0 (unpriced).
|
||||
func (a *pgUsageAccumulator) price(ctx context.Context, p ParsedUsage) float64 {
|
||||
if a.modelID != nil && a.prices != nil {
|
||||
if mp, err := a.prices.PriceForModel(ctx, *a.modelID); err == nil && mp.Found {
|
||||
return (float64(p.PromptTokens)*mp.PromptPerM +
|
||||
float64(p.CompletionTokens)*mp.CompletionPerM +
|
||||
float64(p.ThinkingTokens)*mp.ThinkingPerM) / 1_000_000
|
||||
}
|
||||
}
|
||||
if p.CostUSD != nil {
|
||||
return *p.CostUSD
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (a *pgUsageAccumulator) checkBreach(ctx context.Context, totPrompt, totCompletion, totThinking int64, totCost float64) (BreachReason, bool) {
|
||||
// Wall-clock: enforced here too (the reaper is the periodic backstop).
|
||||
if a.caps.MaxWallClockSeconds != nil {
|
||||
elapsed := a.now().Sub(a.startedAt)
|
||||
if elapsed >= time.Duration(*a.caps.MaxWallClockSeconds)*time.Second {
|
||||
return BreachWallClock, true
|
||||
}
|
||||
}
|
||||
if a.caps.MaxTokens != nil {
|
||||
if totPrompt+totCompletion+totThinking >= *a.caps.MaxTokens {
|
||||
return BreachTokens, true
|
||||
}
|
||||
}
|
||||
if a.caps.MaxCostUSD != nil && totCost >= *a.caps.MaxCostUSD {
|
||||
return BreachCost, true
|
||||
}
|
||||
// Per-project soft cap (advisory kill, slight overshoot tolerated under
|
||||
// concurrency — spec §7 risk). Only queried when a project budget is set.
|
||||
if a.projectBudgetUSD > 0 {
|
||||
if projCost, err := a.repo.SumProjectCostUSD(ctx, a.projectID, a.projectSince); err == nil {
|
||||
if projCost >= a.projectBudgetUSD {
|
||||
return BreachCost, true
|
||||
}
|
||||
} else {
|
||||
a.log.Warn("acp.usage.project_cost_check", "project_id", a.projectID, "err", err.Error())
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
Reference in New Issue
Block a user