refactor(session): simplify pending withdrawal

packages/app/src/context/server-session-v2-reducer.test.ts

@@ -153,4 +153,42 @@ describe("v2 session reducer", () => {
 
     expect(result).toMatchObject({ sessionID: "ses_1", missing: "msg_user", touched: [] })
   })
+
+  test("forgets withdrawn input before a later promotion", () => {
+    const reducer = createV2SessionReducer()
+    reducer.reduce(
+      [],
+      event({
+        ...base,
+        id: "evt_admitted",
+        type: "session.input.admitted",
+        data: {
+          sessionID: "ses_1",
+          inputID: "msg_user",
+          input: { type: "user", delivery: "steer", data: { text: "hello" } },
+        },
+      }),
+    )
+    reducer.reduce(
+      [],
+      event({
+        ...base,
+        id: "evt_withdrawn",
+        type: "session.input.withdrawn",
+        data: { sessionID: "ses_1", inputID: "msg_user" },
+      }),
+    )
+
+    const result = reducer.reduce(
+      [],
+      event({
+        ...base,
+        id: "evt_promoted",
+        type: "session.input.promoted",
+        data: { sessionID: "ses_1", inputID: "msg_user" },
+      }),
+    )
+
+    expect(result).toMatchObject({ sessionID: "ses_1", missing: "msg_user", touched: [] })
+  })
 })

packages/app/src/context/server-session-v2-reducer.ts

@@ -29,6 +29,9 @@ export function createV2SessionReducer() {
       case "session.input.admitted":
         pending.set(key(sessionID, event.data.inputID), event.data.input)
         return result([...source])
+      case "session.input.withdrawn":
+        pending.delete(key(sessionID, event.data.inputID))
+        return
       case "session.input.promoted": {
         const input = pending.get(key(sessionID, event.data.inputID))
         pending.delete(key(sessionID, event.data.inputID))

packages/client/src/effect/api/api.ts

@@ -370,6 +370,15 @@ export type Endpoint5_26Output =
             readonly input: SessionPending.Message
           }
         }
+      | {
+          readonly id: Event.ID
+          readonly created: DateTime.Utc
+          readonly metadata?: { readonly [x: string]: unknown } | undefined
+          readonly type: "session.input.withdrawn"
+          readonly durable: { readonly aggregateID: string; readonly seq: Event.Seq; readonly version: Event.Version }
+          readonly location?: Location.Ref | undefined
+          readonly data: { readonly sessionID: Session.ID; readonly inputID: SessionMessage.ID }
+        }
       | {
           readonly id: Event.ID
           readonly created: DateTime.Utc
@@ -841,6 +850,12 @@ export type Endpoint5_29Input = { readonly sessionID: Session.ID; readonly messa
 export type Endpoint5_29Output = SessionMessage.Info
 export type SessionMessageOperation<E = never> = (input: Endpoint5_29Input) => Effect.Effect<Endpoint5_29Output, E>
 
+export type Endpoint5_30Input = { readonly sessionID: Session.ID; readonly inputID: SessionMessage.ID }
+export type Endpoint5_30Output = boolean
+export type SessionPendingWithdrawOperation<E = never> = (
+  input: Endpoint5_30Input,
+) => Effect.Effect<Endpoint5_30Output, E>
+
 export interface SessionApi<E = never> {
   readonly list: SessionListOperation<E>
   readonly create: SessionCreateOperation<E>
@@ -865,7 +880,10 @@ export interface SessionApi<E = never> {
     readonly commit: SessionRevertCommitOperation<E>
   }
   readonly context: SessionContextOperation<E>
-  readonly pending: { readonly list: SessionPendingListOperation<E> }
+  readonly pending: {
+    readonly list: SessionPendingListOperation<E>
+    readonly withdraw: SessionPendingWithdrawOperation<E>
+  }
   readonly instructions: {
     readonly entry: {
       readonly list: SessionInstructionsEntryListOperation<E>

packages/client/src/effect/generated/client.ts

@@ -76,6 +76,8 @@ import type {
   Endpoint5_28Output,
   Endpoint5_29Input,
   Endpoint5_29Output,
+  Endpoint5_30Input,
+  Endpoint5_30Output,
   Endpoint6_0Input,
   Endpoint6_0Output,
   Endpoint7_0Input,
@@ -550,6 +552,14 @@ const Endpoint5_29 = (raw: RawClient["server.session"]) => (input: Endpoint5_29I
     ),
   )
 
+const Endpoint5_30 = (raw: RawClient["server.session"]) => (input: Endpoint5_30Input) =>
+  preserveEffect<Endpoint5_30Output>()(
+    raw["session.pending.withdraw"]({ params: { sessionID: input["sessionID"], inputID: input["inputID"] } }).pipe(
+      Effect.mapError(mapClientError),
+      Effect.map((value) => value.data),
+    ),
+  )
+
 const adaptGroup5 = (raw: RawClient["server.session"]) => ({
   list: Endpoint5_0(raw),
   create: Endpoint5_1(raw),
@@ -570,7 +580,7 @@ const adaptGroup5 = (raw: RawClient["server.session"]) => ({
   wait: Endpoint5_16(raw),
   revert: { stage: Endpoint5_17(raw), clear: Endpoint5_18(raw), commit: Endpoint5_19(raw) },
   context: Endpoint5_20(raw),
-  pending: { list: Endpoint5_21(raw) },
+  pending: { list: Endpoint5_21(raw), withdraw: Endpoint5_30(raw) },
   instructions: { entry: { list: Endpoint5_22(raw), put: Endpoint5_23(raw), remove: Endpoint5_24(raw) } },
   generate: Endpoint5_25(raw),
   log: Endpoint5_26(raw),

packages/client/src/promise/generated/client.ts

@@ -70,6 +70,8 @@ import type {
   SessionBackgroundOutput,
   SessionMessageInput,
   SessionMessageOutput,
+  SessionPendingWithdrawInput,
+  SessionPendingWithdrawOutput,
   MessageListInput,
   MessageListOutput,
   ModelListInput,
@@ -732,6 +734,17 @@ export function make(options: ClientOptions) {
             },
             requestOptions,
           ).then((value) => value.data),
+        withdraw: (input: SessionPendingWithdrawInput, requestOptions?: RequestOptions) =>
+          request<{ readonly data: SessionPendingWithdrawOutput }>(
+            {
+              method: "POST",
+              path: `/api/session/${encodeURIComponent(input.sessionID)}/pending/${encodeURIComponent(input.inputID)}/withdraw`,
+              successStatus: 200,
+              declaredStatuses: [404, 400, 401],
+              empty: false,
+            },
+            requestOptions,
+          ).then((value) => value.data),
       },
       instructions: {
         entry: {

packages/client/src/promise/generated/types.ts

@@ -637,6 +637,16 @@ export type SessionInputPromoted = {
   data: { sessionID: string; inputID: string }
 }
 
+export type SessionInputWithdrawn = {
+  id: string
+  created: number
+  metadata?: { [x: string]: any }
+  type: "session.input.withdrawn"
+  durable: { aggregateID: string; seq: number; version: 1 }
+  location?: LocationRef
+  data: { sessionID: string; inputID: string }
+}
+
 export type SessionExecutionStarted = {
   id: string
   created: number
@@ -2157,6 +2167,7 @@ export type SessionEventDurable =
   | SessionForked
   | SessionInputPromoted
   | SessionInputAdmitted
+  | SessionInputWithdrawn
   | SessionExecutionStarted
   | SessionExecutionSucceeded
   | SessionExecutionFailed
@@ -2252,6 +2263,7 @@ export type V2Event =
   | SessionForked
   | SessionInputPromoted
   | SessionInputAdmitted
+  | SessionInputWithdrawn
   | SessionExecutionStarted
   | SessionExecutionSucceeded
   | SessionExecutionFailed
@@ -3207,6 +3219,13 @@ export type SessionMessageInput = {
 
 export type SessionMessageOutput = { data: SessionMessageInfo }["data"]
 
+export type SessionPendingWithdrawInput = {
+  readonly sessionID: { readonly sessionID: string; readonly inputID: string }["sessionID"]
+  readonly inputID: { readonly sessionID: string; readonly inputID: string }["inputID"]
+}
+
+export type SessionPendingWithdrawOutput = { data: boolean }["data"]
+
 export type MessageListInput = {
   readonly sessionID: { readonly sessionID: string }["sessionID"]
   readonly limit?: {

packages/client/test/promise.test.ts

@@ -300,7 +300,7 @@ test("session instructions methods use the public HTTP contract", async () => {
   ])
 })
 
-test("session.pending.list uses the public HTTP contract", async () => {
+test("session.pending uses the public HTTP contract", async () => {
   const requests: Array<{ method: string; url: string }> = []
   const pending = [
     {
@@ -317,14 +317,22 @@ test("session.pending.list uses the public HTTP contract", async () => {
     fetch: async (input, init) => {
       const request = input instanceof Request ? input : new Request(input, init)
       requests.push({ method: request.method, url: request.url })
-      return Response.json({ data: pending })
+      return Response.json({ data: request.method === "GET" ? pending : true })
     },
   })
 
   const result = await client.session.pending.list({ sessionID: "ses_test" })
+  const withdrawn = await client.session.pending.withdraw({ sessionID: "ses_test", inputID: "msg_pending" })
 
   expect(result).toEqual(pending)
-  expect(requests).toEqual([{ method: "GET", url: "http://localhost:3000/api/session/ses_test/pending" }])
+  expect(withdrawn).toBe(true)
+  expect(requests).toEqual([
+    { method: "GET", url: "http://localhost:3000/api/session/ses_test/pending" },
+    {
+      method: "POST",
+      url: "http://localhost:3000/api/session/ses_test/pending/msg_pending/withdraw",
+    },
+  ])
 })
 
 test("event.subscribe exposes the Promise event stream wire projection", async () => {

packages/core/src/session.ts

@@ -209,6 +209,10 @@ export interface Interface {
    * unhandled compaction barriers.
    */
   readonly pending: (sessionID: SessionSchema.ID) => Effect.Effect<SessionPending.Info[], NotFoundError>
+  readonly withdraw: (input: {
+    sessionID: SessionSchema.ID
+    inputID: SessionMessage.ID
+  }) => Effect.Effect<boolean, NotFoundError>
   /**
    * Durable, ordered session log read. Replays durable session bus after
    * the exclusive `after` cursor, emits a `Synced` marker at the captured
@@ -540,6 +544,10 @@ const layer = Layer.effect(
         yield* result.get(sessionID)
         return yield* SessionPending.list(db, sessionID)
       }),
+      withdraw: Effect.fn("Session.withdraw")(function* (input) {
+        yield* result.get(input.sessionID)
+        return yield* SessionPending.withdraw(db, bus, input)
+      }),
       log: (input) =>
         Stream.unwrap(
           result

packages/core/src/session/message-updater.ts

@@ -175,6 +175,7 @@ export function update(adapter: Adapter, event: SessionEvent.Event) {
       "session.forked": () => Effect.void,
       "session.input.promoted": () => Effect.void,
       "session.input.admitted": () => Effect.void,
+      "session.input.withdrawn": () => Effect.void,
       "session.execution.started": () => Effect.void,
       "session.execution.succeeded": () => clearCurrentRetry,
       "session.execution.failed": () => clearCurrentRetry,

packages/core/src/session/pending.ts

@@ -1,6 +1,6 @@
 export * as SessionPending from "./pending"
 
-import { and, asc, eq, or } from "drizzle-orm"
+import { and, asc, eq, or, sql } from "drizzle-orm"
 import { DateTime, Effect, Schema } from "effect"
 import {
   Compaction,
@@ -38,10 +38,15 @@ const encodeUser = Schema.encodeSync(UserData)
 const decodeSynthetic = Schema.decodeUnknownSync(SyntheticData)
 const encodeSynthetic = Schema.encodeSync(SyntheticData)
 const decodeAdmittedEvent = Schema.decodeUnknownOption(SessionEvent.InputAdmitted.data)
+const decodeWithdrawnEvent = Schema.decodeUnknownOption(SessionEvent.InputWithdrawn.data)
 const admittedEventType = Bus.versionedType(
   SessionEvent.InputAdmitted.type,
   SessionEvent.InputAdmitted.durable.version,
 )
+const withdrawnEventType = Bus.versionedType(
+  SessionEvent.InputWithdrawn.type,
+  SessionEvent.InputWithdrawn.durable.version,
+)
 const inboxLocks = KeyedMutex.makeUnsafe<SessionSchema.ID>()
 
 export class LifecycleConflict extends Schema.TaggedErrorClass<LifecycleConflict>()(
@@ -103,6 +108,38 @@ export const compaction = Effect.fn("SessionPending.compaction")(function* (
   return entry.type === "compaction" ? entry : undefined
 })
 
+const admittedFromHistory = Effect.fn("SessionPending.admittedFromHistory")(function* (
+  db: DatabaseService,
+  sessionID: SessionSchema.ID,
+  id: SessionMessage.ID,
+) {
+  const row = yield* db
+    .select()
+    .from(EventTable)
+    .where(
+      and(
+        eq(EventTable.aggregate_id, sessionID),
+        eq(EventTable.type, admittedEventType),
+        sql`json_extract(${EventTable.data}, '$.inputID') = ${id}`,
+      ),
+    )
+    .limit(1)
+    .get()
+    .pipe(Effect.orDie)
+  if (row) {
+    const decoded = decodeAdmittedEvent(row.data)
+    if (decoded._tag !== "Some" || decoded.value.inputID !== id)
+      return yield* Effect.die(new LifecycleConflict({ id }))
+    const base = { id, sessionID, timeCreated: DateTime.makeUnsafe(row.created) }
+    return decoded.value.input.type === "user"
+      ? User.make({ ...base, ...decoded.value.input })
+      : Synthetic.make({ ...base, ...decoded.value.input })
+  }
+  // A projected message without an admitted event in this aggregate (for
+  // example fork-copied history) is not a retryable admission.
+  return yield* Effect.die(new LifecycleConflict({ id }))
+})
+
 /**
  * Reconstruct the admitted record for a pending row that was already consumed
  * by promotion. The projected `session_message` row proves promotion happened;
@@ -123,27 +160,30 @@ const promotedFromHistory = Effect.fn("SessionPending.promotedFromHistory")(func
   if (message === undefined) return undefined
   if (message.session_id !== sessionID || (message.type !== "user" && message.type !== "synthetic"))
     return yield* Effect.die(new LifecycleConflict({ id }))
-  const rows = yield* db
+  return yield* admittedFromHistory(db, sessionID, id)
-    .select()
+})
+
+const wasWithdrawn = Effect.fn("SessionPending.wasWithdrawn")(function* (
+  db: DatabaseService,
+  sessionID: SessionSchema.ID,
+  id: SessionMessage.ID,
+) {
+  const row = yield* db
+    .select({ data: EventTable.data })
     .from(EventTable)
-    .where(and(eq(EventTable.aggregate_id, sessionID), eq(EventTable.type, admittedEventType)))
+    .where(
-    .all()
+      and(
+        eq(EventTable.aggregate_id, sessionID),
+        eq(EventTable.type, withdrawnEventType),
+        sql`json_extract(${EventTable.data}, '$.inputID') = ${id}`,
+      ),
+    )
+    .limit(1)
+    .get()
     .pipe(Effect.orDie)
-  for (const row of rows) {
+  if (!row) return false
-    const decoded = decodeAdmittedEvent(row.data)
+  const decoded = decodeWithdrawnEvent(row.data)
-    if (decoded._tag !== "Some" || decoded.value.inputID !== id) continue
+  return decoded._tag === "Some" && decoded.value.inputID === id
-    const base = {
-      id,
-      sessionID,
-      timeCreated: DateTime.makeUnsafe(row.created),
-    }
-    return decoded.value.input.type === "user"
-      ? User.make({ ...base, ...decoded.value.input })
-      : Synthetic.make({ ...base, ...decoded.value.input })
-  }
-  // A projected message without an admitted event in this aggregate (for
-  // example fork-copied history) is not a retryable admission.
-  return yield* Effect.die(new LifecycleConflict({ id }))
 })
 
 export const admit = Effect.fn("SessionPending.admit")(function* (
@@ -162,6 +202,8 @@ export const admit = Effect.fn("SessionPending.admit")(function* (
   }
   const promoted = yield* promotedFromHistory(db, request.sessionID, request.id)
   if (promoted !== undefined) return promoted
+  if (yield* wasWithdrawn(db, request.sessionID, request.id))
+    return yield* admittedFromHistory(db, request.sessionID, request.id)
   return yield* bus
     .publish(SessionEvent.InputAdmitted, {
       inputID: request.id,
@@ -309,14 +351,13 @@ export const projectCompactionAdmitted = Effect.fn("SessionPending.projectCompac
  * message insert inside the same event transaction; the deleted row is what
  * makes the table pending-only.
  */
-export const projectPromoted = Effect.fn("SessionPending.projectPromoted")(function* (
+const consumeInput = Effect.fn("SessionPending.consumeInput")(function* (
   db: DatabaseService,
   input: {
     readonly id: SessionMessage.ID
     readonly sessionID: SessionSchema.ID
   },
 ) {
-  if (yield* compaction(db, input.sessionID)) return yield* Effect.die(new LifecycleConflict({ id: input.id }))
   const deleted = yield* db
     .delete(SessionPendingTable)
     .where(and(eq(SessionPendingTable.id, input.id), eq(SessionPendingTable.session_id, input.sessionID)))
@@ -329,6 +370,27 @@ export const projectPromoted = Effect.fn("SessionPending.projectPromoted")(funct
   return stored
 })
 
+export const projectPromoted = Effect.fn("SessionPending.projectPromoted")(function* (
+  db: DatabaseService,
+  input: {
+    readonly id: SessionMessage.ID
+    readonly sessionID: SessionSchema.ID
+  },
+) {
+  if (yield* compaction(db, input.sessionID)) return yield* Effect.die(new LifecycleConflict({ id: input.id }))
+  return yield* consumeInput(db, input)
+})
+
+export const projectWithdrawn = Effect.fn("SessionPending.projectWithdrawn")(function* (
+  db: DatabaseService,
+  input: {
+    readonly id: SessionMessage.ID
+    readonly sessionID: SessionSchema.ID
+  },
+) {
+  return yield* consumeInput(db, input)
+})
+
 export const settleCompaction = Effect.fn("SessionPending.settleCompaction")(function* (
   db: DatabaseService,
   input: { readonly sessionID: SessionSchema.ID },
@@ -406,6 +468,30 @@ export const equivalent = (
   return false
 }
 
+export const withdraw = Effect.fn("SessionPending.withdraw")(function* (
+  db: DatabaseService,
+  bus: Bus.Interface,
+  input: { readonly sessionID: SessionSchema.ID; readonly inputID: SessionMessage.ID },
+) {
+  return yield* inboxLocks.withLock(input.sessionID)(
+    Effect.gen(function* () {
+      const pending = yield* find(db, input.inputID)
+      if (!pending) return yield* wasWithdrawn(db, input.sessionID, input.inputID)
+      if (pending.sessionID !== input.sessionID || pending.type === "compaction") return false
+      yield* bus
+        .publish(SessionEvent.InputWithdrawn, input)
+        .pipe(
+          Effect.catchDefect((defect) =>
+            wasWithdrawn(db, input.sessionID, input.inputID).pipe(
+              Effect.flatMap((withdrawn) => (withdrawn ? Effect.void : Effect.die(defect))),
+            ),
+          ),
+        )
+      return true
+    }),
+  )
+})
+
 const publish = Effect.fn("SessionPending.publish")(function* (
   db: DatabaseService,
   bus: Bus.Interface,

packages/core/src/session/projector.ts

@@ -665,6 +665,12 @@ const layer = Layer.effectDiscard(
           .pipe(Effect.orDie)
       }),
     )
+    yield* bus.project(SessionEvent.InputWithdrawn, (event) =>
+      SessionPending.projectWithdrawn(db, {
+        id: event.data.inputID,
+        sessionID: event.data.sessionID,
+      }),
+    )
     yield* bus.project(SessionEvent.Compaction.Admitted, (event) =>
       Effect.gen(function* () {
         if (event.durable === undefined)

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

@@ -996,6 +996,52 @@ describe("Session.pending", () => {
     }),
   )
 
+  it.effect("withdraws an interrupted input before promotion without resurrecting exact retries", () =>
+    Effect.gen(function* () {
+      yield* setup
+      const session = yield* Session.Service
+
+      const admitted = yield* session.prompt({
+        id: SessionMessage.ID.make("msg_withdrawn"),
+        sessionID,
+        text: "Withdraw me",
+        resume: false,
+      })
+
+      yield* session.interrupt(sessionID)
+      expect(yield* session.withdraw({ sessionID, inputID: admitted.id })).toBe(true)
+      expect(yield* session.pending(sessionID)).toEqual([])
+      expect(yield* session.messages({ sessionID })).toEqual([])
+      expect(yield* session.withdraw({ sessionID, inputID: admitted.id })).toBe(true)
+
+      const retried = yield* session.prompt({
+        id: admitted.id,
+        sessionID,
+        text: "Withdraw me",
+        resume: false,
+      })
+      expect(retried.id).toBe(admitted.id)
+      expect(yield* session.pending(sessionID)).toEqual([])
+      expect(yield* eventCount(Bus.versionedType(SessionEvent.InputAdmitted.type, 1))).toBe(1)
+      expect(yield* eventCount(Bus.versionedType(SessionEvent.InputWithdrawn.type, 1))).toBe(1)
+    }),
+  )
+
+  it.effect("leaves promoted input for revert when withdrawal loses the race", () =>
+    Effect.gen(function* () {
+      yield* setup
+      const session = yield* Session.Service
+      const bus = yield* Bus.Service
+      const { db } = yield* Database.Service
+      const admitted = yield* session.prompt({ sessionID, text: "Promote me", resume: false })
+
+      yield* SessionPending.promote(db, bus, sessionID, "steer")
+
+      expect(yield* session.withdraw({ sessionID, inputID: admitted.id })).toBe(false)
+      expect(yield* session.messages({ sessionID })).toMatchObject([{ id: admitted.id, type: "user" }])
+    }),
+  )
+
   it.effect("lists an unhandled compaction barrier until it settles", () =>
     Effect.gen(function* () {
       yield* setup

packages/protocol/src/groups/session.ts

@@ -628,6 +628,22 @@ export const makeSessionGroup = <I extends HttpApiMiddleware.AnyId, S>(sessionLo
           }),
         ),
     )
+    .add(
+      HttpApiEndpoint.post("session.pending.withdraw", "/api/session/:sessionID/pending/:inputID/withdraw", {
+        params: { sessionID: Session.ID, inputID: SessionMessage.ID },
+        success: Schema.Struct({ data: Schema.Boolean }),
+        error: SessionNotFoundError,
+      })
+        .middleware(sessionLocationMiddleware)
+        .annotateMerge(
+          OpenApi.annotations({
+            identifier: "v2.session.pending.withdraw",
+            summary: "Withdraw pending session input",
+            description:
+              "Withdraw one admitted input before promotion. Returns true when the input was withdrawn or had already been withdrawn, and false when it is no longer pending so callers can fall back to reverting projected history.",
+          }),
+        ),
+    )
     .annotateMerge(
       OpenApi.annotations({
         title: "session",

packages/schema/src/session-event.ts

@@ -155,6 +155,16 @@ export const InputAdmitted = Event.durable({
 })
 export type InputAdmitted = typeof InputAdmitted.Type
 
+export const InputWithdrawn = Event.durable({
+  type: "session.input.withdrawn",
+  ...options,
+  schema: {
+    ...Base,
+    inputID: SessionMessage.ID,
+  },
+})
+export type InputWithdrawn = typeof InputWithdrawn.Type
+
 export namespace Execution {
   export const Started = Event.durable({ type: "session.execution.started", ...options, schema: Base })
   export type Started = typeof Started.Type
@@ -556,6 +566,7 @@ export const Definitions = Event.inventory(
   Forked,
   InputPromoted,
   InputAdmitted,
+  InputWithdrawn,
   Execution.Started,
   Execution.Succeeded,
   Execution.Failed,

packages/schema/test/event-manifest.test.ts

@@ -108,6 +108,7 @@ describe("public event manifest", () => {
         "session.forked.2",
         "session.input.promoted.1",
         "session.input.admitted.1",
+        "session.input.withdrawn.1",
         "session.execution.started.1",
         "session.execution.succeeded.1",
         "session.execution.failed.1",

packages/server/src/handlers/session.ts

@@ -608,6 +608,23 @@ export const SessionHandler = HttpApiBuilder.group(Api, "server.session", (handl
           }
         }),
       )
+      .handle(
+        "session.pending.withdraw",
+        Effect.fn(function* (ctx) {
+          return {
+            data: yield* session.withdraw(ctx.params).pipe(
+              Effect.catchTag("Session.NotFoundError", (error) =>
+                Effect.fail(
+                  new SessionNotFoundError({
+                    sessionID: error.sessionID,
+                    message: `Session not found: ${error.sessionID}`,
+                  }),
+                ),
+              ),
+            ),
+          }
+        }),
+      )
       .handle(
         "session.instructions.entry.list",
         Effect.fn(function* (ctx) {

packages/tui/src/context/data.tsx

@@ -188,6 +188,12 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
         index.set(item.id, messages.length)
         messages.push(item)
       },
+      reindex(messages: SessionMessageInfo[], index: Map<string, number>, start: number) {
+        for (let position = start; position < messages.length; position++) {
+          const item = messages[position]
+          if (item) index.set(item.id, position)
+        }
+      },
       activeAssistant(messages: SessionMessageInfo[]) {
         const item = messages.findLast((item) => item.type === "assistant" && !item.time.completed)
         return item?.type === "assistant" ? item : undefined
@@ -395,8 +401,7 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
             existing.time.created = event.created
             draft.splice(position, 1)
             draft.push(existing)
-            index.clear()
+            message.reindex(draft, index, position)
-            draft.forEach((message, indexValue) => index.set(message.id, indexValue))
           })
           setStore(
             "session",
@@ -438,6 +443,25 @@ export const { use: useData, provider: DataProvider } = createSimpleContext({
             )
           })
           break
+        case "session.input.withdrawn": {
+          removePending(event.data.sessionID, event.data.inputID)
+          if (store.session.input[event.data.sessionID]?.includes(event.data.inputID))
+            message.update(event.data.sessionID, (draft, index) => {
+              const position = index.get(event.data.inputID)
+              if (position === undefined) return
+              draft.splice(position, 1)
+              index.delete(event.data.inputID)
+              message.reindex(draft, index, position)
+            })
+          if (store.session.input[event.data.sessionID]?.includes(event.data.inputID))
+            setStore(
+              "session",
+              "input",
+              event.data.sessionID,
+              (store.session.input[event.data.sessionID] ?? []).filter((id) => id !== event.data.inputID),
+            )
+          break
+        }
         case "session.instructions.updated":
           const instructions = event.metadata?.instructions
           if (

packages/tui/src/mini/stream-v2.subagent.ts

@@ -643,6 +643,10 @@ export function createSubagentTracker(input: SubagentTrackerInput): SubagentTrac
       if (event.data.input.type === "user") child.prompts.set(event.data.inputID, event.data.input.data.text)
       return
     }
+    if (event.type === "session.input.withdrawn") {
+      child.prompts.delete(event.data.inputID)
+      return
+    }
     if (event.type === "session.input.promoted") {
       const prompt = child.prompts.get(event.data.inputID)
       if (prompt === undefined) return

packages/tui/src/mini/stream-v2.transport.ts

@@ -911,6 +911,12 @@ export async function createSessionTransport(input: StreamInput): Promise<Sessio
       })
       return
     }
+    if (event.type === "session.input.withdrawn") {
+      state.admitted.delete(event.data.inputID)
+      state.pending.delete(event.data.inputID)
+      syncPending()
+      return
+    }
     if (event.type === "session.input.promoted") {
       const waiting = state.wait?.messageID === event.data.inputID
       if (state.wait) promoteWait(state.wait, true, event.data.inputID)

packages/tui/src/routes/session/dialog-message.tsx

@@ -8,6 +8,7 @@ import { errorMessage } from "../../util/error"
 import { DialogFork } from "./dialog-fork"
 import type { PromptInfo } from "../../prompt/history"
 import { projectedPromptInput } from "../../prompt/codec"
+import { undoMessage } from "./undo"
 
 export function DialogMessage(props: {
   messageID: string
@@ -42,9 +43,10 @@ export function DialogMessage(props: {
                 pasted: [],
               })
             }
-            void client.api.session.revert
+            void undoMessage(client.api, {
-              .stage({ sessionID: props.sessionID, messageID: props.messageID })
+              sessionID: props.sessionID,
-              .catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
+              messageID: props.messageID,
+            }).catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
             dialog.clear()
           },
         },

packages/tui/src/routes/session/index.tsx

@@ -64,6 +64,7 @@ import { useToast } from "../../ui/toast"
 import stripAnsi from "strip-ansi"
 import { usePromptRef } from "../../context/prompt"
 import { projectedPromptInput } from "../../prompt/codec"
+import { undoMessage } from "./undo"
 import { useEpilogue } from "../../context/epilogue"
 import { normalizePath } from "../../util/path"
 import { PermissionPrompt } from "./permission"
@@ -590,9 +591,10 @@ export function Session() {
           dialog.clear()
           return
         }
-        void client.api.session.revert
+        void undoMessage(client.api, {
-          .stage({ sessionID: route.sessionID, messageID: message.id })
+          sessionID: route.sessionID,
-          .catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
+          messageID: message.id,
+        }).catch((error) => toast.show({ message: errorMessage(error), variant: "error", duration: 5000 }))
         prompt()?.set({
           ...projectedPromptInput(message),
           pasted: [],

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

@@ -0,0 +1,11 @@
+import type { OpenCodeClient } from "@opencode-ai/client"
+
+export function undoMessage(
+  client: OpenCodeClient,
+  input: { readonly sessionID: string; readonly messageID: string },
+) {
+  const revert = () => client.session.revert.stage(input).then(() => undefined)
+  return client.session.pending
+    .withdraw({ sessionID: input.sessionID, inputID: input.messageID })
+    .then((withdrawn) => (withdrawn ? undefined : revert()))
+}

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

@@ -853,6 +853,65 @@ test("completes exploration when a queued prompt is promoted", async () => {
   }
 })
 
+test("removes optimistic input when it is withdrawn", async () => {
+  const events = createEventStream()
+  const sessionID = "session-withdrawal"
+  const calls = createFetch((url) => {
+    if (url.pathname === `/api/session/${sessionID}/message`) return json({ data: [], cursor: {} })
+  }, events)
+  let data!: ReturnType<typeof useData>
+  let client!: ReturnType<typeof useClient>
+
+  function Probe() {
+    data = useData()
+    client = useClient()
+    return <box />
+  }
+
+  const app = await testRender(() => (
+    <TestTuiContexts>
+      <ClientProvider api={createApi(calls.fetch)}>
+        <ProjectProvider>
+          <DataProvider>
+            <Probe />
+          </DataProvider>
+        </ProjectProvider>
+      </ClientProvider>
+    </TestTuiContexts>
+  ))
+
+  try {
+    await wait(() => client.connection.status() === "connected")
+    emitEvent(events, {
+      id: "evt_prompt_admitted",
+      created: 1,
+      type: "session.input.admitted",
+      durable: durable(sessionID, 1),
+      data: {
+        sessionID,
+        inputID: "message-user",
+        input: { type: "user", data: { text: "Never mind" }, delivery: "steer" },
+      },
+    })
+    await wait(() => data.session.message.get(sessionID, "message-user") !== undefined)
+    expect(data.session.input.has(sessionID, "message-user")).toBe(true)
+
+    emitEvent(events, {
+      id: "evt_prompt_withdrawn",
+      created: 2,
+      type: "session.input.withdrawn",
+      durable: durable(sessionID, 2),
+      data: { sessionID, inputID: "message-user" },
+    })
+
+    await wait(() => data.session.message.get(sessionID, "message-user") === undefined)
+    expect(data.session.input.has(sessionID, "message-user")).toBe(false)
+    expect(data.session.pending.list(sessionID)).toEqual([])
+  } finally {
+    app.renderer.destroy()
+  }
+})
+
 test("classifies live tool rows independently of their call ID", async () => {
   const events = createEventStream()
   const sessionID = "session-tool-call-id"

packages/tui/test/cli/tui/undo.test.ts

@@ -0,0 +1,26 @@
+import { expect, test } from "bun:test"
+import { OpenCode } from "@opencode-ai/client"
+import { undoMessage } from "../../../src/routes/session/undo"
+
+test.each([
+  { withdrawn: true, expected: ["withdraw"] },
+  { withdrawn: false, expected: ["withdraw", "revert"] },
+])("routes undo for withdrawn=$withdrawn", async ({ withdrawn, expected }) => {
+  const calls: string[] = []
+  const client = OpenCode.make({
+    baseUrl: "http://localhost:3000",
+    fetch: Object.assign(
+      async (input: URL | RequestInfo, init?: BunFetchRequestInit | RequestInit) => {
+        const request = input instanceof Request ? input : new Request(input, init)
+        const operation = request.url.endsWith("/withdraw") ? "withdraw" : "revert"
+        calls.push(operation)
+        return Response.json({ data: operation === "withdraw" ? withdrawn : { messageID: "msg_user" } })
+      },
+      { preconnect: fetch.preconnect },
+    ),
+  })
+
+  await undoMessage(client, { sessionID: "ses_test", messageID: "msg_user" })
+
+  expect(calls).toEqual([...expected])
+})