git.iko.soy/crush
on main help sign in
fix(agent): correct fallback usage accounting
2e9c650Greg Slepak
internal/agent/agent.go+2−4
compare with main
unifiedsplit
@@ -1166,11 +1166,9 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,
1166 }
1166 }
1167
1167
1168 session.Cost += cost
1168 session.Cost += cost
1169− if usage.OutputTokens != 0 {
1169+ if !usageIsZero(usage) {
1170 session.CompletionTokens = usage.OutputTokens
1170 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
1174 }
1172 }
1175}
1173}
1176
1174
internal/agent/usage_fallback.go+6−2
@@ -78,9 +78,13 @@ func estimateStepCompletionTokens(step fantasy.StepResult) int64 {
78 case *fantasy.ToolCallContent:
78 case *fantasy.ToolCallContent:
79 tokens += estimateToolCallTokens(c.ToolName, c.Input)
79 tokens += estimateToolCallTokens(c.ToolName, c.Input)
80 case fantasy.ToolResultContent:
80 case fantasy.ToolResultContent:
81+ if c.ProviderExecuted {
81− tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
82+ tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
83+ }
82 case *fantasy.ToolResultContent:
84 case *fantasy.ToolResultContent:
85+ if c.ProviderExecuted {
83− tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
86+ tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
87+ }
84 }
88 }
85 }
89 }
86 return tokens
90 return tokens
internal/agent/usage_fallback_test.go+63
@@ -145,6 +145,51 @@ func TestFallbackStepUsageEstimatesToolResults(t *testing.T) {
145 require.Equal(t, usage.InputTokens, usage.TotalTokens)
145 require.Equal(t, usage.InputTokens, usage.TotalTokens)
146}
146}
147
147
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+
148func TestFallbackStepUsageReturnsZeroWithoutContent(t *testing.T) {
193func TestFallbackStepUsageReturnsZeroWithoutContent(t *testing.T) {
149 t.Parallel()
194 t.Parallel()
150
195
@@ -187,6 +232,24 @@ func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) {
187 require.Equal(t, int64(456), currentSession.CompletionTokens)
232 require.Equal(t, int64(456), currentSession.CompletionTokens)
188}
233}
189
234
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+
190func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {
253func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {
191 t.Parallel()
254 t.Parallel()
192
255