git.iko.soy/crush
on main help sign in
fix(agent): correct fallback usage accounting
commit2e9c65057b8ca5abc8aa54b21856b869c733d52c authorGreg Slepak <contact@taoeffect.com> dateTue 12 May 2026, 15:56 (5 months ago) parent6ed8852
3 files changed+71−6
compare with main
unifiedsplit
internal/agent/agent.go+2−4 internal/agent/usage_fallback.go+6−2 internal/agent/usage_fallback_test.go+63
internal/agent/agent.go+2−4
@@ -1166,11 +1166,9 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
11661166 }
11671167
11681168 session.Cost += cost
1169− if usage.OutputTokens != 0 {
1169+ if !usageIsZero(usage) {
11701170 session.CompletionTokens = usage.OutputTokens
1171− }
1172− if promptTokens := usage.InputTokens + usage.CacheReadTokens; promptTokens != 0 {
1173− session.PromptTokens = promptTokens
1171+ session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
11741172 }
11751173}
11761174
internal/agent/usage_fallback.go+6−2
@@ -78,9 +78,13 @@ func estimateStepCompletionTokens(step fantasy.StepResult) int64 {
7878 case *fantasy.ToolCallContent:
7979 tokens += estimateToolCallTokens(c.ToolName, c.Input)
8080 case fantasy.ToolResultContent:
81− tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
81+ if c.ProviderExecuted {
82+ tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
83+ }
8284 case *fantasy.ToolResultContent:
83− tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
85+ if c.ProviderExecuted {
86+ tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
87+ }
8488 }
8589 }
8690 return tokens
internal/agent/usage_fallback_test.go+63
@@ -145,6 +145,51 @@ func TestFallbackStepUsageEstimatesToolResults(t *testing.T) {
145145 require.Equal(t, usage.InputTokens, usage.TotalTokens)
146146}
147147
148+func TestFallbackStepUsageSkipsClientToolResultsAsOutput(t *testing.T) {
149+ t.Parallel()
150+
151+ step := fantasy.StepResult{
152+ Response: fantasy.Response{
153+ Content: fantasy.ResponseContent{
154+ fantasy.ToolResultContent{
155+ ToolCallID: "tool-call-1",
156+ ToolName: "bash",
157+ Result: fantasy.ToolResultOutputContentText{
158+ Text: "large client-executed payload that should not count as model output tokens",
159+ },
160+ },
161+ },
162+ },
163+ }
164+
165+ usage, estimated := fallbackStepUsage(nil, step)
166+ require.False(t, estimated)
167+ require.Zero(t, usage.OutputTokens)
168+}
169+
170+func TestFallbackStepUsageCountsProviderToolResultsAsOutput(t *testing.T) {
171+ t.Parallel()
172+
173+ step := fantasy.StepResult{
174+ Response: fantasy.Response{
175+ Content: fantasy.ResponseContent{
176+ fantasy.ToolResultContent{
177+ ToolCallID: "tool-call-1",
178+ ToolName: "web_search",
179+ ProviderExecuted: true,
180+ ClientMetadata: "provider metadata",
181+ Result: fantasy.ToolResultOutputContentText{Text: "provider-executed result"},
182+ },
183+ },
184+ },
185+ }
186+
187+ usage, estimated := fallbackStepUsage(nil, step)
188+ require.True(t, estimated)
189+ require.Positive(t, usage.OutputTokens)
190+ require.Equal(t, usage.OutputTokens, usage.TotalTokens)
191+}
192+
148193func TestFallbackStepUsageReturnsZeroWithoutContent(t *testing.T) {
149194 t.Parallel()
150195
@@ -187,6 +232,24 @@ func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) {
187232 require.Equal(t, int64(456), currentSession.CompletionTokens)
188233}
189234
235+func TestUpdateSessionUsageReplacesCountersForPartialUsage(t *testing.T) {
236+ t.Parallel()
237+
238+ agent := &sessionAgent{}
239+ currentSession := &session.Session{
240+ ID: "session-id",
241+ PromptTokens: 123,
242+ CompletionTokens: 456,
243+ }
244+ model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
245+ usage := fantasy.Usage{InputTokens: 789}
246+
247+ agent.updateSessionUsage(model, currentSession, usage, nil, false)
248+
249+ require.Equal(t, int64(789), currentSession.PromptTokens)
250+ require.Zero(t, currentSession.CompletionTokens)
251+}
252+
190253func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {
191254 t.Parallel()
192255