agentic-coding-workflow/internal/acp/usage.go

// 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
}