diff --git a/desktop/src/components/chat/MessageList.tsx b/desktop/src/components/chat/MessageList.tsx index ed262bd5..b2b2bcfa 100644 --- a/desktop/src/components/chat/MessageList.tsx +++ b/desktop/src/components/chat/MessageList.tsx @@ -1,4 +1,4 @@ -import { useRef, useEffect, useMemo, memo, useState, useCallback, useLayoutEffect, type ReactNode } from 'react' +import { useRef, useEffect, useMemo, memo, useState, useCallback, useDeferredValue, useLayoutEffect, type ReactNode } from 'react' import { ArrowDown, BookMarked, Bot, CheckCircle2, ChevronDown, ChevronRight, CircleStop, FileStack, LoaderCircle, MessageCircle, Settings, Target, XCircle } from 'lucide-react' import { ApiError } from '../../api/client' import { sessionsApi, type SessionTurnCheckpoint } from '../../api/sessions' @@ -24,6 +24,11 @@ import { CurrentTurnChangeCard } from './CurrentTurnChangeCard' import type { AgentTaskNotification, UIMessage } from '../../types/chat' import { ConfirmDialog } from '../shared/ConfirmDialog' import { clearWindowSelection, getSelectionPopoverPosition, useSelectionPopoverDismiss } from '../../hooks/useSelectionPopoverDismiss' +import { + getHeightsForSession, + getMetricsForSession, + type VirtualRenderItemMetric, +} from './virtualHeightCache' type ToolCall = Extract type ToolResult = Extract @@ -831,12 +836,6 @@ type VirtualTranscriptWindow = { items: VirtualTranscriptItem[] } -type VirtualRenderItemMetric = { - signature: string - contentWeight: number - estimatedHeight: number -} - const sessionScrollSnapshots = new Map() function isNearScrollBottom(element: HTMLElement) { @@ -1233,8 +1232,12 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { (chatState === 'thinking' && Boolean(activeThinkingId)) const scrollContainerRef = useRef(null) const scrollContentRef = useRef(null) - const virtualItemHeightsRef = useRef(new Map()) - const virtualItemMetricCacheRef = useRef(new Map()) + const virtualItemHeightsRef = useRef>( + resolvedSessionId ? getHeightsForSession(resolvedSessionId) : new Map(), + ) + const virtualItemMetricCacheRef = useRef>( + resolvedSessionId ? getMetricsForSession(resolvedSessionId) : new Map(), + ) const pendingMeasuredHeightsRef = useRef(false) const measureFlushFrameRef = useRef(null) const lastAutoScrollAtRef = useRef(0) @@ -1393,8 +1396,12 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { const snapshot = resolvedSessionId ? sessionScrollSnapshots.get(resolvedSessionId) : undefined shouldAutoScrollRef.current = snapshot?.wasAtBottom ?? true lastSessionIdRef.current = resolvedSessionId - virtualItemHeightsRef.current.clear() - virtualItemMetricCacheRef.current.clear() + virtualItemHeightsRef.current = resolvedSessionId + ? getHeightsForSession(resolvedSessionId) + : new Map() + virtualItemMetricCacheRef.current = resolvedSessionId + ? getMetricsForSession(resolvedSessionId) + : new Map() pendingMeasuredHeightsRef.current = false if (measureFlushFrameRef.current !== null) { cancelAnimationFrame(measureFlushFrameRef.current) @@ -1412,7 +1419,29 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { viewportHeight: container.clientHeight || current.viewportHeight || VIRTUAL_DEFAULT_VIEWPORT_HEIGHT, })) setShowJumpToLatest(true) + } else if (container) { + // Switch to a session we were at the bottom of (or first visit): write + // the bottom sentinel without going through scrollToBottom's read path, + // so we never force a layout flush during the switch's commit. + ignoreProgrammaticScrollUntilRef.current = performance.now() + 250 + ignoreProgrammaticScrollTopRef.current = null + lastAutoScrollAtRef.current = performance.now() + shouldAutoScrollRef.current = true + setScrollToBottomWithoutLayoutRead(container, 'auto') + setVirtualViewport((current) => ({ + scrollTop: SCROLL_BOTTOM_SENTINEL, + viewportHeight: container.clientHeight || current.viewportHeight || VIRTUAL_DEFAULT_VIEWPORT_HEIGHT, + })) + setShowJumpToLatest(false) + if (resolvedSessionId) { + sessionScrollSnapshots.set(resolvedSessionId, { + scrollTop: container.scrollTop, + wasAtBottom: true, + }) + } } else { + // No container yet (initial mount before ref settles): fall back to the + // existing scrollToBottom path which is safe pre-mount. scrollToBottom('auto') } } @@ -1465,11 +1494,21 @@ export function MessageList({ sessionId, compact = false }: MessageListProps = { () => buildRenderModel(messages, activeAskUserQuestionToolUseId), [activeAskUserQuestionToolUseId, messages], ) + // Defer the per-message branchable / completed-turn computations so the first + // commit on tab switch can render the virtualization window without doing two + // additional O(N) walks synchronously. They re-run in a low-priority render + // once the initial frame is painted. + const deferredMessages = useDeferredValue(messages) const branchableMessageTargets = useMemo( - () => branchActionsDisabled ? new Map() : getBranchableMessageTargets(messages), - [branchActionsDisabled, messages], + () => branchActionsDisabled + ? new Map() + : getBranchableMessageTargets(deferredMessages), + [branchActionsDisabled, deferredMessages], + ) + const completedTurnTargets = useMemo( + () => getCompletedTurnTargets(deferredMessages), + [deferredMessages], ) - const completedTurnTargets = useMemo(() => getCompletedTurnTargets(messages), [messages]) const latestCompletedTurnId = completedTurnTargets.length > 0 ? completedTurnTargets[completedTurnTargets.length - 1]?.messageId ?? null diff --git a/desktop/src/components/chat/virtualHeightCache.test.ts b/desktop/src/components/chat/virtualHeightCache.test.ts new file mode 100644 index 00000000..c411cf21 --- /dev/null +++ b/desktop/src/components/chat/virtualHeightCache.test.ts @@ -0,0 +1,61 @@ +import { beforeEach, describe, it, expect } from 'vitest' +import { + __virtualHeightCacheInternals, + dropSession, + getHeightsForSession, + getMetricsForSession, +} from './virtualHeightCache' + +describe('virtualHeightCache', () => { + beforeEach(() => { + __virtualHeightCacheInternals.reset() + }) + + it('returns the same map instance across calls for the same session', () => { + const a = getHeightsForSession('session-a') + a.set('msg-1', 100) + const aAgain = getHeightsForSession('session-a') + expect(aAgain).toBe(a) + expect(aAgain.get('msg-1')).toBe(100) + }) + + it('isolates per-session entries', () => { + getHeightsForSession('s1').set('x', 42) + getHeightsForSession('s2').set('x', 99) + expect(getHeightsForSession('s1').get('x')).toBe(42) + expect(getHeightsForSession('s2').get('x')).toBe(99) + }) + + it('evicts the least recently used session beyond the LRU bound', () => { + for (let i = 0; i < 18; i++) { + getHeightsForSession(`session-${i}`).set('marker', i) + } + expect(__virtualHeightCacheInternals.size()).toBe(16) + // session-0 should have been evicted (oldest, never re-touched) + const restored = getHeightsForSession('session-0') + expect(restored.size).toBe(0) + }) + + it('dropSession removes both height and metric maps', () => { + getHeightsForSession('s').set('a', 1) + getMetricsForSession('s').set('a', { signature: 'sig', contentWeight: 10, estimatedHeight: 20 }) + + dropSession('s') + + expect(getHeightsForSession('s').size).toBe(0) + expect(getMetricsForSession('s').size).toBe(0) + }) + + it('switching back to a previously visited session preserves measurements', () => { + const original = getHeightsForSession('long-session') + original.set('item-1', 240) + original.set('item-2', 360) + + // simulate switching to another session, then back + getHeightsForSession('other-session').set('item-1', 100) + const restored = getHeightsForSession('long-session') + + expect(restored.get('item-1')).toBe(240) + expect(restored.get('item-2')).toBe(360) + }) +}) diff --git a/desktop/src/components/chat/virtualHeightCache.ts b/desktop/src/components/chat/virtualHeightCache.ts new file mode 100644 index 00000000..941d6e32 --- /dev/null +++ b/desktop/src/components/chat/virtualHeightCache.ts @@ -0,0 +1,68 @@ +// Module-level caches keyed by sessionId so switching tabs back to a previously +// rendered long transcript can skip estimate→measure thrash. Bounded by an LRU +// across sessions to avoid unbounded growth across long-running sessions. + +export type VirtualRenderItemMetric = { + signature: string + contentWeight: number + estimatedHeight: number +} + +const MAX_TRACKED_SESSIONS = 16 + +const sessionHeightCache = new Map>() +const sessionMetricCache = new Map>() + +function touchSession(sessionId: string, map: Map>) { + // Reinsert to move to LRU tail. + const existing = map.get(sessionId) + if (existing) { + map.delete(sessionId) + map.set(sessionId, existing) + } +} + +function evictSessionsBeyondLimit(): void { + while (sessionHeightCache.size > MAX_TRACKED_SESSIONS) { + const oldest = sessionHeightCache.keys().next().value + if (typeof oldest !== 'string') break + sessionHeightCache.delete(oldest) + sessionMetricCache.delete(oldest) + } +} + +export function getHeightsForSession(sessionId: string): Map { + let heights = sessionHeightCache.get(sessionId) + if (!heights) { + heights = new Map() + sessionHeightCache.set(sessionId, heights) + evictSessionsBeyondLimit() + } else { + touchSession(sessionId, sessionHeightCache as Map>) + } + return heights +} + +export function getMetricsForSession(sessionId: string): Map { + let metrics = sessionMetricCache.get(sessionId) + if (!metrics) { + metrics = new Map() + sessionMetricCache.set(sessionId, metrics) + } else { + touchSession(sessionId, sessionMetricCache as Map>) + } + return metrics +} + +export function dropSession(sessionId: string): void { + sessionHeightCache.delete(sessionId) + sessionMetricCache.delete(sessionId) +} + +export const __virtualHeightCacheInternals = { + size: () => sessionHeightCache.size, + reset: () => { + sessionHeightCache.clear() + sessionMetricCache.clear() + }, +} diff --git a/desktop/src/stores/tabStore.ts b/desktop/src/stores/tabStore.ts index 1d4d7ad9..aff96cad 100644 --- a/desktop/src/stores/tabStore.ts +++ b/desktop/src/stores/tabStore.ts @@ -1,5 +1,6 @@ import { create } from 'zustand' import { sessionsApi } from '../api/sessions' +import { dropSession as dropVirtualHeightSession } from '../components/chat/virtualHeightCache' const TAB_STORAGE_KEY = 'cc-haha-open-tabs' @@ -108,6 +109,7 @@ export const useTabStore = create((set, get) => ({ set({ tabs: newTabs, activeTabId: newActiveId }) get().saveTabs() + dropVirtualHeightSession(sessionId) }, setActiveTab: (sessionId) => {