diff --git a/pkg/tui/components/sidebar/context_percent_test.go b/pkg/tui/components/sidebar/context_percent_test.go index 502733830..7b9a21e90 100644 --- a/pkg/tui/components/sidebar/context_percent_test.go +++ b/pkg/tui/components/sidebar/context_percent_test.go @@ -121,3 +121,92 @@ func TestContextPercent_FallbackToSingleSession(t *testing.T) { assert.Equal(t, "10%", m.contextPercent()) } + +// TestContextPercent_StaleSessionIDAfterSubAgent verifies that contextPercent() +// returns the correct value throughout a transfer_task round-trip. +// +// After a sub-agent's stream stops, the parent's AgentInfo event restores +// currentAgent to the parent while currentSessionID still references the child +// session. The sidebar must detect this stale ID and fall back to an +// agent-name lookup so the displayed context % matches the parent. +func TestContextPercent_StaleSessionIDAfterSubAgent(t *testing.T) { + t.Parallel() + + m := newTestSidebar() + + // Parent starts. + m.setAgent("root") + m.startStream("parent-session", "root") + m.recordUsage("parent-session", "root", 30000, 100000) + assert.Equal(t, "30%", m.contextPercent(), "parent at 30%%") + + // --- transfer_task to "developer" --- + m.setAgent("developer") + m.startStream("child-session-1", "developer") + m.recordUsage("child-session-1", "developer", 10000, 200000) + assert.Equal(t, "5%", m.contextPercent(), "developer sub-agent at 5%%") + + m.stopStream() + m.setAgent("root") // parent restored + + // Key assertion: stale currentSessionID must not cause a wrong lookup. + assert.Equal(t, "30%", m.contextPercent(), + "after sub-agent returns, context %% must reflect the parent (30%%), not the child (5%%)") + + // --- transfer_task to "researcher" (second round-trip) --- + m.setAgent("researcher") + m.startStream("child-session-2", "researcher") + m.recordUsage("child-session-2", "researcher", 80000, 100000) + assert.Equal(t, "80%", m.contextPercent(), "researcher sub-agent at 80%%") + + m.stopStream() + m.setAgent("root") // parent restored again + + assert.Equal(t, "30%", m.contextPercent(), + "after second sub-agent returns, context %% must still reflect the parent (30%%)") + + // Parent resumes with a new stream iteration. + m.startStream("parent-session", "root") + m.recordUsage("parent-session", "root", 40000, 100000) + assert.Equal(t, "40%", m.contextPercent(), "parent resumes at 40%%") +} + +// testSidebar wraps *model with helpers that mirror the sidebar field mutations +// performed by Update() for each runtime event — without touching the global +// spinner/animation coordinator, which would leak state across test runs. +type testSidebar struct { + *model +} + +func newTestSidebar() *testSidebar { + sess := session.New() + return &testSidebar{ + model: New(service.NewSessionState(sess)).(*model), + } +} + +func (s *testSidebar) setAgent(name string) { + s.currentAgent = name +} + +func (s *testSidebar) startStream(sessionID, agentName string) { + s.workingAgent = agentName + s.currentSessionID = sessionID +} + +func (s *testSidebar) stopStream() { + s.workingAgent = "" +} + +func (s *testSidebar) recordUsage(sessionID, agentName string, contextLen, contextLimit int64) { + s.SetTokenUsage(&runtime.TokenUsageEvent{ + SessionID: sessionID, + AgentContext: runtime.AgentContext{AgentName: agentName}, + Usage: &runtime.Usage{ + InputTokens: contextLen / 2, + OutputTokens: contextLen / 2, + ContextLength: contextLen, + ContextLimit: contextLimit, + }, + }) +} diff --git a/pkg/tui/components/sidebar/sidebar.go b/pkg/tui/components/sidebar/sidebar.go index f1e8f6cbd..3d6be8a57 100644 --- a/pkg/tui/components/sidebar/sidebar.go +++ b/pkg/tui/components/sidebar/sidebar.go @@ -487,12 +487,18 @@ func formatCost(cost float64) string { } // currentSessionUsage returns the usage snapshot for the current agent's session. -// It uses a 3-tier lookup: session ID (most reliable) → agent name → single-session fallback. +// It uses a 3-tier lookup: session ID → agent name → single-session fallback. func (m *model) currentSessionUsage() (*runtime.Usage, bool) { - // Direct lookup by current session ID (most reliable, no map iteration ambiguity). + // Direct lookup by current session ID, skipping when the session belongs + // to a different agent (stale after a sub-agent's stream stops while + // currentAgent has already been restored to the parent). if m.currentSessionID != "" { - if usage, ok := m.sessionUsage[m.currentSessionID]; ok { - return usage, true + owner := m.sessionAgent[m.currentSessionID] + stale := owner != "" && m.currentAgent != "" && owner != m.currentAgent + if !stale { + if usage, ok := m.sessionUsage[m.currentSessionID]; ok { + return usage, true + } } }