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
3 changed files with 100 additions and 13 deletions
+9 -4
View File
@@ -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)
}
}
+38 -9
View File
@@ -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)
+53
View File
@@ -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)
}
}