fix(app): reconcile session pages with concurrent events (#34042)

packages/app/src/context/server-sdk.test.ts

@@ -1,5 +1,5 @@
 import { describe, expect, test } from "bun:test"
-import { coalesceServerEvents, resumeStreamAfterPageShow } from "./server-sdk"
+import { coalesceServerEvents, enqueueServerEvent, resumeStreamAfterPageShow } from "./server-sdk"
 import type { Event } from "@opencode-ai/sdk/v2/client"
 
 describe("resumeStreamAfterPageShow", () => {
@@ -15,19 +15,23 @@ describe("resumeStreamAfterPageShow", () => {
 })
 
 describe("coalesceServerEvents", () => {
-  const delta = (value: string, field = "text") => ({
+  const delta = (value: string, field = "text", partID = "part") => ({
     directory: "/repo",
     payload: {
       type: "message.part.delta",
-      properties: { messageID: "msg", partID: "part", field, delta: value },
+      properties: { messageID: "msg", partID, field, delta: value },
     } as Event,
   })
 
   test("merges adjacent deltas for the same field", () => {
-    const result = coalesceServerEvents([delta("hello "), delta("world")])
+    const first = delta("hello ")
+    const second = delta("world")
+    first.payload.id = "first"
+    second.payload.id = "second"
+    const result = coalesceServerEvents([first, second])
 
     expect(result).toHaveLength(1)
-    expect(result[0]?.payload).toMatchObject({ properties: { delta: "hello world" } })
+    expect(result[0]?.payload).toMatchObject({ id: "second", properties: { delta: "hello world" } })
   })
 
   test("preserves event boundaries and distinct fields", () => {
@@ -45,9 +49,113 @@ describe("coalesceServerEvents", () => {
     ])
   })
 
-  test("drops stale deltas", () => {
-    const result = coalesceServerEvents([delta("stale")], new Set(["/repo:msg:part"]))
+  test("preserves event ID order across interleaved deltas", () => {
+    const first = delta("a")
+    const other = delta("b", "text", "other")
+    const last = delta("c")
+    first.payload.id = "1"
+    other.payload.id = "2"
+    last.payload.id = "3"
 
-    expect(result).toEqual([])
+    const result = coalesceServerEvents([first, other, last])
+
+    expect(result.map((event) => event.payload.id)).toEqual(["1", "2", "3"])
+  })
+
+})
+
+describe("enqueueServerEvent", () => {
+  const partUpdated = (text: string) =>
+    ({
+      type: "message.part.updated",
+      properties: {
+        sessionID: "session",
+        part: { id: "part", sessionID: "session", messageID: "message", type: "text", text },
+      },
+    }) as Event
+
+  test("preserves part updates across message remove and re-add barriers", () => {
+    const events: Array<{ directory: string; payload: Event }> = []
+    const enqueue = (payload: Event) => enqueueServerEvent(events, { directory: "/repo", payload })
+
+    enqueue(partUpdated("old"))
+    enqueue({ type: "message.removed", properties: { sessionID: "session", messageID: "message" } } as Event)
+    enqueue({
+      type: "message.updated",
+      properties: {
+        sessionID: "session",
+        info: {
+          id: "message",
+          sessionID: "session",
+          role: "user",
+          time: { created: 1 },
+          agent: "build",
+          model: { providerID: "provider", modelID: "model" },
+        },
+      },
+    } as Event)
+    enqueue(partUpdated("new"))
+
+    expect(events.map((event) => event.payload.type)).toEqual([
+      "message.part.updated",
+      "message.removed",
+      "message.updated",
+      "message.part.updated",
+    ])
+  })
+
+  test("preserves deltas after a replacement snapshot", () => {
+    const events: Array<{ directory: string; payload: Event }> = []
+    const enqueue = (payload: Event) => enqueueServerEvent(events, { directory: "/repo", payload })
+
+    enqueue(partUpdated("a"))
+    enqueue(partUpdated("ab"))
+    enqueue({
+      type: "message.part.delta",
+      properties: { sessionID: "session", messageID: "message", partID: "part", field: "text", delta: "c" },
+    } as Event)
+
+    const result = coalesceServerEvents(events)
+    expect(result.map((event) => event.payload.type)).toEqual(["message.part.updated", "message.part.delta"])
+    expect(result[0]?.payload).toMatchObject({ properties: { part: { text: "ab" } } })
+    expect(result[1]?.payload).toMatchObject({ properties: { delta: "c" } })
+  })
+
+  test("preserves updates after session deletion", () => {
+    const events: Array<{ directory: string; payload: Event }> = []
+    const enqueue = (payload: Event) => enqueueServerEvent(events, { directory: "/repo", payload })
+
+    enqueue(partUpdated("old"))
+    enqueue({
+      type: "session.deleted",
+      properties: { sessionID: "session", info: { id: "session" } },
+    } as Event)
+    enqueue(partUpdated("new"))
+
+    expect(events.map((event) => event.payload.type)).toEqual([
+      "message.part.updated",
+      "session.deleted",
+      "message.part.updated",
+    ])
+  })
+
+  test("does not coalesce edge-triggered session statuses", () => {
+    const events: Array<{ directory: string; payload: Event }> = []
+    const enqueue = (status: "retry" | "busy") =>
+      enqueueServerEvent(events, {
+        directory: "/repo",
+        payload: {
+          type: "session.status",
+          properties: {
+            sessionID: "session",
+            status: status === "retry" ? { type: "retry", attempt: 1, message: "retry", next: 1 } : { type: "busy" },
+          },
+        } as Event,
+      })
+
+    enqueue("retry")
+    enqueue("busy")
+
+    expect(events).toHaveLength(2)
   })
 })

packages/app/src/context/server-sdk.tsx

@@ -17,34 +17,59 @@ const isAbortError = (error: unknown) =>
 const isStreamClosed = (error: unknown, signal?: AbortSignal) => isAbortError(error) || signal?.aborted === true
 type QueuedServerEvent = { directory: string; payload: Event }
 
-const deltaKey = (directory: string, messageID: string, partID: string) => `${directory}:${messageID}:${partID}`
+const coalescedKey = (event: QueuedServerEvent) => {
+  if (event.payload.type === "lsp.updated") return `lsp.updated:${event.directory}`
+  if (event.payload.type === "message.part.updated") {
+    const part = event.payload.properties.part
+    return `message.part.updated:${event.directory}:${part.messageID}:${part.id}`
+  }
+  return undefined
+}
+
+export function enqueueServerEvent(
+  queue: QueuedServerEvent[],
+  event: QueuedServerEvent,
+) {
+  const key = coalescedKey(event)
+  const previous = queue[queue.length - 1]
+  if (key && previous && coalescedKey(previous) === key) {
+    queue[queue.length - 1] = event
+    return false
+  }
+  queue.push(event)
+  return true
+}
 
-export function coalesceServerEvents(events: QueuedServerEvent[], stale?: Set<string>) {
+export function coalesceServerEvents(events: QueuedServerEvent[]) {
   const output: QueuedServerEvent[] = []
-  const deltas = new Map<string, number>()
   events.forEach((event) => {
-    if (stale && event.payload.type === "message.part.delta") {
-      const props = event.payload.properties
-      if (stale.has(deltaKey(event.directory, props.messageID, props.partID))) return
-    }
     if (event.payload.type !== "message.part.delta") {
-      deltas.clear()
       output.push(event)
       return
     }
     const props = event.payload.properties
-    const id = `${deltaKey(event.directory, props.messageID, props.partID)}:${props.field}`
-    const index = deltas.get(id)
-    const existing = index === undefined ? undefined : output[index]
-    if (!existing || existing.payload.type !== "message.part.delta") {
-      deltas.set(id, output.length)
+    const previous = output[output.length - 1]
+    if (
+      !previous ||
+      previous.payload.type !== "message.part.delta" ||
+      previous.directory !== event.directory ||
+      previous.payload.properties.messageID !== props.messageID ||
+      previous.payload.properties.partID !== props.partID ||
+      previous.payload.properties.field !== props.field
+    ) {
       output.push({
         directory: event.directory,
         payload: { ...event.payload, properties: { ...props } },
       })
       return
     }
-    existing.payload.properties.delta += props.delta
+    output[output.length - 1] = {
+      directory: event.directory,
+      payload: {
+        ...event.payload,
+        properties: { ...props, delta: previous.payload.properties.delta + props.delta },
+      },
+    }
   })
   return output
 }
@@ -85,20 +110,9 @@ function createServerSdkContextBase(server: ServerConnection.Any, scope: ServerS
 
   let queue: Queued[] = []
   let buffer: Queued[] = []
-  const coalesced = new Map<string, number>()
-  const staleDeltas = new Set<string>()
   let timer: ReturnType<typeof setTimeout> | undefined
   let last = 0
 
-  const key = (directory: string, payload: Event) => {
-    if (payload.type === "session.status") return `session.status:${directory}:${payload.properties.sessionID}`
-    if (payload.type === "lsp.updated") return `lsp.updated:${directory}`
-    if (payload.type === "message.part.updated") {
-      const part = payload.properties.part
-      return `message.part.updated:${directory}:${part.messageID}:${part.id}`
-    }
-  }
-
   const flush = () => {
     if (timer) clearTimeout(timer)
     timer = undefined
@@ -106,15 +120,12 @@ function createServerSdkContextBase(server: ServerConnection.Any, scope: ServerS
     if (queue.length === 0) return
 
     const events = queue
-    const skip = staleDeltas.size > 0 ? new Set(staleDeltas) : undefined
     queue = buffer
     buffer = events
     queue.length = 0
-    coalesced.clear()
-    staleDeltas.clear()
 
     last = Date.now()
-    const output = coalesceServerEvents(events, skip)
+    const output = coalesceServerEvents(events)
     batch(() => {
       output.forEach((event) => emitter.emit(event.directory, event.payload))
     })
@@ -184,28 +195,11 @@ function createServerSdkContextBase(server: ServerConnection.Any, scope: ServerS
           for await (const event of events.stream) {
             resetHeartbeat()
             streamErrorLogged = false
-            const directory = event.directory ?? "global"
-            if (event.payload.type === "sync") {
-              continue
-            }
-
-            const payload = event.payload as Event
-
-            const k = key(directory, payload)
-            if (k) {
-              const i = coalesced.get(k)
-              if (i !== undefined) {
-                queue[i] = { directory, payload }
-                if (payload.type === "message.part.updated") {
-                  const part = payload.properties.part
-                  staleDeltas.add(deltaKey(directory, part.messageID, part.id))
-                }
-                continue
-              }
-              coalesced.set(k, queue.length)
+            if (event.payload.type !== "sync") {
+              const directory = event.directory ?? "global"
+              const payload = event.payload as Event
+              if (enqueueServerEvent(queue, { directory, payload })) schedule()
             }
-            queue.push({ directory, payload })
-            schedule()
 
             if (Date.now() - yielded < STREAM_YIELD_MS) continue
             yielded = Date.now()

packages/app/src/context/server-session.ts

@@ -18,44 +18,70 @@ import { rootSession } from "@/utils/session-route"
 import { dropSessionCaches, pickSessionCacheEvictions, SESSION_CACHE_LIMIT } from "./global-sync/session-cache"
 
 const cmp = (a: string, b: string) => (a < b ? -1 : a > b ? 1 : 0)
+const cmpMessage = (a: Message, b: Message) => a.time.created - b.time.created || cmp(a.id, b.id)
 const SKIP_PARTS = new Set(["patch", "step-start", "step-finish"])
 const initialMessagePageSize = 2
 const historyMessagePageSize = 200
 const sessionInfoLimit = 2_048
+const emptyIDs: ReadonlySet<string> = new Set()
 
 type OptimisticItem = {
   message: Message
   parts: Part[]
+  confirmedParts?: Part[]
+  confirmedMessage?: boolean
 }
 
-const hasParts = (parts: Part[] | undefined, want: Part[]) => {
-  if (!parts) return want.length === 0
-  return want.every((part) => Binary.search(parts, part.id, (item) => item.id).found)
+type MessagePage = {
+  session: Message[]
+  part: { id: string; part: Part[] }[]
+  cursor?: string
+  complete: boolean
+}
+
+// Most markers describe the current HTTP attempt; deltaParts persists non-durable stream state across retries.
+type MessageLoadState = {
+  touchedMessages: Set<string>
+  removedMessages: Set<string>
+  retainedMessages: Set<string>
+  touchedParts: Map<string, Set<string>>
+  deltaParts: Map<string, Set<string>>
+  carriedDeltaParts: Map<string, Set<string>>
+  removedParts: Map<string, Set<string>>
+  optimisticParts: Map<string, Set<string>>
+  orphanParents: Set<string>
+  clearedMessageParts: Set<string>
 }
 
 function mergeOptimisticPage(
-  page: { session: Message[]; part: { id: string; part: Part[] }[]; cursor?: string; complete: boolean },
+  page: MessagePage,
   items: OptimisticItem[],
 ) {
-  if (items.length === 0) return { ...page, confirmed: [] as string[] }
+  if (items.length === 0) return { ...page, observed: [] as { messageID: string; parts: Part[] }[] }
   const session = [...page.session]
   const part = new Map(page.part.map((item) => [item.id, item.part]))
-  const confirmed: string[] = []
+  const observed: { messageID: string; parts: Part[] }[] = []
   for (const item of items) {
     const result = Binary.search(session, item.message.id, (message) => message.id)
     if (!result.found) session.splice(result.index, 0, item.message)
     const current = part.get(item.message.id)
-    if (result.found && hasParts(current, item.parts)) {
-      confirmed.push(item.message.id)
-      continue
-    }
-    part.set(item.message.id, merge(current ?? [], item.parts))
+    const confirmed = result.found
+      ? item.parts.filter((part) => Binary.search(current ?? [], part.id, (value) => value.id).found)
+      : []
+    if (result.found) observed.push({ messageID: item.message.id, parts: confirmed })
+    part.set(
+      item.message.id,
+      merge(
+        result.found ? (current ?? []) : merge(item.confirmedParts ?? [], current ?? []),
+        item.parts.filter((part) => !confirmed.includes(part)),
+      ),
+    )
   }
   return {
     ...page,
     session,
     part: [...part.entries()].sort((a, b) => cmp(a[0], b[0])).map(([id, parts]) => ({ id, part: parts })),
-    confirmed,
+    observed,
   }
 }
 
@@ -75,7 +101,38 @@ function merge<T extends { id: string }>(a: readonly T[], b: readonly T[]) {
   return [...items.values()].sort((x, y) => cmp(x.id, y.id))
 }
 
-export function createServerSession(client: OpencodeClient) {
+function reconcileFetched<T extends { id: string }>(
+  fetched: T[],
+  current: readonly T[],
+  options: {
+    touched?: ReadonlySet<string>
+    retained?: ReadonlySet<string>
+    preserveUnfetched?: boolean | ((item: T) => boolean)
+  } = {},
+) {
+  const result = new Map(fetched.map((item) => [item.id, item]))
+  const live = new Map(current.map((item) => [item.id, item]))
+  if (options.preserveUnfetched) {
+    for (const item of current) {
+      if (!result.has(item.id) && (options.preserveUnfetched === true || options.preserveUnfetched(item)))
+        result.set(item.id, item)
+    }
+  }
+  for (const id of options.retained ?? emptyIDs) {
+    if (result.has(id)) continue
+    const item = live.get(id)
+    if (item) result.set(id, item)
+  }
+  // Events observed while the request is pending are the freshest client state for those identities.
+  for (const id of options.touched ?? emptyIDs) {
+    const item = live.get(id)
+    if (item) result.set(id, item)
+    if (!item) result.delete(id)
+  }
+  return [...result.values()].sort((a, b) => cmp(a.id, b.id))
+}
+
+export function createServerSession(client: OpencodeClient, options?: { retry?: typeof retry }) {
   const [data, setData] = createStore({
     info: {} as Record<string, Session | undefined>,
     session_status: {} as Record<string, SessionStatus>,
@@ -95,10 +152,29 @@ export function createServerSession(client: OpencodeClient) {
   const inflightDiff = new Map<string, Promise<void>>()
   const inflightTodo = new Map<string, Promise<void>>()
   const optimistic = new Map<string, Map<string, OptimisticItem>>()
+  const messageLoads = new Map<string, MessageLoadState>()
+  const pendingParts = new Map<string, Map<string, Set<string>>>()
+  const orphanParts = new Map<string, Set<string>>()
+  const removedMessages = new Map<string, Set<string>>()
+  const deltaBases = new Map<string, { base: string; sessionID: string }>()
+  const deleteMessageParts = (cache: { part: Record<string, Part[] | undefined>; part_text_accum_delta: Record<string, string | undefined> }, messageID: string) => {
+    for (const part of cache.part[messageID] ?? []) {
+      delete cache.part_text_accum_delta[part.id]
+      deltaBases.delete(part.id)
+    }
+    delete cache.part[messageID]
+  }
   const seen = new Set<string>()
   const infoSeen = new Set<string>()
   const pinned = new Map<string, number>()
-  const generations = new Map<string, number>()
+  const generations = new Map<string, object>()
+  const generation = (sessionID: string) => {
+    const current = generations.get(sessionID)
+    if (current) return current
+    const created = {}
+    generations.set(sessionID, created)
+    return created
+  }
   const [meta, setMeta] = createStore({
     limit: {} as Record<string, number | undefined>,
     cursor: {} as Record<string, string | undefined>,
@@ -115,6 +191,11 @@ export function createServerSession(client: OpencodeClient) {
       const preserve = new Set([
         ...pinned.keys(),
         ...requests.keys(),
+        ...inflight.keys(),
+        ...inflightDiff.keys(),
+        ...inflightTodo.keys(),
+        ...messageLoads.keys(),
+        ...optimistic.keys(),
         ...Object.entries(data.permission)
           .filter(([, items]) => items.length > 0)
           .map(([sessionID]) => sessionID),
@@ -138,6 +219,7 @@ export function createServerSession(client: OpencodeClient) {
         if (!preserve.has(sessionID)) stale.push(sessionID)
       }
       stale.forEach((sessionID) => infoSeen.delete(sessionID))
+      stale.forEach((sessionID) => generations.delete(sessionID))
       setData(
         "info",
         produce((draft) => stale.forEach((sessionID) => delete draft[sessionID])),
@@ -151,20 +233,29 @@ export function createServerSession(client: OpencodeClient) {
     if (cached && !options?.force) return Promise.resolve(cached)
     const pending = requests.get(sessionID)
     if (pending) return pending
-    const generation = generations.get(sessionID) ?? 0
+    const active = generation(sessionID)
     const request = client.session.get({ sessionID }).then((result) => {
       if (!result.data) throw new Error(`Session not found: ${sessionID}`)
-      if ((generations.get(sessionID) ?? 0) !== generation) return result.data
+      if (generations.get(sessionID) !== active) return result.data
       return remember(result.data)
     })
     requests.set(sessionID, request)
+    const cleanup = () => {
+      if (requests.get(sessionID) === request) requests.delete(sessionID)
+      if (
+        generations.get(sessionID) === active &&
+        !data.info[sessionID] &&
+        !requests.has(sessionID) &&
+        !messageLoads.has(sessionID) &&
+        !inflight.has(sessionID) &&
+        !inflightDiff.has(sessionID) &&
+        !inflightTodo.has(sessionID)
+      )
+        generations.delete(sessionID)
+    }
     void request.then(
-      () => {
-        if (requests.get(sessionID) === request) requests.delete(sessionID)
-      },
-      () => {
-        if (requests.get(sessionID) === request) requests.delete(sessionID)
-      },
+      cleanup,
+      cleanup,
     )
     return request
   }
@@ -195,15 +286,120 @@ export function createServerSession(client: OpencodeClient) {
     if (items.size === 0) optimistic.delete(sessionID)
   }
 
+  const clearOptimisticPart = (sessionID: string, messageID: string, partID: string) => {
+    const items = optimistic.get(sessionID)
+    const item = items?.get(messageID)
+    if (!items || !item) return
+    const parts = item.parts.filter((part) => part.id !== partID)
+    const confirmedParts = item.confirmedParts?.filter((part) => part.id !== partID)
+    if (parts.length === 0) {
+      clearOptimistic(sessionID, messageID)
+      return
+    }
+    items.set(messageID, { ...item, parts, confirmedParts, confirmedMessage: true })
+  }
+
+  const confirmOptimisticPart = (sessionID: string, messageID: string, part: Part) => {
+    const items = optimistic.get(sessionID)
+    const item = items?.get(messageID)
+    if (!items || !item) return
+    const parts = item.parts.filter((value) => value.id !== part.id)
+    if (parts.length === 0) {
+      clearOptimistic(sessionID, messageID)
+      return
+    }
+    items.set(messageID, {
+      ...item,
+      parts,
+      confirmedParts: merge(item.confirmedParts ?? [], [part]),
+      confirmedMessage: true,
+    })
+  }
+
+  const confirmOptimistic = (sessionID: string, messageID: string, confirmedParts: Part[]) => {
+    const items = optimistic.get(sessionID)
+    const item = items?.get(messageID)
+    if (!items || !item) return
+    const confirmed = new Set(confirmedParts.map((part) => part.id))
+    const parts = item.parts.filter((part) => !confirmed.has(part.id))
+    if (parts.length === 0) {
+      clearOptimistic(sessionID, messageID)
+      return
+    }
+    items.set(messageID, {
+      ...item,
+      parts,
+      confirmedParts: merge(item.confirmedParts ?? [], confirmedParts),
+      confirmedMessage: true,
+    })
+  }
+
+  const trackPartChange = (sessionID: string, messageID: string, partID: string) => {
+    const load = messageLoads.get(sessionID)
+    if (!load) return
+    // A part event keeps an existing parent when the fetched page omits it without overriding fetched metadata.
+    const messages = data.message[sessionID]
+    if (messages && Binary.search(messages, messageID, (message) => message.id).found) load.retainedMessages.add(messageID)
+    const parts = load.touchedParts.get(messageID)
+    if (parts) {
+      parts.add(partID)
+      return
+    }
+    load.touchedParts.set(messageID, new Set([partID]))
+  }
+
+  const resetMessageLoad = (sessionID: string, load: MessageLoadState) => {
+    load.touchedMessages.clear()
+    load.retainedMessages.clear()
+    load.touchedParts.clear()
+    load.carriedDeltaParts.clear()
+    load.clearedMessageParts.clear()
+    for (const messageID of load.removedMessages) {
+      load.touchedMessages.add(messageID)
+      load.clearedMessageParts.add(messageID)
+    }
+    for (const [messageID, parts] of load.deltaParts) {
+      load.touchedParts.set(messageID, new Set(parts))
+      load.carriedDeltaParts.set(messageID, new Set(parts))
+      const messages = data.message[sessionID]
+      if (messages && Binary.search(messages, messageID, (message) => message.id).found)
+        load.retainedMessages.add(messageID)
+    }
+    for (const [messageID, parts] of load.removedParts) {
+      const touched = load.touchedParts.get(messageID) ?? new Set<string>()
+      parts.forEach((partID) => touched.add(partID))
+      load.touchedParts.set(messageID, touched)
+      const messages = data.message[sessionID]
+      if (messages && Binary.search(messages, messageID, (message) => message.id).found)
+        load.retainedMessages.add(messageID)
+    }
+    for (const [messageID, parts] of load.optimisticParts) {
+      load.removedMessages.delete(messageID)
+      load.clearedMessageParts.add(messageID)
+      load.touchedMessages.add(messageID)
+      const touched = load.touchedParts.get(messageID) ?? new Set<string>()
+      parts.forEach((partID) => touched.add(partID))
+      load.touchedParts.set(messageID, touched)
+    }
+  }
+
   const evict = (sessionIDs: string[]) => {
     if (sessionIDs.length === 0) return
+    const evicted = new Set(sessionIDs)
+    for (const [partID, item] of deltaBases) {
+      if (evicted.has(item.sessionID)) deltaBases.delete(partID)
+    }
     sessionIDs.forEach((sessionID) => {
-      generations.set(sessionID, (generations.get(sessionID) ?? 0) + 1)
+      generations.delete(sessionID)
       clearOptimistic(sessionID)
       requests.delete(sessionID)
       inflight.delete(sessionID)
       inflightDiff.delete(sessionID)
       inflightTodo.delete(sessionID)
+      messageLoads.delete(sessionID)
+      pendingParts.delete(sessionID)
+      orphanParts.delete(sessionID)
+      removedMessages.delete(sessionID)
     })
     setData(
       produce((draft) => {
@@ -230,6 +426,7 @@ export function createServerSession(client: OpencodeClient) {
       ...inflight.keys(),
       ...inflightDiff.keys(),
       ...inflightTodo.keys(),
+      ...messageLoads.keys(),
       ...optimistic.keys(),
       ...Object.entries(data.permission)
         .filter(([, items]) => items.length > 0)
@@ -247,8 +444,11 @@ export function createServerSession(client: OpencodeClient) {
       pickSessionCacheEvictions({ seen, keep: sessionID, limit: SESSION_CACHE_LIMIT, preserve: protectedSessions() }),
     )
 
-  const fetchMessages = async (sessionID: string, limit: number, before?: string) => {
-    const response = await retry(() => client.session.messages({ sessionID, limit, before }))
+  const fetchMessages = async (sessionID: string, limit: number, before?: string, onAttempt?: () => void) => {
+    const response = await (options?.retry ?? retry)(() => {
+      onAttempt?.()
+      return client.session.messages({ sessionID, limit, before })
+    })
     const items = (response.data ?? []).filter((item) => !!item?.info?.id)
     return {
       session: items.map((item) => cleanMessage(item.info)).sort((a, b) => cmp(a.id, b.id)),
@@ -261,30 +461,164 @@ export function createServerSession(client: OpencodeClient) {
     }
   }
 
+  const replaceMessages = (sessionID: string, messages: Message[]) => {
+    const messageIDs = new Set(messages.map((message) => message.id))
+    const dropped = (data.message[sessionID] ?? []).filter((message) => !messageIDs.has(message.id))
+    setData("message", sessionID, reconcile(messages, { key: "id" }))
+    setData(
+      produce((draft) => {
+        for (const message of dropped) deleteMessageParts(draft, message.id)
+      }),
+    )
+    return messageIDs
+  }
+
+  const replaceParts = (
+    sessionID: string,
+    items: MessagePage["part"],
+    messageIDs: Set<string>,
+    load?: MessageLoadState,
+  ) => {
+    for (const item of items) {
+      if (!messageIDs.has(item.id)) continue
+      const fetched = load?.clearedMessageParts.has(item.id)
+        ? []
+        : item.part.filter((part) => !SKIP_PARTS.has(part.type))
+      const fetchedIDs = new Set(fetched.map((part) => part.id))
+      const pending = pendingParts.get(sessionID)?.get(item.id)
+      const touched = new Set([...(load?.touchedParts.get(item.id) ?? []), ...(pending ?? [])])
+      for (const part of fetched) {
+        const accumulated = data.part_text_accum_delta[part.id]
+        const base = deltaBases.get(part.id)?.base
+        const preserveDelta =
+          base !== undefined &&
+          accumulated !== undefined &&
+          "text" in part &&
+          typeof part.text === "string" &&
+          part.text.startsWith(base) &&
+          accumulated.startsWith(part.text) &&
+          accumulated !== part.text
+        if (preserveDelta) touched.add(part.id)
+        if (load?.carriedDeltaParts.get(item.id)?.has(part.id) && !preserveDelta) touched.delete(part.id)
+      }
+      for (const partID of load?.carriedDeltaParts.get(item.id) ?? []) {
+        if (!fetchedIDs.has(partID)) touched.delete(partID)
+      }
+      const parts = reconcileFetched(
+        fetched,
+        data.part[item.id] ?? [],
+        { touched },
+      )
+      if (!parts.length) {
+        orphanParts.get(sessionID)?.delete(item.id)
+        setData(produce((draft) => deleteMessageParts(draft, item.id)))
+        continue
+      }
+      const partIDs = new Set(parts.map((part) => part.id))
+      setData(
+        "part_text_accum_delta",
+        produce((draft) => {
+          for (const part of data.part[item.id] ?? []) {
+            if (!partIDs.has(part.id) || !touched.has(part.id)) {
+              delete draft[part.id]
+              deltaBases.delete(part.id)
+            }
+          }
+        }),
+      )
+      setData("part", item.id, reconcile(parts, { key: "id" }))
+      orphanParts.get(sessionID)?.delete(item.id)
+    }
+  }
+
+  const applyMessagePage = (
+    sessionID: string,
+    page: MessagePage,
+    load: MessageLoadState | undefined,
+    preserveUnfetched: boolean | ((message: Message) => boolean),
+    cleanupOrphans: boolean,
+  ) => {
+    const merged = mergeOptimisticPage(page, [...(optimistic.get(sessionID)?.values() ?? [])])
+    merged.observed.forEach((item) => {
+      if (!load?.clearedMessageParts.has(item.messageID)) confirmOptimistic(sessionID, item.messageID, item.parts)
+    })
+    const touchedMessages = new Set([
+      ...(load?.touchedMessages ?? []),
+      ...(removedMessages.get(sessionID) ?? []),
+    ])
+    const messages = reconcileFetched(
+      merged.session,
+      data.message[sessionID] ?? [],
+      {
+        touched: touchedMessages,
+        retained: load?.retainedMessages,
+        preserveUnfetched,
+      },
+    )
+    batch(() => {
+      const messageIDs = replaceMessages(sessionID, messages)
+      replaceParts(sessionID, merged.part, messageIDs, load)
+      const orphans = orphanParts.get(sessionID)
+      if (cleanupOrphans && page.complete && orphans) {
+        for (const messageID of orphans) {
+          if (!messageIDs.has(messageID)) setData(produce((draft) => deleteMessageParts(draft, messageID)))
+        }
+        orphanParts.delete(sessionID)
+      }
+      setMeta("limit", sessionID, messages.length)
+      setMeta("cursor", sessionID, merged.cursor)
+      setMeta("complete", sessionID, merged.complete)
+      setMeta("at", sessionID, Date.now())
+    })
+  }
+
   const loadMessages = async (sessionID: string, limit: number, before?: string, mode?: "replace" | "prepend") => {
     if (meta.loading[sessionID]) return
-    const generation = generations.get(sessionID) ?? 0
+    const active = generation(sessionID)
+    const load: MessageLoadState = {
+      touchedMessages: new Set(),
+      removedMessages: new Set(),
+      retainedMessages: new Set(),
+      touchedParts: new Map(),
+      deltaParts: new Map(),
+      carriedDeltaParts: new Map(),
+      removedParts: new Map(),
+      optimisticParts: new Map(),
+      orphanParents: new Set(),
+      clearedMessageParts: new Set(),
+    }
+    messageLoads.set(sessionID, load)
     setMeta("loading", sessionID, true)
-    await fetchMessages(sessionID, limit, before)
+    let applied = false
+    await fetchMessages(sessionID, limit, before, () => resetMessageLoad(sessionID, load))
       .then((page) => {
-        if ((generations.get(sessionID) ?? 0) !== generation) return
-        const next = mergeOptimisticPage(page, [...(optimistic.get(sessionID)?.values() ?? [])])
-        next.confirmed.forEach((messageID) => clearOptimistic(sessionID, messageID))
-        const messages = mode === "prepend" ? merge(data.message[sessionID] ?? [], next.session) : next.session
-        batch(() => {
-          setData("message", sessionID, reconcile(messages, { key: "id" }))
-          for (const item of next.part) {
-            const parts = item.part.filter((part) => !SKIP_PARTS.has(part.type))
-            if (parts.length) setData("part", item.id, reconcile(parts, { key: "id" }))
-          }
-          setMeta("limit", sessionID, messages.length)
-          setMeta("cursor", sessionID, next.cursor)
-          setMeta("complete", sessionID, next.complete)
-          setMeta("at", sessionID, Date.now())
-        })
+        if (generations.get(sessionID) !== active) return
+        const first = page.session.reduce<Message | undefined>(
+          (oldest, message) => (!oldest || cmpMessage(message, oldest) < 0 ? message : oldest),
+          undefined,
+        )
+        const preserveUnfetched =
+          mode === "prepend" || (!page.complete && (!first || ((message: Message) => cmpMessage(message, first) < 0)))
+        applyMessagePage(
+          sessionID,
+          page,
+          messageLoads.get(sessionID) === load ? load : undefined,
+          preserveUnfetched,
+          mode !== "prepend",
+        )
+        applied = true
       })
       .finally(() => {
-        if ((generations.get(sessionID) ?? 0) === generation) setMeta("loading", sessionID, false)
+        if (!applied && generations.get(sessionID) === active && messageLoads.get(sessionID) === load) {
+          for (const messageID of load.orphanParents) {
+            if (!orphanParts.get(sessionID)?.has(messageID)) continue
+            setData(produce((draft) => deleteMessageParts(draft, messageID)))
+            orphanParts.get(sessionID)?.delete(messageID)
+          }
+          if (orphanParts.get(sessionID)?.size === 0) orphanParts.delete(sessionID)
+        }
+        if (messageLoads.get(sessionID) === load) messageLoads.delete(sessionID)
+        if (generations.get(sessionID) === active) setMeta("loading", sessionID, false)
       })
   }
 
@@ -339,7 +673,13 @@ export function createServerSession(client: OpencodeClient) {
     const eventID = eventSessionID(event)
     if (eventID) {
       touch(eventID)
-      if (!data.info[eventID]) void resolve(eventID).catch(() => {})
+      if (
+        !data.info[eventID] &&
+        event.type !== "session.created" &&
+        event.type !== "session.updated" &&
+        event.type !== "session.deleted"
+      )
+        void resolve(eventID).catch(() => {})
     }
     switch (event.type) {
       case "session.created":
@@ -378,6 +718,21 @@ export function createServerSession(client: OpencodeClient) {
       }
       case "message.updated": {
         const info = cleanMessage((event.properties as { info: Message }).info)
+        const load = messageLoads.get(info.sessionID)
+        load?.touchedMessages.add(info.id)
+        load?.removedMessages.delete(info.id)
+        const items = optimistic.get(info.sessionID)
+        const item = items?.get(info.id)
+        if (items && item) {
+          if (item.parts.length === 0) clearOptimistic(info.sessionID, info.id)
+          if (item.parts.length > 0) items.set(info.id, { ...item, confirmedMessage: true })
+        }
+        const orphans = orphanParts.get(info.sessionID)
+        orphans?.delete(info.id)
+        if (orphans?.size === 0) orphanParts.delete(info.sessionID)
+        const removedMessagesForSession = removedMessages.get(info.sessionID)
+        removedMessagesForSession?.delete(info.id)
+        if (removedMessagesForSession?.size === 0) removedMessages.delete(info.sessionID)
         const messages = data.message[info.sessionID]
         if (!messages) {
           setData("message", info.sessionID, [info])
@@ -395,6 +750,20 @@ export function createServerSession(client: OpencodeClient) {
       }
       case "message.removed": {
         const props = event.properties as { sessionID: string; messageID: string }
+        const load = messageLoads.get(props.sessionID)
+        load?.touchedMessages.add(props.messageID)
+        load?.removedMessages.add(props.messageID)
+        load?.clearedMessageParts.add(props.messageID)
+        load?.deltaParts.delete(props.messageID)
+        load?.carriedDeltaParts.delete(props.messageID)
+        load?.removedParts.delete(props.messageID)
+        load?.optimisticParts.delete(props.messageID)
+        pendingParts.get(props.sessionID)?.delete(props.messageID)
+        if (pendingParts.get(props.sessionID)?.size === 0) pendingParts.delete(props.sessionID)
+        const removedMessagesForSession = removedMessages.get(props.sessionID) ?? new Set<string>()
+        removedMessagesForSession.add(props.messageID)
+        removedMessages.set(props.sessionID, removedMessagesForSession)
+        clearOptimistic(props.sessionID, props.messageID)
         setData(
           produce((draft) => {
             const messages = draft.message[props.sessionID]
@@ -402,8 +771,7 @@ export function createServerSession(client: OpencodeClient) {
               const result = Binary.search(messages, props.messageID, (message) => message.id)
               if (result.found) messages.splice(result.index, 1)
             }
-            for (const part of draft.part[props.messageID] ?? []) delete draft.part_text_accum_delta[part.id]
-            delete draft.part[props.messageID]
+            deleteMessageParts(draft, props.messageID)
           }),
         )
         return
@@ -411,6 +779,40 @@ export function createServerSession(client: OpencodeClient) {
       case "message.part.updated": {
         const part = (event.properties as { part: Part }).part
         if (SKIP_PARTS.has(part.type)) return
+        const messages = data.message[part.sessionID]
+        const load = messageLoads.get(part.sessionID)
+        const missing = !messages || !Binary.search(messages, part.messageID, (message) => message.id).found
+        // Outside a page load, accepting a part without its ordered parent event would create an unbounded orphan.
+        if (
+          missing &&
+          (!load || load.clearedMessageParts.has(part.messageID) || removedMessages.get(part.sessionID)?.has(part.messageID))
+        )
+          return
+        if (missing) {
+          const orphans = orphanParts.get(part.sessionID) ?? new Set<string>()
+          orphans.add(part.messageID)
+          orphanParts.set(part.sessionID, orphans)
+          load?.orphanParents.add(part.messageID)
+        }
+        const deltas = load?.deltaParts.get(part.messageID)
+        deltas?.delete(part.id)
+        if (deltas?.size === 0) load?.deltaParts.delete(part.messageID)
+        const carried = load?.carriedDeltaParts.get(part.messageID)
+        carried?.delete(part.id)
+        if (carried?.size === 0) load?.carriedDeltaParts.delete(part.messageID)
+        const removed = load?.removedParts.get(part.messageID)
+        removed?.delete(part.id)
+        if (removed?.size === 0) load?.removedParts.delete(part.messageID)
+        const pending = pendingParts.get(part.sessionID)?.get(part.messageID)
+        pending?.delete(part.id)
+        if (pending?.size === 0) pendingParts.get(part.sessionID)?.delete(part.messageID)
+        if (pendingParts.get(part.sessionID)?.size === 0) pendingParts.delete(part.sessionID)
+        const optimistic = load?.optimisticParts.get(part.messageID)
+        optimistic?.delete(part.id)
+        if (optimistic?.size === 0) load?.optimisticParts.delete(part.messageID)
+        deltaBases.delete(part.id)
+        trackPartChange(part.sessionID, part.messageID, part.id)
+        confirmOptimisticPart(part.sessionID, part.messageID, part)
         setData(
           "part_text_accum_delta",
           produce((draft) => void delete draft[part.id]),
@@ -431,10 +833,34 @@ export function createServerSession(client: OpencodeClient) {
         return
       }
       case "message.part.removed": {
-        const props = event.properties as { messageID: string; partID: string }
+        const props = event.properties as { sessionID: string; messageID: string; partID: string }
+        // Part removal is event-only on the server, so its tombstone lasts until a later update or eviction.
+        const pending = pendingParts.get(props.sessionID) ?? new Map<string, Set<string>>()
+        const parts = pending.get(props.messageID) ?? new Set<string>()
+        parts.add(props.partID)
+        pending.set(props.messageID, parts)
+        pendingParts.set(props.sessionID, pending)
+        const deltas = messageLoads.get(props.sessionID)?.deltaParts.get(props.messageID)
+        deltas?.delete(props.partID)
+        if (deltas?.size === 0) messageLoads.get(props.sessionID)?.deltaParts.delete(props.messageID)
+        const load = messageLoads.get(props.sessionID)
+        const carried = load?.carriedDeltaParts.get(props.messageID)
+        carried?.delete(props.partID)
+        if (carried?.size === 0) load?.carriedDeltaParts.delete(props.messageID)
+        if (load) {
+          const parts = load.removedParts.get(props.messageID) ?? new Set<string>()
+          parts.add(props.partID)
+          load.removedParts.set(props.messageID, parts)
+          const optimistic = load.optimisticParts.get(props.messageID)
+          optimistic?.delete(props.partID)
+          if (optimistic?.size === 0) load.optimisticParts.delete(props.messageID)
+        }
+        trackPartChange(props.sessionID, props.messageID, props.partID)
+        clearOptimisticPart(props.sessionID, props.messageID, props.partID)
         setData(
           produce((draft) => {
             delete draft.part_text_accum_delta[props.partID]
+            deltaBases.delete(props.partID)
             const parts = draft.part[props.messageID]
             if (!parts) return
             const result = Binary.search(parts, props.partID, (part) => part.id)
@@ -445,13 +871,31 @@ export function createServerSession(client: OpencodeClient) {
         return
       }
       case "message.part.delta": {
-        const props = event.properties as { messageID: string; partID: string; field: string; delta: string }
+        const props = event.properties as {
+          sessionID: string
+          messageID: string
+          partID: string
+          field: string
+          delta: string
+        }
         const parts = data.part[props.messageID]
         if (!parts) return
         const result = Binary.search(parts, props.partID, (part) => part.id)
         if (!result.found) return
+        trackPartChange(props.sessionID, props.messageID, props.partID)
+        const load = messageLoads.get(props.sessionID)
+        if (load) {
+          const parts = load.deltaParts.get(props.messageID) ?? new Set<string>()
+          parts.add(props.partID)
+          load.deltaParts.set(props.messageID, parts)
+          const carried = load.carriedDeltaParts.get(props.messageID)
+          carried?.delete(props.partID)
+          if (carried?.size === 0) load.carriedDeltaParts.delete(props.messageID)
+        }
         const field = props.field as keyof (typeof parts)[number]
         const current = parts[result.index]?.[field]
+        if (!deltaBases.has(props.partID) && typeof current === "string")
+          deltaBases.set(props.partID, { base: current, sessionID: props.sessionID })
         setData(
           "part_text_accum_delta",
           props.partID,
@@ -559,32 +1003,67 @@ export function createServerSession(client: OpencodeClient) {
     },
     optimistic: {
       add(input: { sessionID: string; message: Message; parts: Part[] }) {
+        const parts = input.parts.filter((part) => !!part?.id && !SKIP_PARTS.has(part.type)).sort((a, b) => cmp(a.id, b.id))
+        const load = messageLoads.get(input.sessionID)
+        if (load?.clearedMessageParts.has(input.message.id)) {
+          const touched = load.touchedParts.get(input.message.id) ?? new Set<string>()
+          parts.forEach((part) => touched.add(part.id))
+          load.touchedParts.set(input.message.id, touched)
+        }
+        if (load) {
+          load.removedMessages.delete(input.message.id)
+          load.optimisticParts.set(input.message.id, new Set(parts.map((part) => part.id)))
+        }
         const items = optimistic.get(input.sessionID)
-        if (items) items.set(input.message.id, input)
-        if (!items) optimistic.set(input.sessionID, new Map([[input.message.id, input]]))
+        const removedMessagesForSession = removedMessages.get(input.sessionID)
+        removedMessagesForSession?.delete(input.message.id)
+        if (removedMessagesForSession?.size === 0) removedMessages.delete(input.sessionID)
+        if (items) items.set(input.message.id, { ...input, parts, confirmedParts: [] })
+        if (!items) optimistic.set(input.sessionID, new Map([[input.message.id, { ...input, parts, confirmedParts: [] }]]))
         setData("message", input.sessionID, (messages = []) => merge(messages, [input.message]))
         setData(
-          "part",
-          input.message.id,
-          input.parts.filter((part) => !!part?.id).sort((a, b) => cmp(a.id, b.id)),
+          "part_text_accum_delta",
+          produce((draft) => {
+            for (const part of [...(data.part[input.message.id] ?? []), ...parts]) {
+              delete draft[part.id]
+              deltaBases.delete(part.id)
+            }
+          }),
         )
+        setData("part", input.message.id, parts)
       },
       remove(input: { sessionID: string; messageID: string }) {
+        const item = optimistic.get(input.sessionID)?.get(input.messageID)
+        if (!item) return
+        messageLoads.get(input.sessionID)?.optimisticParts.delete(input.messageID)
         clearOptimistic(input.sessionID, input.messageID)
+        if (item.confirmedMessage) {
+          const partIDs = new Set(item.parts.map((part) => part.id))
+          setData(
+            produce((draft) => {
+              for (const part of item.parts) {
+                delete draft.part_text_accum_delta[part.id]
+                deltaBases.delete(part.id)
+              }
+              const parts = draft.part[input.messageID]
+              if (!parts) return
+              draft.part[input.messageID] = parts.filter((part) => !partIDs.has(part.id))
+              if (draft.part[input.messageID]?.length === 0) delete draft.part[input.messageID]
+            }),
+          )
+          return
+        }
         setData("message", input.sessionID, (messages) => messages?.filter((message) => message.id !== input.messageID))
-        setData(
-          "part",
-          produce((draft) => void delete draft[input.messageID]),
-        )
+        setData(produce((draft) => deleteMessageParts(draft, input.messageID)))
       },
     },
     diff(sessionID: string, options?: { force?: boolean }) {
       touch(sessionID)
       if (data.session_diff[sessionID] !== undefined && !options?.force) return Promise.resolve()
       return runInflight(inflightDiff, sessionID, () => {
-        const generation = generations.get(sessionID) ?? 0
+        const active = generation(sessionID)
         return retry(() => client.session.diff({ sessionID })).then((result) => {
-          if ((generations.get(sessionID) ?? 0) !== generation) return
+          if (generations.get(sessionID) !== active) return
           setData("session_diff", sessionID, reconcile(cleanDiffs(result.data), { key: "file" }))
         })
       })
@@ -593,9 +1072,9 @@ export function createServerSession(client: OpencodeClient) {
       touch(sessionID)
       if (data.todo[sessionID] !== undefined && !options?.force) return Promise.resolve()
       return runInflight(inflightTodo, sessionID, () => {
-        const generation = generations.get(sessionID) ?? 0
+        const active = generation(sessionID)
         return retry(() => client.session.todo({ sessionID })).then((result) => {
-          if ((generations.get(sessionID) ?? 0) !== generation) return
+          if (generations.get(sessionID) !== active) return
           setData("todo", sessionID, reconcile(result.data ?? [], { key: "id" }))
         })
       })