diff --git a/pkg/agent/pipeline_llm.go b/pkg/agent/pipeline_llm.go index 685647713..b4c87634e 100644 --- a/pkg/agent/pipeline_llm.go +++ b/pkg/agent/pipeline_llm.go @@ -495,11 +495,16 @@ func (p *Pipeline) CallLLM( } } - // Save finishReason to turnState for SubTurn truncation detection - if innerTS := turnStateFromContext(ctx); innerTS != nil { - innerTS.SetLastFinishReason(exec.response.FinishReason) + // Save finishReason and usage on the turn state. Use ts directly (the + // authoritative turn state for this call) rather than a context lookup: + // the raw ctx passed to CallLLM is not seeded with turnState (only turnCtx + // is), so turnStateFromContext(ctx) returns nil here and silently dropped + // both the finish reason and the per-turn token usage. ts is also exactly + // what the streaming publisher reads via GetLastUsage at finalize. + if ts != nil { + ts.SetLastFinishReason(exec.response.FinishReason) if exec.response.Usage != nil { - innerTS.SetLastUsage(exec.response.Usage) + ts.SetLastUsage(exec.response.Usage) } } diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index d7c6553a2..bea8154f8 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -699,18 +699,34 @@ func setStreamerModelName(streamer any, modelName string) { setter.SetModelName(modelName) } +type turnUsageStreamer interface { + SetTurnUsage(inputTokens, outputTokens int) +} + +// setStreamerTurnUsage forwards real per-turn token usage to a streamer that +// supports it, transparently unwrapping the manager's streamer wrappers. +func setStreamerTurnUsage(streamer any, inputTokens, outputTokens int) { + setter, ok := streamer.(turnUsageStreamer) + if !ok { + return + } + setter.SetTurnUsage(inputTokens, outputTokens) +} + // splitMarkerStreamer turns accumulated streaming text containing // MessageSplitMarker into separate channel stream messages. type splitMarkerStreamer struct { - mu sync.Mutex - current bus.Streamer - reasoning bus.ReasoningStreamer - begin func(context.Context) (bus.Streamer, error) - completedParts int - finalized bool - onFinalize func(context.Context, string) - clearMarker func() - modelName string + mu sync.Mutex + current bus.Streamer + reasoning bus.ReasoningStreamer + begin func(context.Context) (bus.Streamer, error) + completedParts int + finalized bool + onFinalize func(context.Context, string) + clearMarker func() + modelName string + turnInputTokens int + turnOutputTokens int } func (s *splitMarkerStreamer) Update(ctx context.Context, content string) error { @@ -761,6 +777,14 @@ func (s *splitMarkerStreamer) SetModelName(modelName string) { setStreamerModelName(s.reasoning, s.modelName) } +func (s *splitMarkerStreamer) SetTurnUsage(inputTokens, outputTokens int) { + s.mu.Lock() + defer s.mu.Unlock() + s.turnInputTokens = inputTokens + s.turnOutputTokens = outputTokens + setStreamerTurnUsage(s.current, s.turnInputTokens, s.turnOutputTokens) +} + func (s *splitMarkerStreamer) Cancel(ctx context.Context) { s.mu.Lock() defer s.mu.Unlock() @@ -840,6 +864,7 @@ func (s *splitMarkerStreamer) ensureCurrentLocked(ctx context.Context) error { } s.current = streamer setStreamerModelName(s.current, s.modelName) + setStreamerTurnUsage(s.current, s.turnInputTokens, s.turnOutputTokens) return nil } @@ -928,6 +953,10 @@ func (s *finalizeHookStreamer) SetModelName(modelName string) { setStreamerModelName(s.Streamer, strings.TrimSpace(modelName)) } +func (s *finalizeHookStreamer) SetTurnUsage(inputTokens, outputTokens int) { + setStreamerTurnUsage(s.Streamer, inputTokens, outputTokens) +} + func (s *finalizeHookStreamer) runFinalizeHook(ctx context.Context, content string) { if s.onFinalize != nil { s.onFinalize(ctx, content) diff --git a/pkg/channels/manager_test.go b/pkg/channels/manager_test.go index a7ceac67d..533d49ecb 100644 --- a/pkg/channels/manager_test.go +++ b/pkg/channels/manager_test.go @@ -3383,3 +3383,56 @@ func TestManager_SendPlaceholder(t *testing.T) { t.Error("expected SendPlaceholder to fail for unknown channel") } } + +// turnUsageTrackingStreamer is a mockStreamer that records SetTurnUsage calls, +// used to verify the manager's streamer wrappers forward per-turn token usage +// to the inner streamer (regression: the wrappers previously dropped it because +// SetTurnUsage is not part of the bus.Streamer interface). +type turnUsageTrackingStreamer struct { + mockStreamer + inputTokens int + outputTokens int + usageCalls int +} + +func (m *turnUsageTrackingStreamer) SetTurnUsage(inputTokens, outputTokens int) { + m.usageCalls++ + m.inputTokens = inputTokens + m.outputTokens = outputTokens +} + +func TestFinalizeHookStreamerForwardsTurnUsage(t *testing.T) { + inner := &turnUsageTrackingStreamer{} + wrapper := &finalizeHookStreamer{Streamer: inner} + + setter, ok := any(wrapper).(turnUsageStreamer) + if !ok { + t.Fatal("finalizeHookStreamer does not satisfy turnUsageStreamer") + } + setter.SetTurnUsage(1234, 567) + + if inner.usageCalls != 1 { + t.Fatalf("inner SetTurnUsage calls = %d, want 1", inner.usageCalls) + } + if inner.inputTokens != 1234 || inner.outputTokens != 567 { + t.Errorf("inner usage = (%d, %d), want (1234, 567)", inner.inputTokens, inner.outputTokens) + } +} + +func TestSplitMarkerStreamerForwardsTurnUsage(t *testing.T) { + inner := &turnUsageTrackingStreamer{} + wrapper := &splitMarkerStreamer{current: inner} + + setter, ok := any(wrapper).(turnUsageStreamer) + if !ok { + t.Fatal("splitMarkerStreamer does not satisfy turnUsageStreamer") + } + setter.SetTurnUsage(1234, 567) + + if inner.usageCalls != 1 { + t.Fatalf("inner SetTurnUsage calls = %d, want 1", inner.usageCalls) + } + if inner.inputTokens != 1234 || inner.outputTokens != 567 { + t.Errorf("inner usage = (%d, %d), want (1234, 567)", inner.inputTokens, inner.outputTokens) + } +}