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

package acp

import (
	"context"
	"encoding/json"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func msgWithResult(t *testing.T, result string) *Message {
	t.Helper()
	return &Message{JSONRPC: "2.0", ID: json.RawMessage(`"client-prompt-1"`), Result: json.RawMessage(result)}
}

func msgNotification(t *testing.T, method, params string) *Message {
	t.Helper()
	return &Message{JSONRPC: "2.0", Method: method, Params: json.RawMessage(params)}
}

func TestExtractUsage(t *testing.T) {
	tests := []struct {
		name       string
		msg        *Message
		wantHas    bool
		wantPrompt int
		wantComp   int
		wantThink  int
		wantCost   *float64
	}{
		{
			name:       "snake_case prompt response",
			msg:        msgWithResult(t, `{"stopReason":"end_turn","usage":{"input_tokens":120,"output_tokens":48}}`),
			wantHas:    true,
			wantPrompt: 120,
			wantComp:   48,
		},
		{
			name:       "camelCase v2 usage with thoughtTokens + cached",
			msg:        msgWithResult(t, `{"stopReason":"end_turn","usage":{"inputTokens":200,"outputTokens":90,"thoughtTokens":15,"cachedReadTokens":50}}`),
			wantHas:    true,
			wantPrompt: 200,
			wantComp:   90,
			wantThink:  15,
		},
		{
			name:     "usage_update notification with cost only",
			msg:      msgNotification(t, "session/update", `{"update":{"sessionUpdate":"usage_update","used":1000,"size":200000,"cost":{"amount":0.0123,"currency":"USD"}}}`),
			wantHas:  true,
			wantCost: f64ptr(0.0123),
		},
		{
			name:    "usage_update without cost -> no usage (gauge ignored)",
			msg:     msgNotification(t, "session/update", `{"update":{"sessionUpdate":"usage_update","used":1000,"size":200000}}`),
			wantHas: false,
		},
		{
			name:    "ordinary agent_message_chunk -> no usage",
			msg:     msgNotification(t, "session/update", `{"update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"hi"}}}`),
			wantHas: false,
		},
		{
			name:    "malformed result -> no panic, no usage",
			msg:     msgWithResult(t, `{"usage":not-json`),
			wantHas: false,
		},
		{
			name:    "result without usage object -> no usage",
			msg:     msgWithResult(t, `{"stopReason":"end_turn"}`),
			wantHas: false,
		},
		{
			name:    "nil message -> no usage",
			msg:     nil,
			wantHas: false,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			got := ExtractUsage(tc.msg)
			assert.Equal(t, tc.wantHas, got.HasUsage)
			if !tc.wantHas {
				return
			}
			assert.Equal(t, tc.wantPrompt, got.PromptTokens)
			assert.Equal(t, tc.wantComp, got.CompletionTokens)
			assert.Equal(t, tc.wantThink, got.ThinkingTokens)
			if tc.wantCost == nil {
				assert.Nil(t, got.CostUSD)
			} else {
				require.NotNil(t, got.CostUSD)
				assert.InDelta(t, *tc.wantCost, *got.CostUSD, 1e-9)
			}
		})
	}
}

func f64ptr(f float64) *float64 { return &f }
func i64ptr(i int64) *int64     { return &i }

// ===== usageAccumulator =====

type fakeUsageRepo struct {
	ledger      []*SessionUsageRecord
	dedup       map[int64]struct{}
	totPrompt   int64
	totComp     int64
	totThink    int64
	totCost     float64
	projectCost float64
}

func (f *fakeUsageRepo) InsertSessionUsage(_ context.Context, r *SessionUsageRecord) (bool, error) {
	if r.SourceEventID != nil {
		if f.dedup == nil {
			f.dedup = map[int64]struct{}{}
		}
		if _, ok := f.dedup[*r.SourceEventID]; ok {
			return false, nil
		}
		f.dedup[*r.SourceEventID] = struct{}{}
	}
	cp := *r
	f.ledger = append(f.ledger, &cp)
	return true, nil
}

func (f *fakeUsageRepo) AddSessionUsageTotals(_ context.Context, _ uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error) {
	f.totPrompt += dp
	f.totComp += dc
	f.totThink += dt
	f.totCost += dCost
	return f.totPrompt, f.totComp, f.totThink, f.totCost, nil
}

func (f *fakeUsageRepo) SumProjectCostUSD(_ context.Context, _ uuid.UUID, _ time.Time) (float64, error) {
	return f.projectCost, nil
}

type fakePriceLookup struct {
	price ModelPrice
}

func (f fakePriceLookup) PriceForModel(_ context.Context, _ uuid.UUID) (ModelPrice, error) {
	return f.price, nil
}

func TestAccumulator_PricesViaModelLookup(t *testing.T) {
	repo := &fakeUsageRepo{}
	mid := uuid.New()
	acc := newUsageAccumulator(AccumulatorParams{
		Repo:    repo,
		Prices:  fakePriceLookup{price: ModelPrice{PromptPerM: 3, CompletionPerM: 15, ThinkingPerM: 0, Found: true}},
		ModelID: &mid,
		Caps:    BudgetCaps{},
	})
	_, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 1_000_000, CompletionTokens: 1_000_000, HasUsage: true}, 1)
	assert.False(t, breached)
	require.Len(t, repo.ledger, 1)
	// 3 USD (prompt) + 15 USD (completion) = 18.
	assert.InDelta(t, 18.0, repo.ledger[0].CostUSD, 1e-6)
	assert.InDelta(t, 18.0, repo.totCost, 1e-6)
}

func TestAccumulator_FallsBackToAgentCost(t *testing.T) {
	repo := &fakeUsageRepo{}
	acc := newUsageAccumulator(AccumulatorParams{Repo: repo}) // no model price
	cost := 0.5
	_, _ = acc.Observe(context.Background(), ParsedUsage{CostUSD: &cost, HasUsage: true}, 1)
	require.Len(t, repo.ledger, 1)
	assert.InDelta(t, 0.5, repo.ledger[0].CostUSD, 1e-9)
}

func TestAccumulator_SkipsWhenNoUsage(t *testing.T) {
	repo := &fakeUsageRepo{}
	acc := newUsageAccumulator(AccumulatorParams{Repo: repo})
	_, breached := acc.Observe(context.Background(), ParsedUsage{HasUsage: false}, 1)
	assert.False(t, breached)
	assert.Empty(t, repo.ledger)
}

func TestAccumulator_DedupBySourceEventID(t *testing.T) {
	repo := &fakeUsageRepo{}
	mid := uuid.New()
	acc := newUsageAccumulator(AccumulatorParams{
		Repo:    repo,
		Prices:  fakePriceLookup{price: ModelPrice{CompletionPerM: 10, Found: true}},
		ModelID: &mid,
	})
	p := ParsedUsage{CompletionTokens: 1000, HasUsage: true}
	_, _ = acc.Observe(context.Background(), p, 7)
	_, _ = acc.Observe(context.Background(), p, 7) // same event id -> no-op
	assert.Len(t, repo.ledger, 1)
	assert.InDelta(t, 0.01, repo.totCost, 1e-9)
}

func TestAccumulator_BreachCost(t *testing.T) {
	repo := &fakeUsageRepo{}
	mid := uuid.New()
	acc := newUsageAccumulator(AccumulatorParams{
		Repo:    repo,
		Prices:  fakePriceLookup{price: ModelPrice{CompletionPerM: 1000, Found: true}},
		ModelID: &mid,
		Caps:    BudgetCaps{MaxCostUSD: f64ptr(0.5)},
	})
	// 1000 completion tokens * 1000/M = 1.0 USD >= 0.5 cap.
	reason, breached := acc.Observe(context.Background(), ParsedUsage{CompletionTokens: 1000, HasUsage: true}, 1)
	assert.True(t, breached)
	assert.Equal(t, BreachCost, reason)
}

func TestAccumulator_BreachTokens(t *testing.T) {
	repo := &fakeUsageRepo{}
	acc := newUsageAccumulator(AccumulatorParams{
		Repo: repo,
		Caps: BudgetCaps{MaxTokens: i64ptr(1000)},
	})
	reason, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 600, CompletionTokens: 600, HasUsage: true}, 1)
	assert.True(t, breached)
	assert.Equal(t, BreachTokens, reason)
}

func TestAccumulator_BreachWallClock(t *testing.T) {
	repo := &fakeUsageRepo{}
	start := time.Now().Add(-2 * time.Hour)
	acc := newUsageAccumulator(AccumulatorParams{
		Repo:      repo,
		StartedAt: start,
		Caps:      BudgetCaps{MaxWallClockSeconds: i32ptr(60)}, // 1 minute cap, 2h elapsed
	})
	reason, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 1, HasUsage: true}, 1)
	assert.True(t, breached)
	assert.Equal(t, BreachWallClock, reason)
}

func TestAccumulator_BreachProjectBudget(t *testing.T) {
	repo := &fakeUsageRepo{projectCost: 100} // already over the project cap
	acc := newUsageAccumulator(AccumulatorParams{
		Repo:             repo,
		ProjectBudgetUSD: 50,
	})
	reason, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 1, HasUsage: true}, 1)
	assert.True(t, breached)
	assert.Equal(t, BreachCost, reason)
}

func TestAccumulator_NoBreachUnderCaps(t *testing.T) {
	repo := &fakeUsageRepo{}
	acc := newUsageAccumulator(AccumulatorParams{
		Repo: repo,
		Caps: BudgetCaps{MaxTokens: i64ptr(1_000_000), MaxCostUSD: f64ptr(100)},
	})
	_, breached := acc.Observe(context.Background(), ParsedUsage{PromptTokens: 10, CompletionTokens: 10, HasUsage: true}, 1)
	assert.False(t, breached)
}

func i32ptr(i int32) *int32 { return &i }