From 06dcf3f221466b239ca272ee6ab342941ace5fad Mon Sep 17 00:00:00 2001 From: Dax Raad Date: Mon, 6 Jul 2026 13:08:25 -0400 Subject: [PATCH] fix(tui): remove committed undo messages --- packages/core/src/session/projector.ts | 6 +-- packages/core/test/session-projector.test.ts | 9 +++- packages/core/test/session-prompt.test.ts | 3 +- packages/tui/src/context/data.tsx | 8 +++ packages/tui/test/cli/tui/data.test.tsx | 54 ++++++++++++++++++++ 5 files changed, 74 insertions(+), 6 deletions(-) diff --git a/packages/core/src/session/projector.ts b/packages/core/src/session/projector.ts index f265aa6fc9..3cbd40c296 100644 --- a/packages/core/src/session/projector.ts +++ b/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() diff --git a/packages/core/test/session-projector.test.ts b/packages/core/test/session-projector.test.ts index 1d77303d43..bb53c13405 100644 --- a/packages/core/test/session-projector.test.ts +++ b/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() }), diff --git a/packages/core/test/session-prompt.test.ts b/packages/core/test/session-prompt.test.ts index e7d3f5189f..9ef5b5d822 100644 --- a/packages/core/test/session-prompt.test.ts +++ b/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() }), ) diff --git a/packages/tui/src/context/data.tsx b/packages/tui/src/context/data.tsx index aa27572ac4..ae6626be39 100644 --- a/packages/tui/src/context/data.tsx +++ b/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 diff --git a/packages/tui/test/cli/tui/data.test.tsx b/packages/tui/test/cli/tui/data.test.tsx index 1820e58f29..0e43377657 100644 --- a/packages/tui/test/cli/tui/data.test.tsx +++ b/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 + + function Probe() { + data = useData() + return + } + + const app = await testRender(() => ( + + + + + + + + + + )) + + 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 | undefined