diff --git a/ui/src/pages/chat/chat-history.test.ts b/ui/src/pages/chat/chat-history.test.ts index a28575915bad..94f4815e1dc4 100644 --- a/ui/src/pages/chat/chat-history.test.ts +++ b/ui/src/pages/chat/chat-history.test.ts @@ -339,6 +339,46 @@ describe("rewindChatHistory", () => { expect(result).toBeNull(); expect(state.handleChatDraftChange).not.toHaveBeenCalled(); }); + + it("reconciles committed rewind history without overwriting a replacement draft", async () => { + let resolveRewind!: (result: { editorText?: string }) => void; + const rewind = new Promise<{ editorText?: string }>((resolve) => { + resolveRewind = resolve; + }); + const canonical = { role: "assistant", content: "canonical history after rewind" }; + const state = createState({ messages: [canonical] }) as TestState & { + handleChatDraftChange: ReturnType; + sessions: { rewind: ReturnType }; + }; + state.chatMessages = [{ role: "assistant", content: "stale replacement history" }]; + state.handleChatDraftChange = vi.fn((next: string) => { + state.chatMessage = next; + }); + state.sessions = { + rewind: vi.fn(() => rewind), + setModelOverride: vi.fn(), + }; + + const pending = rewindChatHistory(state as never, "user-entry"); + state.connected = false; + state.connectionEpoch += 1; + state.connected = true; + state.connectionEpoch += 1; + state.chatMessage = "new connection draft"; + state.chatAttachments = [ + { id: "new", mimeType: "image/jpeg", dataUrl: "data:image/jpeg;base64,bmV3" }, + ]; + resolveRewind({ editorText: "stale rewind draft" }); + + const result = await pending; + expect(state.chatMessage).toBe("new connection draft"); + expect(state.chatAttachments).toEqual([ + { id: "new", mimeType: "image/jpeg", dataUrl: "data:image/jpeg;base64,bmV3" }, + ]); + expect(state.chatMessages).toEqual([canonical]); + expect(state.handleChatDraftChange).not.toHaveBeenCalled(); + expect(result).toBeNull(); + }); }); describe("switchChatHistoryBranch", () => { @@ -487,6 +527,37 @@ describe("switchChatHistoryBranch", () => { expect(request).toHaveBeenCalledTimes(2); expect(state.chatMessages).toEqual([selected]); }); + + it("reconciles a committed branch switch after a same-client reconnect", async () => { + let resolveSwitch!: () => void; + const switched = new Promise((resolve) => { + resolveSwitch = () => resolve({}); + }); + const selected = { role: "assistant", content: "selected branch after reconnect" }; + const state = createState({ messages: [selected] }) as TestState & { + sessions: { + listBranches: ReturnType; + switchBranch: ReturnType; + }; + }; + state.chatMessages = [{ role: "assistant", content: "stale branch after reconnect" }]; + state.sessions = { + listBranches: vi.fn().mockResolvedValue([]), + switchBranch: vi.fn(() => switched), + setModelOverride: vi.fn(), + }; + + const pending = switchChatHistoryBranch(state as never, "stale-leaf"); + state.connected = false; + state.connectionEpoch += 1; + state.connected = true; + state.connectionEpoch += 1; + resolveSwitch(); + + await expect(pending).resolves.toBe(false); + expect(state.chatMessages).toEqual([selected]); + expect(state.sessions.listBranches).toHaveBeenCalledWith(state.sessionKey, expect.any(Object)); + }); }); describe("canonical history snapshot projection", () => { diff --git a/ui/src/pages/chat/chat-history.ts b/ui/src/pages/chat/chat-history.ts index ce96ebc37a2e..1eb2c2941a25 100644 --- a/ui/src/pages/chat/chat-history.ts +++ b/ui/src/pages/chat/chat-history.ts @@ -1421,6 +1421,12 @@ export async function rewindChatHistory( } const sessionKey = state.sessionKey; const agentParams = scopedAgentParamsForSession(state, sessionKey); + const client = state.client; + const connectionEpoch = state.connectionEpoch; + const connectionIsCurrent = () => + state.connected && state.client === client && state.connectionEpoch === connectionEpoch; + const viewMatches = () => visibleSessionMatches(state, sessionKey, agentParams.agentId); + const viewIsCurrent = () => connectionIsCurrent() && viewMatches(); try { const result = await state.sessions.rewind(sessionKey, entryId, agentParams); const editorText = result.editorText ?? ""; @@ -1430,16 +1436,18 @@ export async function rewindChatHistory( agentId: agentParams.agentId, }); } - persistChatComposerState(state, sessionKey, { - agentId: agentParams.agentId, - draft: editorText, - }); - if (!visibleSessionMatches(state, sessionKey, agentParams.agentId)) { + if (connectionIsCurrent()) { + persistChatComposerState(state, sessionKey, { + agentId: agentParams.agentId, + draft: editorText, + }); + } + if (!viewMatches()) { return null; } resetChatHistoryProjection(state, agentParams.agentId); await Promise.all([loadChatHistory(state), loadChatBranches(state)]); - if (!visibleSessionMatches(state, sessionKey, agentParams.agentId)) { + if (!viewIsCurrent()) { return null; } // Restored images intentionally stay in this tab's memory; persisted composer drafts remain @@ -1451,8 +1459,10 @@ export async function rewindChatHistory( state.handleChatDraftChange(editorText); return result; } catch (error) { - setChatError(state, formatUiError(error)); - scheduleChatScroll(state); + if (viewIsCurrent()) { + setChatError(state, formatUiError(error)); + scheduleChatScroll(state); + } return null; } } @@ -1466,6 +1476,12 @@ export async function switchChatHistoryBranch( } const sessionKey = state.sessionKey; const agentParams = scopedAgentParamsForSession(state, sessionKey); + const client = state.client; + const connectionEpoch = state.connectionEpoch; + const connectionIsCurrent = () => + state.connected && state.client === client && state.connectionEpoch === connectionEpoch; + const viewMatches = () => visibleSessionMatches(state, sessionKey, agentParams.agentId); + const viewIsCurrent = () => connectionIsCurrent() && viewMatches(); try { await state.sessions.switchBranch(sessionKey, leafEntryId, agentParams); if (state.chatMessagesBySession) { @@ -1474,15 +1490,17 @@ export async function switchChatHistoryBranch( agentId: agentParams.agentId, }); } - if (!visibleSessionMatches(state, sessionKey, agentParams.agentId)) { + if (!viewMatches()) { return false; } resetChatHistoryProjection(state, agentParams.agentId); await Promise.all([loadChatHistory(state), loadChatBranches(state)]); - return visibleSessionMatches(state, sessionKey, agentParams.agentId); + return viewIsCurrent(); } catch (error) { - setChatError(state, formatUiError(error)); - scheduleChatScroll(state); + if (viewIsCurrent()) { + setChatError(state, formatUiError(error)); + scheduleChatScroll(state); + } return false; } } diff --git a/ui/src/pages/chat/chat-pane-history.ts b/ui/src/pages/chat/chat-pane-history.ts index 3af0655aa9dd..16ec15a1d37e 100644 --- a/ui/src/pages/chat/chat-pane-history.ts +++ b/ui/src/pages/chat/chat-pane-history.ts @@ -488,16 +488,20 @@ export abstract class ChatPaneHistory extends ChatPaneReplyNavigation { } protected async forkFromMessage(entryId: string): Promise { - const state = this.state; - if (!state) { + const scope = this.captureConnectionScope(); + if (!scope) { return; } + const state = scope.state; const sourceKey = state.sessionKey; const agentParams = scopedAgentParamsForSession(state, sourceKey); try { const result = await state.sessions.forkAtMessage(sourceKey, entryId, agentParams); const editorText = result.editorText ?? ""; - if (this.state !== state || !visibleSessionMatches(state, sourceKey, agentParams.agentId)) { + if ( + !this.isConnectionScopeCurrent(scope) || + !visibleSessionMatches(state, sourceKey, agentParams.agentId) + ) { return; } if (this.onPaneSessionChange?.(this.paneId, result.sessionKey) === false) { @@ -512,6 +516,12 @@ export abstract class ChatPaneHistory extends ChatPaneReplyNavigation { draft: editorText, }); } catch (error) { + if ( + !this.isConnectionScopeCurrent(scope) || + !visibleSessionMatches(state, sourceKey, agentParams.agentId) + ) { + return; + } state.lastError = formatUiError(error); state.chatError = state.lastError; state.requestUpdate?.(); diff --git a/ui/src/pages/chat/chat-pane.message-cut.test.ts b/ui/src/pages/chat/chat-pane.message-cut.test.ts index 271586a8994f..4fff94c6bf4c 100644 --- a/ui/src/pages/chat/chat-pane.message-cut.test.ts +++ b/ui/src/pages/chat/chat-pane.message-cut.test.ts @@ -132,4 +132,44 @@ describe("chat pane message cuts", () => { expect(state.sessionKey).toBe("global"); expect(state.assistantAgentId).toBe("work"); }); + + it("does not navigate to a fork that finishes after a same-client reconnect", async () => { + const forked = createDeferred<{ sessionKey: string; editorText?: string }>(); + const sessions = { + forkAtMessage: vi.fn(() => forked.promise), + } as unknown as SessionCapability; + const client = {} as GatewayBrowserClient; + const { pane, state } = createTestChatPane({ client, sessions }); + const navigate = vi.fn(); + pane.onPaneSessionChange = navigate; + + const pending = pane.forkFromMessage("user-entry"); + pane.connectionGeneration += 1; + state.connectionEpoch = pane.connectionGeneration; + forked.resolve({ sessionKey: "agent:main:forked", editorText: "stale draft" }); + + await pending; + expect(navigate).not.toHaveBeenCalled(); + expect(consumePaneSessionHandoff(pane.context, pane.paneId, "agent:main:forked")).toBeNull(); + }); + + it("does not paint a stale fork error after the selected session changes", async () => { + let rejectFork!: (error: Error) => void; + const forked = new Promise((_resolve, reject) => { + rejectFork = reject; + }); + const sessions = { + forkAtMessage: vi.fn(() => forked), + } as unknown as SessionCapability; + const client = {} as GatewayBrowserClient; + const { pane, state } = createTestChatPane({ client, sessions }); + + const pending = pane.forkFromMessage("user-entry"); + state.sessionKey = "agent:main:replacement"; + rejectFork(new Error("stale fork failed")); + + await pending; + expect(state.lastError).toBeNull(); + expect(state.chatError).toBeNull(); + }); });