fix(desktop): more fine-grained state updates

4 files changed, 126 insertions, 66 deletions

diff --git a/packages/app/src/context/global-sync.tsx b/packages/app/src/context/global-sync.tsx
index a496d59cb..0296dfad3 100644
--- a/packages/app/src/context/global-sync.tsx
+++ b/packages/app/src/context/global-sync.tsx
@@ -71,21 +71,19 @@ function createGlobalSync() {
     project: Project[]
     provider: ProviderListResponse
     provider_auth: ProviderAuthResponse
-    children: Record<string, State>
   }>({
     ready: false,
     path: { state: "", config: "", worktree: "", directory: "", home: "" },
     project: [],
     provider: { all: [], connected: [], default: {} },
     provider_auth: {},
-    children: {},
   })
 
   const children: Record<string, ReturnType<typeof createStore<State>>> = {}
   function child(directory: string) {
     if (!directory) console.error("No directory provided")
     if (!children[directory]) {
-      setGlobalStore("children", directory, {
+      children[directory] = createStore<State>({
         project: "",
         provider: { all: [], connected: [], default: {} },
         config: {},
@@ -105,7 +103,6 @@ function createGlobalSync() {
         message: {},
         part: {},
       })
-      children[directory] = createStore(globalStore.children[directory])
       bootstrapInstance(directory)
     }
     return children[directory]
diff --git a/packages/app/src/context/sync.tsx b/packages/app/src/context/sync.tsx
index 941b8b629..dad035fb3 100644
--- a/packages/app/src/context/sync.tsx
+++ b/packages/app/src/context/sync.tsx
@@ -1,5 +1,5 @@
-import { produce } from "solid-js/store"
-import { createMemo } from "solid-js"
+import { batch, createMemo } from "solid-js"
+import { produce, reconcile } from "solid-js/store"
 import { Binary } from "@opencode-ai/util/binary"
 import { retry } from "@opencode-ai/util/retry"
 import { createSimpleContext } from "@opencode-ai/ui/context"
@@ -67,22 +67,46 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
             retry(() => sdk.client.session.todo({ sessionID })),
             retry(() => sdk.client.session.diff({ sessionID })),
           ])
-          setStore(
-            produce((draft) => {
-              const match = Binary.search(draft.session, sessionID, (s) => s.id)
-              if (match.found) draft.session[match.index] = session.data!
-              if (!match.found) draft.session.splice(match.index, 0, session.data!)
-              draft.todo[sessionID] = todo.data ?? []
-              draft.message[sessionID] = messages
-                .data!.map((x) => x.info)
-                .slice()
-                .sort((a, b) => a.id.localeCompare(b.id))
-              for (const message of messages.data!) {
-                draft.part[message.info.id] = message.parts.slice().sort((a, b) => a.id.localeCompare(b.id))
-              }
-              draft.session_diff[sessionID] = diff.data ?? []
-            }),
-          )
+
+          batch(() => {
+            setStore(
+              "session",
+              produce((draft) => {
+                const match = Binary.search(draft, sessionID, (s) => s.id)
+                if (match.found) {
+                  draft[match.index] = session.data!
+                  return
+                }
+                draft.splice(match.index, 0, session.data!)
+              }),
+            )
+
+            setStore("todo", sessionID, reconcile(todo.data ?? []))
+            setStore(
+              "message",
+              sessionID,
+              reconcile(
+                (messages.data ?? [])
+                  .map((x) => x.info)
+                  .slice()
+                  .sort((a, b) => a.id.localeCompare(b.id)),
+                { key: "id" },
+              ),
+            )
+
+            for (const message of messages.data ?? []) {
+              setStore(
+                "part",
+                message.info.id,
+                reconcile(
+                  message.parts.slice().sort((a, b) => a.id.localeCompare(b.id)),
+                  { key: "id" },
+                ),
+              )
+            }
+
+            setStore("session_diff", sessionID, reconcile(diff.data ?? [], { key: "file" }))
+          })
         },
         fetch: async (count = 10) => {
           setStore("limit", (x) => x + count)
diff --git a/packages/app/src/pages/session.tsx b/packages/app/src/pages/session.tsx
index ec55acf1e..6bc39daca 100644
--- a/packages/app/src/pages/session.tsx
+++ b/packages/app/src/pages/session.tsx
@@ -79,11 +79,7 @@ export default function Page() {
   const info = createMemo(() => (params.id ? sync.session.get(params.id) : undefined))
   const revertMessageID = createMemo(() => info()?.revert?.messageID)
   const messages = createMemo(() => (params.id ? (sync.data.message[params.id] ?? []) : []))
-  const userMessages = createMemo(() =>
-    messages()
-      .filter((m) => m.role === "user")
-      .sort((a, b) => a.id.localeCompare(b.id)),
-  )
+  const userMessages = createMemo(() => messages().filter((m) => m.role === "user"))
   const visibleUserMessages = createMemo(() => {
     const revert = revertMessageID()
     if (!revert) return userMessages()
@@ -587,6 +583,7 @@ export default function Page() {
             <SessionTurn
               sessionID={params.id!}
               messageID={message.id}
+              lastUserMessageID={lastUserMessage()?.id}
               stepsExpanded={store.mobileStepsExpanded[message.id] ?? false}
               onStepsExpandedToggle={() => setStore("mobileStepsExpanded", message.id, (x) => !x)}
               onUserInteracted={() => setStore("userInteracted", true)}
@@ -643,6 +640,7 @@ export default function Page() {
             <SessionTurn
               sessionID={params.id!}
               messageID={activeMessage()!.id}
+              lastUserMessageID={lastUserMessage()?.id}
               stepsExpanded={store.stepsExpanded}
               onStepsExpandedToggle={() => setStore("stepsExpanded", (x) => !x)}
               onUserInteracted={() => setStore("userInteracted", true)}
diff --git a/packages/ui/src/components/session-turn.tsx b/packages/ui/src/components/session-turn.tsx
index 8cb426387..8e9aa9f23 100644
--- a/packages/ui/src/components/session-turn.tsx
+++ b/packages/ui/src/components/session-turn.tsx
@@ -3,6 +3,7 @@ import { useData } from "../context"
 import { useDiffComponent } from "../context/diff"
 import { getDirectory, getFilename } from "@opencode-ai/util/path"
 import { checksum } from "@opencode-ai/util/encode"
+import { Binary } from "@opencode-ai/util/binary"
 import { createEffect, createMemo, For, Match, onCleanup, ParentProps, Show, Switch } from "solid-js"
 import { createResizeObserver } from "@solid-primitives/resize-observer"
 import { DiffChanges } from "./diff-changes"
@@ -62,24 +63,29 @@ function computeStatusFromPart(part: PartType | undefined): string | undefined {
 
 function AssistantMessageItem(props: {
   message: AssistantMessage
-  summary: string | undefined
-  response: string | undefined
-  lastTextPartId: string | undefined
-  working: boolean
+  responsePartId: string | undefined
+  hideResponsePart: boolean
 }) {
   const data = useData()
   const msgParts = createMemo(() => data.store.part[props.message.id] ?? [])
-  const lastTextPart = createMemo(() =>
-    msgParts()
-      .filter((p) => p?.type === "text")
-      .at(-1),
-  )
+  const lastTextPart = createMemo(() => {
+    const parts = msgParts()
+    for (let i = parts.length - 1; i >= 0; i--) {
+      const part = parts[i]
+      if (part?.type === "text") return part as TextPart
+    }
+    return undefined
+  })
 
   const filteredParts = createMemo(() => {
-    if (!props.working && !props.summary && props.response && props.lastTextPartId === lastTextPart()?.id) {
-      return msgParts().filter((p) => p?.id !== lastTextPart()?.id)
-    }
-    return msgParts()
+    const parts = msgParts()
+    if (!props.hideResponsePart) return parts
+
+    const responsePartId = props.responsePartId
+    if (!responsePartId) return parts
+    if (responsePartId !== lastTextPart()?.id) return parts
+
+    return parts.filter((part) => part?.id !== responsePartId)
   })
 
   return <Message message={props.message} parts={filteredParts()} />
@@ -89,6 +95,7 @@ export function SessionTurn(
   props: ParentProps<{
     sessionID: string
     messageID: string
+    lastUserMessageID?: string
     stepsExpanded?: boolean
     onStepsExpandedToggle?: () => void
     onUserInteracted?: () => void
@@ -103,14 +110,30 @@ export function SessionTurn(
   const diffComponent = useDiffComponent()
 
   const allMessages = createMemo(() => data.store.message[props.sessionID] ?? [])
-  const userMessages = createMemo(() =>
-    allMessages()
-      .filter((m) => m.role === "user")
-      .sort((a, b) => a.id.localeCompare(b.id)),
-  )
 
-  const message = createMemo(() => userMessages().find((m) => m.id === props.messageID))
-  const isLastUserMessage = createMemo(() => message()?.id === userMessages().at(-1)?.id)
+  const message = createMemo(() => {
+    const messages = allMessages()
+    const result = Binary.search(messages, props.messageID, (m) => m.id)
+    if (!result.found) return undefined
+
+    const msg = messages[result.index]
+    if (msg.role !== "user") return undefined
+
+    return msg
+  })
+
+  const lastUserMessageID = createMemo(() => {
+    if (props.lastUserMessageID) return props.lastUserMessageID
+
+    const messages = allMessages()
+    for (let i = messages.length - 1; i >= 0; i--) {
+      const msg = messages[i]
+      if (msg?.role === "user") return msg.id
+    }
+    return undefined
+  })
+
+  const isLastUserMessage = createMemo(() => props.messageID === lastUserMessageID())
 
   const parts = createMemo(() => {
     const msg = message()
@@ -118,10 +141,29 @@ export function SessionTurn(
     return data.store.part[msg.id] ?? []
   })
 
+  const messageIndex = createMemo(() => {
+    const messages = allMessages()
+    const result = Binary.search(messages, props.messageID, (m) => m.id)
+    if (!result.found) return -1
+    return result.index
+  })
+
   const assistantMessages = createMemo(() => {
     const msg = message()
     if (!msg) return [] as AssistantMessage[]
-    return allMessages().filter((m) => m.role === "assistant" && m.parentID === msg.id) as AssistantMessage[]
+
+    const messages = allMessages()
+    const index = messageIndex()
+    if (index < 0) return [] as AssistantMessage[]
+
+    const result: AssistantMessage[] = []
+    for (let i = index + 1; i < messages.length; i++) {
+      const item = messages[i]
+      if (!item) continue
+      if (item.role === "user") break
+      if (item.role === "assistant" && item.parentID === msg.id) result.push(item as AssistantMessage)
+    }
+    return result
   })
 
   const lastAssistantMessage = createMemo(() => assistantMessages().at(-1))
@@ -152,16 +194,18 @@ export function SessionTurn(
   })
 
   const permissionParts = createMemo(() => {
-    const result: { part: ToolPart; message: AssistantMessage }[] = []
     const permissions = data.store.permission?.[props.sessionID] ?? []
-    if (!permissions.length) return result
+    if (!permissions.length) return [] as { part: ToolPart; message: AssistantMessage }[]
 
-    for (const m of assistantMessages()) {
-      const msgParts = data.store.part[m.id] ?? []
-      for (const p of msgParts) {
-        if (p?.type === "tool" && permissions.some((perm) => perm.callID === (p as ToolPart).callID)) {
-          result.push({ part: p as ToolPart, message: m })
-        }
+    const ids = new Set(permissions.map((perm) => perm.callID))
+    const result: { part: ToolPart; message: AssistantMessage }[] = []
+
+    for (const message of assistantMessages()) {
+      const parts = data.store.part[message.id] ?? []
+      for (const part of parts) {
+        if (part?.type !== "tool") continue
+        const tool = part as ToolPart
+        if (ids.has(tool.callID)) result.push({ part: tool, message })
       }
     }
     return result
@@ -245,12 +289,11 @@ export function SessionTurn(
     return s
   })
 
-  const summary = () => message()?.summary?.body
-  const response = () => {
-    const part = lastTextPart()
-    return part?.type === "text" ? (part as TextPart).text : undefined
-  }
-  const hasDiffs = () => message()?.summary?.diffs?.length
+  const summary = createMemo(() => message()?.summary?.body)
+  const response = createMemo(() => lastTextPart()?.text)
+  const responsePartId = createMemo(() => lastTextPart()?.id)
+  const hasDiffs = createMemo(() => message()?.summary?.diffs?.length)
+  const hideResponsePart = createMemo(() => !working() && !summary() && !!responsePartId())
 
   function duration() {
     const msg = message()
@@ -477,10 +520,8 @@ export function SessionTurn(
                           {(assistantMessage) => (
                             <AssistantMessageItem
                               message={assistantMessage}
-                              summary={summary()}
-                              response={response()}
-                              lastTextPartId={lastTextPart()?.id}
-                              working={working()}
+                              responsePartId={responsePartId()}
+                              hideResponsePart={hideResponsePart()}
                             />
                           )}
                         </For>