mirror of
https://github.com/mark3labs/kit.git
synced 2026-06-17 21:08:43 +00:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| b68b3dd0bf | |||
| 48521bf76d | |||
| 16df3a738c | |||
| 9d0b8c8cef |
+64
-31
@@ -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) ---
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user