fix(tui): serialize session hydration

Author: kitlangton

packages/tui/src/context/data.tsx

@@ -121,14 +121,14 @@ function createSync() {
   }
 }
 
-// Apply local admissions and promotions that arrived while the server snapshot was loading.
-function reconcilePendingInputs(baseline: string[], current: string[], server: string[]) {
-  const result = new Set(server)
-  const before = new Set(baseline)
-  const now = new Set(current)
-  baseline.filter((id) => !now.has(id)).forEach((id) => result.delete(id))
-  current.filter((id) => !before.has(id)).forEach((id) => result.add(id))
-  return [...result]
+// Apply admissions and promotions that arrived while the pending snapshot was loading.
+function reconcilePending(baseline: SessionPendingInfo[], current: SessionPendingInfo[], server: SessionPendingInfo[]) {
+  const result = new Map(server.map((item) => [item.id, item]))
+  const before = new Set(baseline.map((item) => item.id))
+  const now = new Map(current.map((item) => [item.id, item]))
+  baseline.filter((item) => !now.has(item.id)).forEach((item) => result.delete(item.id))
+  current.filter((item) => !before.has(item.id)).forEach((item) => result.set(item.id, item))
+  return [...result.values()]
 }
 
 export const { use: useData, provider: DataProvider } = createSimpleContext({
@@ -961,7 +961,9 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           },
           sync(sessionID: string) {
             return sync.run(`session.pending:${sessionID}`, async () => {
-              const pending = await client.api.session.pending.list({ sessionID })
+              const baseline = [...(store.session.pending[sessionID] ?? [])]
+              const server = await client.api.session.pending.list({ sessionID })
+              const pending = reconcilePending(baseline, store.session.pending[sessionID] ?? [], server)
               setStore("session", "pending", sessionID, reconcile(pending))
               setStore(
                 "session",
@@ -995,32 +997,17 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           },
           sync(sessionID: string) {
             return sync.run(`session.message:${sessionID}`, async () => {
-              const localInputs = store.session.input[sessionID] ?? []
-              const [projected, pending] = await Promise.all([
-                client.api.message.list({ sessionID, limit: 200, order: "desc" }),
-                client.api.session.pending.list({ sessionID }),
-              ])
-              const pendingMessages = pending.map(message.fromPending)
+              const localInputs = [...(store.session.input[sessionID] ?? [])]
+              const pendingMessages = [...(store.session.pending[sessionID] ?? [])].map(message.fromPending)
+              const projected = await client.api.message.list({ sessionID, limit: 200, order: "desc" })
               const next = projected.data.toReversed()
               const index = new Map(next.map((message, index) => [message.id, index]))
-              const localInputIDs = new Set(localInputs)
+              const localInputIDs = new Set([...localInputs, ...(store.session.input[sessionID] ?? [])])
               store.session.message[sessionID]
                 ?.filter((message) => localInputIDs.has(message.id))
                 .forEach((item) => message.append(next, index, item))
               pendingMessages.forEach((item) => message.append(next, index, item))
               messageIndex.set(sessionID, index)
-              setStore(
-                "session",
-                "input",
-                sessionID,
-                reconcilePendingInputs(
-                  localInputs,
-                  store.session.input[sessionID] ?? [],
-                  pendingMessages.flatMap((message) =>
-                    message.type === "user" || message.type === "synthetic" ? [message.id] : [],
-                  ),
-                ),
-              )
               setStore("session", "message", sessionID, reconcile(next))
             })
           },

packages/tui/src/routes/session/rows.ts

@@ -73,14 +73,16 @@ export function createSessionRows(sessionID: Accessor<string>) {
     on([sessionID, () => client.connection.status()], ([id, status]) => {
       if (status !== "connected") return
       setRows(reconcile(reduce()))
-      void data.session.pending.sync(id).catch(() => undefined)
-      void data.session.message.sync(id).then(
-        () => {
-          if (sessionID() !== id) return
-          setRows(reconcile(reduce()))
-        },
-        () => undefined,
-      )
+      void data.session.pending
+        .sync(id)
+        .then(() => data.session.message.sync(id))
+        .then(
+          () => {
+            if (sessionID() !== id) return
+            setRows(reconcile(reduce()))
+          },
+          () => undefined,
+        )
     }),
   )
 

packages/tui/test/cli/tui/data.test.tsx

@@ -2468,7 +2468,7 @@ test("preserves admitted prompts when hydration races with promotion", async ()
     ])
     expect(sync.session.input.list(sessionID)).toEqual([messageID])
 
-    const refresh = sync.session.message.sync(sessionID)
+    const refresh = sync.session.pending.sync(sessionID).then(() => sync.session.message.sync(sessionID))
     await requested.promise
     emitEvent(events, {
       id: "evt_prompted_1",
@@ -2504,31 +2504,21 @@ test("preserves admitted prompts when hydration races with promotion", async ()
   }
 })
 
-test("preserves an admitted prompt when hydration races with promotion", async () => {
+test("preserves admissions that arrive while pending work hydrates", async () => {
   const events = createEventStream()
-  const sessionID = "session-hydration-race"
-  const messageID = "msg_initial_user"
-  let resolveMessages!: (response: Response) => void
-  let requestMessages!: () => void
-  const requested = new Promise<void>((resolve) => {
-    requestMessages = resolve
-  })
+  const sessionID = "session-pending-race"
+  const messageID = "msg_late_user"
+  const requested = Promise.withResolvers<void>()
+  const response = Promise.withResolvers<Response>()
   const calls = createFetch((url) => {
-    if (url.pathname !== `/api/session/${sessionID}/message`) return
-    requestMessages()
-    return new Promise<Response>((resolve) => {
-      resolveMessages = resolve
-    })
+    if (url.pathname !== `/api/session/${sessionID}/pending`) return
+    requested.resolve()
+    return response.promise
   }, events)
   let data!: ReturnType<typeof useData>
-  let ready!: () => void
-  const mounted = new Promise<void>((resolve) => {
-    ready = resolve
-  })
 
   function Probe() {
     data = useData()
-    onMount(ready)
     return <box />
   }
 
@@ -2545,129 +2535,32 @@ test("preserves an admitted prompt when hydration races with promotion", async (
   ))
 
   try {
-    await mounted
-    emitEvent(events, {
-      id: "evt_initial_admitted",
-      created: 1,
-      type: "session.input.admitted",
-      durable: durable(sessionID),
-      data: {
-        sessionID,
-        inputID: messageID,
-        input: { type: "user", data: { text: "Initial prompt" }, delivery: "steer" },
-      },
-    })
-    await wait(() => data.session.message.get(sessionID, messageID)?.type === "user")
-
-    const refresh = data.session.message.sync(sessionID)
-    await requested
-    emitEvent(events, {
-      id: "evt_initial_promoted",
-      created: 2,
-      type: "session.input.promoted",
-      durable: durable(sessionID, 1),
-      data: { sessionID, inputID: messageID },
-    })
-    await wait(() => data.session.input.list(sessionID).length === 0)
-
-    resolveMessages(json({ data: [], cursor: {} }))
-    await refresh
-
-    expect(data.session.message.get(sessionID, messageID)).toMatchObject({
-      id: messageID,
-      type: "user",
-      text: "Initial prompt",
-      time: { created: 2 },
-    })
-  } finally {
-    app.renderer.destroy()
-  }
-})
-
-/*
-test("preserves an admitted prompt when hydration races with promotion", async () => {
-  const events = createEventStream()
-  const sessionID = "session-hydration-race"
-  const messageID = "msg_initial_user"
-  let resolveMessages!: (response: Response) => void
-  let requestMessages!: () => void
-  const requested = new Promise<void>((resolve) => {
-    requestMessages = resolve
-  })
-  const calls = createFetch((url) => {
-    if (url.pathname !== `/api/session/${sessionID}/message`) return
-    requestMessages()
-    return new Promise<Response>((resolve) => {
-      resolveMessages = resolve
-    })
-  }, events)
-  let data!: ReturnType<typeof useData>
-  let ready!: () => void
-  const mounted = new Promise<void>((resolve) => {
-    ready = resolve
-  })
-
-  function Probe() {
-    data = useData()
-    onMount(ready)
-    return <box />
-  }
-
-  const app = await testRender(() => (
-    <TestTuiContexts>
-      <SDKProvider client={createClient(calls.fetch)} api={createApi(calls.fetch)}>
-        <ProjectProvider>
-          <DataProvider>
-            <Probe />
-          </DataProvider>
-        </ProjectProvider>
-      </SDKProvider>
-    </TestTuiContexts>
-  ))
-
-  try {
-    await mounted
+    const sync = data.session.pending.sync(sessionID)
+    await requested.promise
     emitEvent(events, {
-      id: "evt_initial_admitted",
+      id: "evt_late_admitted",
       created: 1,
       type: "session.input.admitted",
       durable: durable(sessionID),
       data: {
         sessionID,
         inputID: messageID,
-        input: { type: "user", data: { text: "Initial prompt" }, delivery: "steer" },
+        input: { type: "user", data: { text: "late" }, delivery: "steer" },
       },
     })
-    await wait(() => data.session.message.get(sessionID, messageID)?.type === "user")
-
-    const refresh = data.session.message.refresh(sessionID)
-    await requested
-    emitEvent(events, {
-      id: "evt_initial_promoted",
-      created: 2,
-      type: "session.input.promoted",
-      durable: durable(sessionID, 1),
-      data: { sessionID, inputID: messageID },
-    })
-    await wait(() => data.session.input.list(sessionID).length === 0)
-
-    resolveMessages(json({ data: [], cursor: {} }))
-    await refresh
+    await wait(() => data.session.pending.list(sessionID).length === 1)
+    response.resolve(json({ data: [] }))
+    await sync
 
-    expect(data.session.message.get(sessionID, messageID)).toMatchObject({
-      id: messageID,
-      type: "user",
-      text: "Initial prompt",
-      time: { created: 2 },
-    })
+    expect(data.session.pending.list(sessionID).map((item) => item.id)).toEqual([messageID])
+    expect(data.session.input.list(sessionID)).toEqual([messageID])
+    expect(data.session.message.get(sessionID, messageID)).toMatchObject({ text: "late" })
   } finally {
     app.renderer.destroy()
   }
 })
 
 test("skips initial instruction state and projects later updates with their message ID", async () => {
-*/
-test("projects live instruction updates with their message ID", async () => {
   const events = createEventStream()
   const calls = createFetch(undefined, events)
   let sync!: ReturnType<typeof useData>