From 052c742fe74e5948c08c635bc82cb5ebbf1cc350 Mon Sep 17 00:00:00 2001
From: Andy Lo-A-Foe <andy.loafoe@gmail.com>
Date: Mon, 22 Jun 2026 13:07:48 +0200
Subject: [PATCH] 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 587f889a..b11befae 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 00000000..ae2d9749
--- /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 0c809430..ac92913d 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"