mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-07-28 01:27:58 +00:00
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:
@@ -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
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user