fix(ui): discard reverted branch on resend

packages/ui/src/sync/DOCUMENTATION.md

@@ -200,6 +200,7 @@ Rules:
 1. If an action mutates session list membership or visible session metadata, update `useGlobalSessionsStore` there.
 2. If an action targets a session by ID, resolve the **session's own directory**. Do not assume the current directory is correct.
 3. `session-ui-store.ts` should delegate to `session-actions.ts` for these mutations instead of duplicating SDK calls.
+4. Sending after a revert commits the new branch optimistically: remove the reverted tail and marker before inserting the new message, and restore both if the send is rejected.
 
 Examples of global-store updates performed in `session-actions.ts`:
 

packages/ui/src/sync/session-actions.test.ts

@@ -617,6 +617,85 @@ describe("optimisticSend target directory", () => {
     expect(currentStore.getState().session_status["session-new"]).toBe(undefined)
   })
 
+  test("commits the new branch locally when sending after a revert", async () => {
+    const retainedMessage = { id: "msg_1", role: "user", sessionID: "session-reverted" } as Message
+    const revertedMessage = { id: "msg_2", role: "user", sessionID: "session-reverted" } as Message
+    const targetStore = createStore({}, {
+      session: [{ id: "session-reverted", revert: { messageID: "msg_2" } } as Session],
+      message: { "session-reverted": [retainedMessage, revertedMessage] },
+      part: { msg_2: [{ id: "part_2", type: "text", text: "old branch" } as Part] },
+    })
+    const childStores = createChildStores([["/target/project", targetStore]])
+    let optimisticMessage: Message | null = null
+
+    const { optimisticSend, setActionRefs, setOptimisticRefs } = await import("./session-actions")
+    setActionRefs(mockSdk as unknown as OpencodeClient, childStores, () => "/target/project")
+    setOptimisticRefs(
+      (input) => {
+        optimisticMessage = input.message
+        targetStore.setState((state) => ({
+          message: { ...state.message, [input.sessionID]: [...(state.message[input.sessionID] ?? []), input.message] },
+          part: { ...state.part, [input.message.id]: input.parts },
+        }))
+      },
+      () => {},
+    )
+
+    await optimisticSend({
+      sessionId: "session-reverted",
+      directory: "/target/project",
+      content: "new branch",
+      providerID: "provider",
+      modelID: "model",
+      send: async () => {},
+    })
+
+    expect(targetStore.getState().session[0].revert).toBe(undefined)
+    expect(targetStore.getState().message["session-reverted"].map((message) => message.id)).toEqual([
+      "msg_1",
+      (optimisticMessage as unknown as Message).id,
+    ])
+    expect(targetStore.getState().part.msg_2).toBe(undefined)
+  })
+
+  test("restores the reverted branch when sending fails", async () => {
+    const retainedMessage = { id: "msg_1", role: "user", sessionID: "session-reverted" } as Message
+    const revertedMessage = { id: "msg_2", role: "user", sessionID: "session-reverted" } as Message
+    const revertedPart = { id: "part_2", type: "text", text: "old branch" } as Part
+    const targetStore = createStore({}, {
+      session: [{ id: "session-reverted", revert: { messageID: "msg_2" } } as Session],
+      message: { "session-reverted": [retainedMessage, revertedMessage] },
+      part: { msg_2: [revertedPart] },
+    })
+    const childStores = createChildStores([["/target/project", targetStore]])
+
+    const { optimisticSend, setActionRefs, setOptimisticRefs } = await import("./session-actions")
+    setActionRefs(mockSdk as unknown as OpencodeClient, childStores, () => "/target/project")
+    setOptimisticRefs(
+      (input) => targetStore.setState((state) => ({
+        message: { ...state.message, [input.sessionID]: [...(state.message[input.sessionID] ?? []), input.message] },
+        part: { ...state.part, [input.message.id]: input.parts },
+      })),
+      (input) => targetStore.setState((state) => ({
+        message: { ...state.message, [input.sessionID]: (state.message[input.sessionID] ?? []).filter((message) => message.id !== input.messageID) },
+        part: Object.fromEntries(Object.entries(state.part).filter(([messageID]) => messageID !== input.messageID)),
+      })),
+    )
+
+    await expect(optimisticSend({
+      sessionId: "session-reverted",
+      directory: "/target/project",
+      content: "new branch",
+      providerID: "provider",
+      modelID: "model",
+      send: async () => { throw new Error("rejected") },
+    })).rejects.toThrow("rejected")
+
+    expect(targetStore.getState().session[0].revert?.messageID).toBe("msg_2")
+    expect(targetStore.getState().message["session-reverted"]).toEqual([retainedMessage, revertedMessage])
+    expect(targetStore.getState().part.msg_2).toEqual([revertedPart])
+  })
+
   test("allows callers to block final send when runtime changes after optimistic insert", async () => {
     const targetStore = createStore({})
     const childStores = createChildStores([["/target/project", targetStore]])

packages/ui/src/sync/session-actions.ts

@@ -866,6 +866,29 @@ export async function optimisticSend(input: {
 
   const targetDirectory = input.directory ?? dir()
   const store = targetDirectory ? dirStoreForDirectory(targetDirectory) : dirStore()
+  const stateBeforeSend = store.getState()
+  const sessionBeforeSend = stateBeforeSend.session.find((session) => session.id === input.sessionId)
+  const revertMessageID = sessionBeforeSend?.revert?.messageID
+  const revertedMessages = revertMessageID
+    ? (stateBeforeSend.message[input.sessionId] ?? []).filter((message) => message.id >= revertMessageID)
+    : []
+  const revertedParts = new Map(
+    revertedMessages.map((message) => [message.id, stateBeforeSend.part[message.id] ?? []] as const),
+  )
+
+  if (revertMessageID) {
+    const session = stateBeforeSend.session.map((candidate) => (
+      candidate.id === input.sessionId ? { ...candidate, revert: undefined } as Session : candidate
+    ))
+    const message = {
+      ...stateBeforeSend.message,
+      [input.sessionId]: (stateBeforeSend.message[input.sessionId] ?? []).filter((candidate) => candidate.id < revertMessageID),
+    }
+    const part = { ...stateBeforeSend.part }
+    for (const revertedMessage of revertedMessages) delete part[revertedMessage.id]
+    store.setState({ session, message, part })
+  }
+
   const messageID = ascendingId("msg")
   input.onMessageID?.(messageID)
   const textPartId = ascendingId("prt")
@@ -934,10 +957,32 @@ export async function optimisticSend(input: {
       directory: targetDirectory,
       messageID,
     })
-    const s = store.getState()
+    const rollbackState = store.getState()
+    let session = rollbackState.session
+    let message = rollbackState.message
+    let part = rollbackState.part
+
+    if (revertMessageID) {
+      session = rollbackState.session.map((candidate) => (
+        candidate.id === input.sessionId ? { ...candidate, revert: sessionBeforeSend?.revert } as Session : candidate
+      ))
+      message = {
+        ...rollbackState.message,
+        [input.sessionId]: [...(rollbackState.message[input.sessionId] ?? []), ...revertedMessages]
+          .sort((a, b) => (a.id < b.id ? -1 : a.id > b.id ? 1 : 0)),
+      }
+      part = { ...rollbackState.part }
+      for (const [revertedMessageID, parts] of revertedParts) {
+        part[revertedMessageID] = parts
+      }
+    }
+
     store.setState({
+      session,
+      message,
+      part,
       session_status: {
-        ...s.session_status,
+        ...rollbackState.session_status,
         [input.sessionId]: { type: "idle" as const },
       },
     })