From d5b024e174c83ba19b03fd92fa05112b0907fc51 Mon Sep 17 00:00:00 2001 From: Diao Shengjia <104132148+GoDiao@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:02:06 +0800 Subject: [PATCH] perf(desktop): subscribe components to the session fields they use ActiveSession, ChatInput, ModelSelector, StreamingIndicator, MessageList, PermissionDialog and AskUserQuestion select only the fields or derived values they render instead of the whole session object or store. Each tool group gets only its own results, children and agent notifications, with references kept while they are unchanged, so finished groups no longer re-render while a reply streams. Part of #153 --- .../src/components/chat/AskUserQuestion.tsx | 14 +- desktop/src/components/chat/ChatInput.tsx | 34 ++- .../src/components/chat/MessageList.test.tsx | 106 ++++++- desktop/src/components/chat/MessageList.tsx | 126 ++++++-- .../src/components/chat/PermissionDialog.tsx | 7 +- .../components/chat/StreamingIndicator.tsx | 21 +- .../chat/streamingRenderCounts.test.tsx | 289 ++++++++++++++++++ .../chat/streamingRenderScenario.ts | 142 +++++++++ .../src/components/controls/ModelSelector.tsx | 8 +- .../pages/ActiveSession.renderCounts.test.tsx | 117 +++++++ desktop/src/pages/ActiveSession.tsx | 30 +- desktop/src/stores/chatStore.ts | 12 + 12 files changed, 829 insertions(+), 77 deletions(-) create mode 100644 desktop/src/components/chat/streamingRenderCounts.test.tsx create mode 100644 desktop/src/components/chat/streamingRenderScenario.ts create mode 100644 desktop/src/pages/ActiveSession.renderCounts.test.tsx diff --git a/desktop/src/components/chat/AskUserQuestion.tsx b/desktop/src/components/chat/AskUserQuestion.tsx index a93610ff..03b4c2c8 100644 --- a/desktop/src/components/chat/AskUserQuestion.tsx +++ b/desktop/src/components/chat/AskUserQuestion.tsx @@ -61,10 +61,13 @@ function getSelectedAnswer(question: Question, selected: string[] | undefined) { } export function AskUserQuestion({ sessionId, toolUseId, input, result }: Props) { - const { respondToPermission } = useChatStore() + const respondToPermission = useChatStore((s) => s.respondToPermission) const activeTabId = useTabStore((s) => s.activeTabId) const targetSessionId = sessionId ?? activeTabId - const pendingPermission = useChatStore((s) => targetSessionId ? s.sessions[targetSessionId]?.pendingPermission : undefined) + const pendingRequestId = useChatStore((s) => { + const pending = targetSessionId ? s.sessions[targetSessionId]?.pendingPermission : undefined + return pending?.toolUseId === toolUseId ? pending.requestId : null + }) const t = useTranslation() const questions = parseInput(input) const inputObject = (input && typeof input === 'object') ? input as Record : {} @@ -88,7 +91,6 @@ export function AskUserQuestion({ sessionId, toolUseId, input, result }: Props) const hasStructuredAnswers = Object.keys(resultAnswers).length > 0 const hasTerminalResult = hasStructuredAnswers || resultText.length > 0 - const pendingRequest = pendingPermission?.toolUseId === toolUseId ? pendingPermission : null const answeredText = useMemo(() => { if (hasStructuredAnswers) { return questions @@ -169,7 +171,7 @@ export function AskUserQuestion({ sessionId, toolUseId, input, result }: Props) const response = parts.join('; ') if (!response) return - if (!targetSessionId || !pendingRequest) return + if (!targetSessionId || !pendingRequestId) return const answers = questions.reduce>((acc, question, index) => { const freeText = freeTexts[index]?.trim() @@ -183,7 +185,7 @@ export function AskUserQuestion({ sessionId, toolUseId, input, result }: Props) }, {}) setHasSubmitted(true) - respondToPermission(targetSessionId, pendingRequest.requestId, true, { + respondToPermission(targetSessionId, pendingRequestId, true, { updatedInput: { ...inputObject, answers, @@ -384,7 +386,7 @@ export function AskUserQuestion({ sessionId, toolUseId, input, result }: Props) + ), +})) + +vi.mock('../controls/ModelSelector', () => ({ + ModelSelector: () => , +})) + +vi.mock('./ToolCallGroup', async (importOriginal) => { + const actual = await importOriginal() + const { memo } = await import('react') + const Real = actual.ToolCallGroup + return { + ...actual, + ToolCallGroup: memo(function CountedToolCallGroup(props: Parameters[0]) { + const id = props.toolCalls[0]!.id + counts.toolGroups[id] = (counts.toolGroups[id] ?? 0) + 1 + counts.toolGroupProps[id] = props + return + }), + } +}) + +import { ChatInput } from './ChatInput' +import { StreamingIndicator } from './StreamingIndicator' +import { MessageList } from './MessageList' +import { PermissionDialog } from './PermissionDialog' +import { AskUserQuestion } from './AskUserQuestion' +import { useChatStore, useSessionRuntimeSelection, type PerSessionState } from '../../stores/chatStore' +import { useSettingsStore } from '../../stores/settingsStore' +import { useTabStore } from '../../stores/tabStore' +import { + HISTORY_GROUP_ID, + LIVE_GROUP_ID, + SCENARIO_SESSION_ID, + SCENARIO_UPDATE_COUNT, + installScenarioSession, + playStreamingReply, +} from './streamingRenderScenario' + +function Counted({ id, children }: { id: string; children: ReactNode }) { + return ( + { counts.commits[id] = (counts.commits[id] ?? 0) + 1 }}> + {children} + + ) +} + +function resetCounts() { + counts.commits = {} + counts.toolGroups = {} + counts.toolGroupProps = {} +} + +function patchSession(update: Partial) { + act(() => { + useChatStore.setState((state) => ({ + sessions: { + ...state.sessions, + [SCENARIO_SESSION_ID]: { ...state.sessions[SCENARIO_SESSION_ID]!, ...update }, + }, + })) + }) +} + +describe('render counts during one streaming reply', () => { + const initialChatState = useChatStore.getInitialState() + const initialTabState = useTabStore.getInitialState() + + beforeEach(() => { + useSettingsStore.setState({ locale: 'en' }) + useChatStore.setState(initialChatState, true) + useTabStore.setState(initialTabState, true) + installScenarioSession() + }) + + afterEach(() => { + cleanup() + }) + + it('re-renders each component only when the fields it shows change', () => { + render( + <> + + + + + + + + + + , + ) + resetCounts() + + playStreamingReply() + + expect({ + ...counts.commits, + historyToolGroup: counts.toolGroups[HISTORY_GROUP_ID] ?? 0, + liveToolGroup: counts.toolGroups[LIVE_GROUP_ID] ?? 0, + }).toEqual({ + ChatInput: 8, + StreamingIndicator: 14, + MessageList: 62, + PermissionDialog: 2, + historyToolGroup: 0, + liveToolGroup: 4, + }) + expect(SCENARIO_UPDATE_COUNT).toBe(31) + expect(counts.toolGroupProps[LIVE_GROUP_ID]!.resultMap.get('tu-live')?.content).toBe('all passed') + }) +}) + +describe('components still update on the fields they use', () => { + const initialChatState = useChatStore.getInitialState() + const initialTabState = useTabStore.getInitialState() + + beforeEach(() => { + useSettingsStore.setState({ locale: 'en' }) + useChatStore.setState(initialChatState, true) + useTabStore.setState(initialTabState, true) + installScenarioSession() + resetCounts() + }) + + afterEach(() => { + cleanup() + }) + + it('ChatInput follows permissionMode, the effective mode and composerPrefill', () => { + render() + + patchSession({ permissionMode: 'plan' }) + expect(screen.getByTestId('permission-mode')).toHaveAttribute('data-value', 'plan') + + const session = useChatStore.getState().sessions[SCENARIO_SESSION_ID]! + patchSession({ + mirror: { ...session.mirror!, permission: { desired: 'plan', effective: 'acceptEdits' } }, + }) + expect(screen.getByTestId('permission-mode')).toHaveAttribute('data-effective', 'acceptEdits') + + patchSession({ composerPrefill: { text: 'try again', nonce: 1 } }) + expect(screen.getByRole('textbox')).toHaveValue('try again') + }) + + it('ChatInput follows slashCommands and chatState', () => { + render() + + patchSession({ slashCommands: [{ name: 'deploy-now', description: 'Ship it' }] }) + fireEvent.change(screen.getByRole('textbox'), { target: { value: '/deploy', selectionStart: 7 } }) + expect(screen.getAllByRole('button').some((button) => button.textContent?.startsWith('/deploy-now'))).toBe(true) + + expect(screen.queryByTitle('Stop generation (Cmd+.)')).not.toBeInTheDocument() + patchSession({ chatState: 'streaming' }) + expect(screen.getByTitle('Stop generation (Cmd+.)')).toBeInTheDocument() + }) + + it('StreamingIndicator follows statusVerb, tokenUsage and apiRetry', () => { + patchSession({ chatState: 'thinking' }) + render() + + patchSession({ statusVerb: 'Pondering' }) + expect(screen.getByText('Pondering...')).toBeInTheDocument() + + patchSession({ tokenUsage: { input_tokens: 10, output_tokens: 42 } }) + expect(screen.getByText(/↓ 42/)).toBeInTheDocument() + + patchSession({ + apiRetry: { + attempt: 1, + maxRetries: 3, + retryDelayMs: 1000, + errorStatus: 529, + receivedAt: Date.now(), + } as PerSessionState['apiRetry'], + }) + expect(screen.getByTestId('api-retry-indicator')).toBeInTheDocument() + }) + + it('PermissionDialog follows its pending request', () => { + render() + expect(screen.queryByText('Allow')).not.toBeInTheDocument() + + patchSession({ pendingPermission: { requestId: 'perm-1', toolName: 'Bash', input: { command: 'ls' } } }) + expect(screen.getByText('Allow')).toBeInTheDocument() + + patchSession({ pendingPermission: null }) + expect(screen.queryByText('Allow')).not.toBeInTheDocument() + }) + + it('AskUserQuestion follows its pending request', () => { + render( + , + ) + fireEvent.click(screen.getByRole('button', { name: /^No$/ })) + expect(screen.getByRole('button', { name: /submit/i })).toBeDisabled() + + patchSession({ + pendingPermission: { requestId: 'perm-q', toolName: 'AskUserQuestion', toolUseId: 'tu-question', input: {} }, + }) + expect(screen.getByRole('button', { name: /submit/i })).toBeEnabled() + }) + + it('the runtime selection keeps its reference until the desired runtime changes', () => { + const { result } = renderHook(() => useSessionRuntimeSelection(SCENARIO_SESSION_ID)) + const session = useChatStore.getState().sessions[SCENARIO_SESSION_ID]! + patchSession({ + mirror: { + ...session.mirror!, + runtime: { desired: { providerId: 'p1', modelId: 'm1', effort: 'high' }, effective: null }, + } as PerSessionState['mirror'], + }) + const first = result.current + expect(first).toEqual({ providerId: 'p1', modelId: 'm1', effort: 'high' }) + + patchSession({ streamingText: 'unrelated', mirror: { ...useChatStore.getState().sessions[SCENARIO_SESSION_ID]!.mirror! } }) + expect(result.current).toBe(first) + + const mirror = useChatStore.getState().sessions[SCENARIO_SESSION_ID]!.mirror! + patchSession({ + mirror: { ...mirror, runtime: { desired: { providerId: 'p1', modelId: 'm2', effort: 'high' }, effective: null } } as PerSessionState['mirror'], + }) + expect(result.current).toEqual({ providerId: 'p1', modelId: 'm2', effort: 'high' }) + }) + + it('a tool group re-renders when its own result arrives', () => { + render() + const session = useChatStore.getState().sessions[SCENARIO_SESSION_ID]! + patchSession({ + messages: [ + ...session.messages, + { id: LIVE_GROUP_ID, type: 'tool_use', toolName: 'Bash', toolUseId: 'tu-live', input: { command: 'ls' }, timestamp: 1 }, + ], + }) + expect(counts.toolGroupProps[LIVE_GROUP_ID]!.resultMap.size).toBe(0) + resetCounts() + + patchSession({ + messages: [ + ...useChatStore.getState().sessions[SCENARIO_SESSION_ID]!.messages, + { id: 'r-live', type: 'tool_result', toolUseId: 'tu-live', content: 'files', isError: false, timestamp: 2 }, + ], + }) + expect(counts.toolGroups[LIVE_GROUP_ID]).toBe(1) + expect(counts.toolGroupProps[LIVE_GROUP_ID]!.resultMap.get('tu-live')?.content).toBe('files') + expect(counts.toolGroups[HISTORY_GROUP_ID]).toBeUndefined() + }) +}) diff --git a/desktop/src/components/chat/streamingRenderScenario.ts b/desktop/src/components/chat/streamingRenderScenario.ts new file mode 100644 index 00000000..ec25f893 --- /dev/null +++ b/desktop/src/components/chat/streamingRenderScenario.ts @@ -0,0 +1,142 @@ +import { act } from 'react' +import { useChatStore, type PerSessionState } from '../../stores/chatStore' +import { useTabStore } from '../../stores/tabStore' +import type { SessionState, UIMessage } from '../../types/chat' + +export const SCENARIO_SESSION_ID = 'render-count-session' + +const T0 = 1_000_000 + +function mirror(turn: SessionState['turn']): SessionState { + return { + id: SCENARIO_SESSION_ID, + version: 1, + permission: { desired: 'default', effective: 'default' }, + runtime: { desired: null, effective: null }, + title: { text: null, source: null }, + workspace: { workDir: '/repo', projectRoot: '/repo', workDirExists: true }, + process: { phase: 'running', generation: 1 }, + turn, + pendingPermissions: [], + slashCommands: [], + activeGoal: null, + backgroundTasks: [], + team: null, + clientCount: 1, + } +} + +const IDLE_TURN = { phase: 'idle' } as SessionState['turn'] +const RUNNING_TURN = { phase: 'running', activity: 'thinking', startedAt: T0 } as SessionState['turn'] + +export const HISTORY_GROUP_ID = 'history-tool' +export const LIVE_GROUP_ID = 'live-tool' + +const historyMessages: UIMessage[] = [ + { id: 'u1', type: 'user_text', content: 'Read the readme', timestamp: T0 - 100 }, + { id: 'a1', type: 'assistant_text', content: 'Reading it.', timestamp: T0 - 90 }, + { id: HISTORY_GROUP_ID, type: 'tool_use', toolName: 'Read', toolUseId: 'tu-history', input: { file_path: '/repo/README.md' }, timestamp: T0 - 80 }, + { id: 'r-history', type: 'tool_result', toolUseId: 'tu-history', content: 'readme body', isError: false, timestamp: T0 - 70 }, + { id: 'a2', type: 'assistant_text', content: 'Done.', timestamp: T0 - 60 }, +] + +export function scenarioInitialSession(): PerSessionState { + return { + messages: historyMessages, + chatState: 'idle', + connectionState: 'connected', + streamingText: '', + streamingToolInput: '', + activeToolUseId: null, + activeToolName: null, + activeThinkingId: null, + pendingPermission: null, + tokenUsage: { input_tokens: 0, output_tokens: 0 }, + statusVerb: '', + apiRetry: null, + slashCommands: [{ name: 'help', description: 'Show help' }], + agentTaskNotifications: {}, + backgroundAgentTasks: {}, + activeGoal: null, + composerPrefill: null, + composerInsertion: null, + composerDraft: null, + mirror: mirror(IDLE_TURN), + } +} + +export function installScenarioSession() { + useChatStore.setState({ sessions: { [SCENARIO_SESSION_ID]: scenarioInitialSession() } }) + useTabStore.setState({ + activeTabId: SCENARIO_SESSION_ID, + tabs: [{ sessionId: SCENARIO_SESSION_ID, title: 'Render counts', type: 'session', status: 'idle' }], + }) +} + +function patch(update: (session: PerSessionState) => Partial) { + act(() => { + useChatStore.setState((state) => { + const session = state.sessions[SCENARIO_SESSION_ID]! + return { sessions: { ...state.sessions, [SCENARIO_SESSION_ID]: { ...session, ...update(session) } } } + }) + }) +} + +function append(...messages: UIMessage[]) { + return (session: PerSessionState) => ({ messages: [...session.messages, ...messages] }) +} + +export function playStreamingReply() { + let tokens = 0 + const usage = () => { + tokens += 25 + return { tokenUsage: { input_tokens: 1200, output_tokens: tokens } } + } + + patch((s) => ({ + ...append({ id: 'u2', type: 'user_text', content: 'Run the tests', timestamp: T0 })(s), + chatState: 'thinking', + statusVerb: 'Thinking', + mirror: mirror(RUNNING_TURN), + })) + patch(usage) + for (let i = 0; i < 10; i++) { + patch((s) => ({ chatState: 'streaming', streamingText: `${s.streamingText}chunk ${i} ` })) + } + for (let i = 0; i < 3; i++) patch(usage) + patch((s) => ({ + ...append( + { id: 'a3', type: 'assistant_text', content: s.streamingText, timestamp: T0 + 10 }, + { id: LIVE_GROUP_ID, type: 'tool_use', toolName: 'Bash', toolUseId: 'tu-live', input: { command: 'bun test' }, timestamp: T0 + 11 }, + )(s), + streamingText: '', + chatState: 'tool_executing', + activeToolUseId: 'tu-live', + activeToolName: 'Bash', + })) + patch(() => ({ + chatState: 'permission_pending', + pendingPermission: { requestId: 'perm-1', toolName: 'Bash', toolUseId: 'tu-live', input: { command: 'bun test' } }, + })) + patch(() => ({ chatState: 'tool_executing', pendingPermission: null })) + patch((s) => ({ + ...append({ id: 'r-live', type: 'tool_result', toolUseId: 'tu-live', content: 'all passed', isError: false, timestamp: T0 + 20 })(s), + activeToolUseId: null, + activeToolName: null, + chatState: 'thinking', + })) + patch(usage) + for (let i = 0; i < 10; i++) { + patch((s) => ({ chatState: 'streaming', streamingText: `${s.streamingText}more ${i} ` })) + } + patch((s) => ({ + ...append({ id: 'a4', type: 'assistant_text', content: s.streamingText, timestamp: T0 + 30 })(s), + streamingText: '', + chatState: 'idle', + statusVerb: '', + mirror: mirror(IDLE_TURN), + ...usage(), + })) +} + +export const SCENARIO_UPDATE_COUNT = 31 diff --git a/desktop/src/components/controls/ModelSelector.tsx b/desktop/src/components/controls/ModelSelector.tsx index 61f4e8c1..832cbac2 100644 --- a/desktop/src/components/controls/ModelSelector.tsx +++ b/desktop/src/components/controls/ModelSelector.tsx @@ -2,7 +2,7 @@ import { useCallback, useEffect, useId, useLayoutEffect, useMemo, useRef, useSta import { createPortal } from 'react-dom' import { OFFICIAL_DEFAULT_MODEL_ID } from '../../constants/modelCatalog' import { useTranslation } from '../../i18n' -import { sessionRuntimeSelection, useChatStore } from '../../stores/chatStore' +import { useChatStore, useSessionRuntimeSelection } from '../../stores/chatStore' import { useProviderStore } from '../../stores/providerStore' import { DRAFT_RUNTIME_SELECTION_KEY, useSessionRuntimeStore } from '../../stores/sessionRuntimeStore' import { useSettingsStore } from '../../stores/settingsStore' @@ -149,10 +149,10 @@ export function ModelSelector({ const draftRuntimeSelection = useSessionRuntimeStore((state) => runtimeKey === DRAFT_RUNTIME_SELECTION_KEY ? state.selections[runtimeKey] : undefined, ) - const runtimeSession = useChatStore((state) => - runtimeKey && runtimeKey !== DRAFT_RUNTIME_SELECTION_KEY ? state.sessions[runtimeKey] : undefined, + const sessionSelection = useSessionRuntimeSelection( + runtimeKey && runtimeKey !== DRAFT_RUNTIME_SELECTION_KEY ? runtimeKey : null, ) - const runtimeSelection = draftRuntimeSelection ?? sessionRuntimeSelection(runtimeSession) + const runtimeSelection = draftRuntimeSelection ?? sessionSelection const [open, setOpen] = useState(false) const [dropdownPosition, setDropdownPosition] = useState(null) const ref = useRef(null) diff --git a/desktop/src/pages/ActiveSession.renderCounts.test.tsx b/desktop/src/pages/ActiveSession.renderCounts.test.tsx new file mode 100644 index 00000000..b3405856 --- /dev/null +++ b/desktop/src/pages/ActiveSession.renderCounts.test.tsx @@ -0,0 +1,117 @@ +import { Profiler } from 'react' +import { afterEach, describe, expect, it, vi } from 'vitest' +import { act, cleanup, render, screen } from '@testing-library/react' +import '@testing-library/jest-dom' + +vi.mock('../api/websocket', () => ({ + wsManager: { + connect: vi.fn(), + disconnect: vi.fn(), + onMessage: vi.fn(() => () => {}), + clearHandlers: vi.fn(), + send: vi.fn(), + }, +})) + +vi.mock('../api/sessions', () => ({ + sessionsApi: new Proxy({}, { get: () => () => new Promise(() => {}) }), +})) + +vi.mock('../hooks/useMobileViewport', () => ({ + useMobileViewport: () => false, +})) + +vi.mock('../components/chat/MessageList', () => ({ + MessageList: () =>
, +})) + +vi.mock('../components/chat/ChatInput', () => ({ + ChatInput: () =>
, +})) + +vi.mock('../components/teams/TeamStatusBar', () => ({ + TeamStatusBar: () =>
, +})) + +vi.mock('../components/chat/SessionTaskBar', () => ({ + SessionTaskBar: () =>
, +})) + +vi.mock('../components/workspace/WorkspacePanel', () => ({ + WorkspacePanel: () =>
, +})) + +import { ActiveSession } from './ActiveSession' +import { useChatStore, type PerSessionState } from '../stores/chatStore' +import { useCLITaskStore } from '../stores/cliTaskStore' +import { useSettingsStore } from '../stores/settingsStore' +import { useTabStore } from '../stores/tabStore' +import { + SCENARIO_SESSION_ID, + installScenarioSession, + playStreamingReply, +} from '../components/chat/streamingRenderScenario' + +afterEach(() => { + cleanup() + useChatStore.setState(useChatStore.getInitialState(), true) + useTabStore.setState(useTabStore.getInitialState(), true) +}) + +describe('ActiveSession render count during one streaming reply', () => { + it('re-renders only when the page-level fields change', () => { + useCLITaskStore.setState({ fetchSessionTasks: vi.fn(async () => {}) }) + installScenarioSession() + let commits = 0 + render( + { commits += 1 }}> + + , + ) + commits = 0 + + playStreamingReply() + + expect(commits).toBe(2) + }) +}) + +function patchSession(update: Partial) { + act(() => { + useChatStore.setState((state) => ({ + sessions: { + ...state.sessions, + [SCENARIO_SESSION_ID]: { ...state.sessions[SCENARIO_SESSION_ID]!, ...update }, + }, + })) + }) +} + +describe('ActiveSession still updates on the fields it uses', () => { + it('follows the running state, background tasks, the goal and the first message', () => { + useSettingsStore.setState({ locale: 'en' }) + useCLITaskStore.setState({ fetchSessionTasks: vi.fn(async () => {}) }) + installScenarioSession() + patchSession({ messages: [] }) + render() + + expect(screen.queryByTestId('message-list')).not.toBeInTheDocument() + patchSession({ streamingText: 'Hello' }) + expect(screen.getByTestId('message-list')).toBeInTheDocument() + + expect(screen.queryByTestId('session-header-running')).not.toBeInTheDocument() + patchSession({ chatState: 'thinking' }) + expect(screen.getByTestId('session-header-running')).toBeInTheDocument() + patchSession({ chatState: 'idle' }) + expect(screen.queryByTestId('session-header-running')).not.toBeInTheDocument() + + patchSession({ backgroundAgentTasks: { t1: { taskId: 't1', status: 'running', startedAt: 1, updatedAt: 1 } } }) + expect(screen.getByTestId('session-header-running')).toBeInTheDocument() + patchSession({ backgroundAgentTasks: { t1: { taskId: 't1', status: 'completed', startedAt: 1, updatedAt: 2 } } }) + expect(screen.queryByTestId('session-header-running')).not.toBeInTheDocument() + + expect(screen.queryByText('Active goal')).not.toBeInTheDocument() + patchSession({ activeGoal: { action: 'created', objective: 'Ship the release', updatedAt: 1 } }) + expect(screen.getByText('Active goal')).toBeInTheDocument() + }) +}) diff --git a/desktop/src/pages/ActiveSession.tsx b/desktop/src/pages/ActiveSession.tsx index aa8aac5c..4c1b165a 100644 --- a/desktop/src/pages/ActiveSession.tsx +++ b/desktop/src/pages/ActiveSession.tsx @@ -6,7 +6,7 @@ import { type TabType, } from '../stores/tabStore' import { useSessionStore } from '../stores/sessionStore' -import { useChatStore } from '../stores/chatStore' +import { useChatStore, useSessionField } from '../stores/chatStore' import { useCLITaskStore } from '../stores/cliTaskStore' import { useTeamStore } from '../stores/teamStore' import { useUIStore } from '../stores/uiStore' @@ -239,14 +239,23 @@ export function ActiveSession() { const activeTabType = useTabStore((s) => s.tabs.find((tab) => tab.sessionId === s.activeTabId)?.type ?? null) const sessions = useSessionStore((s) => s.sessions) const connectToSession = useChatStore((s) => s.connectToSession) - const sessionState = useChatStore((s) => activeTabId ? s.sessions[activeTabId] : undefined) + const isRunning = useChatStore((s) => { + const chatState = activeTabId ? s.sessions[activeTabId]?.chatState : undefined + return chatState !== undefined && chatState !== 'idle' + }) + const hasRunningBackgroundTasks = useChatStore((s) => { + const tasks = activeTabId ? s.sessions[activeTabId]?.backgroundAgentTasks : undefined + return tasks ? Object.values(tasks).some((task) => task.status === 'running') : false + }) + const hasLoadedTranscript = useChatStore((s) => { + const sessionState = activeTabId ? s.sessions[activeTabId] : undefined + return Boolean(sessionState && (sessionState.messages.length > 0 || sessionState.streamingText)) + }) + const activeGoal = useSessionField(activeTabId, 'activeGoal') ?? null const fetchSessionTasks = useCLITaskStore((s) => s.fetchSessionTasks) const trackedTaskSessionId = useCLITaskStore((s) => s.sessionId) const hasIncompleteTasks = useCLITaskStore((s) => s.tasks.some((task) => task.status !== 'completed')) const hasRunningTasks = useCLITaskStore((s) => s.tasks.some((task) => task.status === 'in_progress')) - const chatState = sessionState?.chatState ?? 'idle' - const hasRunningBackgroundTasks = Object.values(sessionState?.backgroundAgentTasks ?? {}) - .some((task) => task.status === 'running') const session = sessions.find((s) => s.id === activeTabId) const memberInfo = useTeamStore((s) => activeTabId ? s.getMemberBySessionId(activeTabId) : null) @@ -299,7 +308,7 @@ export function ActiveSession() { if (!activeTabId || isMemberSession) return const shouldPollTasks = - chatState !== 'idle' || + isRunning || (trackedTaskSessionId === activeTabId && hasIncompleteTasks) if (!shouldPollTasks) return @@ -314,19 +323,16 @@ export function ActiveSession() { }, [ activeTabId, isMemberSession, - chatState, + isRunning, trackedTaskSessionId, hasIncompleteTasks, fetchSessionTasks, ]) const t = useTranslation() - const messages = sessionState?.messages ?? [] - const streamingText = sessionState?.streamingText ?? '' - const activeGoal = sessionState?.activeGoal ?? null - const isEmpty = messages.length === 0 && !streamingText && (session?.messageCount ?? 0) === 0 + const isEmpty = !hasLoadedTranscript && (session?.messageCount ?? 0) === 0 - const isActive = chatState !== 'idle' || + const isActive = isRunning || (trackedTaskSessionId === activeTabId && hasRunningTasks) || hasRunningBackgroundTasks diff --git a/desktop/src/stores/chatStore.ts b/desktop/src/stores/chatStore.ts index ffc44088..b67010a2 100644 --- a/desktop/src/stores/chatStore.ts +++ b/desktop/src/stores/chatStore.ts @@ -1,4 +1,5 @@ import { create } from 'zustand' +import { useShallow } from 'zustand/react/shallow' import { wsManager } from '../api/websocket' import { sessionsApi } from '../api/sessions' import { useTeamStore } from './teamStore' @@ -693,6 +694,17 @@ export function sessionRuntimeSelection(session: PerSessionState | undefined): R } } +export function useSessionField( + sessionId: string | null | undefined, + key: K, +): PerSessionState[K] | undefined { + return useChatStore((state) => (sessionId ? state.sessions[sessionId]?.[key] : undefined)) +} + +export function useSessionRuntimeSelection(sessionId: string | null | undefined): RuntimeSelection | undefined { + return useChatStore(useShallow((state) => (sessionId ? sessionRuntimeSelection(state.sessions[sessionId]) : undefined))) +} + function runtimeConfigMessage(selection: RuntimeSelection) { return { type: 'set_runtime_config' as const,