From aeb704367cdb4a569fda98a18f24bfc632ad7d9e Mon Sep 17 00:00:00 2001 From: Ed Zynda Date: Thu, 23 Apr 2026 12:56:00 +0300 Subject: [PATCH] feat(app): update token counts and context fill after every step - Set context tokens per-step in recordStepUsage instead of waiting for turn completion; each step re-sends the full conversation so the reported usage monotonically increases - Add UsageUpdatedEvent to trigger a TUI re-render after each step so the status bar reflects updated tokens, cost, and context % even during gaps between streaming chunks - Update test to expect per-step context token updates --- internal/app/app.go | 33 ++++++++++++++++++++++++++------- internal/app/app_test.go | 18 +++++++++++------- internal/app/events.go | 6 ++++++ internal/ui/model.go | 6 ++++++ 4 files changed, 49 insertions(+), 14 deletions(-) diff --git a/internal/app/app.go b/internal/app/app.go index 20ea9e39..bc4b3912 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -923,7 +923,7 @@ func (a *App) subscribeSDKEvents(sendFn func(tea.Msg), stepUsageSeen *atomic.Boo case kit.SteerConsumedEvent: sendFn(SteerConsumedEvent{}) case kit.StepUsageEvent: - a.recordStepUsage(ev, stepUsageSeen) + a.recordStepUsage(ev, stepUsageSeen, sendFn) case kit.PasswordPromptEvent: // Convert SDK PasswordPromptEvent to app PasswordPromptEvent // The TUI will handle this and send the response back @@ -1241,7 +1241,16 @@ func (a *App) PrintBlockFromExtension(opts extensions.PrintBlockOpts) { // recordStepUsage applies token/cost usage reported for a completed step. // Step usage events arrive even when a turn is later cancelled, so this keeps // the usage widget accurate on all stop paths. -func (a *App) recordStepUsage(ev kit.StepUsageEvent, stepUsageSeen *atomic.Bool) { +// +// Both session totals (cost, token counts) and the context window fill level +// are updated here so the status bar reflects progress after every LLM call, +// not just at the end of the full turn. Context fill monotonically increases +// across steps because each step re-sends the entire conversation plus any +// new tool results, so the numbers only go up. +// +// sendFn is called with a UsageUpdatedEvent to trigger a TUI re-render so +// the updated values are visible immediately. +func (a *App) recordStepUsage(ev kit.StepUsageEvent, stepUsageSeen *atomic.Bool, sendFn func(tea.Msg)) { hasUsage := ev.InputTokens > 0 || ev.OutputTokens > 0 || ev.CacheReadTokens > 0 || ev.CacheWriteTokens > 0 if a.opts.Debug { log.Printf("[DEBUG] recordStepUsage: hasUsage=%v input=%d output=%d cacheRead=%d cacheWrite=%d", @@ -1262,11 +1271,21 @@ func (a *App) recordStepUsage(ev kit.StepUsageEvent, stepUsageSeen *atomic.Bool) int(ev.CacheReadTokens), int(ev.CacheWriteTokens), ) - // NOTE: We do NOT call SetContextTokens here. Context fill is set once - // at turn completion via updateUsageFromTurnResult, which sums all token - // categories (Input + CacheRead + CacheCreate + Output) from FinalUsage. - // Per-step context tokens would cause the display to jump around during - // multi-step tool calls. + // Update context window fill from this step's usage. Each step sends + // the full conversation to the LLM, so the reported token counts + // represent the actual context utilization at that point. + contextFill := int(ev.InputTokens) + int(ev.CacheReadTokens) + int(ev.CacheWriteTokens) + int(ev.OutputTokens) + if contextFill > 0 { + if a.opts.Debug { + log.Printf("[DEBUG] recordStepUsage: SetContextTokens=%d (Input=%d + CacheRead=%d + CacheWrite=%d + Output=%d)", + contextFill, ev.InputTokens, ev.CacheReadTokens, ev.CacheWriteTokens, ev.OutputTokens) + } + a.opts.UsageTracker.SetContextTokens(contextFill) + } + // Notify the TUI so it re-renders the status bar with updated values. + if sendFn != nil { + sendFn(UsageUpdatedEvent{}) + } } // updateUsageFromTurnResult records token usage from an SDK TurnResult into the diff --git a/internal/app/app_test.go b/internal/app/app_test.go index 3202af6b..e0f2217a 100644 --- a/internal/app/app_test.go +++ b/internal/app/app_test.go @@ -534,9 +534,9 @@ func TestQueueLength_reflects(t *testing.T) { } // TestRecordStepUsage_updatesTracker verifies that per-step usage updates are -// recorded immediately for cost tracking. Context tokens are NOT updated here -// (only via updateUsageFromTurnResult) to avoid display jumps during multi-step -// tool calls. +// recorded immediately for cost tracking. Context tokens are also updated so +// the status bar reflects context fill after every LLM call in a multi-step +// turn, not just at the end. func TestRecordStepUsage_updatesTracker(t *testing.T) { usage := &usageUpdaterStub{} app := New(Options{UsageTracker: usage}, nil) @@ -547,7 +547,7 @@ func TestRecordStepUsage_updatesTracker(t *testing.T) { OutputTokens: 45, CacheReadTokens: 5, CacheWriteTokens: 2, - }, nil) + }, nil, nil) usage.mu.Lock() defer usage.mu.Unlock() @@ -559,9 +559,13 @@ func TestRecordStepUsage_updatesTracker(t *testing.T) { t.Fatalf("unexpected usage update payload: in=%d out=%d cache_read=%d cache_write=%d", usage.lastUpdateInput, usage.lastUpdateOutput, usage.lastUpdateCacheRead, usage.lastUpdateCacheWrite) } - // Context tokens should NOT be updated by recordStepUsage (only by updateUsageFromTurnResult) - if usage.contextCalls != 0 { - t.Fatalf("expected 0 context token updates from recordStepUsage, got %d", usage.contextCalls) + // Context tokens should now be updated per-step (Input + CacheRead + CacheWrite + Output). + if usage.contextCalls != 1 { + t.Fatalf("expected 1 context token update from recordStepUsage, got %d", usage.contextCalls) + } + expectedContext := 120 + 45 + 5 + 2 + if usage.lastContextTokens != expectedContext { + t.Fatalf("expected context tokens %d, got %d", expectedContext, usage.lastContextTokens) } } diff --git a/internal/app/events.go b/internal/app/events.go index e56226d4..71923106 100644 --- a/internal/app/events.go +++ b/internal/app/events.go @@ -210,6 +210,12 @@ type ModelChangedEvent struct { ModelName string } +// UsageUpdatedEvent is sent after each completed LLM step to notify the TUI +// that token counts and costs have changed. The UsageTracker is updated +// in-place before this event is sent; the TUI just needs to re-render to +// reflect the new values in the status bar. +type UsageUpdatedEvent struct{} + // WidgetUpdateEvent is sent when an extension adds, updates, or removes a // widget via ctx.SetWidget or ctx.RemoveWidget. The TUI re-reads widget state // from its WidgetProvider on the next render cycle. diff --git a/internal/ui/model.go b/internal/ui/model.go index cfc02d01..5118bd88 100644 --- a/internal/ui/model.go +++ b/internal/ui/model.go @@ -2076,6 +2076,12 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.providerName = msg.ProviderName m.modelName = msg.ModelName + case app.UsageUpdatedEvent: + // Token usage was updated after a completed LLM step. No state + // changes needed — the UsageTracker was already mutated in-place. + // Returning from Update() triggers View() which re-renders the + // status bar with the latest token counts, cost, and context %. + case app.WidgetUpdateEvent: // Extension widget changed — recalculate height distribution so the // stream region accounts for widget space. View() will read the