From fe0ef0d1dace44f9bec3e6caff1ba6cdbb977dc1 Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Sat, 25 Jul 2026 01:01:21 +0300 Subject: [PATCH] fix(ui): discard reverted branch on resend --- packages/ui/src/sync/DOCUMENTATION.md | 1 + packages/ui/src/sync/session-actions.test.ts | 79 ++++++++++++++++++++ packages/ui/src/sync/session-actions.ts | 49 +++++++++++- 3 files changed, 127 insertions(+), 2 deletions(-) diff --git a/packages/ui/src/sync/DOCUMENTATION.md b/packages/ui/src/sync/DOCUMENTATION.md index 34eb5889..c72ebdfb 100644 --- a/packages/ui/src/sync/DOCUMENTATION.md +++ b/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`: diff --git a/packages/ui/src/sync/session-actions.test.ts b/packages/ui/src/sync/session-actions.test.ts index ddcd39e8..5f1d5b41 100644 --- a/packages/ui/src/sync/session-actions.test.ts +++ b/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]]) diff --git a/packages/ui/src/sync/session-actions.ts b/packages/ui/src/sync/session-actions.ts index 1826a1d9..b3357368 100644 --- a/packages/ui/src/sync/session-actions.ts +++ b/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 }, }, })