fix(ui): preserve new composer state across reconnects (#128620)

Prevent delayed rewind, branch-switch, and fork completions from overwriting replacement chat state after reconnects while preserving canonical history reconciliation.

Refs #128617.
Reviewed-by: @shakkernerd
Co-authored-by: Shakker <165377636+shakkernerd@users.noreply.github.com>

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<typeof vi.fn>;
+      sessions: { rewind: ReturnType<typeof vi.fn> };
+    };
+    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<object>((resolve) => {
+      resolveSwitch = () => resolve({});
+    });
+    const selected = { role: "assistant", content: "selected branch after reconnect" };
+    const state = createState({ messages: [selected] }) as TestState & {
+      sessions: {
+        listBranches: ReturnType<typeof vi.fn>;
+        switchBranch: ReturnType<typeof vi.fn>;
+      };
+    };
+    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", () => {

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;
   }
 }

ui/src/pages/chat/chat-pane-history.ts

@@ -488,16 +488,20 @@ export abstract class ChatPaneHistory extends ChatPaneReplyNavigation {
   }
 
   protected async forkFromMessage(entryId: string): Promise<void> {
-    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?.();

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<never>((_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();
+  });
 });