Skip to content

Commit 0378171

Browse files
riturajFiriturajFi
committed
token count fix
1 parent 7b8e38d commit 0378171

File tree

3 files changed

+162
-179
lines changed

3 files changed

+162
-179
lines changed

‎pkg/runtime/runtime.go‎

Lines changed: 87 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -536,18 +536,18 @@ func (r *LocalRuntime) Run(ctx context.Context, sess *session.Session) ([]sessio
536536
}
537537

538538
func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStream, a *agent.Agent, agentTools []tools.Tool, sess *session.Session, m *modelsdev.Model, events chan Event) (streamResult, error) {
539-
defer stream.Close()
540-
541-
var fullContent strings.Builder
542-
var fullReasoningContent strings.Builder
543-
var thinkingSignature string
544-
var toolCalls []tools.ToolCall
545-
// Track which tool call indices we've already emitted partial events for
546-
emittedPartialEvents := make(map[string]bool)
547-
// Track the latest observed self usage and cost for this streaming call.
548-
// We will add these to lifetime totals once the call completes.
549-
var lastSelfInput, lastSelfOutput int
550-
var lastCallCost float64
539+
defer stream.Close()
540+
541+
var fullContent strings.Builder
542+
var fullReasoningContent strings.Builder
543+
var thinkingSignature string
544+
var toolCalls []tools.ToolCall
545+
// Track which tool call indices we've already emitted partial events for
546+
emittedPartialEvents := make(map[string]bool)
547+
// Track the latest observed self usage and cost for this streaming call.
548+
// We will add these to lifetime totals once the call completes.
549+
var lastSelfInput, lastSelfOutput int
550+
var lastCallCost float64
551551

552552
for {
553553
response, err := stream.Recv()
@@ -558,31 +558,31 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
558558
return streamResult{Stopped: true}, fmt.Errorf("error receiving from stream: %w", err)
559559
}
560560

561-
if response.Usage != nil {
562-
childInputTotal, childOutputTotal := childTokenTotals(sess)
561+
if response.Usage != nil {
562+
childInputTotal, childOutputTotal := childTokenTotals(sess)
563563

564-
selfInput := response.Usage.InputTokens + response.Usage.CachedInputTokens
565-
selfOutput := response.Usage.OutputTokens + response.Usage.CachedOutputTokens + response.Usage.ReasoningTokens
564+
selfInput := response.Usage.InputTokens + response.Usage.CachedInputTokens
565+
selfOutput := response.Usage.OutputTokens + response.Usage.CachedOutputTokens + response.Usage.ReasoningTokens
566566

567567
var callCost float64
568-
if m != nil {
569-
callCost = (float64(response.Usage.InputTokens)*m.Cost.Input +
570-
float64(response.Usage.OutputTokens+response.Usage.ReasoningTokens)*m.Cost.Output +
571-
float64(response.Usage.CachedInputTokens)*m.Cost.CacheRead +
572-
float64(response.Usage.CachedOutputTokens)*m.Cost.CacheWrite) / 1e6
573-
sess.Cost += callCost
574-
}
575-
576-
sess.SelfCost = callCost
577-
// Update per-call snapshot fields used for UI live view and compaction heuristics
578-
sess.SelfInputTokens = selfInput
579-
sess.SelfOutputTokens = selfOutput
580-
sess.InputTokens = childInputTotal + selfInput
581-
sess.OutputTokens = childOutputTotal + selfOutput
582-
// Remember latest self snapshot for adding to lifetime totals when the call finishes
583-
lastSelfInput = selfInput
584-
lastSelfOutput = selfOutput
585-
lastCallCost = callCost
568+
if m != nil {
569+
callCost = (float64(response.Usage.InputTokens)*m.Cost.Input +
570+
float64(response.Usage.OutputTokens+response.Usage.ReasoningTokens)*m.Cost.Output +
571+
float64(response.Usage.CachedInputTokens)*m.Cost.CacheRead +
572+
float64(response.Usage.CachedOutputTokens)*m.Cost.CacheWrite) / 1e6
573+
sess.Cost += callCost
574+
}
575+
576+
sess.SelfCost = callCost
577+
// Update per-call snapshot fields used for UI live view and compaction heuristics
578+
sess.SelfInputTokens = selfInput
579+
sess.SelfOutputTokens = selfOutput
580+
sess.InputTokens = childInputTotal + selfInput
581+
sess.OutputTokens = childOutputTotal + selfOutput
582+
// Remember latest self snapshot for adding to lifetime totals when the call finishes
583+
lastSelfInput = selfInput
584+
lastSelfOutput = selfOutput
585+
lastCallCost = callCost
586586

587587
modelName := "unknown"
588588
if m != nil {
@@ -591,27 +591,27 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
591591
telemetry.RecordTokenUsage(ctx, modelName, int64(response.Usage.InputTokens), int64(response.Usage.OutputTokens+response.Usage.ReasoningTokens), sess.Cost)
592592
}
593593

594-
if len(response.Choices) == 0 {
595-
continue
596-
}
597-
choice := response.Choices[0]
598-
if choice.FinishReason == chat.FinishReasonStop || choice.FinishReason == chat.FinishReasonLength {
599-
// On finish, fold this call's final self usage into lifetime totals.
600-
if lastSelfInput > 0 || lastSelfOutput > 0 {
601-
sess.TotalInputTokens += lastSelfInput
602-
sess.TotalOutputTokens += lastSelfOutput
603-
sess.SelfTotalInputTokens += lastSelfInput
604-
sess.SelfTotalOutputTokens += lastSelfOutput
605-
sess.SelfTotalCost += lastCallCost
606-
}
607-
return streamResult{
608-
Calls: toolCalls,
609-
Content: fullContent.String(),
610-
ReasoningContent: fullReasoningContent.String(),
611-
ThinkingSignature: thinkingSignature,
612-
Stopped: true,
613-
}, nil
614-
}
594+
if len(response.Choices) == 0 {
595+
continue
596+
}
597+
choice := response.Choices[0]
598+
if choice.FinishReason == chat.FinishReasonStop || choice.FinishReason == chat.FinishReasonLength {
599+
// On finish, fold this call's final self usage into lifetime totals.
600+
if lastSelfInput > 0 || lastSelfOutput > 0 {
601+
sess.TotalInputTokens += lastSelfInput
602+
sess.TotalOutputTokens += lastSelfOutput
603+
sess.SelfTotalInputTokens += lastSelfInput
604+
sess.SelfTotalOutputTokens += lastSelfOutput
605+
sess.SelfTotalCost += lastCallCost
606+
}
607+
return streamResult{
608+
Calls: toolCalls,
609+
Content: fullContent.String(),
610+
ReasoningContent: fullReasoningContent.String(),
611+
ThinkingSignature: thinkingSignature,
612+
Stopped: true,
613+
}, nil
614+
}
615615

616616
// Handle tool calls
617617
if len(choice.Delta.ToolCalls) > 0 {
@@ -704,30 +704,30 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
704704

705705
// buildInclusiveUsageSnapshot captures the session's current inclusive usage in the shared event format.
706706
func buildInclusiveUsageSnapshot(sess *session.Session, contextLimit int) *Usage {
707-
// Build a live lifetime-inclusive snapshot: lifetime totals + current in-flight self snapshot
708-
liveInput := sess.TotalInputTokens + sess.SelfInputTokens
709-
liveOutput := sess.TotalOutputTokens + sess.SelfOutputTokens
710-
return &Usage{
711-
ContextLength: liveInput + liveOutput,
712-
ContextLimit: contextLimit,
713-
InputTokens: liveInput,
714-
OutputTokens: liveOutput,
715-
Cost: sess.Cost,
716-
}
707+
// Build a live lifetime-inclusive snapshot: lifetime totals + current in-flight self snapshot
708+
liveInput := sess.TotalInputTokens + sess.SelfInputTokens
709+
liveOutput := sess.TotalOutputTokens + sess.SelfOutputTokens
710+
return &Usage{
711+
ContextLength: liveInput + liveOutput,
712+
ContextLimit: contextLimit,
713+
InputTokens: liveInput,
714+
OutputTokens: liveOutput,
715+
Cost: sess.Cost,
716+
}
717717
}
718718

719719
func buildSelfUsageSnapshot(sess *session.Session, contextLimit int) *Usage {
720-
// Build a live self-only lifetime snapshot: cumulative self totals + current in-flight self snapshot
721-
liveInput := sess.SelfTotalInputTokens + sess.SelfInputTokens
722-
liveOutput := sess.SelfTotalOutputTokens + sess.SelfOutputTokens
723-
liveCost := sess.SelfTotalCost + sess.SelfCost
724-
return &Usage{
725-
ContextLength: liveInput + liveOutput,
726-
ContextLimit: contextLimit,
727-
InputTokens: liveInput,
728-
OutputTokens: liveOutput,
729-
Cost: liveCost,
730-
}
720+
// Build a live self-only lifetime snapshot: cumulative self totals + current in-flight self snapshot
721+
liveInput := sess.SelfTotalInputTokens + sess.SelfInputTokens
722+
liveOutput := sess.SelfTotalOutputTokens + sess.SelfOutputTokens
723+
liveCost := sess.SelfTotalCost + sess.SelfCost
724+
return &Usage{
725+
ContextLength: liveInput + liveOutput,
726+
ContextLimit: contextLimit,
727+
InputTokens: liveInput,
728+
OutputTokens: liveOutput,
729+
Cost: liveCost,
730+
}
731731
}
732732

733733
func childTokenTotals(sess *session.Session) (int, int) {
@@ -1064,17 +1064,17 @@ func (r *LocalRuntime) handleTaskTransfer(ctx context.Context, sess *session.Ses
10641064

10651065
sess.ToolsApproved = s.ToolsApproved
10661066

1067-
sess.Cost += s.Cost
1068-
// Mirror cost behavior: once the child finishes, fold its token usage into the parent.
1069-
// Keep per-turn inclusive fields for compaction heuristics, and separately track lifetime totals.
1070-
childInputTotal, childOutputTotal := childTokenTotals(sess)
1071-
childInputTotal += s.InputTokens
1072-
childOutputTotal += s.OutputTokens
1073-
sess.InputTokens = childInputTotal + sess.SelfInputTokens
1074-
sess.OutputTokens = childOutputTotal + sess.SelfOutputTokens
1075-
// Lifetime cumulative totals include completed child session totals
1076-
sess.TotalInputTokens += s.TotalInputTokens
1077-
sess.TotalOutputTokens += s.TotalOutputTokens
1067+
sess.Cost += s.Cost
1068+
// Mirror cost behavior: once the child finishes, fold its token usage into the parent.
1069+
// Keep per-turn inclusive fields for compaction heuristics, and separately track lifetime totals.
1070+
childInputTotal, childOutputTotal := childTokenTotals(sess)
1071+
childInputTotal += s.InputTokens
1072+
childOutputTotal += s.OutputTokens
1073+
sess.InputTokens = childInputTotal + sess.SelfInputTokens
1074+
sess.OutputTokens = childOutputTotal + sess.SelfOutputTokens
1075+
// Lifetime cumulative totals include completed child session totals
1076+
sess.TotalInputTokens += s.TotalInputTokens
1077+
sess.TotalOutputTokens += s.TotalOutputTokens
10781078

10791079
sess.AddSubSession(s)
10801080

‎pkg/session/session.go‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -71,15 +71,15 @@ type Session struct {
7171
// Lifetime/cumulative token counters across the entire session (self + children).
7272
TotalInputTokens int `json:"total_input_tokens,omitempty"`
7373
TotalOutputTokens int `json:"total_output_tokens,omitempty"`
74-
// Self* fields track the most recent provider-reported usage for this session only (no children).
75-
SelfInputTokens int `json:"self_input_tokens"`
76-
SelfOutputTokens int `json:"self_output_tokens"`
77-
SelfCost float64 `json:"self_cost"`
78-
79-
// SelfTotal* fields track the cumulative self-only usage for this session across all passes.
80-
SelfTotalInputTokens int `json:"self_total_input_tokens,omitempty"`
81-
SelfTotalOutputTokens int `json:"self_total_output_tokens,omitempty"`
82-
SelfTotalCost float64 `json:"self_total_cost,omitempty"`
74+
// Self* fields track the most recent provider-reported usage for this session only (no children).
75+
SelfInputTokens int `json:"self_input_tokens"`
76+
SelfOutputTokens int `json:"self_output_tokens"`
77+
SelfCost float64 `json:"self_cost"`
78+
79+
// SelfTotal* fields track the cumulative self-only usage for this session across all passes.
80+
SelfTotalInputTokens int `json:"self_total_input_tokens,omitempty"`
81+
SelfTotalOutputTokens int `json:"self_total_output_tokens,omitempty"`
82+
SelfTotalCost float64 `json:"self_total_cost,omitempty"`
8383
}
8484

8585
// Message is a message from an agent

0 commit comments

Comments
 (0)