git.iko.soy/crush
on main help sign in
fix(agent): harden fallback usage accounting
74e6e37Greg Slepak
internal/agent/agent.go+20−7
compare with main
unifiedsplit
@@ -760,7 +760,9 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
760 // Just in case, get just the last usage info.
760 // Just in case, get just the last usage info.
761 usage := resp.Response.Usage
761 usage := resp.Response.Usage
762 currentSession.SummaryMessageID = summaryMessage.ID
762 currentSession.SummaryMessageID = summaryMessage.ID
763+ if completionTokens := summaryCompletionTokens(usage, summaryMessage); completionTokens != 0 {
763− currentSession.CompletionTokens = usage.OutputTokens
764+ currentSession.CompletionTokens = completionTokens
765+ }
764 currentSession.PromptTokens = 0
766 currentSession.PromptTokens = 0
765 _, err = a.sessions.Save(genCtx, currentSession)
767 _, err = a.sessions.Save(genCtx, currentSession)
766 if err != nil {
768 if err != nil {
@@ -1145,11 +1147,9 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
1145 modelConfig.CostPer1MIn/1e6*float64(usage.InputTokens) +
1147 modelConfig.CostPer1MIn/1e6*float64(usage.InputTokens) +
1146 modelConfig.CostPer1MOut/1e6*float64(usage.OutputTokens)
1148 modelConfig.CostPer1MOut/1e6*float64(usage.OutputTokens)
1147
1149
1148− eventCost := cost
1149− if estimated {
1150+ if !estimated {
1150− eventCost = 0
1151+ a.eventTokensUsed(session.ID, model, usage, cost)
1151 }
1152 }
1152− a.eventTokensUsed(session.ID, model, usage, eventCost)
1153
1153
1154 if estimated {
1154 if estimated {
1155 cost = 0
1155 cost = 0
@@ -1166,10 +1166,23 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
1166 }
1166 }
1167
1167
1168 session.Cost += cost
1168 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 {
1170 session.CompletionTokens = usage.OutputTokens
1174 session.CompletionTokens = usage.OutputTokens
1171− session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
1172 }
1175 }
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())
1173}
1186}
1174
1187
1175func (a *sessionAgent) Cancel(sessionID string) {
1188func (a *sessionAgent) Cancel(sessionID string) {
internal/agent/usage_fallback_test.go+73−2
@@ -6,6 +6,7 @@ import (
6
6
7 "charm.land/catwalk/pkg/catwalk"
7 "charm.land/catwalk/pkg/catwalk"
8 "charm.land/fantasy"
8 "charm.land/fantasy"
9+ "github.com/charmbracelet/crush/internal/message"
9 "github.com/charmbracelet/crush/internal/session"
10 "github.com/charmbracelet/crush/internal/session"
10 "github.com/stretchr/testify/require"
11 "github.com/stretchr/testify/require"
11)
12)
@@ -232,7 +233,7 @@ func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) {
232 require.Equal(t, int64(456), currentSession.CompletionTokens)
233 require.Equal(t, int64(456), currentSession.CompletionTokens)
233}
234}
234
235
235−func TestUpdateSessionUsageReplacesCountersForPartialUsage(t *testing.T) {
236+func TestUpdateSessionUsagePreservesOmittedCountersForPartialUsage(t *testing.T) {
236 t.Parallel()
237 t.Parallel()
237
238
238 agent := &sessionAgent{}
239 agent := &sessionAgent{}
@@ -247,7 +248,77 @@ func TestUpdateSessionUsageReplacesCountersForPartialUsage(t *testing.T) {
247 agent.updateSessionUsage(model, currentSession, usage, nil, false)
248 agent.updateSessionUsage(model, currentSession, usage, nil, false)
248
249
249 require.Equal(t, int64(789), currentSession.PromptTokens)
250 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{}))
251}
322}
252
323
253func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {
324func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {