Просмотр исходного кода

fix(tui): preserve admitted messages during hydration

Kit Langton 1 месяц назад
Родитель
Сommit
00ad2bed67

+ 96 - 57
packages/tui/src/context/data.tsx

@@ -1,7 +1,7 @@
 // Client data layer: apply server events and cache API reads into a Solid store.
-// Prefer straightforward projection. Do not add generation counters, stale-response
-// merges, live/history overlays, or other race machinery here—last write wins.
-// Reconnect invalidates cached reads; active UI owners decide what to sync again.
+// Prefer straightforward projection. API reads replace cached state, except admitted
+// inputs survive history refresh until the server projects them. Reconnect invalidates
+// cached reads; active UI owners decide what to sync again.
 
 import type {
   AgentInfo,
@@ -71,7 +71,6 @@ type Store = {
     active: Record<string, DataSessionStatus>
     message: Record<string, SessionMessageInfo[]>
     pending: Record<string, SessionPendingInfo[]>
-    input: Record<string, string[]>
     permission: Record<string, PermissionV2Request[]>
     // Pending forms keyed by owner: a session ID or the temporary "global" elicitation sentinel.
     form: Record<string, FormWithLocation[]>
@@ -82,6 +81,11 @@ type Store = {
   location: Record<string, LocationData>
 }
 
+type PendingOperation =
+  | { type: "admitted"; item: SessionPendingInfo }
+  | { type: "promoted"; inputID: string }
+  | { type: "reverted"; to: string }
+
 function locationKey(location: LocationRef) {
   return JSON.stringify([location.directory, location.workspaceID])
 }
@@ -131,7 +135,6 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
         active: {},
         message: {},
         pending: {},
-        input: {},
         permission: {},
         form: {},
       },
@@ -147,18 +150,21 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
     })
     const messageIndex = new Map<string, Map<string, number>>()
     const sync = createSync()
+    const pendingOperations = new Map<string, PendingOperation[]>()
 
     function setSessionActive(sessionID: string, status: DataSessionStatus) {
       setStore("session", "active", sessionID, status)
     }
 
     function addPending(item: SessionPendingInfo) {
+      pendingOperations.get(item.sessionID)?.push({ type: "admitted", item })
       if (store.session.pending[item.sessionID]?.some((pending) => pending.id === item.id)) return
       setStore("session", "pending", item.sessionID, [...(store.session.pending[item.sessionID] ?? []), item])
     }
 
     function removePending(sessionID: string, inputID?: string) {
       if (!inputID) return
+      pendingOperations.get(sessionID)?.push({ type: "promoted", inputID })
       setStore(
         "session",
         "pending",
@@ -167,6 +173,36 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
       )
     }
 
+    function pendingInputs(sessionID: string) {
+      return (store.session.pending[sessionID] ?? []).filter((item) => item.type !== "compaction")
+    }
+
+    function syncPending(sessionID: string) {
+      return sync.run(`session.pending:${sessionID}`, async () => {
+        const operations = pendingOperations.get(sessionID) ?? []
+        pendingOperations.set(sessionID, operations)
+        try {
+          const pending = new Map((await client.api.session.pending.list({ sessionID })).map((item) => [item.id, item]))
+          operations.forEach((operation) => {
+            if (operation.type === "admitted") {
+              pending.set(operation.item.id, operation.item)
+              return
+            }
+            if (operation.type === "promoted") {
+              pending.delete(operation.inputID)
+              return
+            }
+            pending.forEach((_, id) => {
+              if (id >= operation.to) pending.delete(id)
+            })
+          })
+          setStore("session", "pending", sessionID, reconcile([...pending.values()]))
+        } finally {
+          if (pendingOperations.get(sessionID) === operations) pendingOperations.delete(sessionID)
+        }
+      })
+    }
+
     const message = {
       update(sessionID: string, fn: (messages: SessionMessageInfo[], index: Map<string, number>) => void) {
         setStore(
@@ -199,6 +235,31 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
         const item = messages.findLast((item) => item.type === "compaction" && item.status === "running")
         return item?.type === "compaction" ? item : undefined
       },
+      fromPending(item: SessionPendingInfo): SessionMessageInfo {
+        if (item.type === "user")
+          return {
+            id: item.id,
+            type: "user",
+            ...item.data,
+            time: { created: item.timeCreated },
+          }
+        if (item.type === "synthetic")
+          return {
+            id: item.id,
+            type: "synthetic",
+            ...item.data,
+            time: { created: item.timeCreated },
+          }
+        return {
+          id: item.id,
+          type: "compaction",
+          status: "running",
+          reason: "manual",
+          summary: "",
+          recent: "",
+          time: { created: item.timeCreated },
+        }
+      },
       latestTool(assistant: SessionMessageAssistant | undefined, callID?: string) {
         return assistant?.content.findLast(
           (item): item is SessionMessageAssistantTool =>
@@ -269,6 +330,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
 
     function removeSession(sessionID: string) {
       messageIndex.delete(sessionID)
+      pendingOperations.delete(sessionID)
       sync.invalidate(`session:${sessionID}`)
       sync.invalidate(`session.pending:${sessionID}`)
       sync.invalidate(`session.message:${sessionID}`)
@@ -281,7 +343,6 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           delete draft.active[sessionID]
           delete draft.message[sessionID]
           delete draft.pending[sessionID]
-          delete draft.input[sessionID]
           delete draft.permission[sessionID]
           delete draft.form[sessionID]
           for (const [rootID, family] of Object.entries(draft.family)) {
@@ -374,59 +435,35 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           }
           break
         case "session.input.promoted": {
+          const pending = store.session.pending[event.data.sessionID]?.some((item) => item.id === event.data.inputID)
           removePending(event.data.sessionID, event.data.inputID)
           message.update(event.data.sessionID, (draft, index) => {
             const position = index.get(event.data.inputID)
             if (position === undefined) return
             const existing = draft[position]
-            if (!existing || !store.session.input[event.data.sessionID]?.includes(event.data.inputID)) return
+            if (!existing || !pending) return
             existing.time.created = event.created
             draft.splice(position, 1)
             draft.push(existing)
             index.clear()
             draft.forEach((message, indexValue) => index.set(message.id, indexValue))
           })
-          setStore(
-            "session",
-            "input",
-            event.data.sessionID,
-            (store.session.input[event.data.sessionID] ?? []).filter((id) => id !== event.data.inputID),
-          )
           break
         }
-        case "session.input.admitted":
-          addPending({
+        case "session.input.admitted": {
+          const pending: SessionPendingInfo = {
             id: event.data.inputID,
             sessionID: event.data.sessionID,
             admittedSeq: event.durable.seq,
             timeCreated: event.created,
             ...event.data.input,
-          })
-          if (!store.session.input[event.data.sessionID]?.includes(event.data.inputID))
-            setStore("session", "input", event.data.sessionID, [
-              ...(store.session.input[event.data.sessionID] ?? []),
-              event.data.inputID,
-            ])
+          }
+          addPending(pending)
           message.update(event.data.sessionID, (draft, index) => {
-            message.append(
-              draft,
-              index,
-              event.data.input.type === "user"
-                ? {
-                    id: event.data.inputID,
-                    type: "user",
-                    ...event.data.input.data,
-                    time: { created: event.created },
-                  }
-                : {
-                    id: event.data.inputID,
-                    type: "synthetic",
-                    ...event.data.input.data,
-                    time: { created: event.created },
-                  },
-            )
+            message.append(draft, index, message.fromPending(pending))
           })
           break
+        }
         case "session.instructions.updated":
           const instructions = event.metadata?.instructions
           if (
@@ -736,11 +773,12 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           if (store.session.info[event.data.sessionID]) {
             setStore("session", "info", event.data.sessionID, "revert", undefined)
           }
+          pendingOperations.get(event.data.sessionID)?.push({ type: "reverted", to: event.data.to })
           setStore(
             "session",
-            "input",
+            "pending",
             event.data.sessionID,
-            (store.session.input[event.data.sessionID] ?? []).filter((id) => id < event.data.to),
+            (store.session.pending[event.data.sessionID] ?? []).filter((item) => item.id < event.data.to),
           )
           message.update(event.data.sessionID, (draft, index) => {
             const position = draft.findIndex((item) => item.id >= event.data.to)
@@ -914,10 +952,10 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
         },
         input: {
           list(sessionID: string) {
-            return store.session.input[sessionID] ?? []
+            return pendingInputs(sessionID).map((item) => item.id)
           },
           has(sessionID: string, inputID: string) {
-            return store.session.input[sessionID]?.includes(inputID) ?? false
+            return pendingInputs(sessionID).some((item) => item.id === inputID)
           },
         },
         pending: {
@@ -925,16 +963,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
             return store.session.pending[sessionID] ?? []
           },
           sync(sessionID: string) {
-            return sync.run(`session.pending:${sessionID}`, async () => {
-              const pending = await client.api.session.pending.list({ sessionID })
-              setStore("session", "pending", sessionID, reconcile(pending))
-              setStore(
-                "session",
-                "input",
-                sessionID,
-                reconcile(pending.filter((item) => item.type !== "compaction").map((item) => item.id)),
-              )
-            })
+            return syncPending(sessionID)
           },
           invalidate(sessionID: string) {
             sync.invalidate(`session.pending:${sessionID}`)
@@ -960,11 +989,21 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
           },
           sync(sessionID: string) {
             return sync.run(`session.message:${sessionID}`, async () => {
-              const messages = (
-                await client.api.message.list({ sessionID, limit: 200, order: "desc" })
-              ).data.toReversed()
-              messageIndex.set(sessionID, new Map(messages.map((message, index) => [message.id, index])))
-              setStore("session", "message", sessionID, reconcile(messages))
+              await syncPending(sessionID)
+              const localInputs = pendingInputs(sessionID).map((item) => item.id)
+              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, ...pendingInputs(sessionID).map((item) => item.id)])
+              localInputIDs.forEach((messageID) => {
+                const position = messageIndex.get(sessionID)?.get(messageID)
+                const item = position === undefined ? undefined : store.session.message[sessionID]?.[position]
+                if (item) message.append(next, index, item)
+              })
+              pendingMessages.forEach((item) => message.append(next, index, item))
+              messageIndex.set(sessionID, index)
+              setStore("session", "message", sessionID, reconcile(next))
             })
           },
           invalidate(sessionID: string) {

+ 1 - 8
packages/tui/src/routes/session/rows.ts

@@ -73,14 +73,7 @@ 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.message.sync(id).catch(() => undefined)
     }),
   )
 

+ 128 - 12
packages/tui/test/cli/tui/data.test.tsx

@@ -2377,16 +2377,40 @@ test("settles pending tools when a live failure arrives", async () => {
   }
 })
 
-test("renders admitted prompts immediately and tracks them until promoted", async () => {
+test("preserves admitted prompts when hydration races with promotion", async () => {
   const events = createEventStream()
   const sessionID = "session-1"
   const messageID = "msg_user_1"
+  const queuedID = "msg_user_2"
+  const requested = Promise.withResolvers<void>()
+  const response = Promise.withResolvers<Response>()
   const calls = createFetch((url) => {
-    if (url.pathname === `/api/session/${sessionID}/message`)
+    if (url.pathname === `/api/session/${sessionID}/pending`)
       return json({
-        data: [{ id: messageID, type: "user", text: "hello", time: { created: 0 } }],
-        cursor: {},
+        data: [
+          {
+            admittedSeq: 0,
+            id: messageID,
+            sessionID,
+            timeCreated: 0,
+            type: "user",
+            data: { text: "hello" },
+            delivery: "steer",
+          },
+          {
+            admittedSeq: 1,
+            id: queuedID,
+            sessionID,
+            timeCreated: 1,
+            type: "user",
+            data: { text: "queued" },
+            delivery: "queue",
+          },
+        ],
       })
+    if (url.pathname !== `/api/session/${sessionID}/message`) return
+    requested.resolve()
+    return response.promise
   }, events)
   let sync!: ReturnType<typeof useData>
   let ready!: () => void
@@ -2444,12 +2468,11 @@ test("renders admitted prompts immediately and tracks them until promoted", asyn
     ])
     expect(sync.session.input.list(sessionID)).toEqual([messageID])
 
-    await sync.session.message.sync(sessionID)
-    expect(sync.session.message.list(sessionID)?.[0]?.metadata).toBeUndefined()
-
+    const refresh = sync.session.message.sync(sessionID)
+    await requested.promise
     emitEvent(events, {
       id: "evt_prompted_1",
-      created: 0,
+      created: 1,
       type: "session.input.promoted",
       durable: durable(sessionID, 1),
       data: {
@@ -2461,14 +2484,17 @@ test("renders admitted prompts immediately and tracks them until promoted", asyn
     await wait(() => received.at(-1) === "session.input.promoted")
     expect(received.slice(-2)).toEqual(["session.input.admitted", "session.input.promoted"])
     unsubscribe()
-    const message = sync.session.message.list(sessionID)?.[0]
+    response.resolve(json({ data: [], cursor: {} }))
+    await refresh
+
+    const message = sync.session.message.get(sessionID, messageID)
     expect(message?.type).toBe("user")
     if (message?.type !== "user") return
     expect(message).toMatchObject({ id: messageID, text: "hello" })
     expect(message.metadata).toBeUndefined()
-    expect(sync.session.pending.list(sessionID)).toEqual([])
-    expect(sync.session.input.list(sessionID)).toEqual([])
-    expect(sync.session.message.list(sessionID).map((message) => message.id)).toEqual([messageID])
+    expect(sync.session.input.list(sessionID)).toEqual([queuedID])
+    expect(sync.session.message.list(sessionID).map((message) => message.id)).toEqual([messageID, queuedID])
+    expect(sync.session.message.get(sessionID, queuedID)).toMatchObject({ id: queuedID, text: "queued" })
     expect(sync.session.message.list("missing")).toEqual([])
     expect(sync.session.message.get(sessionID, messageID)).toBe(message)
     expect(sync.session.message.get(sessionID, "missing")).toBeUndefined()
@@ -2478,6 +2504,96 @@ test("renders admitted prompts immediately and tracks them until promoted", asyn
   }
 })
 
+test("reconciles admissions and promotions that arrive while pending work hydrates", async () => {
+  const events = createEventStream()
+  const sessionID = "session-pending-race"
+  const messageID = "msg_late_user"
+  const promotedID = "msg_promoted_user"
+  const requested = Promise.withResolvers<void>()
+  const response = Promise.withResolvers<Response>()
+  const calls = createFetch((url) => {
+    if (url.pathname !== `/api/session/${sessionID}/pending`) return
+    requested.resolve()
+    return response.promise
+  }, events)
+  let data!: ReturnType<typeof useData>
+
+  function Probe() {
+    data = useData()
+    return <box />
+  }
+
+  const app = await testRender(() => (
+    <TestTuiContexts>
+      <ClientProvider api={createApi(calls.fetch)}>
+        <ProjectProvider>
+          <DataProvider>
+            <Probe />
+          </DataProvider>
+        </ProjectProvider>
+      </ClientProvider>
+    </TestTuiContexts>
+  ))
+
+  try {
+    const sync = data.session.pending.sync(sessionID)
+    await requested.promise
+    emitEvent(events, {
+      id: "evt_late_admitted",
+      created: 1,
+      type: "session.input.admitted",
+      durable: durable(sessionID),
+      data: {
+        sessionID,
+        inputID: messageID,
+        input: { type: "user", data: { text: "late" }, delivery: "steer" },
+      },
+    })
+    emitEvent(events, {
+      id: "evt_promoted_admitted",
+      created: 2,
+      type: "session.input.admitted",
+      durable: durable(sessionID, 1),
+      data: {
+        sessionID,
+        inputID: promotedID,
+        input: { type: "user", data: { text: "promoted" }, delivery: "steer" },
+      },
+    })
+    emitEvent(events, {
+      id: "evt_promoted",
+      created: 3,
+      type: "session.input.promoted",
+      durable: durable(sessionID, 2),
+      data: { sessionID, inputID: promotedID },
+    })
+    await wait(() => data.session.pending.list(sessionID).length === 1)
+    response.resolve(
+      json({
+        data: [
+          {
+            admittedSeq: 1,
+            id: promotedID,
+            sessionID,
+            timeCreated: 2,
+            type: "user",
+            data: { text: "promoted" },
+            delivery: "steer",
+          },
+        ],
+      }),
+    )
+    await sync
+
+    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" })
+    expect(data.session.message.get(sessionID, promotedID)).toMatchObject({ text: "promoted" })
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
 test("skips initial instruction state and projects later updates with their message ID", async () => {
   const events = createEventStream()
   const calls = createFetch(undefined, events)

+ 1 - 0
packages/tui/test/fixture/tui-client.ts

@@ -109,6 +109,7 @@ export function createFetch(override?: FetchHandler, events?: ReturnType<typeof
       })
     if (url.pathname === "/api/session") return json({ data: [], cursor: {} })
     if (url.pathname === "/api/session/active") return json({ data: {} })
+    if (/^\/api\/session\/[^/]+\/pending$/.test(url.pathname)) return json({ data: [] })
     if (url.pathname === "/api/permission/request")
       return json({ location: { directory, project: { id: "proj_test", directory: worktree } }, data: [] })
     if (url.pathname === "/api/form/request")