refactor: chatStore supports per-session isolated state

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

desktop/src/stores/chatStore.ts

@@ -11,7 +11,7 @@ import type { AttachmentRef, ChatState, UIAttachment, UIMessage, ServerMessage,
 
 type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'reconnecting'
 
-type ChatStore = {
+export type PerSessionState = {
   messages: UIMessage[]
   chatState: ChatState
   connectionState: ConnectionState
@@ -29,31 +29,11 @@ type ChatStore = {
   tokenUsage: TokenUsage
   elapsedSeconds: number
   statusVerb: string
-  connectedSessionId: string | null
   slashCommands: Array<{ name: string; description: string }>
-
+  elapsedTimer: ReturnType<typeof setInterval> | null
-  // Actions
-  connectToSession: (sessionId: string) => void
-  disconnectSession: () => void
-  sendMessage: (content: string, attachments?: AttachmentRef[]) => void
-  respondToPermission: (requestId: string, allowed: boolean, rule?: string) => void
-  setSessionPermissionMode: (mode: PermissionMode) => void
-  stopGeneration: () => void
-  loadHistory: (sessionId: string) => Promise<void>
-  clearMessages: () => void
-  handleServerMessage: (msg: ServerMessage) => void
 }
 
-const TASK_TOOL_NAMES = new Set(['TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList', 'TodoWrite'])
+const DEFAULT_SESSION_STATE: PerSessionState = {
-
-/** Track tool_use IDs for task-related tools, so we can refresh on tool_result */
-const pendingTaskToolUseIds = new Set<string>()
-
-let msgCounter = 0
-const nextId = () => `msg-${++msgCounter}-${Date.now()}`
-let elapsedTimer: ReturnType<typeof setInterval> | null = null
-
-export const useChatStore = create<ChatStore>((set, get) => ({
   messages: [],
   chatState: 'idle',
   connectionState: 'disconnected',
@@ -66,97 +46,127 @@ export const useChatStore = create<ChatStore>((set, get) => ({
   tokenUsage: { input_tokens: 0, output_tokens: 0 },
   elapsedSeconds: 0,
   statusVerb: '',
-  connectedSessionId: null,
   slashCommands: [],
+  elapsedTimer: null,
+}
 
-  connectToSession: (sessionId: string) => {
+function createDefaultSessionState(): PerSessionState {
-    const current = get().connectedSessionId
+  return { ...DEFAULT_SESSION_STATE, messages: [], tokenUsage: { input_tokens: 0, output_tokens: 0 } }
-    if (current === sessionId) return
+}
 
-    // Disconnect previous
+type ChatStore = {
-    if (current) wsManager.disconnect()
+  sessions: Record<string, PerSessionState>
 
-    set({
+  getSession: (sessionId: string) => PerSessionState
-      connectedSessionId: sessionId,
+  connectToSession: (sessionId: string) => void
-      connectionState: 'connecting',
+  disconnectSession: (sessionId: string) => void
-      messages: [],
+  sendMessage: (sessionId: string, content: string, attachments?: AttachmentRef[]) => void
-      chatState: 'idle',
+  respondToPermission: (sessionId: string, requestId: string, allowed: boolean, rule?: string) => void
-      streamingText: '',
+  setSessionPermissionMode: (sessionId: string, mode: PermissionMode) => void
-      streamingToolInput: '',
+  stopGeneration: (sessionId: string) => void
-      activeToolUseId: null,
+  loadHistory: (sessionId: string) => Promise<void>
-      activeToolName: null,
+  clearMessages: (sessionId: string) => void
-      activeThinkingId: null,
+  handleServerMessage: (sessionId: string, msg: ServerMessage) => void
-      pendingPermission: null,
+}
-      elapsedSeconds: 0,
-      slashCommands: [],
-    })
 
-    // Clear all previous handlers before registering new ones
+const TASK_TOOL_NAMES = new Set(['TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList', 'TodoWrite'])
-    // This prevents handler accumulation causing message duplication
+const pendingTaskToolUseIds = new Set<string>()
-    wsManager.clearHandlers()
+
+let msgCounter = 0
+const nextId = () => `msg-${++msgCounter}-${Date.now()}`
+
+/** Helper: immutably update a specific session within the sessions record */
+function updateSessionIn(
+  sessions: Record<string, PerSessionState>,
+  sessionId: string,
+  updater: (s: PerSessionState) => Partial<PerSessionState>,
+): Record<string, PerSessionState> {
+  const session = sessions[sessionId]
+  if (!session) return sessions
+  return { ...sessions, [sessionId]: { ...session, ...updater(session) } }
+}
+
+export const useChatStore = create<ChatStore>((set, get) => ({
+  sessions: {},
+
+  getSession: (sessionId) => get().sessions[sessionId] ?? createDefaultSessionState(),
+
+  connectToSession: (sessionId) => {
+    const existing = get().sessions[sessionId]
+    if (existing && existing.connectionState !== 'disconnected') return
+
+    set((s) => ({
+      sessions: {
+        ...s.sessions,
+        [sessionId]: {
+          ...createDefaultSessionState(),
+          connectionState: 'connecting',
+          messages: existing?.messages ?? [],
+        },
+      },
+    }))
+
+    wsManager.clearHandlers(sessionId)
     wsManager.connect(sessionId)
-    wsManager.onMessage((msg) => {
+    wsManager.onMessage(sessionId, (msg) => {
       if (msg.type === 'connected') {
-        set({ connectionState: 'connected' })
+        set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ connectionState: 'connected' })) }))
       }
-      get().handleServerMessage(msg)
+      get().handleServerMessage(sessionId, msg)
     })
 
-    // Load history and tasks
     get().loadHistory(sessionId)
     useCLITaskStore.getState().fetchSessionTasks(sessionId)
     sessionsApi.getSlashCommands(sessionId)
       .then(({ commands }) => {
-        if (get().connectedSessionId === sessionId) {
+        if (get().sessions[sessionId]) {
-          set({ slashCommands: commands })
+          set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ slashCommands: commands })) }))
         }
       })
       .catch(() => {
-        if (get().connectedSessionId === sessionId) {
+        if (get().sessions[sessionId]) {
-          set({ slashCommands: [] })
+          set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ slashCommands: [] })) }))
         }
       })
   },
 
-  disconnectSession: () => {
+  disconnectSession: (sessionId) => {
-    wsManager.disconnect()
+    const session = get().sessions[sessionId]
-    if (elapsedTimer) clearInterval(elapsedTimer)
+    if (session?.elapsedTimer) clearInterval(session.elapsedTimer)
-    useCLITaskStore.getState().clearTasks()
+    wsManager.disconnect(sessionId)
-    set({ connectedSessionId: null, chatState: 'idle' })
+    set((s) => {
+      const { [sessionId]: _, ...rest } = s.sessions
+      return { sessions: rest }
+    })
   },
 
-  sendMessage: (content: string, attachments?: AttachmentRef[]) => {
+  sendMessage: (sessionId, content, attachments?) => {
     const userFacingContent = content.trim()
     const uiAttachments: UIAttachment[] | undefined =
       attachments && attachments.length > 0
-        ? attachments.map((attachment) => ({
+        ? attachments.map((a) => ({
-            type: attachment.type,
+            type: a.type,
-            name: attachment.name || attachment.path || attachment.mimeType || attachment.type,
+            name: a.name || a.path || a.mimeType || a.type,
-            data: attachment.data,
+            data: a.data,
-            mimeType: attachment.mimeType,
+            mimeType: a.mimeType,
           }))
         : undefined
 
-    // If all tasks are completed, inline the task summary before the new user message
     const taskStore = useCLITaskStore.getState()
     const allTasksDone = taskStore.tasks.length > 0 && taskStore.tasks.every((t) => t.status === 'completed')
 
-    // Add user message to UI (with optional task summary before it)
     set((s) => {
-      const newMessages = [...s.messages]
+      const session = s.sessions[sessionId]
+      if (!session) return s
+
+      const newMessages = [...session.messages]
       if (allTasksDone) {
         newMessages.push({
           id: nextId(),
           type: 'task_summary',
-          tasks: taskStore.tasks.map((t) => ({
+          tasks: taskStore.tasks.map((t) => ({ id: t.id, subject: t.subject, status: t.status, activeForm: t.activeForm })),
-            id: t.id,
-            subject: t.subject,
-            status: t.status,
-            activeForm: t.activeForm,
-          })),
           timestamp: Date.now(),
         })
-        // Clear sticky task bar since we inlined the summary
         taskStore.clearTasks()
       }
       newMessages.push({
@@ -166,64 +176,66 @@ export const useChatStore = create<ChatStore>((set, get) => ({
         attachments: uiAttachments,
         timestamp: Date.now(),
       })
+
+      if (session.elapsedTimer) clearInterval(session.elapsedTimer)
+
+      const timer = setInterval(() => {
+        set((st) => ({ sessions: updateSessionIn(st.sessions, sessionId, (sess) => ({ elapsedSeconds: sess.elapsedSeconds + 1 })) }))
+      }, 1000)
+
       return {
-        messages: newMessages,
+        sessions: {
-        chatState: 'thinking',
+          ...s.sessions,
-        elapsedSeconds: 0,
+          [sessionId]: {
-        streamingText: '',
+            ...session,
-        statusVerb: randomSpinnerVerb(),
+            messages: newMessages,
+            chatState: 'thinking',
+            elapsedSeconds: 0,
+            streamingText: '',
+            statusVerb: randomSpinnerVerb(),
+            elapsedTimer: timer,
+          },
+        },
       }
     })
 
-    // Start elapsed timer
+    wsManager.send(sessionId, { type: 'user_message', content, attachments })
-    if (elapsedTimer) clearInterval(elapsedTimer)
-    elapsedTimer = setInterval(() => {
-      set((s) => ({ elapsedSeconds: s.elapsedSeconds + 1 }))
-    }, 1000)
-
-    wsManager.send({ type: 'user_message', content, attachments })
   },
 
-  respondToPermission: (requestId: string, allowed: boolean, rule?: string) => {
+  respondToPermission: (sessionId, requestId, allowed, rule?) => {
-    wsManager.send({ type: 'permission_response', requestId, allowed, ...(rule ? { rule } : {}) })
+    wsManager.send(sessionId, { type: 'permission_response', requestId, allowed, ...(rule ? { rule } : {}) })
-    set({ pendingPermission: null, chatState: allowed ? 'tool_executing' : 'idle' })
+    set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ pendingPermission: null, chatState: allowed ? 'tool_executing' : 'idle' })) }))
   },
 
-  setSessionPermissionMode: (mode) => {
+  setSessionPermissionMode: (sessionId, mode) => {
-    if (!get().connectedSessionId) return
+    if (!get().sessions[sessionId]) return
-    wsManager.send({ type: 'set_permission_mode', mode })
+    wsManager.send(sessionId, { type: 'set_permission_mode', mode })
   },
 
-  stopGeneration: () => {
+  stopGeneration: (sessionId) => {
-    wsManager.send({ type: 'stop_generation' })
+    wsManager.send(sessionId, { type: 'stop_generation' })
-    if (elapsedTimer) clearInterval(elapsedTimer)
+    set((s) => {
-    set({ chatState: 'idle' })
+      const session = s.sessions[sessionId]
+      if (!session) return s
+      if (session.elapsedTimer) clearInterval(session.elapsedTimer)
+      return { sessions: { ...s.sessions, [sessionId]: { ...session, chatState: 'idle', elapsedTimer: null } } }
+    })
   },
 
-  loadHistory: async (sessionId: string) => {
+  loadHistory: async (sessionId) => {
     try {
       const { messages } = await sessionsApi.getMessages(sessionId)
       const uiMessages = mapHistoryMessagesToUiMessages(messages)
-
       set((state) => {
-        if (state.connectedSessionId !== sessionId || state.messages.length > 0) {
+        const session = state.sessions[sessionId]
-          return state
+        if (!session || session.messages.length > 0) return state
-        }
+        return { sessions: updateSessionIn(state.sessions, sessionId, () => ({ messages: uiMessages })) }
-        return { ...state, messages: uiMessages }
       })
-
-      // Extract the last TodoWrite input from history so TaskBar shows for V1 sessions
       const lastTodos = extractLastTodoWriteFromHistory(messages)
       if (lastTodos && lastTodos.length > 0) {
         const taskStore = useCLITaskStore.getState()
-        // Only set if V2 task fetch didn't already populate tasks
+        if (taskStore.tasks.length === 0) taskStore.setTasksFromTodos(lastTodos)
-        if (taskStore.tasks.length === 0) {
-          taskStore.setTasksFromTodos(lastTodos)
-        }
       }
-
-      // For both V1 and V2: if all tasks completed and the user already
-      // continued chatting, suppress the sticky TaskBar on page refresh.
       if (hasUserMessagesAfterTaskCompletion(messages)) {
         useCLITaskStore.getState().markCompletedAndDismissed()
       }
@@ -232,89 +244,77 @@ export const useChatStore = create<ChatStore>((set, get) => ({
     }
   },
 
-  clearMessages: () => set({ messages: [], streamingText: '', chatState: 'idle' }),
+  clearMessages: (sessionId) => {
+    set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, () => ({ messages: [], streamingText: '', chatState: 'idle' })) }))
+  },
+
+  handleServerMessage: (sessionId, msg) => {
+    const update = (updater: (session: PerSessionState) => Partial<PerSessionState>) => {
+      set((s) => ({ sessions: updateSessionIn(s.sessions, sessionId, updater) }))
+    }
 
-  handleServerMessage: (msg: ServerMessage) => {
     switch (msg.type) {
       case 'connected':
         break
 
       case 'status':
-        set({
+        update((session) => ({
           chatState: msg.state,
-          // Only override statusVerb if the server sends something other than
-          // the generic 'Thinking' — otherwise keep the random verb we picked
-          // in sendMessage so the user sees fun loading text.
           ...(msg.verb && msg.verb !== 'Thinking' ? { statusVerb: msg.verb } : {}),
-          ...(msg.tokens ? { tokenUsage: { ...get().tokenUsage, output_tokens: msg.tokens } } : {}),
+          ...(msg.tokens ? { tokenUsage: { ...session.tokenUsage, output_tokens: msg.tokens } } : {}),
           ...(msg.state === 'idle' ? { activeThinkingId: null, statusVerb: '' } : {}),
-        })
+        }))
-        if (msg.state === 'idle' && elapsedTimer) {
+        if (msg.state === 'idle') {
-          clearInterval(elapsedTimer)
+          const session = get().sessions[sessionId]
-          elapsedTimer = null
+          if (session?.elapsedTimer) {
+            clearInterval(session.elapsedTimer)
+            update(() => ({ elapsedTimer: null }))
+          }
         }
         break
 
       case 'content_start': {
-        // Flush any accumulated streamingText as assistant_text BEFORE
+        const session = get().sessions[sessionId]
-        // switching to the next block. This preserves intermediate text
+        if (!session) break
-        // segments between tool calls (e.g. "项目已经有 node_modules，直接启动").
+        const pendingText = session.streamingText.trim()
-        const pendingText = get().streamingText.trim()
         if (pendingText) {
-          set((s) => ({
+          update((s) => ({
-            messages: [...s.messages, {
+            messages: [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }],
-              id: nextId(),
-              type: 'assistant_text',
-              content: pendingText,
-              timestamp: Date.now(),
-            }],
             streamingText: '',
           }))
         }
-
         if (msg.blockType === 'text') {
-          set({ streamingText: '', chatState: 'streaming', activeThinkingId: null })
+          update(() => ({ streamingText: '', chatState: 'streaming', activeThinkingId: null }))
         } else if (msg.blockType === 'tool_use') {
-          set({
+          update(() => ({
             activeToolUseId: msg.toolUseId ?? null,
             activeToolName: msg.toolName ?? null,
             streamingToolInput: '',
             chatState: 'tool_executing',
             activeThinkingId: null,
-          })
+          }))
         }
         break
       }
 
       case 'content_delta':
-        if (msg.text !== undefined) {
+        if (msg.text !== undefined) update((s) => ({ streamingText: s.streamingText + msg.text }))
-          set((s) => ({ streamingText: s.streamingText + msg.text }))
+        if (msg.toolInput !== undefined) update((s) => ({ streamingToolInput: s.streamingToolInput + msg.toolInput }))
-        }
-        if (msg.toolInput !== undefined) {
-          set((s) => ({ streamingToolInput: s.streamingToolInput + msg.toolInput }))
-        }
         break
 
       case 'thinking':
-        // Merge consecutive thinking deltas into one message.
+        update((s) => {
-        // Also flush any pending streamingText first — otherwise the text
-        // becomes invisible because MessageList only renders streamingText
-        // when chatState === 'streaming', and we're about to set it to 'thinking'.
-        set((s) => {
           const pendingText = s.streamingText.trim()
           const base = pendingText
             ? [...s.messages, { id: nextId(), type: 'assistant_text' as const, content: pendingText, timestamp: Date.now() }]
             : s.messages
-
           const last = base[base.length - 1]
           if (last && last.type === 'thinking') {
-            // Append to existing thinking message
             const updated = [...base]
             updated[updated.length - 1] = { ...last, content: last.content + msg.text }
             return { messages: updated, chatState: 'thinking', activeThinkingId: last.id, streamingText: '' }
           }
           const id = nextId()
-          // Create new thinking message
           return {
             messages: [...base, { id, type: 'thinking', content: msg.text, timestamp: Date.now() }],
             chatState: 'thinking',
@@ -325,49 +325,33 @@ export const useChatStore = create<ChatStore>((set, get) => ({
         break
 
       case 'tool_use_complete': {
-        const toolName = msg.toolName || get().activeToolName || 'unknown'
+        const session = get().sessions[sessionId]
-        set((s) => ({
+        const toolName = msg.toolName || session?.activeToolName || 'unknown'
+        update((s) => ({
           messages: [...s.messages, {
-            id: nextId(),
+            id: nextId(), type: 'tool_use', toolName,
-            type: 'tool_use',
-            toolName,
             toolUseId: msg.toolUseId || s.activeToolUseId || '',
-            input: msg.input,
+            input: msg.input, timestamp: Date.now(), parentToolUseId: msg.parentToolUseId,
-            timestamp: Date.now(),
-            parentToolUseId: msg.parentToolUseId,
           }],
-          activeToolUseId: null,
+          activeToolUseId: null, activeToolName: null, activeThinkingId: null, streamingToolInput: '',
-          activeToolName: null,
-          activeThinkingId: null,
-          streamingToolInput: '',
         }))
-        // TodoWrite: input contains the full todo list — update tasks immediately
         if (toolName === 'TodoWrite' && Array.isArray((msg.input as any)?.todos)) {
           useCLITaskStore.getState().setTasksFromTodos((msg.input as any).todos)
         } else if (TASK_TOOL_NAMES.has(toolName)) {
-          // V2 task tools — refresh will happen on tool_result
+          const useId = msg.toolUseId || session?.activeToolUseId
-          // when the tool has actually finished executing and written to disk
-          const useId = msg.toolUseId || get().activeToolUseId
           if (useId) pendingTaskToolUseIds.add(useId)
         }
         break
       }
 
       case 'tool_result':
-        set((s) => ({
+        update((s) => ({
           messages: [...s.messages, {
-            id: nextId(),
+            id: nextId(), type: 'tool_result', toolUseId: msg.toolUseId,
-            type: 'tool_result',
+            content: msg.content, isError: msg.isError, timestamp: Date.now(), parentToolUseId: msg.parentToolUseId,
-            toolUseId: msg.toolUseId,
-            content: msg.content,
-            isError: msg.isError,
-            timestamp: Date.now(),
-            parentToolUseId: msg.parentToolUseId,
           }],
-          chatState: 'thinking',
+          chatState: 'thinking', activeThinkingId: null,
-          activeThinkingId: null,
         }))
-        // Refresh tasks after a task tool has finished executing
         if (pendingTaskToolUseIds.has(msg.toolUseId)) {
           pendingTaskToolUseIds.delete(msg.toolUseId)
           useCLITaskStore.getState().refreshTasks()
@@ -375,230 +359,116 @@ export const useChatStore = create<ChatStore>((set, get) => ({
         break
 
       case 'permission_request':
-        set({
+        update((s) => ({
-          pendingPermission: {
+          pendingPermission: { requestId: msg.requestId, toolName: msg.toolName, input: msg.input, description: msg.description },
-            requestId: msg.requestId,
-            toolName: msg.toolName,
-            input: msg.input,
-            description: msg.description,
-          },
           chatState: 'permission_pending',
           activeThinkingId: null,
-        })
-        set((s) => ({
           messages: [...s.messages, {
-            id: nextId(),
+            id: nextId(), type: 'permission_request', requestId: msg.requestId,
-            type: 'permission_request',
+            toolName: msg.toolName, input: msg.input, description: msg.description, timestamp: Date.now(),
-            requestId: msg.requestId,
-            toolName: msg.toolName,
-            input: msg.input,
-            description: msg.description,
-            timestamp: Date.now(),
           }],
         }))
         break
 
       case 'message_complete': {
-        // Flush streaming text as a message
+        const session = get().sessions[sessionId]
-        const text = get().streamingText
+        if (!session) break
+        const text = session.streamingText
         if (text) {
-          set((s) => ({
+          update((s) => ({
             messages: [...s.messages, { id: nextId(), type: 'assistant_text', content: text, timestamp: Date.now() }],
             streamingText: '',
           }))
         }
-        set({ tokenUsage: msg.usage, chatState: 'idle', activeThinkingId: null })
+        if (session.elapsedTimer) clearInterval(session.elapsedTimer)
-        if (elapsedTimer) {
+        update(() => ({ tokenUsage: msg.usage, chatState: 'idle', activeThinkingId: null, elapsedTimer: null }))
-          clearInterval(elapsedTimer)
-          elapsedTimer = null
-        }
         break
       }
 
       case 'error':
-        set((s) => ({
+        update((s) => ({
           messages: [...s.messages, { id: nextId(), type: 'error', message: msg.message, code: msg.code, timestamp: Date.now() }],
-          chatState: 'idle',
+          chatState: 'idle', activeThinkingId: null,
-          activeThinkingId: null,
         }))
-        if (elapsedTimer) {
+        {
-          clearInterval(elapsedTimer)
+          const session = get().sessions[sessionId]
-          elapsedTimer = null
+          if (session?.elapsedTimer) {
+            clearInterval(session.elapsedTimer)
+            update(() => ({ elapsedTimer: null }))
+          }
         }
         break
 
       case 'team_created':
         useTeamStore.getState().handleTeamCreated(msg.teamName)
         break
-
       case 'team_update':
         useTeamStore.getState().handleTeamUpdate(msg.teamName, msg.members)
         break
-
       case 'team_deleted':
         useTeamStore.getState().handleTeamDeleted(msg.teamName)
         break
-
       case 'task_update':
         break
-
       case 'session_title_updated':
         useSessionStore.getState().updateSessionTitle(msg.sessionId, msg.title)
         break
-
       case 'system_notification':
-        // Cache slash commands from CLI init
         if (msg.subtype === 'slash_commands' && Array.isArray(msg.data)) {
-          set({ slashCommands: msg.data as Array<{ name: string; description: string }> })
+          update(() => ({ slashCommands: msg.data as Array<{ name: string; description: string }> }))
         }
         break
-
       case 'pong':
         break
     }
   },
 }))
 
-type AssistantHistoryBlock = {
+// ─── History mapping helpers (unchanged from original) ─────────
-  type: string
-  text?: string
-  thinking?: string
-  name?: string
-  id?: string
-  input?: unknown
-}
 
-type UserHistoryBlock = {
+type AssistantHistoryBlock = { type: string; text?: string; thinking?: string; name?: string; id?: string; input?: unknown }
-  type: string
+type UserHistoryBlock = { type: string; text?: string; tool_use_id?: string; content?: unknown; is_error?: boolean; source?: { data?: string }; mimeType?: string; media_type?: string; name?: string }
-  text?: string
-  tool_use_id?: string
-  content?: unknown
-  is_error?: boolean
-  source?: { data?: string }
-  mimeType?: string
-  media_type?: string
-  name?: string
-}
 
 export function mapHistoryMessagesToUiMessages(messages: MessageEntry[]): UIMessage[] {
   const uiMessages: UIMessage[] = []
-
   for (const msg of messages) {
     const timestamp = new Date(msg.timestamp).getTime()
-
     if (msg.type === 'user' && typeof msg.content === 'string') {
-      uiMessages.push({
+      uiMessages.push({ id: msg.id || nextId(), type: 'user_text', content: msg.content, timestamp })
-        id: msg.id || nextId(),
-        type: 'user_text',
-        content: msg.content,
-        timestamp,
-      })
       continue
     }
-
     if (msg.type === 'assistant' && typeof msg.content === 'string') {
-      uiMessages.push({
+      uiMessages.push({ id: msg.id || nextId(), type: 'assistant_text', content: msg.content, timestamp, model: msg.model })
-        id: msg.id || nextId(),
-        type: 'assistant_text',
-        content: msg.content,
-        timestamp,
-        model: msg.model,
-      })
       continue
     }
-
-    // Server marks assistant messages containing tool_use blocks as type 'tool_use',
-    // but the content array structure is the same as 'assistant' — handle both.
     if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
       for (const block of msg.content as AssistantHistoryBlock[]) {
-        if (block.type === 'thinking' && block.thinking) {
+        if (block.type === 'thinking' && block.thinking) uiMessages.push({ id: nextId(), type: 'thinking', content: block.thinking, timestamp })
-          uiMessages.push({
+        else if (block.type === 'text' && block.text) uiMessages.push({ id: nextId(), type: 'assistant_text', content: block.text, timestamp, model: msg.model })
-            id: nextId(),
+        else if (block.type === 'tool_use') uiMessages.push({ id: nextId(), type: 'tool_use', toolName: block.name ?? 'unknown', toolUseId: block.id ?? '', input: block.input, timestamp, parentToolUseId: msg.parentToolUseId })
-            type: 'thinking',
-            content: block.thinking,
-            timestamp,
-          })
-        } else if (block.type === 'text' && block.text) {
-          uiMessages.push({
-            id: nextId(),
-            type: 'assistant_text',
-            content: block.text,
-            timestamp,
-            model: msg.model,
-          })
-        } else if (block.type === 'tool_use') {
-          uiMessages.push({
-            id: nextId(),
-            type: 'tool_use',
-            toolName: block.name ?? 'unknown',
-            toolUseId: block.id ?? '',
-            input: block.input,
-            timestamp,
-            parentToolUseId: msg.parentToolUseId,
-          })
-        }
       }
       continue
     }
-
-    // Server marks user messages containing tool_result blocks as type 'tool_result',
-    // but the content array structure is the same as 'user' — handle both.
     if ((msg.type === 'user' || msg.type === 'tool_result') && Array.isArray(msg.content)) {
       const textParts: string[] = []
       const attachments: UIAttachment[] = []
-
       for (const block of msg.content as UserHistoryBlock[]) {
-        if (block.type === 'text' && block.text) {
+        if (block.type === 'text' && block.text) textParts.push(block.text)
-          textParts.push(block.text)
+        else if (block.type === 'image') attachments.push({ type: 'image', name: block.name || 'image', data: block.source?.data, mimeType: block.mimeType || block.media_type })
-        } else if (block.type === 'image') {
+        else if (block.type === 'file') attachments.push({ type: 'file', name: block.name || 'file' })
-          attachments.push({
+        else if (block.type === 'tool_result') uiMessages.push({ id: nextId(), type: 'tool_result', toolUseId: block.tool_use_id ?? '', content: block.content, isError: !!block.is_error, timestamp, parentToolUseId: msg.parentToolUseId })
-            type: 'image',
-            name: block.name || 'image',
-            data: block.source?.data,
-            mimeType: block.mimeType || block.media_type,
-          })
-        } else if (block.type === 'file') {
-          attachments.push({
-            type: 'file',
-            name: block.name || 'file',
-          })
-        } else if (block.type === 'tool_result') {
-          uiMessages.push({
-            id: nextId(),
-            type: 'tool_result',
-            toolUseId: block.tool_use_id ?? '',
-            content: block.content,
-            isError: !!block.is_error,
-            timestamp,
-            parentToolUseId: msg.parentToolUseId,
-          })
-        }
       }
-
       if (textParts.length > 0 || attachments.length > 0) {
-        uiMessages.push({
+        uiMessages.push({ id: nextId(), type: 'user_text', content: textParts.join('\n'), attachments: attachments.length > 0 ? attachments : undefined, timestamp })
-          id: nextId(),
-          type: 'user_text',
-          content: textParts.join('\n'),
-          attachments: attachments.length > 0 ? attachments : undefined,
-          timestamp,
-        })
       }
     }
   }
-
   return uiMessages
 }
 
-/** Scan history messages for the last TodoWrite tool_use and return the todos array.
+function extractLastTodoWriteFromHistory(messages: MessageEntry[]): Array<{ content: string; status: string; activeForm?: string }> | null {
- *  Returns null if all todos were completed and the user continued chatting after. */
-function extractLastTodoWriteFromHistory(
-  messages: MessageEntry[],
-): Array<{ content: string; status: string; activeForm?: string }> | null {
   let foundIndex = -1
   let todos: Array<{ content: string; status: string; activeForm?: string }> | null = null
-
-  // Walk backwards to find the most recent TodoWrite
   for (let i = messages.length - 1; i >= 0; i--) {
     const msg = messages[i]!
     if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
@@ -617,45 +487,28 @@ function extractLastTodoWriteFromHistory(
       if (todos) break
     }
   }
-
   if (!todos) return null
-
-  // If all todos are completed and there are user messages after, the tasks
-  // were already inlined — don't re-show the sticky bar.
   const allDone = todos.every((t) => t.status === 'completed')
   if (allDone) {
     for (let i = foundIndex + 1; i < messages.length; i++) {
-      const msg = messages[i]!
+      if (messages[i]!.type === 'user' && messages[i]!.content) return null
-      if (msg.type === 'user' && msg.content) return null
     }
   }
-
   return todos
 }
 
 const TASK_RELATED_TOOL_NAMES = new Set(['TodoWrite', 'TaskCreate', 'TaskUpdate', 'TaskGet', 'TaskList'])
 
-/** Check if there are user messages after the last task-related tool call.
- *  Used to determine if completed tasks should be shown as sticky or inline. */
 function hasUserMessagesAfterTaskCompletion(messages: MessageEntry[]): boolean {
-  // Find the index of the last task-related tool call
   let lastTaskIndex = -1
   for (let i = messages.length - 1; i >= 0; i--) {
     const msg = messages[i]!
     if ((msg.type === 'assistant' || msg.type === 'tool_use') && Array.isArray(msg.content)) {
       const blocks = msg.content as AssistantHistoryBlock[]
-      if (blocks.some((b) => b.type === 'tool_use' && TASK_RELATED_TOOL_NAMES.has(b.name ?? ''))) {
+      if (blocks.some((b) => b.type === 'tool_use' && TASK_RELATED_TOOL_NAMES.has(b.name ?? ''))) { lastTaskIndex = i; break }
-        lastTaskIndex = i
-        break
-      }
     }
   }
-
   if (lastTaskIndex < 0) return false
-
+  for (let i = lastTaskIndex + 1; i < messages.length; i++) { if (messages[i]!.type === 'user') return true }
-  // Check for user messages after the last task tool
-  for (let i = lastTaskIndex + 1; i < messages.length; i++) {
-    if (messages[i]!.type === 'user') return true
-  }
   return false
 }