perf(desktop): subscribe components to the session fields they use

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<string, unknown> : {}
@@ -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<Record<string, string>>((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)
         <Button
           variant="primary"
           size="xs"
-          disabled={!allAnswered || !pendingRequest}
+          disabled={!allAnswered || !pendingRequestId}
           onClick={handleSubmit}
         >
           {t('question.submit')}

desktop/src/components/chat/ChatInput.tsx

@@ -1,6 +1,12 @@
 import { useState, useRef, useEffect, useCallback, useMemo } from 'react'
 import { useTranslation } from '../../i18n'
-import { sessionRuntimeSelection, useChatStore } from '../../stores/chatStore'
+import {
+  sessionRuntimeSelection,
+  useChatStore,
+  useSessionField,
+  useSessionRuntimeSelection,
+  type PerSessionState,
+} from '../../stores/chatStore'
 import { useTabStore } from '../../stores/tabStore'
 import { useUIStore } from '../../stores/uiStore'
 import { useSessionStore } from '../../stores/sessionStore'
@@ -52,6 +58,7 @@ type ChatInputProps = {
 }
 
 const EMPTY_WORKSPACE_REFERENCES: WorkspaceChatReference[] = []
+const EMPTY_SLASH_COMMANDS: PerSessionState['slashCommands'] = []
 
 function workspaceReferenceToAttachment(reference: WorkspaceChatReference): Attachment {
   return {
@@ -124,24 +131,29 @@ export function ChatInput({ variant = 'default', compact = false }: ChatInputPro
       return next
     })
   }, [])
-  const { sendMessage, stopGeneration, clearComposerInsertion, setSessionPermissionMode } = useChatStore()
+  const sendMessage = useChatStore((s) => s.sendMessage)
+  const stopGeneration = useChatStore((s) => s.stopGeneration)
+  const clearComposerInsertion = useChatStore((s) => s.clearComposerInsertion)
+  const setSessionPermissionMode = useChatStore((s) => s.setSessionPermissionMode)
   const activeTabId = useTabStore((s) => s.activeTabId)
-  const sessionState = useChatStore((s) => activeTabId ? s.sessions[activeTabId] : undefined)
-  const chatState = sessionState?.chatState ?? 'idle'
-  const slashCommands = sessionState?.slashCommands ?? []
-  const composerPrefill = sessionState?.composerPrefill ?? null
-  const composerInsertion = sessionState?.composerInsertion ?? null
+  const chatState = useSessionField(activeTabId, 'chatState') ?? 'idle'
+  const slashCommands = useSessionField(activeTabId, 'slashCommands') ?? EMPTY_SLASH_COMMANDS
+  const composerPrefill = useSessionField(activeTabId, 'composerPrefill') ?? null
+  const composerInsertion = useSessionField(activeTabId, 'composerInsertion') ?? null
   const defaultPermissionMode = useSettingsStore((s) => s.permissionMode)
-  const permissionMode = sessionState?.permissionMode ?? defaultPermissionMode
-  const runtimeSelection = sessionRuntimeSelection(sessionState)
+  const permissionMode = useSessionField(activeTabId, 'permissionMode') ?? defaultPermissionMode
+  const effectivePermissionMode = useChatStore((s) =>
+    activeTabId ? s.sessions[activeTabId]?.mirror?.permission.effective ?? undefined : undefined,
+  )
+  const runtimeSelection = useSessionRuntimeSelection(activeTabId)
   const currentModel = useSettingsStore((state) => state.currentModel)
   const runtimeSelectionKey = runtimeSelection
     ? `${runtimeSelection.providerId ?? 'none'}:${runtimeSelection.modelId}`
     : undefined
   const providerMissing = useProviderMissing(runtimeSelection)
   const runtimeModelLabel = runtimeSelection?.modelId || currentModel?.name || currentModel?.id
   const activeSession = useSessionStore((state) => activeTabId ? state.sessions.find((session) => session.id === activeTabId) ?? null : null)
-  const loadedMessageCount = sessionState?.messages?.length ?? 0
+  const loadedMessageCount = useChatStore((s) => activeTabId ? s.sessions[activeTabId]?.messages?.length ?? 0 : 0)
   const messageCount = Math.max(loadedMessageCount, activeSession?.messageCount ?? 0)
   const memberInfo = useTeamStore((s) => activeTabId ? s.getMemberBySessionId(activeTabId) : null)
   const [gitInfo, setGitInfo] = useState<GitInfo | null>(null)
@@ -1093,7 +1105,7 @@ export function ChatInput({ variant = 'default', compact = false }: ChatInputPro
                     <PermissionModeSelector
                       compact={useCompactControls}
                       value={permissionMode}
-                      effectiveValue={(sessionState?.mirror?.permission.effective ?? undefined) as typeof permissionMode | undefined}
+                      effectiveValue={effectivePermissionMode as typeof permissionMode | undefined}
                       onChange={(mode) => {
                         if (activeTabId) setSessionPermissionMode(activeTabId, mode)
                       }}

desktop/src/components/chat/MessageList.test.tsx

@@ -1,6 +1,6 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { act, fireEvent, render, screen, waitFor, within } from '@testing-library/react'
-import { MessageList, buildRenderModel } from './MessageList'
+import { MessageList, buildRenderModel, buildToolGroupResults } from './MessageList'
 import { relativizeWorkspacePath } from './CurrentTurnChangeCard'
 import { sessionsApi } from '../../api/sessions'
 import { useChatStore } from '../../stores/chatStore'
@@ -9,7 +9,7 @@ import { useSettingsStore } from '../../stores/settingsStore'
 import { useSessionStore } from '../../stores/sessionStore'
 import { useTabStore } from '../../stores/tabStore'
 import { useUIStore } from '../../stores/uiStore'
-import type { UIMessage } from '../../types/chat'
+import type { AgentTaskNotification, UIMessage } from '../../types/chat'
 import type { PerSessionState } from '../../stores/chatStore'
 
 const ACTIVE_TAB = 'active-tab'
@@ -3600,3 +3600,105 @@ describe('MessageList nested tool calls', () => {
     ).toBeTruthy()
   })
 })
+
+describe('buildToolGroupResults', () => {
+  const read = (id: string, parentToolUseId?: string): UIMessage => ({
+    id: `use-${id}`,
+    type: 'tool_use',
+    toolName: id.startsWith('agent') ? 'Agent' : 'Read',
+    toolUseId: id,
+    input: {},
+    timestamp: 1,
+    ...(parentToolUseId ? { parentToolUseId } : {}),
+  })
+  const result = (id: string, parentToolUseId?: string): UIMessage => ({
+    id: `result-${id}`,
+    type: 'tool_result',
+    toolUseId: id,
+    content: `${id} done`,
+    isError: false,
+    timestamp: 2,
+    ...(parentToolUseId ? { parentToolUseId } : {}),
+  })
+  const note = (id: string, status: AgentTaskNotification['status'] = 'completed'): AgentTaskNotification => ({
+    taskId: `task-${id}`,
+    toolUseId: id,
+    status,
+    summary: id,
+  })
+  const build = (
+    messages: UIMessage[],
+    previous?: ReturnType<typeof buildToolGroupResults>,
+    notifications: Record<string, AgentTaskNotification> = {},
+  ) => {
+    const model = buildRenderModel(messages)
+    return buildToolGroupResults(model.renderItems, model.toolResultMap, model.childToolCallsByParent, notifications, previous)
+  }
+
+  it('gives each group only its own results and children', () => {
+    const groups = build([
+      read('a'),
+      result('a'),
+      { id: 'text', type: 'assistant_text', content: 'between', timestamp: 3 },
+      read('agent-1'),
+      read('child', 'agent-1'),
+      result('child', 'agent-1'),
+    ])
+
+    expect([...groups.get('group-use-a')!.resultMap.keys()]).toEqual(['a'])
+    const agentGroup = groups.get('group-use-agent-1')!
+    expect([...agentGroup.resultMap.keys()]).toEqual(['child'])
+    expect(agentGroup.childToolCallsByParent.get('agent-1')?.map((call) => call.id)).toEqual(['use-child'])
+  })
+
+  it('keeps references for groups whose results did not change', () => {
+    const history = [read('a'), result('a'), { id: 'text', type: 'assistant_text', content: 'next', timestamp: 3 } as UIMessage]
+    const useB = read('b')
+    const before = build([...history, useB])
+    const after = build([...history, useB, result('b')], before)
+
+    expect(after.get('group-use-a')).toEqual(before.get('group-use-a'))
+    expect(after.get('group-use-a')!.toolCalls).toBe(before.get('group-use-a')!.toolCalls)
+    expect(after.get('group-use-a')!.resultMap).toBe(before.get('group-use-a')!.resultMap)
+    expect(after.get('group-use-b')!.toolCalls).toBe(before.get('group-use-b')!.toolCalls)
+    expect(after.get('group-use-b')!.resultMap).not.toBe(before.get('group-use-b')!.resultMap)
+    expect(after.get('group-use-b')!.resultMap.get('b')).toMatchObject({ content: 'b done' })
+  })
+
+  it('gives a group only its own agent notifications and replaces them only when they change', () => {
+    const messages: UIMessage[] = [
+      read('agent-1'),
+      { id: 'text', type: 'assistant_text', content: 'between', timestamp: 3 },
+      read('agent-2'),
+    ]
+    const failed = note('agent-1', 'failed')
+    const before = build(messages, undefined, { 'agent-1': failed })
+    expect(Object.keys(before.get('group-use-agent-1')!.agentTaskNotifications)).toEqual(['agent-1'])
+    expect(before.get('group-use-agent-2')!.agentTaskNotifications).toEqual({})
+
+    const same = build(messages, before, { 'agent-1': failed })
+    expect(same.get('group-use-agent-1')!.agentTaskNotifications).toBe(before.get('group-use-agent-1')!.agentTaskNotifications)
+
+    const changed = build(messages, before, { 'agent-1': note('agent-1') })
+    expect(changed.get('group-use-agent-1')!.agentTaskNotifications).not.toBe(before.get('group-use-agent-1')!.agentTaskNotifications)
+    expect(changed.get('group-use-agent-1')!.agentTaskNotifications['agent-1']?.status).toBe('completed')
+    expect(changed.get('group-use-agent-2')!.agentTaskNotifications).toBe(before.get('group-use-agent-2')!.agentTaskNotifications)
+  })
+
+  it('includes grandchildren and keeps the child map stable', () => {
+    const messages = [
+      read('agent-1'),
+      read('agent-child', 'agent-1'),
+      read('leaf', 'agent-child'),
+      result('leaf', 'agent-child'),
+    ]
+    const before = build(messages, undefined, { 'agent-child': note('agent-child') })
+    const group = before.get('group-use-agent-1')!
+    expect(group.resultMap.has('leaf')).toBe(true)
+    expect(group.agentTaskNotifications['agent-child']).toBeDefined()
+    expect([...group.childToolCallsByParent.keys()].sort()).toEqual(['agent-1', 'agent-child'])
+
+    const again = build(messages, before, group.agentTaskNotifications)
+    expect(again.get('group-use-agent-1')!.childToolCallsByParent).toBe(group.childToolCallsByParent)
+  })
+})