Compare commits

...

4 Commits

7 changed files with 399 additions and 103 deletions
+64 -31
View File
@@ -5,6 +5,7 @@ import (
"fmt"
"os"
"sync"
"sync/atomic"
tea "charm.land/bubbletea/v2"
"charm.land/fantasy"
@@ -598,9 +599,10 @@ func (a *App) executeStep(ctx context.Context, prompt string, eventFn func(tea.M
}
}
// Subscribe to SDK events for TUI rendering. The subscription is
// temporary — it lives only for the duration of this step.
unsub := a.subscribeSDKEvents(sendFn)
// Subscribe to SDK events for TUI rendering and per-step usage updates.
// The subscription is temporary — it lives only for the duration of this step.
var sawStepUsage atomic.Bool
unsub := a.subscribeSDKEvents(sendFn, &sawStepUsage)
defer unsub()
// Show spinner while the agent works.
@@ -620,8 +622,9 @@ func (a *App) executeStep(ctx context.Context, prompt string, eventFn func(tea.M
// Sync in-memory store with the SDK's authoritative conversation.
a.store.Replace(result.Messages)
// Update usage tracker.
a.updateUsageFromTurnResult(result, prompt)
// Update usage tracker. If per-step usage was already recorded from
// StepUsageEvent callbacks, avoid double-counting totals.
a.updateUsageFromTurnResult(result, prompt, sawStepUsage.Load())
return result, nil
}
@@ -645,9 +648,10 @@ func (a *App) executeBatch(ctx context.Context, items []queueItem, eventFn func(
}
}
// Subscribe to SDK events for TUI rendering. The subscription is
// temporary — it lives only for the duration of this step.
unsub := a.subscribeSDKEvents(sendFn)
// Subscribe to SDK events for TUI rendering and per-step usage updates.
// The subscription is temporary — it lives only for the duration of this step.
var sawStepUsage atomic.Bool
unsub := a.subscribeSDKEvents(sendFn, &sawStepUsage)
defer unsub()
// Show spinner while the agent works.
@@ -702,8 +706,10 @@ func (a *App) executeBatch(ctx context.Context, items []queueItem, eventFn func(
// Sync in-memory store with the SDK's authoritative conversation.
a.store.Replace(result.Messages)
// Update usage tracker (using last item's prompt for tracking).
a.updateUsageFromTurnResult(result, items[len(items)-1].Prompt)
// Update usage tracker (using last item's prompt for fallback estimation).
// If per-step usage was already recorded from StepUsageEvent callbacks,
// avoid double-counting totals.
a.updateUsageFromTurnResult(result, items[len(items)-1].Prompt, sawStepUsage.Load())
return result, nil
}
@@ -720,9 +726,10 @@ func (a *App) sendEvent(msg tea.Msg) {
}
// subscribeSDKEvents registers temporary SDK event subscribers that convert
// SDK events to tea.Msg events and dispatch them via sendFn. Returns an
// unsubscribe function that removes all listeners.
func (a *App) subscribeSDKEvents(sendFn func(tea.Msg)) func() {
// SDK events to tea.Msg events and dispatch them via sendFn. When stepUsageSeen
// is provided, it is set to true after any non-zero StepUsageEvent is observed.
// Returns an unsubscribe function that removes all listeners.
func (a *App) subscribeSDKEvents(sendFn func(tea.Msg), stepUsageSeen *atomic.Bool) func() {
k := a.opts.Kit
var unsubs []func()
@@ -756,6 +763,8 @@ func (a *App) subscribeSDKEvents(sendFn func(tea.Msg)) func() {
})
case kit.SteerConsumedEvent:
sendFn(SteerConsumedEvent{})
case kit.StepUsageEvent:
a.recordStepUsage(ev, stepUsageSeen)
}
}))
@@ -925,32 +934,56 @@ 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) {
hasUsage := ev.InputTokens > 0 || ev.OutputTokens > 0 || ev.CacheReadTokens > 0 || ev.CacheWriteTokens > 0
if !hasUsage {
return
}
if stepUsageSeen != nil {
stepUsageSeen.Store(true)
}
if a.opts.UsageTracker == nil {
return
}
a.opts.UsageTracker.UpdateUsage(
int(ev.InputTokens),
int(ev.OutputTokens),
int(ev.CacheReadTokens),
int(ev.CacheWriteTokens),
)
// Keep context fill reasonably fresh during long/partial turns.
a.opts.UsageTracker.SetContextTokens(int(ev.InputTokens + ev.OutputTokens))
}
// updateUsageFromTurnResult records token usage from an SDK TurnResult into the
// configured UsageTracker. Called once per turn after the turn completes.
//
// Cost/token accumulation uses TotalUsage (sum across all tool-calling steps in
// the turn). Context-window fill uses FinalUsage.InputTokens only — that is the
// number of tokens sent to the model on the last API call, which equals the
// actual context window occupation (all accumulated messages + tool results).
// OutputTokens are not added here because they are the response length, not
// context fill.
func (a *App) updateUsageFromTurnResult(result *kit.TurnResult, userPrompt string) {
// When sawStepUsage is true, totals were already accumulated incrementally via
// StepUsageEvent callbacks; in that case this method only updates context fill.
// Otherwise it falls back to TotalUsage (or estimation) to keep costs/tokens
// visible for providers/modes that don't emit per-step usage.
func (a *App) updateUsageFromTurnResult(result *kit.TurnResult, userPrompt string, sawStepUsage bool) {
if a.opts.UsageTracker == nil || result == nil {
return
}
// --- Accumulate cost/token totals for the session ---
if result.TotalUsage != nil && result.TotalUsage.InputTokens > 0 {
a.opts.UsageTracker.UpdateUsage(
int(result.TotalUsage.InputTokens),
int(result.TotalUsage.OutputTokens),
int(result.TotalUsage.CacheReadTokens),
int(result.TotalUsage.CacheCreationTokens),
)
} else {
// Provider didn't report token counts — fall back to character-based
// estimates so the footer shows something rather than nothing.
a.opts.UsageTracker.EstimateAndUpdateUsage(userPrompt, result.Response)
if !sawStepUsage {
if result.TotalUsage != nil && result.TotalUsage.InputTokens > 0 {
a.opts.UsageTracker.UpdateUsage(
int(result.TotalUsage.InputTokens),
int(result.TotalUsage.OutputTokens),
int(result.TotalUsage.CacheReadTokens),
int(result.TotalUsage.CacheCreationTokens),
)
} else {
// Provider didn't report token counts — fall back to character-based
// estimates so the footer shows something rather than nothing.
a.opts.UsageTracker.EstimateAndUpdateUsage(userPrompt, result.Response)
}
}
// --- Context window fill (drives the % bar) ---
+107
View File
@@ -7,6 +7,8 @@ import (
"testing"
"time"
"charm.land/fantasy"
kit "github.com/mark3labs/kit/pkg/kit"
)
@@ -14,6 +16,47 @@ import (
// Helpers
// --------------------------------------------------------------------------
type usageUpdaterStub struct {
mu sync.Mutex
updateCalls int
estimateCalls int
contextCalls int
lastUpdateInput int
lastUpdateOutput int
lastUpdateCacheRead int
lastUpdateCacheWrite int
lastContextTokens int
lastEstimateInput string
lastEstimateOutput string
}
func (s *usageUpdaterStub) UpdateUsage(inputTokens, outputTokens, cacheReadTokens, cacheWriteTokens int) {
s.mu.Lock()
defer s.mu.Unlock()
s.updateCalls++
s.lastUpdateInput = inputTokens
s.lastUpdateOutput = outputTokens
s.lastUpdateCacheRead = cacheReadTokens
s.lastUpdateCacheWrite = cacheWriteTokens
}
func (s *usageUpdaterStub) EstimateAndUpdateUsage(inputText, outputText string) {
s.mu.Lock()
defer s.mu.Unlock()
s.estimateCalls++
s.lastEstimateInput = inputText
s.lastEstimateOutput = outputText
}
func (s *usageUpdaterStub) SetContextTokens(tokens int) {
s.mu.Lock()
defer s.mu.Unlock()
s.contextCalls++
s.lastContextTokens = tokens
}
// turnResult builds a minimal TurnResult with response text t.
func turnResult(t string) *kit.TurnResult {
return &kit.TurnResult{Response: t}
@@ -489,3 +532,67 @@ func TestQueueLength_reflects(t *testing.T) {
t.Fatalf("expected 3, got %d", got)
}
}
// TestRecordStepUsage_updatesTracker verifies that per-step usage updates are
// recorded immediately (including context tokens) for stop-path correctness.
func TestRecordStepUsage_updatesTracker(t *testing.T) {
usage := &usageUpdaterStub{}
app := New(Options{UsageTracker: usage}, nil)
defer app.Close()
app.recordStepUsage(kit.StepUsageEvent{
InputTokens: 120,
OutputTokens: 45,
CacheReadTokens: 5,
CacheWriteTokens: 2,
}, nil)
usage.mu.Lock()
defer usage.mu.Unlock()
if usage.updateCalls != 1 {
t.Fatalf("expected 1 update call, got %d", usage.updateCalls)
}
if usage.lastUpdateInput != 120 || usage.lastUpdateOutput != 45 || usage.lastUpdateCacheRead != 5 || usage.lastUpdateCacheWrite != 2 {
t.Fatalf("unexpected usage update payload: in=%d out=%d cache_read=%d cache_write=%d",
usage.lastUpdateInput, usage.lastUpdateOutput, usage.lastUpdateCacheRead, usage.lastUpdateCacheWrite)
}
if usage.contextCalls != 1 {
t.Fatalf("expected 1 context token update, got %d", usage.contextCalls)
}
if usage.lastContextTokens != 165 {
t.Fatalf("expected context tokens 165, got %d", usage.lastContextTokens)
}
}
// TestUpdateUsageFromTurnResult_skipsTotalsWhenStepUsageSeen ensures we avoid
// double-counting totals once StepUsageEvent-based updates were already applied.
func TestUpdateUsageFromTurnResult_skipsTotalsWhenStepUsageSeen(t *testing.T) {
usage := &usageUpdaterStub{}
app := New(Options{UsageTracker: usage}, nil)
defer app.Close()
app.updateUsageFromTurnResult(&kit.TurnResult{
Response: "ok",
TotalUsage: &fantasy.Usage{
InputTokens: 999,
OutputTokens: 111,
CacheReadTokens: 7,
CacheCreationTokens: 3,
},
FinalUsage: &fantasy.Usage{InputTokens: 456},
}, "prompt", true)
usage.mu.Lock()
defer usage.mu.Unlock()
if usage.updateCalls != 0 {
t.Fatalf("expected no total usage update when sawStepUsage=true, got %d", usage.updateCalls)
}
if usage.estimateCalls != 0 {
t.Fatalf("expected no estimate update when sawStepUsage=true, got %d", usage.estimateCalls)
}
if usage.contextCalls != 1 || usage.lastContextTokens != 456 {
t.Fatalf("expected final context tokens=456, got calls=%d tokens=%d", usage.contextCalls, usage.lastContextTokens)
}
}
+79 -10
View File
@@ -349,7 +349,7 @@ func TestStreamComponent_SpinnerKeepsRunningDuringStreaming(t *testing.T) {
c = sendStreamMsg(c, app.StreamChunkEvent{Content: "hello"})
// Flush pending chunks (simulates the 16ms tick firing).
c = sendStreamMsg(c, streamFlushTickMsg{})
c = sendStreamMsg(c, streamFlushTickMsg{generation: c.flushGeneration})
if !c.spinning {
t.Fatal("expected spinning=true after first chunk")
@@ -376,7 +376,7 @@ func TestStreamComponent_ChunkAccumulation(t *testing.T) {
}
// Flush pending chunks (simulates the 16ms tick firing).
c = sendStreamMsg(c, streamFlushTickMsg{})
c = sendStreamMsg(c, streamFlushTickMsg{generation: c.flushGeneration})
got := c.streamContent.String()
want := "Hello, world!"
@@ -396,6 +396,7 @@ func TestStreamComponent_ToolExecution_IsStarting_ShowsSpinner(t *testing.T) {
c := newTestStream()
_, cmd := c.Update(app.ToolExecutionEvent{
ToolCallID: "call-exec-1",
ToolName: "exec_tool",
IsStarting: true,
})
@@ -403,8 +404,9 @@ func TestStreamComponent_ToolExecution_IsStarting_ShowsSpinner(t *testing.T) {
if !c.spinning {
t.Fatal("expected spinning=true during tool execution")
}
if len(c.activeTools) != 1 || !strings.Contains(c.activeTools[0], "exec_tool") {
t.Fatalf("expected activeTools to contain tool name, got %v", c.activeTools)
tools := c.activeToolDisplays()
if len(tools) != 1 || !strings.Contains(tools[0], "exec_tool") {
t.Fatalf("expected activeTools to contain tool name, got %v", tools)
}
if cmd == nil {
t.Fatal("expected tick cmd from ToolExecutionEvent{IsStarting:true}")
@@ -418,11 +420,13 @@ func TestStreamComponent_ToolExecution_NotStarting_KeepsSpinning(t *testing.T) {
c = sendStreamMsg(c, app.SpinnerEvent{Show: true})
// Simulate a tool starting
c = sendStreamMsg(c, app.ToolExecutionEvent{
ToolCallID: "call-some-1",
ToolName: "some_tool",
IsStarting: true,
})
c = sendStreamMsg(c, app.ToolExecutionEvent{
ToolCallID: "call-some-1",
ToolName: "some_tool",
IsStarting: false,
})
@@ -440,9 +444,9 @@ func TestStreamComponent_ParallelToolExecution(t *testing.T) {
c := newTestStream()
// Start three tools in parallel
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolName: "read", IsStarting: true})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolName: "grep", IsStarting: true})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolName: "find", IsStarting: true})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-read", ToolName: "read", IsStarting: true})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-grep", ToolName: "grep", IsStarting: true})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-find", ToolName: "find", IsStarting: true})
if len(c.activeTools) != 3 {
t.Fatalf("expected 3 active tools, got %d: %v", len(c.activeTools), c.activeTools)
@@ -455,19 +459,44 @@ func TestStreamComponent_ParallelToolExecution(t *testing.T) {
}
// Finish one tool
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolName: "grep", IsStarting: false})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-grep", ToolName: "grep", IsStarting: false})
if len(c.activeTools) != 2 {
t.Fatalf("expected 2 active tools after one finished, got %d: %v", len(c.activeTools), c.activeTools)
}
// Finish remaining tools
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolName: "read", IsStarting: false})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolName: "find", IsStarting: false})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-read", ToolName: "read", IsStarting: false})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-find", ToolName: "find", IsStarting: false})
if len(c.activeTools) != 0 {
t.Fatalf("expected 0 active tools after all finished, got %d: %v", len(c.activeTools), c.activeTools)
}
}
// TestStreamComponent_ParallelSameToolName_UsesToolCallID verifies finishing one
// tool call does not remove another concurrent call with the same tool name.
func TestStreamComponent_ParallelSameToolName_UsesToolCallID(t *testing.T) {
c := newTestStream()
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-read-1", ToolName: "read", IsStarting: true})
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-read-2", ToolName: "read", IsStarting: true})
tools := c.activeToolDisplays()
if len(tools) != 2 {
t.Fatalf("expected 2 active read calls, got %d (%v)", len(tools), tools)
}
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-read-1", ToolName: "read", IsStarting: false})
tools = c.activeToolDisplays()
if len(tools) != 1 {
t.Fatalf("expected 1 active read call after finishing one ID, got %d (%v)", len(tools), tools)
}
c = sendStreamMsg(c, app.ToolExecutionEvent{ToolCallID: "call-read-2", ToolName: "read", IsStarting: false})
if len(c.activeToolDisplays()) != 0 {
t.Fatalf("expected no active tools after finishing both IDs, got %v", c.activeToolDisplays())
}
}
// --------------------------------------------------------------------------
// TestStreamComponent_GetRenderedContent verifies the method returns rendered
// text when content is accumulated, and empty string when not.
@@ -621,3 +650,43 @@ func TestStreamComponent_StaleTick_Discarded(t *testing.T) {
t.Fatal("current-gen tick should reschedule")
}
}
// TestStreamComponent_StaleFlushTick_Discarded verifies that flush ticks from a
// previous generation (e.g. pre-Reset) are ignored.
func TestStreamComponent_StaleFlushTick_Discarded(t *testing.T) {
c := newTestStream()
// Start a pending flush and capture its generation.
c = sendStreamMsg(c, app.StreamChunkEvent{Content: "old"})
staleGen := c.flushGeneration
if !c.flushPending {
t.Fatal("precondition: expected flushPending=true after first chunk")
}
// Reset should invalidate in-flight flush ticks.
c.Reset()
if c.flushGeneration == staleGen {
t.Fatal("expected flushGeneration to change after Reset")
}
// New content in a new generation.
c = sendStreamMsg(c, app.StreamChunkEvent{Content: "new"})
if got := c.pendingStream.String(); got != "new" {
t.Fatalf("expected pendingStream='new', got %q", got)
}
// Stale flush tick should be ignored.
c = sendStreamMsg(c, streamFlushTickMsg{generation: staleGen})
if got := c.pendingStream.String(); got != "new" {
t.Fatalf("stale flush tick should not commit pending stream, got %q", got)
}
// Current generation flush should commit.
c = sendStreamMsg(c, streamFlushTickMsg{generation: c.flushGeneration})
if got := c.pendingStream.String(); got != "" {
t.Fatalf("expected pendingStream empty after current flush, got %q", got)
}
if got := c.streamContent.String(); got != "new" {
t.Fatalf("expected streamContent='new' after current flush, got %q", got)
}
}
+2 -11
View File
@@ -8,7 +8,6 @@ import (
"os"
"os/exec"
"strings"
"sync"
"time"
tea "charm.land/bubbletea/v2"
@@ -584,8 +583,8 @@ type AppModel struct {
streamingBashStderr []string
// streamingBashMaxLines caps how many lines to accumulate to prevent memory issues.
streamingBashMaxLines int
// streamingMu protects the streaming bash output fields from concurrent access.
streamingMu sync.RWMutex
// streaming bash fields are only mutated/read from the Bubble Tea event loop
// (Update/View), so no mutex is required here.
// streamingBashCommand holds the command being executed for display as a header.
streamingBashCommand string
}
@@ -1378,9 +1377,7 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
Command string `json:"command"`
}
if err := json.Unmarshal([]byte(msg.ToolArgs), &args); err == nil && args.Command != "" {
m.streamingMu.Lock()
m.streamingBashCommand = args.Command
m.streamingMu.Unlock()
}
}
@@ -1395,11 +1392,9 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
// Buffer tool result for scrollback.
m.printToolResult(msg)
// Clear streaming bash output since tool completed.
m.streamingMu.Lock()
m.streamingBashOutput = nil
m.streamingBashStderr = nil
m.streamingBashCommand = ""
m.streamingMu.Unlock()
// Start spinner again while waiting for the next LLM response.
if m.stream != nil {
_, cmd := m.stream.Update(app.SpinnerEvent{Show: true})
@@ -1408,7 +1403,6 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case app.ToolOutputEvent:
// Accumulate streaming bash output for display.
m.streamingMu.Lock()
if msg.IsStderr {
m.streamingBashStderr = append(m.streamingBashStderr, msg.Chunk)
// Cap stderr lines to prevent memory issues.
@@ -1422,7 +1416,6 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.streamingBashOutput = m.streamingBashOutput[len(m.streamingBashOutput)-m.streamingBashMaxLines:]
}
}
m.streamingMu.Unlock()
case app.ToolCallContentEvent:
// In streaming mode this text was already delivered via StreamChunkEvents
@@ -1847,13 +1840,11 @@ func (m *AppModel) renderStream() string {
// Lines are truncated to the terminal width and capped to maxBashLines to prevent
// long-running commands from blowing up the TUI layout.
func (m *AppModel) renderStreamingBashOutput(theme Theme) string {
m.streamingMu.RLock()
stdoutLines := make([]string, len(m.streamingBashOutput))
copy(stdoutLines, m.streamingBashOutput)
stderrLines := make([]string, len(m.streamingBashStderr))
copy(stderrLines, m.streamingBashStderr)
command := m.streamingBashCommand
m.streamingMu.RUnlock()
if len(stdoutLines) == 0 && len(stderrLines) == 0 {
return ""
+88 -47
View File
@@ -79,7 +79,12 @@ func streamSpinnerTickCmd(generation uint64) tea.Cmd {
// streamFlushTickMsg fires when it's time to commit pending chunks to the
// main content builders and trigger a re-render. This coalesces rapid
// streaming chunks into fewer expensive markdown re-renders.
type streamFlushTickMsg struct{}
//
// generation ties the tick to the pending flush session that created it so
// stale ticks from a prior Reset() are discarded.
type streamFlushTickMsg struct {
generation uint64
}
// streamFlushInterval is the coalescing window for stream chunks. Chunks
// arriving within this window are batched into a single render pass.
@@ -89,9 +94,9 @@ const streamFlushInterval = 16 * time.Millisecond
// streamFlushTickCmd returns a tea.Cmd that fires streamFlushTickMsg after
// the coalescing interval.
func streamFlushTickCmd() tea.Cmd {
func streamFlushTickCmd(generation uint64) tea.Cmd {
return tea.Tick(streamFlushInterval, func(_ time.Time) tea.Msg {
return streamFlushTickMsg{}
return streamFlushTickMsg{generation: generation}
})
}
@@ -149,9 +154,11 @@ type StreamComponent struct {
// spinnerFrame is the current frame index.
spinnerFrame int
// activeTools tracks the names of tools currently executing in parallel.
// When multiple tools run concurrently, all are displayed in the spinner.
activeTools []string
// activeTools maps ToolCallID -> display label for currently running tools.
activeTools map[string]string
// activeToolOrder preserves deterministic display order for active tools.
activeToolOrder []string
// streamContent holds committed streaming text (flushed from pending).
streamContent strings.Builder
@@ -172,6 +179,10 @@ type StreamComponent struct {
// the same coalescing window.
flushPending bool
// flushGeneration is incremented when stream state resets so stale flush
// ticks from a previous step can be discarded.
flushGeneration uint64
// renderCache holds the last rendered output string. Reused by View()
// between flush ticks to avoid redundant markdown re-parsing.
renderCache string
@@ -190,14 +201,8 @@ type StreamComponent struct {
// reasoningDuration holds the total reasoning time, frozen when streaming text begins.
reasoningDuration time.Duration
// messageRenderer renders assistant messages in standard mode.
messageRenderer *MessageRenderer
// compactRenderer renders assistant messages in compact mode.
compactRenderer *CompactRenderer
// compactMode selects which renderer to use.
compactMode bool
// renderer renders streaming assistant text in either compact or standard mode.
renderer Renderer
// modelName is displayed in the streaming text header.
modelName string
@@ -218,13 +223,19 @@ func NewStreamComponent(compactMode bool, width int, modelName string) *StreamCo
if width == 0 {
width = 80
}
var renderer Renderer
if compactMode {
renderer = NewCompactRenderer(width, false)
} else {
renderer = newMessageRenderer(width, false)
}
return &StreamComponent{
spinnerFrames: knightRiderFrames(),
compactMode: compactMode,
modelName: modelName,
messageRenderer: newMessageRenderer(width, false),
compactRenderer: NewCompactRenderer(width, false),
width: width,
spinnerFrames: knightRiderFrames(),
modelName: modelName,
renderer: renderer,
width: width,
}
}
@@ -251,11 +262,13 @@ func (s *StreamComponent) Reset() {
s.spinnerGeneration++ // invalidate any in-flight tick commands
s.spinnerFrame = 0
s.activeTools = nil
s.activeToolOrder = nil
s.streamContent.Reset()
s.reasoningContent.Reset()
s.pendingStream.Reset()
s.pendingReasoning.Reset()
s.flushPending = false
s.flushGeneration++
s.renderCache = ""
s.renderDirty = false
s.timestamp = time.Time{}
@@ -323,8 +336,9 @@ func (s *StreamComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case tea.WindowSizeMsg:
s.width = msg.Width
s.messageRenderer.SetWidth(s.width)
s.compactRenderer.SetWidth(s.width)
if s.renderer != nil {
s.renderer.SetWidth(s.width)
}
// Invalidate render cache — width change affects wrapping/styling.
s.renderCache = ""
s.renderDirty = true
@@ -360,6 +374,9 @@ func (s *StreamComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
}
case streamFlushTickMsg:
if msg.generation != s.flushGeneration {
break
}
s.flushPending = false
s.commitPending()
@@ -374,7 +391,7 @@ func (s *StreamComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
s.pendingReasoning.WriteString(msg.Delta)
if !s.flushPending {
s.flushPending = true
return s, streamFlushTickCmd()
return s, streamFlushTickCmd(s.flushGeneration)
}
case app.StreamChunkEvent:
@@ -389,14 +406,25 @@ func (s *StreamComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
s.pendingStream.WriteString(msg.Content)
if !s.flushPending {
s.flushPending = true
return s, streamFlushTickCmd()
return s, streamFlushTickCmd(s.flushGeneration)
}
case app.ToolExecutionEvent:
toolID := msg.ToolCallID
if toolID == "" {
// Defensive fallback for older/third-party emitters that may omit
// ToolCallID. Best-effort only: same-name+args concurrent calls can
// still collide without a stable ID.
toolID = fmt.Sprintf("%s|%s", msg.ToolName, msg.ToolArgs)
}
if msg.IsStarting {
// Add tool to active list for parallel execution display.
toolDisplay := formatToolExecutionMessage(msg.ToolName, msg.ToolArgs)
s.activeTools = append(s.activeTools, toolDisplay)
if s.activeTools == nil {
s.activeTools = make(map[string]string)
}
if _, exists := s.activeTools[toolID]; !exists {
s.activeToolOrder = append(s.activeToolOrder, toolID)
}
s.activeTools[toolID] = formatToolExecutionMessage(msg.ToolName)
s.spinnerFrame = 0
if !s.spinning {
s.phase = streamPhaseActive
@@ -405,9 +433,10 @@ func (s *StreamComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
return s, streamSpinnerTickCmd(s.spinnerGeneration)
}
} else {
// Tool finished — remove from active list but keep spinning if others remain.
toolDisplay := formatToolExecutionMessage(msg.ToolName, msg.ToolArgs)
s.activeTools = removeFromSlice(s.activeTools, toolDisplay)
if s.activeTools != nil {
delete(s.activeTools, toolID)
}
s.activeToolOrder = removeToolID(s.activeToolOrder, toolID)
}
}
@@ -568,7 +597,8 @@ func (s *StreamComponent) SpinnerView() string {
return ""
}
frame := s.spinnerFrames[s.spinnerFrame%len(s.spinnerFrames)]
if len(s.activeTools) == 0 {
tools := s.activeToolDisplays()
if len(tools) == 0 {
return " " + frame
}
theme := GetTheme()
@@ -578,10 +608,10 @@ func (s *StreamComponent) SpinnerView() string {
// Format active tools list
var toolsMsg string
if len(s.activeTools) == 1 {
toolsMsg = s.activeTools[0]
if len(tools) == 1 {
toolsMsg = tools[0]
} else {
toolsMsg = "Running: " + strings.Join(s.activeTools, ", ")
toolsMsg = "Running: " + strings.Join(tools, ", ")
}
return " " + frame + " " + msgStyle.Render(toolsMsg)
}
@@ -593,28 +623,39 @@ func (s *StreamComponent) renderStreamingText(text string) string {
if ts.IsZero() {
ts = time.Now()
}
if s.compactMode {
msg := s.compactRenderer.RenderAssistantMessage(text, ts, s.modelName)
return msg.Content
if s.renderer == nil {
return text
}
msg := s.messageRenderer.RenderAssistantMessage(text, ts, s.modelName)
msg := s.renderer.RenderAssistantMessage(text, ts, s.modelName)
return msg.Content
}
// removeFromSlice removes the first occurrence of a string from a slice.
func removeFromSlice(slice []string, s string) []string {
for i, v := range slice {
if v == s {
return append(slice[:i], slice[i+1:]...)
func (s *StreamComponent) activeToolDisplays() []string {
if len(s.activeTools) == 0 {
return nil
}
out := make([]string, 0, len(s.activeToolOrder))
for _, id := range s.activeToolOrder {
if display, ok := s.activeTools[id]; ok {
out = append(out, display)
}
}
return slice
return out
}
// removeToolID removes the first occurrence of a tool ID from a slice.
func removeToolID(ids []string, id string) []string {
for i, v := range ids {
if v == id {
return append(ids[:i], ids[i+1:]...)
}
}
return ids
}
// formatToolExecutionMessage creates a descriptive spinner message for tool execution.
// For spawn_subagent, it shows simply as "Subagent" with optional task preview.
func formatToolExecutionMessage(toolName, toolArgs string) string {
// For spawn_subagent, it shows simply as "Subagent".
func formatToolExecutionMessage(toolName string) string {
if toolName == "spawn_subagent" {
return "Subagent"
}
-4
View File
@@ -151,10 +151,6 @@ func (ut *UsageTracker) RenderUsageInfo() string {
ut.mu.RLock()
defer ut.mu.RUnlock()
if ut.sessionStats.RequestCount == 0 {
return ""
}
baseStyle := lipgloss.NewStyle()
// Display the current context window token count (from the last API call),
+59
View File
@@ -67,3 +67,62 @@ func TestUsageTracker_RenderUsageInfo_OAuth(t *testing.T) {
t.Errorf("Expected regular rendered output to show actual cost, got: %s", regularRendered)
}
}
func TestUsageTracker_RenderUsageInfo_StartupState(t *testing.T) {
// Create a mock model info with costs and context limit
modelInfo := &models.ModelInfo{
ID: "claude-3-5-sonnet-20241022",
Name: "Claude 3.5 Sonnet v2",
Cost: models.Cost{
Input: 3.0,
Output: 15.0,
},
Limit: models.Limit{
Context: 200000,
Output: 8192,
},
}
// Test startup state (no requests made yet) - Regular API key
regularTracker := NewUsageTracker(modelInfo, "anthropic", 80, false)
rendered := stripAnsi(regularTracker.RenderUsageInfo())
// Should NOT return empty string on startup
if rendered == "" {
t.Errorf("Expected non-empty output on startup, got empty string")
}
// Should show 0 tokens
if !strings.Contains(rendered, "Tokens: 0") {
t.Errorf("Expected 'Tokens: 0' on startup, got: %s", rendered)
}
// Should NOT show percentage when tokens are 0
if strings.Contains(rendered, "(%") {
t.Errorf("Expected no percentage on startup with 0 tokens, got: %s", rendered)
}
// Should show $0.0000 cost for regular API key
if !strings.Contains(rendered, "Cost: $0.0000") {
t.Errorf("Expected 'Cost: $0.0000' on startup, got: %s", rendered)
}
// Test startup state (no requests made yet) - OAuth
oauthTracker := NewUsageTracker(modelInfo, "anthropic", 80, true)
oauthRendered := stripAnsi(oauthTracker.RenderUsageInfo())
// Should NOT return empty string on startup
if oauthRendered == "" {
t.Errorf("Expected non-empty output on startup for OAuth, got empty string")
}
// Should show 0 tokens for OAuth
if !strings.Contains(oauthRendered, "Tokens: 0") {
t.Errorf("Expected 'Tokens: 0' on startup for OAuth, got: %s", oauthRendered)
}
// Should show $0.00 cost for OAuth
if !strings.Contains(oauthRendered, "Cost: $0.00") {
t.Errorf("Expected 'Cost: $0.00' on startup for OAuth, got: %s", oauthRendered)
}
}