fix(desktop): scroll jank in session turn and review

2 files changed, 74 insertions, 55 deletions

diff --git a/packages/app/src/context/global-sync.tsx b/packages/app/src/context/global-sync.tsx
index 085c5067c..10607b1d2 100644
--- a/packages/app/src/context/global-sync.tsx
+++ b/packages/app/src/context/global-sync.tsx
@@ -215,7 +215,7 @@ function createGlobalSync() {
         break
       }
       case "session.diff":
-        setStore("session_diff", event.properties.sessionID, reconcile(event.properties.diff))
+        setStore("session_diff", event.properties.sessionID, reconcile(event.properties.diff, { key: "file" }))
         break
       case "todo.updated":
         setStore("todo", event.properties.sessionID, reconcile(event.properties.todos))
diff --git a/packages/ui/src/components/session-turn.tsx b/packages/ui/src/components/session-turn.tsx
index a160f4ba3..a0368b0d4 100644
--- a/packages/ui/src/components/session-turn.tsx
+++ b/packages/ui/src/components/session-turn.tsx
@@ -124,77 +124,96 @@ export function SessionTurn(
     return allMessages().filter((m) => m.role === "assistant" && m.parentID === msg.id) as AssistantMessage[]
   })
 
-  const assistantParts = createMemo(() => {
-    const result: PartType[] = []
-    for (const m of assistantMessages()) {
-      const msgParts = data.store.part[m.id]
-      if (msgParts) {
-        for (const p of msgParts) {
-          if (p) result.push(p)
-        }
-      }
-    }
-    return result
-  })
-
   const lastAssistantMessage = createMemo(() => assistantMessages().at(-1))
 
   const error = createMemo(() => assistantMessages().find((m) => m.error)?.error)
 
   const lastTextPart = createMemo(() => {
-    const ap = assistantParts()
-    for (let i = ap.length - 1; i >= 0; i--) {
-      if (ap[i]?.type === "text") return ap[i]
+    const msgs = assistantMessages()
+    for (let mi = msgs.length - 1; mi >= 0; mi--) {
+      const msgParts = data.store.part[msgs[mi].id] ?? []
+      for (let pi = msgParts.length - 1; pi >= 0; pi--) {
+        const part = msgParts[pi]
+        if (part?.type === "text") return part as TextPart
+      }
     }
     return undefined
   })
 
-  const hasSteps = createMemo(() => assistantParts().some((p) => p?.type === "tool"))
-
-  const isShellMode = createMemo(() => {
-    const p = parts()
-    const ap = assistantParts()
-    if (p.every((part) => part?.type === "text" && part?.synthetic) && ap.length === 1) {
-      const assistantPart = ap[0]
-      if (assistantPart?.type === "tool" && assistantPart?.tool === "bash") return true
+  const hasSteps = createMemo(() => {
+    for (const m of assistantMessages()) {
+      const msgParts = data.store.part[m.id]
+      if (!msgParts) continue
+      for (const p of msgParts) {
+        if (p?.type === "tool") return true
+      }
     }
     return false
   })
 
+  const shellModePart = createMemo(() => {
+    const p = parts()
+    if (!p.every((part) => part?.type === "text" && part?.synthetic)) return
+
+    const msgs = assistantMessages()
+    if (msgs.length !== 1) return
+
+    const msgParts = data.store.part[msgs[0].id] ?? []
+    if (msgParts.length !== 1) return
+
+    const assistantPart = msgParts[0]
+    if (assistantPart?.type === "tool" && assistantPart.tool === "bash") return assistantPart
+  })
+
+  const isShellMode = createMemo(() => !!shellModePart())
+
   const rawStatus = createMemo(() => {
-    const ap = assistantParts()
-    const currentTask = ap.findLast(
-      (p) =>
-        p &&
-        p.type === "tool" &&
-        p.tool === "task" &&
-        p.state &&
-        "metadata" in p.state &&
-        p.state.metadata &&
-        p.state.metadata.sessionId &&
-        p.state.status === "running",
-    ) as ToolPart | undefined
-
-    let resolvedParts = ap
-    if (currentTask?.state && "metadata" in currentTask.state && currentTask.state.metadata?.sessionId) {
-      const taskMessages = data.store.message[currentTask.state.metadata.sessionId as string]?.filter(
-        (m) => m.role === "assistant",
-      )
-      if (taskMessages) {
-        const taskParts: PartType[] = []
-        for (const m of taskMessages) {
-          const msgParts = data.store.part[m.id]
-          if (msgParts) {
-            for (const p of msgParts) {
-              if (p) taskParts.push(p)
-            }
-          }
+    const msgs = assistantMessages()
+    let last: PartType | undefined
+    let currentTask: ToolPart | undefined
+
+    for (let mi = msgs.length - 1; mi >= 0; mi--) {
+      const msgParts = data.store.part[msgs[mi].id] ?? []
+      for (let pi = msgParts.length - 1; pi >= 0; pi--) {
+        const part = msgParts[pi]
+        if (!part) continue
+        if (!last) last = part
+
+        if (
+          part.type === "tool" &&
+          part.tool === "task" &&
+          part.state &&
+          "metadata" in part.state &&
+          part.state.metadata?.sessionId &&
+          part.state.status === "running"
+        ) {
+          currentTask = part as ToolPart
+          break
+        }
+      }
+      if (currentTask) break
+    }
+
+    const taskSessionId =
+      currentTask?.state && "metadata" in currentTask.state
+        ? (currentTask.state.metadata?.sessionId as string | undefined)
+        : undefined
+
+    if (taskSessionId) {
+      const taskMessages = data.store.message[taskSessionId] ?? []
+      for (let mi = taskMessages.length - 1; mi >= 0; mi--) {
+        const msg = taskMessages[mi]
+        if (!msg || msg.role !== "assistant") continue
+
+        const msgParts = data.store.part[msg.id] ?? []
+        for (let pi = msgParts.length - 1; pi >= 0; pi--) {
+          const part = msgParts[pi]
+          if (part) return computeStatusFromPart(part)
         }
-        if (taskParts.length > 0) resolvedParts = taskParts
       }
     }
 
-    return computeStatusFromPart(resolvedParts.at(-1))
+    return computeStatusFromPart(last)
   })
 
   const status = createMemo(
@@ -367,7 +386,7 @@ export function SessionTurn(
               >
                 <Switch>
                   <Match when={isShellMode()}>
-                    <Part part={assistantParts()[0]} message={msg()} defaultOpen />
+                    <Part part={shellModePart()!} message={msg()} defaultOpen />
                   </Match>
                   <Match when={true}>
                     {/* Title (sticky) */}