From 1fb14f211d4e69ba17c8f81ee10c6d615b1f17b2 Mon Sep 17 00:00:00 2001 From: David Gageot Date: Tue, 10 Mar 2026 18:21:02 +0100 Subject: [PATCH] Fix sidebar context % flickering during sub-agent transfers After a sub-agent's stream stops, handleTaskTransfer restores the parent agent via AgentInfo, but currentSessionID still references the child session. currentSessionUsage() used to prioritise currentSessionID unconditionally, so it returned the child's context usage instead of the parent's until the next StreamStarted event reset the session ID. Detect this stale state by checking whether the session's owning agent matches currentAgent. When they differ, skip the session-ID lookup and fall through to the agent-name search, which finds the correct parent session. Add a test that replays two sequential transfer_task round-trips and asserts contextPercent() stays correct at every step. Assisted-By: docker-agent --- .../sidebar/context_percent_test.go | 89 +++++++++++++++++++ pkg/tui/components/sidebar/sidebar.go | 14 ++- 2 files changed, 99 insertions(+), 4 deletions(-) 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 + } } }