git.iko.soy/crush
on main help sign in
fix(agent): harden fallback usage accounting
commit74e6e378e37493440cb01c956ef6de7ebcde38e1 authorGreg Slepak <contact@taoeffect.com> dateTue 12 May 2026, 16:27 (5 months ago) parent2e9c650
2 files changed+93−9
compare with main
unifiedsplit
internal/agent/agent.go+20−7 internal/agent/usage_fallback_test.go+73−2
internal/agent/agent.go+20−7
@@ -760,7 +760,9 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
760760 // Just in case, get just the last usage info.
761761 usage := resp.Response.Usage
762762 currentSession.SummaryMessageID = summaryMessage.ID
763− currentSession.CompletionTokens = usage.OutputTokens
763+ if completionTokens := summaryCompletionTokens(usage, summaryMessage); completionTokens != 0 {
764+ currentSession.CompletionTokens = completionTokens
765+ }
764766 currentSession.PromptTokens = 0
765767 _, err = a.sessions.Save(genCtx, currentSession)
766768 if err != nil {
@@ -1145,11 +1147,9 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
11451147 modelConfig.CostPer1MIn/1e6*float64(usage.InputTokens) +
11461148 modelConfig.CostPer1MOut/1e6*float64(usage.OutputTokens)
11471149
1148− eventCost := cost
1149− if estimated {
1150− eventCost = 0
1150+ if !estimated {
1151+ a.eventTokensUsed(session.ID, model, usage, cost)
11511152 }
1152− a.eventTokensUsed(session.ID, model, usage, eventCost)
11531153
11541154 if estimated {
11551155 cost = 0
@@ -1166,10 +1166,23 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
11661166 }
11671167
11681168 session.Cost += cost
1169− if !usageIsZero(usage) {
1169+ updateSessionTokenCounters(session, usage)
1170+}
1171+
1172+func updateSessionTokenCounters(session *session.Session, usage fantasy.Usage) {
1173+ if usage.OutputTokens != 0 {
11701174 session.CompletionTokens = usage.OutputTokens
1171− session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
11721175 }
1176+ if promptTokens := usage.InputTokens + usage.CacheReadTokens; promptTokens != 0 {
1177+ session.PromptTokens = promptTokens
1178+ }
1179+}
1180+
1181+func summaryCompletionTokens(usage fantasy.Usage, summaryMessage message.Message) int64 {
1182+ if usage.OutputTokens != 0 {
1183+ return usage.OutputTokens
1184+ }
1185+ return approxTokenCount(summaryMessage.Content().Text) + approxTokenCount(summaryMessage.ReasoningContent().String())
11731186}
11741187
11751188func (a *sessionAgent) Cancel(sessionID string) {
internal/agent/usage_fallback_test.go+73−2
@@ -6,6 +6,7 @@ import (
66
77 "charm.land/catwalk/pkg/catwalk"
88 "charm.land/fantasy"
9+ "github.com/charmbracelet/crush/internal/message"
910 "github.com/charmbracelet/crush/internal/session"
1011 "github.com/stretchr/testify/require"
1112)
@@ -232,7 +233,7 @@ func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) {
232233 require.Equal(t, int64(456), currentSession.CompletionTokens)
233234}
234235
235−func TestUpdateSessionUsageReplacesCountersForPartialUsage(t *testing.T) {
236+func TestUpdateSessionUsagePreservesOmittedCountersForPartialUsage(t *testing.T) {
236237 t.Parallel()
237238
238239 agent := &sessionAgent{}
@@ -247,7 +248,77 @@ func TestUpdateSessionUsageReplacesCountersForPartialUsage(t *testing.T) {
247248 agent.updateSessionUsage(model, currentSession, usage, nil, false)
248249
249250 require.Equal(t, int64(789), currentSession.PromptTokens)
250− require.Zero(t, currentSession.CompletionTokens)
251+ require.Equal(t, int64(456), currentSession.CompletionTokens)
252+}
253+
254+func TestUpdateSessionUsagePreservesCountersForTotalOnlyUsage(t *testing.T) {
255+ t.Parallel()
256+
257+ agent := &sessionAgent{}
258+ currentSession := &session.Session{
259+ ID: "session-id",
260+ PromptTokens: 123,
261+ CompletionTokens: 456,
262+ }
263+ model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
264+ usage := fantasy.Usage{TotalTokens: 100}
265+
266+ agent.updateSessionUsage(model, currentSession, usage, nil, false)
267+
268+ require.Equal(t, int64(123), currentSession.PromptTokens)
269+ require.Equal(t, int64(456), currentSession.CompletionTokens)
270+}
271+
272+func TestUpdateSessionUsagePreservesPromptForOutputOnlyUsage(t *testing.T) {
273+ t.Parallel()
274+
275+ agent := &sessionAgent{}
276+ currentSession := &session.Session{
277+ ID: "session-id",
278+ PromptTokens: 123,
279+ CompletionTokens: 456,
280+ }
281+ model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
282+ usage := fantasy.Usage{OutputTokens: 50}
283+
284+ agent.updateSessionUsage(model, currentSession, usage, nil, false)
285+
286+ require.Equal(t, int64(123), currentSession.PromptTokens)
287+ require.Equal(t, int64(50), currentSession.CompletionTokens)
288+}
289+
290+func TestUpdateSessionUsageKeepsCountersForEstimatedZeroUsage(t *testing.T) {
291+ t.Parallel()
292+
293+ agent := &sessionAgent{}
294+ currentSession := &session.Session{
295+ ID: "session-id",
296+ PromptTokens: 123,
297+ CompletionTokens: 456,
298+ Cost: 1.25,
299+ }
300+ model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
301+
302+ agent.updateSessionUsage(model, currentSession, fantasy.Usage{}, nil, true)
303+
304+ require.Equal(t, 1.25, currentSession.Cost)
305+ require.Equal(t, int64(123), currentSession.PromptTokens)
306+ require.Equal(t, int64(456), currentSession.CompletionTokens)
307+}
308+
309+func TestSummaryCompletionTokens(t *testing.T) {
310+ t.Parallel()
311+
312+ summaryMessage := message.Message{
313+ Parts: []message.ContentPart{
314+ message.TextContent{Text: "summary text"},
315+ message.ReasoningContent{Thinking: "reasoning text"},
316+ },
317+ }
318+
319+ require.Equal(t, int64(42), summaryCompletionTokens(fantasy.Usage{OutputTokens: 42}, summaryMessage))
320+ require.Equal(t, approxTokenCount("summary text")+approxTokenCount("reasoning text"), summaryCompletionTokens(fantasy.Usage{}, summaryMessage))
321+ require.Zero(t, summaryCompletionTokens(fantasy.Usage{}, message.Message{}))
251322}
252323
253324func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {