git.iko.soy/crush
on main help sign in
fix(agent): estimate missing streamed usage
6ed8852Greg Slepak
internal/agent/agent.go+30−12
compare with main
unifiedsplit
@@ -261,6 +261,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
261 a.eventPromptSent(call.SessionID)
261 a.eventPromptSent(call.SessionID)
262
262
263 var currentAssistant *message.Message
263 var currentAssistant *message.Message
264+ var stepMessages []fantasy.Message
264 var shouldSummarize bool
265 var shouldSummarize bool
265 // Don't send MaxOutputTokens if 0 — some providers (e.g. LM Studio) reject it
266 // Don't send MaxOutputTokens if 0 — some providers (e.g. LM Studio) reject it
266 var maxOutputTokens *int64
267 var maxOutputTokens *int64
@@ -319,6 +320,10 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
319 prepared.Messages = append([]fantasy.Message{fantasy.NewSystemMessage(promptPrefix)}, prepared.Messages...)
320 prepared.Messages = append([]fantasy.Message{fantasy.NewSystemMessage(promptPrefix)}, prepared.Messages...)
320 }
321 }
321
322
323+ sessionLock.Lock()
324+ stepMessages = cloneFantasyMessages(prepared.Messages)
325+ sessionLock.Unlock()
326+
322 var assistantMsg message.Message
327 var assistantMsg message.Message
323 assistantMsg, err = a.messages.Create(callContext, call.SessionID, message.CreateMessageParams{
328 assistantMsg, err = a.messages.Create(callContext, call.SessionID, message.CreateMessageParams{
324 Role: message.Assistant,
329 Role: message.Assistant,
@@ -444,7 +449,8 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy
444 if getSessionErr != nil {
449 if getSessionErr != nil {
445 return getSessionErr
450 return getSessionErr
446 }
451 }
452+ usage, estimated := fallbackStepUsage(stepMessages, stepResult)
447− a.updateSessionUsage(largeModel, &updatedSession, stepResult.Usage, a.openrouterCost(stepResult.ProviderMetadata))
453+ a.updateSessionUsage(largeModel, &updatedSession, usage, a.openrouterCost(stepResult.ProviderMetadata), estimated)
448 _, sessionErr := a.sessions.Save(ctx, updatedSession)
454 _, sessionErr := a.sessions.Save(ctx, updatedSession)
449 if sessionErr != nil {
455 if sessionErr != nil {
450 return sessionErr
456 return sessionErr
@@ -749,7 +755,7 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
749 }
755 }
750 }
756 }
751
757
752− a.updateSessionUsage(largeModel, &currentSession, resp.TotalUsage, openrouterCost)
758+ a.updateSessionUsage(largeModel, &currentSession, resp.TotalUsage, openrouterCost, false)
753
759
754 // Just in case, get just the last usage info.
760 // Just in case, get just the last usage info.
755 usage := resp.Response.Usage
761 usage := resp.Response.Usage
@@ -1132,28 +1138,40 @@ func (a *sessionAgent) openrouterCost(metadata fantasy.ProviderMetadata) *float6
1132 return &opts.Usage.Cost
1138 return &opts.Usage.Cost
1133}
1139}
1134
1140
1135−func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session, usage fantasy.Usage, overrideCost *float64) {
1141+func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session, usage fantasy.Usage, overrideCost *float64, estimated bool) {
1136 modelConfig := model.CatwalkCfg
1142 modelConfig := model.CatwalkCfg
1137 cost := modelConfig.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
1143 cost := modelConfig.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
1138 modelConfig.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
1144 modelConfig.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
1139 modelConfig.CostPer1MIn/1e6*float64(usage.InputTokens) +
1145 modelConfig.CostPer1MIn/1e6*float64(usage.InputTokens) +
1140 modelConfig.CostPer1MOut/1e6*float64(usage.OutputTokens)
1146 modelConfig.CostPer1MOut/1e6*float64(usage.OutputTokens)
1141
1147
1142− a.eventTokensUsed(session.ID, model, usage, cost)
1148+ eventCost := cost
1143−
1149+ if estimated {
1144− // Use override cost if available (e.g., from OpenRouter).
1150+ eventCost = 0
1145− if overrideCost != nil {
1146− cost = *overrideCost
1147 }
1151 }
1152+ a.eventTokensUsed(session.ID, model, usage, eventCost)
1148
1153
1149− // Skip cost accumulation
1154+ if estimated {
1150− if model.FlatRate {
1151 cost = 0
1155 cost = 0
1156+ } else {
1157+ // Use override cost if available (e.g., from OpenRouter).
1158+ if overrideCost != nil {
1159+ cost = *overrideCost
1160+ }
1161+
1162+ // Skip cost accumulation
1163+ if model.FlatRate {
1164+ cost = 0
1165+ }
1152 }
1166 }
1153
1167
1154 session.Cost += cost
1168 session.Cost += cost
1169+ if usage.OutputTokens != 0 {
1155− session.CompletionTokens = usage.OutputTokens
1170+ session.CompletionTokens = usage.OutputTokens
1171+ }
1156− session.PromptTokens = usage.InputTokens + usage.CacheReadTokens
1172+ if promptTokens := usage.InputTokens + usage.CacheReadTokens; promptTokens != 0 {
1173+ session.PromptTokens = promptTokens
1174+ }
1157}
1175}
1158
1176
1159func (a *sessionAgent) Cancel(sessionID string) {
1177func (a *sessionAgent) Cancel(sessionID string) {
internal/agent/usage_fallback.gonew file+172
@@ -0,0 +1,172 @@
1+package agent
2+
3+import (
4+ "fmt"
5+
6+ "charm.land/fantasy"
7+)
8+
9+func usageIsZero(usage fantasy.Usage) bool {
10+ return usage.InputTokens == 0 &&
11+ usage.OutputTokens == 0 &&
12+ usage.TotalTokens == 0 &&
13+ usage.ReasoningTokens == 0 &&
14+ usage.CacheCreationTokens == 0 &&
15+ usage.CacheReadTokens == 0
16+}
17+
18+func fallbackStepUsage(messages []fantasy.Message, step fantasy.StepResult) (fantasy.Usage, bool) {
19+ if !usageIsZero(step.Usage) {
20+ return step.Usage, false
21+ }
22+
23+ inputTokens := estimateMessageTokens(messages)
24+ outputTokens := estimateStepCompletionTokens(step)
25+ if inputTokens == 0 && outputTokens == 0 {
26+ return fantasy.Usage{}, false
27+ }
28+
29+ return fantasy.Usage{
30+ InputTokens: inputTokens,
31+ OutputTokens: outputTokens,
32+ TotalTokens: inputTokens + outputTokens,
33+ }, true
34+}
35+
36+func cloneFantasyMessages(messages []fantasy.Message) []fantasy.Message {
37+ cloned := make([]fantasy.Message, len(messages))
38+ for i, msg := range messages {
39+ cloned[i] = msg
40+ cloned[i].Content = append([]fantasy.MessagePart(nil), msg.Content...)
41+ }
42+ return cloned
43+}
44+
45+func estimateMessageTokens(messages []fantasy.Message) int64 {
46+ var tokens int64
47+ for _, msg := range messages {
48+ tokens += approxTokenCount(string(msg.Role))
49+ for _, part := range msg.Content {
50+ tokens += estimateMessagePartTokens(part)
51+ }
52+ }
53+ return tokens
54+}
55+
56+func estimateStepCompletionTokens(step fantasy.StepResult) int64 {
57+ var tokens int64
58+ for _, content := range step.Content {
59+ switch c := content.(type) {
60+ case fantasy.TextContent:
61+ tokens += approxTokenCount(c.Text)
62+ case *fantasy.TextContent:
63+ tokens += approxTokenCount(c.Text)
64+ case fantasy.ReasoningContent:
65+ tokens += approxTokenCount(c.Text)
66+ case *fantasy.ReasoningContent:
67+ tokens += approxTokenCount(c.Text)
68+ case fantasy.FileContent:
69+ tokens += estimateGeneratedFileTokens(c)
70+ case *fantasy.FileContent:
71+ tokens += estimateGeneratedFileTokens(*c)
72+ case fantasy.SourceContent:
73+ tokens += estimateSourceTokens(c)
74+ case *fantasy.SourceContent:
75+ tokens += estimateSourceTokens(*c)
76+ case fantasy.ToolCallContent:
77+ tokens += estimateToolCallTokens(c.ToolName, c.Input)
78+ case *fantasy.ToolCallContent:
79+ tokens += estimateToolCallTokens(c.ToolName, c.Input)
80+ case fantasy.ToolResultContent:
81+ tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
82+ case *fantasy.ToolResultContent:
83+ tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result)
84+ }
85+ }
86+ return tokens
87+}
88+
89+func estimateMessagePartTokens(part fantasy.MessagePart) int64 {
90+ switch p := part.(type) {
91+ case fantasy.TextPart:
92+ return approxTokenCount(p.Text)
93+ case *fantasy.TextPart:
94+ return approxTokenCount(p.Text)
95+ case fantasy.ReasoningPart:
96+ return approxTokenCount(p.Text)
97+ case *fantasy.ReasoningPart:
98+ return approxTokenCount(p.Text)
99+ case fantasy.FilePart:
100+ return estimateFilePartTokens(p)
101+ case *fantasy.FilePart:
102+ return estimateFilePartTokens(*p)
103+ case fantasy.ToolCallPart:
104+ return estimateToolCallTokens(p.ToolName, p.Input)
105+ case *fantasy.ToolCallPart:
106+ return estimateToolCallTokens(p.ToolName, p.Input)
107+ case fantasy.ToolResultPart:
108+ return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output)
109+ case *fantasy.ToolResultPart:
110+ return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output)
111+ default:
112+ return 0
113+ }
114+}
115+
116+func estimateToolCallTokens(toolName, input string) int64 {
117+ return approxTokenCount(toolName) + approxTokenCount(input)
118+}
119+
120+func estimateToolResultContentTokens(toolCallID, toolName, metadata string, output fantasy.ToolResultOutputContent) int64 {
121+ tokens := approxTokenCount(toolCallID) + approxTokenCount(toolName) + approxTokenCount(metadata)
122+ switch result := output.(type) {
123+ case fantasy.ToolResultOutputContentText:
124+ tokens += approxTokenCount(result.Text)
125+ case *fantasy.ToolResultOutputContentText:
126+ tokens += approxTokenCount(result.Text)
127+ case fantasy.ToolResultOutputContentError:
128+ if result.Error != nil {
129+ tokens += approxTokenCount(result.Error.Error())
130+ }
131+ case *fantasy.ToolResultOutputContentError:
132+ if result.Error != nil {
133+ tokens += approxTokenCount(result.Error.Error())
134+ }
135+ case fantasy.ToolResultOutputContentMedia:
136+ tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data))
137+ case *fantasy.ToolResultOutputContentMedia:
138+ tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data))
139+ }
140+ return tokens
141+}
142+
143+func estimateFilePartTokens(file fantasy.FilePart) int64 {
144+ return estimateMediaTokens(file.MediaType, file.Filename, len(file.Data))
145+}
146+
147+func estimateGeneratedFileTokens(file fantasy.FileContent) int64 {
148+ return estimateMediaTokens(file.MediaType, "", len(file.Data))
149+}
150+
151+func estimateMediaTokens(mediaType, text string, dataBytes int) int64 {
152+ if dataBytes == 0 {
153+ return approxTokenCount(mediaType) + approxTokenCount(text)
154+ }
155+ return approxTokenCount(fmt.Sprintf("%s %s %d bytes", mediaType, text, dataBytes))
156+}
157+
158+func estimateSourceTokens(source fantasy.SourceContent) int64 {
159+ return approxTokenCount(string(source.SourceType)) +
160+ approxTokenCount(source.ID) +
161+ approxTokenCount(source.URL) +
162+ approxTokenCount(source.Title) +
163+ approxTokenCount(source.MediaType) +
164+ approxTokenCount(source.Filename)
165+}
166+
167+func approxTokenCount(s string) int64 {
168+ if s == "" {
169+ return 0
170+ }
171+ return int64((len(s) + 3) / 4)
172+}
internal/agent/usage_fallback_test.gonew file+203
@@ -0,0 +1,203 @@
1+package agent
2+
3+import (
4+ "errors"
5+ "testing"
6+
7+ "charm.land/catwalk/pkg/catwalk"
8+ "charm.land/fantasy"
9+ "github.com/charmbracelet/crush/internal/session"
10+ "github.com/stretchr/testify/require"
11+)
12+
13+func TestUsageIsZero(t *testing.T) {
14+ t.Parallel()
15+
16+ require.True(t, usageIsZero(fantasy.Usage{}))
17+ require.False(t, usageIsZero(fantasy.Usage{InputTokens: 1}))
18+ require.False(t, usageIsZero(fantasy.Usage{OutputTokens: 1}))
19+ require.False(t, usageIsZero(fantasy.Usage{TotalTokens: 1}))
20+ require.False(t, usageIsZero(fantasy.Usage{ReasoningTokens: 1}))
21+ require.False(t, usageIsZero(fantasy.Usage{CacheCreationTokens: 1}))
22+ require.False(t, usageIsZero(fantasy.Usage{CacheReadTokens: 1}))
23+}
24+
25+func TestFallbackStepUsageKeepsProviderUsage(t *testing.T) {
26+ t.Parallel()
27+
28+ usage := fantasy.Usage{
29+ InputTokens: 10,
30+ OutputTokens: 5,
31+ TotalTokens: 15,
32+ }
33+ step := fantasy.StepResult{
34+ Response: fantasy.Response{Usage: usage},
35+ }
36+
37+ fallbackUsage, estimated := fallbackStepUsage(nil, step)
38+ require.False(t, estimated)
39+ require.Equal(t, usage, fallbackUsage)
40+}
41+
42+func TestFallbackStepUsageEstimatesPromptAndAssistantText(t *testing.T) {
43+ t.Parallel()
44+
45+ messages := []fantasy.Message{
46+ fantasy.NewUserMessage("please explain the implementation details"),
47+ }
48+ step := fantasy.StepResult{
49+ Response: fantasy.Response{
50+ Content: fantasy.ResponseContent{
51+ fantasy.TextContent{Text: "the implementation stores state safely"},
52+ },
53+ },
54+ }
55+
56+ usage, estimated := fallbackStepUsage(messages, step)
57+ require.True(t, estimated)
58+ require.Positive(t, usage.InputTokens)
59+ require.Positive(t, usage.OutputTokens)
60+ require.Equal(t, usage.InputTokens+usage.OutputTokens, usage.TotalTokens)
61+}
62+
63+func TestFallbackStepUsageEstimatesReasoning(t *testing.T) {
64+ t.Parallel()
65+
66+ messages := []fantasy.Message{
67+ {
68+ Role: fantasy.MessageRoleAssistant,
69+ Content: []fantasy.MessagePart{
70+ fantasy.ReasoningPart{Text: "first reason about the request"},
71+ },
72+ },
73+ }
74+ step := fantasy.StepResult{
75+ Response: fantasy.Response{
76+ Content: fantasy.ResponseContent{
77+ fantasy.ReasoningContent{Text: "second reason about the answer"},
78+ },
79+ },
80+ }
81+
82+ usage, estimated := fallbackStepUsage(messages, step)
83+ require.True(t, estimated)
84+ require.Positive(t, usage.InputTokens)
85+ require.Positive(t, usage.OutputTokens)
86+}
87+
88+func TestFallbackStepUsageEstimatesToolCalls(t *testing.T) {
89+ t.Parallel()
90+
91+ step := fantasy.StepResult{
92+ Response: fantasy.Response{
93+ Content: fantasy.ResponseContent{
94+ fantasy.ToolCallContent{
95+ ToolCallID: "tool-call-1",
96+ ToolName: "view",
97+ Input: `{"file_path":"/tmp/example.go"}`,
98+ },
99+ },
100+ },
101+ }
102+
103+ usage, estimated := fallbackStepUsage(nil, step)
104+ require.True(t, estimated)
105+ require.Zero(t, usage.InputTokens)
106+ require.Positive(t, usage.OutputTokens)
107+ require.Equal(t, usage.OutputTokens, usage.TotalTokens)
108+}
109+
110+func TestFallbackStepUsageEstimatesToolResults(t *testing.T) {
111+ t.Parallel()
112+
113+ messages := []fantasy.Message{
114+ {
115+ Role: fantasy.MessageRoleTool,
116+ Content: []fantasy.MessagePart{
117+ fantasy.ToolResultPart{
118+ ToolCallID: "tool-call-1",
119+ Output: fantasy.ToolResultOutputContentText{
120+ Text: "file contents returned by the tool",
121+ },
122+ },
123+ fantasy.ToolResultPart{
124+ ToolCallID: "tool-call-2",
125+ Output: fantasy.ToolResultOutputContentError{
126+ Error: errors.New("permission denied"),
127+ },
128+ },
129+ fantasy.ToolResultPart{
130+ ToolCallID: "tool-call-3",
131+ Output: fantasy.ToolResultOutputContentMedia{
132+ MediaType: "image/png",
133+ Text: "screenshot",
134+ Data: "abc123",
135+ },
136+ },
137+ },
138+ },
139+ }
140+
141+ usage, estimated := fallbackStepUsage(messages, fantasy.StepResult{})
142+ require.True(t, estimated)
143+ require.Positive(t, usage.InputTokens)
144+ require.Zero(t, usage.OutputTokens)
145+ require.Equal(t, usage.InputTokens, usage.TotalTokens)
146+}
147+
148+func TestFallbackStepUsageReturnsZeroWithoutContent(t *testing.T) {
149+ t.Parallel()
150+
151+ usage, estimated := fallbackStepUsage(nil, fantasy.StepResult{})
152+ require.False(t, estimated)
153+ require.True(t, usageIsZero(usage))
154+}
155+
156+func TestUpdateSessionUsageSkipsEstimatedCost(t *testing.T) {
157+ t.Parallel()
158+
159+ agent := &sessionAgent{}
160+ currentSession := &session.Session{ID: "session-id", Cost: 1.25}
161+ model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
162+ usage := fantasy.Usage{InputTokens: 1000, OutputTokens: 2000}
163+
164+ agent.updateSessionUsage(model, currentSession, usage, nil, true)
165+
166+ require.Equal(t, 1.25, currentSession.Cost)
167+ require.Equal(t, int64(1000), currentSession.PromptTokens)
168+ require.Equal(t, int64(2000), currentSession.CompletionTokens)
169+}
170+
171+func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) {
172+ t.Parallel()
173+
174+ agent := &sessionAgent{}
175+ currentSession := &session.Session{
176+ ID: "session-id",
177+ PromptTokens: 123,
178+ CompletionTokens: 456,
179+ Cost: 1.25,
180+ }
181+ model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
182+
183+ agent.updateSessionUsage(model, currentSession, fantasy.Usage{}, nil, false)
184+
185+ require.Equal(t, 1.25, currentSession.Cost)
186+ require.Equal(t, int64(123), currentSession.PromptTokens)
187+ require.Equal(t, int64(456), currentSession.CompletionTokens)
188+}
189+
190+func TestUpdateSessionUsageAddsProviderCost(t *testing.T) {
191+ t.Parallel()
192+
193+ agent := &sessionAgent{}
194+ currentSession := &session.Session{ID: "session-id", Cost: 1.25}
195+ model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}}
196+ usage := fantasy.Usage{InputTokens: 1000, OutputTokens: 2000}
197+
198+ agent.updateSessionUsage(model, currentSession, usage, nil, false)
199+
200+ require.Equal(t, 1.3, currentSession.Cost)
201+ require.Equal(t, int64(1000), currentSession.PromptTokens)
202+ require.Equal(t, int64(2000), currentSession.CompletionTokens)
203+}