fix(pico): deliver per-turn token usage to the streamer

Two bugs prevented the usage block from ever reaching the wire:

1. CallLLM read turnStateFromContext(ctx), but the raw ctx is not seeded
   with the turn state (only turnCtx is), so SetLastUsage/SetLastFinishReason
   were dropped — GetLastUsage() returned nil at finalize. Set them on the
   ts parameter directly, which is also what the streaming publisher reads.

2. The manager wraps the channel streamer in finalizeHookStreamer /
   splitMarkerStreamer, neither of which forwarded SetTurnUsage (it is not
   part of the bus.Streamer interface), so the type assertion in the
   publisher's Finalize failed silently. Mirror the existing SetModelName
   forwarding: add a turnUsageStreamer interface + setStreamerTurnUsage
   helper and SetTurnUsage methods on both wrappers (splitMarker also stores
   and re-applies usage to each freshly-begun part streamer).

Adds regression tests asserting both wrappers forward SetTurnUsage.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Andy Lo-A-Foe 2026-06-22 15:38:46 +02:00
parent 697e94fe8c
commit cc7b4ca86b
No known key found for this signature in database
GPG key ID: C0E4EB79E9E6A23D
3 changed files with 100 additions and 13 deletions

View file

@ -495,11 +495,16 @@ func (p *Pipeline) CallLLM(
} }
} }
// Save finishReason to turnState for SubTurn truncation detection // Save finishReason and usage on the turn state. Use ts directly (the
if innerTS := turnStateFromContext(ctx); innerTS != nil { // authoritative turn state for this call) rather than a context lookup:
innerTS.SetLastFinishReason(exec.response.FinishReason) // 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 { if exec.response.Usage != nil {
innerTS.SetLastUsage(exec.response.Usage) ts.SetLastUsage(exec.response.Usage)
} }
} }

View file

@ -699,6 +699,20 @@ func setStreamerModelName(streamer any, modelName string) {
setter.SetModelName(modelName) 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 // splitMarkerStreamer turns accumulated streaming text containing
// MessageSplitMarker into separate channel stream messages. // MessageSplitMarker into separate channel stream messages.
type splitMarkerStreamer struct { type splitMarkerStreamer struct {
@ -711,6 +725,8 @@ type splitMarkerStreamer struct {
onFinalize func(context.Context, string) onFinalize func(context.Context, string)
clearMarker func() clearMarker func()
modelName string modelName string
turnInputTokens int
turnOutputTokens int
} }
func (s *splitMarkerStreamer) Update(ctx context.Context, content string) error { 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) 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) { func (s *splitMarkerStreamer) Cancel(ctx context.Context) {
s.mu.Lock() s.mu.Lock()
defer s.mu.Unlock() defer s.mu.Unlock()
@ -840,6 +864,7 @@ func (s *splitMarkerStreamer) ensureCurrentLocked(ctx context.Context) error {
} }
s.current = streamer s.current = streamer
setStreamerModelName(s.current, s.modelName) setStreamerModelName(s.current, s.modelName)
setStreamerTurnUsage(s.current, s.turnInputTokens, s.turnOutputTokens)
return nil return nil
} }
@ -928,6 +953,10 @@ func (s *finalizeHookStreamer) SetModelName(modelName string) {
setStreamerModelName(s.Streamer, strings.TrimSpace(modelName)) 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) { func (s *finalizeHookStreamer) runFinalizeHook(ctx context.Context, content string) {
if s.onFinalize != nil { if s.onFinalize != nil {
s.onFinalize(ctx, content) s.onFinalize(ctx, content)

View file

@ -3383,3 +3383,56 @@ func TestManager_SendPlaceholder(t *testing.T) {
t.Error("expected SendPlaceholder to fail for unknown channel") 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)
}
}