fix(tui): remove committed undo messages

packages/core/src/session/projector.ts

@@ -1,6 +1,6 @@
 export * as SessionProjector from "./projector"
 
-import { and, asc, desc, eq, gt, inArray, lt, or, sql } from "drizzle-orm"
+import { and, asc, desc, eq, gt, gte, inArray, lt, or, sql } from "drizzle-orm"
 import { DateTime, Effect, Layer, Schema } from "effect"
 import { Database } from "../database/database"
 import { EventV2 } from "../event"
@@ -698,7 +698,7 @@ const layer = Layer.effectDiscard(
         yield* db
           .delete(SessionMessageTable)
           .where(
-            and(eq(SessionMessageTable.session_id, event.data.sessionID), gt(SessionMessageTable.seq, boundary.seq)),
+            and(eq(SessionMessageTable.session_id, event.data.sessionID), gte(SessionMessageTable.seq, boundary.seq)),
           )
           .run()
           .pipe(Effect.orDie)
@@ -707,7 +707,7 @@ const layer = Layer.effectDiscard(
           .where(
             and(
               eq(SessionInputTable.session_id, event.data.sessionID),
-              or(gt(SessionInputTable.admitted_seq, boundary.seq), gt(SessionInputTable.promoted_seq, boundary.seq)),
+              or(gte(SessionInputTable.admitted_seq, boundary.seq), gte(SessionInputTable.promoted_seq, boundary.seq)),
             ),
           )
           .run()

packages/core/test/session-projector.test.ts

@@ -70,9 +70,14 @@ describe("SessionProjector", () => {
         })
         .run()
       const boundary = SessionMessage.ID.make("msg_boundary")
+      const earlier = SessionMessage.ID.make("msg_earlier")
       yield* db
         .insert(SessionMessageTable)
-        .values([assistantRow(boundary, 1), assistantRow(SessionMessage.ID.make("msg_later"), 2)])
+        .values([
+          assistantRow(earlier, 0),
+          assistantRow(boundary, 1),
+          assistantRow(SessionMessage.ID.make("msg_later"), 2),
+        ])
         .run()
       yield* db
         .insert(SessionContextCheckpointTable)
@@ -100,7 +105,7 @@ describe("SessionProjector", () => {
       })
       expect(
         (yield* db.select({ id: SessionMessageTable.id }).from(SessionMessageTable).all()).map((row) => row.id),
-      ).toEqual([boundary])
+      ).toEqual([earlier])
       // A committed revert resets the context checkpoint so the next turn re-initializes.
       expect(yield* db.select().from(SessionContextCheckpointTable).get().pipe(Effect.orDie)).toBeUndefined()
     }),

packages/core/test/session-prompt.test.ts

@@ -212,7 +212,8 @@ describe("SessionV2.prompt", () => {
         (yield* db.select({ id: SessionMessageTable.id }).from(SessionMessageTable).all().pipe(Effect.orDie)).map(
           (row) => row.id,
         ),
-      ).not.toContain(stale)
+      ).not.toContainAnyValues([boundary.id, stale])
+      expect(yield* SessionInput.find(db, boundary.id)).toBeUndefined()
     }),
   )
 

packages/tui/src/context/data.tsx

@@ -553,9 +553,17 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
             setStore("session", "info", event.data.sessionID, "revert", event.data.revert)
           break
         case "session.revert.cleared":
+          if (store.session.info[event.data.sessionID])
+            setStore("session", "info", event.data.sessionID, "revert", undefined)
+          break
         case "session.revert.committed":
           if (store.session.info[event.data.sessionID])
             setStore("session", "info", event.data.sessionID, "revert", undefined)
+          message.update(event.data.sessionID, (draft, index) => {
+            const position = draft.findIndex((item) => item.id >= event.data.messageID)
+            if (position === -1) return
+            for (const item of draft.splice(position)) index.delete(item.id)
+          })
           break
         case "session.compaction.delta":
           break

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

@@ -284,6 +284,60 @@ test("completes exploration when a queued prompt is promoted", async () => {
   }
 })
 
+test("removes committed revert messages from local state", async () => {
+  const events = createEventStream()
+  const sessionID = "session-revert"
+  const calls = createFetch((url) => {
+    if (url.pathname === `/api/session/${sessionID}/message`) return json({ data: [], cursor: {} })
+  }, events)
+  let data!: ReturnType<typeof useData>
+
+  function Probe() {
+    data = useData()
+    return <box />
+  }
+
+  const app = await testRender(() => (
+    <TestTuiContexts>
+      <SDKProvider client={createClient(calls.fetch)} api={createApi(calls.fetch)}>
+        <ProjectProvider>
+          <DataProvider>
+            <Probe />
+          </DataProvider>
+        </ProjectProvider>
+      </SDKProvider>
+    </TestTuiContexts>
+  ))
+
+  try {
+    for (const [seq, inputID] of ["msg_001", "msg_002", "msg_003"].entries()) {
+      emitEvent(events, {
+        id: EventV2.ID.create(),
+        created: seq,
+        type: "session.prompt.admitted",
+        durable: durable(sessionID, seq),
+        data: { sessionID, inputID, prompt: { text: inputID }, delivery: "steer" },
+      })
+    }
+    await wait(() => data.session.message.ids(sessionID).length === 3)
+
+    emitEvent(events, {
+      id: EventV2.ID.create(),
+      created: 3,
+      type: "session.revert.committed",
+      durable: durable(sessionID, 3),
+      data: { sessionID, messageID: "msg_002" },
+    })
+
+    await wait(() => data.session.message.ids(sessionID).length === 1)
+    expect(data.session.message.ids(sessionID)).toEqual(["msg_001"])
+    expect(data.session.message.get(sessionID, "msg_002")).toBeUndefined()
+    expect(data.session.message.get(sessionID, "msg_003")).toBeUndefined()
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
 test("connectedOnce is false until first connect and persists across disconnect", async () => {
   const encoder = new TextEncoder()
   let stream: ReadableStreamDefaultController<Uint8Array> | undefined