From 052c742fe74e5948c08c635bc82cb5ebbf1cc350 Mon Sep 17 00:00:00 2001 From: Andy Lo-A-Foe Date: Mon, 22 Jun 2026 13:07:48 +0200 Subject: [PATCH 1/3] feat(pico): add SetTurnUsage and usage payload helper --- pkg/channels/pico/pico.go | 27 +++++++++++++++++++++ pkg/channels/pico/pico_usage_test.go | 36 ++++++++++++++++++++++++++++ pkg/channels/pico/protocol.go | 1 + 3 files changed, 64 insertions(+) create mode 100644 pkg/channels/pico/pico_usage_test.go diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index 587f889ab..b11befae1 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -531,6 +531,8 @@ type picoStreamer struct { channel *PicoChannel chatID string modelName string + turnInputTokens int + turnOutputTokens int messageID string reasoningID string throttleInterval time.Duration @@ -553,6 +555,17 @@ func (s *picoStreamer) SetModelName(modelName string) { s.modelName = strings.TrimSpace(modelName) } +// SetTurnUsage records the real per-turn LLM token usage to emit on finalize. +func (s *picoStreamer) SetTurnUsage(inputTokens, outputTokens int) { + if s == nil { + return + } + s.mu.Lock() + defer s.mu.Unlock() + s.turnInputTokens = inputTokens + s.turnOutputTokens = outputTokens +} + func (s *picoStreamer) Update(ctx context.Context, content string) error { s.mu.Lock() defer s.mu.Unlock() @@ -1403,6 +1416,20 @@ func setContextUsagePayload(payload map[string]any, u *bus.ContextUsage) { } } +// setTurnUsagePayload attaches real per-turn LLM token usage to the payload. +// Input and output are kept separate (billed at different rates); total is a +// convenience sum. Omitted entirely when both counts are zero. +func setTurnUsagePayload(payload map[string]any, inputTokens, outputTokens int) { + if inputTokens <= 0 && outputTokens <= 0 { + return + } + payload[PayloadKeyUsage] = map[string]any{ + "input_tokens": inputTokens, + "output_tokens": outputTokens, + "total_tokens": inputTokens + outputTokens, + } +} + func picoToolCallsPayload(msg bus.OutboundMessage) ([]utils.VisibleToolCall, bool) { raw := strings.TrimSpace(msg.Context.Raw[PayloadKeyToolCalls]) if raw == "" { diff --git a/pkg/channels/pico/pico_usage_test.go b/pkg/channels/pico/pico_usage_test.go new file mode 100644 index 000000000..ae2d9749a --- /dev/null +++ b/pkg/channels/pico/pico_usage_test.go @@ -0,0 +1,36 @@ +package pico + +import "testing" + +func TestSetTurnUsagePayload(t *testing.T) { + t.Run("populates usage block when counts present", func(t *testing.T) { + payload := map[string]any{PayloadKeyContent: "hi"} + setTurnUsagePayload(payload, 1234, 567) + + raw, ok := payload[PayloadKeyUsage] + if !ok { + t.Fatalf("expected %q key in payload", PayloadKeyUsage) + } + usage, ok := raw.(map[string]any) + if !ok { + t.Fatalf("usage block is not a map: %T", raw) + } + if usage["input_tokens"] != 1234 { + t.Errorf("input_tokens = %v, want 1234", usage["input_tokens"]) + } + if usage["output_tokens"] != 567 { + t.Errorf("output_tokens = %v, want 567", usage["output_tokens"]) + } + if usage["total_tokens"] != 1801 { + t.Errorf("total_tokens = %v, want 1801", usage["total_tokens"]) + } + }) + + t.Run("omits usage block when both counts zero", func(t *testing.T) { + payload := map[string]any{PayloadKeyContent: "hi"} + setTurnUsagePayload(payload, 0, 0) + if _, ok := payload[PayloadKeyUsage]; ok { + t.Errorf("expected no %q key when counts are zero", PayloadKeyUsage) + } + }) +} diff --git a/pkg/channels/pico/protocol.go b/pkg/channels/pico/protocol.go index 0c809430d..ac92913d4 100644 --- a/pkg/channels/pico/protocol.go +++ b/pkg/channels/pico/protocol.go @@ -28,6 +28,7 @@ const ( PayloadKeyPlaceholder = "placeholder" PayloadKeyToolCalls = "tool_calls" PayloadKeyModelName = "model_name" + PayloadKeyUsage = "usage" MessageKindThought = "thought" MessageKindToolCalls = "tool_calls" From 697e94fe8c198ac7cc5d67724cfbc40be912e5e9 Mon Sep 17 00:00:00 2001 From: Andy Lo-A-Foe Date: Mon, 22 Jun 2026 13:15:35 +0200 Subject: [PATCH 2/3] feat(pico): emit real per-turn token usage on finalized message --- pkg/agent/pipeline_streaming.go | 7 ++++++ pkg/channels/pico/pico.go | 14 ++++++++++- pkg/channels/pico/pico_usage_test.go | 35 +++++++++++++++++++++++++++- 3 files changed, 54 insertions(+), 2 deletions(-) diff --git a/pkg/agent/pipeline_streaming.go b/pkg/agent/pipeline_streaming.go index 115cbffcb..5457e4723 100644 --- a/pkg/agent/pipeline_streaming.go +++ b/pkg/agent/pipeline_streaming.go @@ -54,6 +54,7 @@ func (p *Pipeline) tryConfiguredStreamingLLM( channel: ts.channel, chatID: ts.chatID, modelName: exec.llmModelName, + ts: ts, } logger.DebugCF("agent", "configured streaming enabled", map[string]any{ @@ -376,6 +377,7 @@ type streamingChunkPublisher struct { published bool reasoningPublished bool err error + ts *turnState } func (p *streamingChunkPublisher) Update(ctx context.Context, accumulated string) { @@ -445,6 +447,11 @@ func (p *streamingChunkPublisher) Finalize(ctx context.Context, content string, if setter, ok := p.streamer.(interface{ SetModelName(modelName string) }); ok { setter.SetModelName(p.modelName) } + if usage := p.ts.GetLastUsage(); usage != nil { + if setter, ok := p.streamer.(interface{ SetTurnUsage(in, out int) }); ok { + setter.SetTurnUsage(usage.PromptTokens, usage.CompletionTokens) + } + } var err error if streamer, ok := p.streamer.(bus.ContextUsageStreamer); ok { err = streamer.FinalizeWithContext(ctx, content, contextUsage) diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index b11befae1..a39143e33 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -105,6 +105,8 @@ type PicoChannel struct { cancel context.CancelFunc progress *channels.ToolFeedbackAnimator deleteMessageFn func(context.Context, string, string) error + // broadcastFn lets tests intercept outbound broadcasts. nil → broadcastToSession. + broadcastFn func(chatID string, msg PicoMessage) error } // NewPicoChannel creates a new Pico Protocol channel. @@ -674,8 +676,9 @@ func (s *picoStreamer) sendLocked(ctx context.Context, content string, contextUs payload[PayloadKeyModelName] = s.modelName } setContextUsagePayload(payload, contextUsage) + setTurnUsagePayload(payload, s.turnInputTokens, s.turnOutputTokens) outMsg := newMessage(TypeMessageCreate, payload) - if err := s.channel.broadcastToSession(s.chatID, outMsg); err != nil { + if err := s.channel.broadcast(s.chatID, outMsg); err != nil { return err } } else if content != s.lastContent || contextUsage != nil { @@ -686,6 +689,7 @@ func (s *picoStreamer) sendLocked(ctx context.Context, content string, contextUs if s.modelName != "" { payload[PayloadKeyModelName] = s.modelName } + setTurnUsagePayload(payload, s.turnInputTokens, s.turnOutputTokens) if err := s.channel.editMessagePayload(ctx, s.chatID, s.messageID, payload, contextUsage); err != nil { return err } @@ -945,6 +949,14 @@ func (c *PicoChannel) handleMediaDownload(w http.ResponseWriter, r *http.Request http.ServeContent(w, r, filename, info.ModTime(), file) } +// broadcast routes through broadcastFn when set (tests), else broadcastToSession. +func (c *PicoChannel) broadcast(chatID string, msg PicoMessage) error { + if c.broadcastFn != nil { + return c.broadcastFn(chatID, msg) + } + return c.broadcastToSession(chatID, msg) +} + // broadcastToSession sends a message to all connections with a matching session. func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error { // chatID format: "pico:" diff --git a/pkg/channels/pico/pico_usage_test.go b/pkg/channels/pico/pico_usage_test.go index ae2d9749a..444fb49f6 100644 --- a/pkg/channels/pico/pico_usage_test.go +++ b/pkg/channels/pico/pico_usage_test.go @@ -1,6 +1,9 @@ package pico -import "testing" +import ( + "context" + "testing" +) func TestSetTurnUsagePayload(t *testing.T) { t.Run("populates usage block when counts present", func(t *testing.T) { @@ -34,3 +37,33 @@ func TestSetTurnUsagePayload(t *testing.T) { } }) } + +// newCaptureStreamer returns a streamer whose broadcasts are captured into the +// returned map pointer, so tests need no live websocket. +func newCaptureStreamer() (*picoStreamer, *map[string]any) { + var last map[string]any + ch := &PicoChannel{} + ch.broadcastFn = func(chatID string, msg PicoMessage) error { + last = msg.Payload + return nil + } + s := &picoStreamer{channel: ch, chatID: "c1"} + return s, &last +} + +func TestStreamerEmitsUsageOnFinalize(t *testing.T) { + s, last := newCaptureStreamer() + s.SetTurnUsage(100, 40) + + // sendLocked with empty messageID takes the create branch, which attaches + // usage from the streamer's stored counts. + s.mu.Lock() + err := s.sendLocked(context.Background(), "answer", nil) + s.mu.Unlock() + if err != nil { + t.Fatalf("sendLocked: %v", err) + } + if _, ok := (*last)[PayloadKeyUsage]; !ok { + t.Fatalf("expected usage in payload, got %+v", *last) + } +} From cc7b4ca86b5162a52ef8d97bb83be6c9f9624cf0 Mon Sep 17 00:00:00 2001 From: Andy Lo-A-Foe Date: Mon, 22 Jun 2026 15:38:46 +0200 Subject: [PATCH 3/3] fix(pico): deliver per-turn token usage to the streamer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- pkg/agent/pipeline_llm.go | 13 ++++++--- pkg/channels/manager.go | 47 ++++++++++++++++++++++++++------ pkg/channels/manager_test.go | 53 ++++++++++++++++++++++++++++++++++++ 3 files changed, 100 insertions(+), 13 deletions(-) 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) + } +}