mirror of
https://github.com/mark3labs/kit.git
synced 2026-06-13 19:20:06 +00:00
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:
+26
-7
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user