You've already forked agentic-coding-workflow
bugfix
This commit is contained in:
@@ -0,0 +1,260 @@
|
||||
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 }
|
||||
Reference in New Issue
Block a user