Fix sidebar context % flickering during sub-agent transfers

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,
+		},
+	})
+}

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
+			}
 		}
 	}