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

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 coalesceServerEvents(events: QueuedServerEvent[], stale?: Set<string>) {
+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[]) {
   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,29 +195,12 @@ 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
+            if (event.payload.type !== "sync") {
+              const directory = event.directory ?? "global"
+              const payload = event.payload as Event
+              if (enqueueServerEvent(queue, { directory, payload })) schedule()
             }
 
-            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)
-            }
-            queue.push({ directory, payload })
-            schedule()
-
             if (Date.now() - yielded < STREAM_YIELD_MS) continue
             yielded = Date.now()
             await wait(0)