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
This commit is contained in:
Ed Zynda
2026-04-23 12:56:00 +03:00
parent d2e23295b6
commit aeb704367c
4 changed files with 49 additions and 14 deletions
+26 -7
View File
@@ -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
+11 -7
View File
@@ -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)
}
}
+6
View File
@@ -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.
+6
View File
@@ -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