From f80a78d5a87c63a8f6cb5161179f1b46bd7fe4ae Mon Sep 17 00:00:00 2001
From: Xubin Ren <52506698+Re-bin@users.noreply.github.com>
Date: Mon, 22 Jun 2026 21:39:42 +0800
Subject: [PATCH] fix(webui): preserve fork replies during history refresh

---
 webui/src/components/thread/ThreadShell.tsx | 48 +++++++++++++-
 webui/src/tests/thread-shell.test.tsx       | 69 +++++++++++++++++++++
 2 files changed, 115 insertions(+), 2 deletions(-)

diff --git a/webui/src/components/thread/ThreadShell.tsx b/webui/src/components/thread/ThreadShell.tsx
index 7ac58ce3..97ccfece 100644
--- a/webui/src/components/thread/ThreadShell.tsx
+++ b/webui/src/components/thread/ThreadShell.tsx
@@ -44,7 +44,9 @@ function projectWebuiThreadMessages(messages: UIMessage[]): UIMessage[] {
   return scrubSubagentUiMessages(normalizeLegacyLongTaskMessages(messages));
 }
 
-function sameMessageShape(a: UIMessage, b: UIMessage): boolean {
+type MessageShape = Pick<UIMessage, "role" | "kind" | "content">;
+
+function sameMessageShape(a: MessageShape, b: MessageShape): boolean {
   return (
     a.role === b.role
     && (a.kind ?? "") === (b.kind ?? "")
@@ -52,9 +54,51 @@ function sameMessageShape(a: UIMessage, b: UIMessage): boolean {
   );
 }
 
+function durableMessageShape(message: UIMessage): MessageShape | null {
+  if (message.kind === "trace") return null;
+  if (message.role !== "user" && message.role !== "assistant") return null;
+  if (message.role === "assistant" && !message.content.trim() && !message.media?.length) {
+    return null;
+  }
+  return {
+    role: message.role,
+    kind: message.kind,
+    content: message.content,
+  };
+}
+
+function preservesDurableMessages(current: UIMessage[], snapshot: UIMessage[]): boolean {
+  // Canonical history refreshes can race with live websocket messages after fork/send.
+  // Never accept a refreshed snapshot that drops a user/assistant message already shown.
+  const expected = current
+    .map(durableMessageShape)
+    .filter((message): message is MessageShape => message !== null);
+  if (expected.length === 0) return true;
+  const candidates = snapshot
+    .map(durableMessageShape)
+    .filter((message): message is MessageShape => message !== null);
+
+  let cursor = 0;
+  for (const message of expected) {
+    let found = false;
+    while (cursor < candidates.length) {
+      const candidate = candidates[cursor];
+      cursor += 1;
+      if (sameMessageShape(message, candidate)) {
+        found = true;
+        break;
+      }
+    }
+    if (!found) return false;
+  }
+  return true;
+}
+
 function isStaleThreadSnapshot(current: UIMessage[], snapshot: UIMessage[]): boolean {
-  if (current.length === 0 || snapshot.length >= current.length) return false;
+  if (current.length === 0) return false;
   if (snapshot.length === 0) return true;
+  if (!preservesDurableMessages(current, snapshot)) return true;
+  if (snapshot.length >= current.length) return false;
   return snapshot.every((message, index) => sameMessageShape(current[index], message));
 }
 
diff --git a/webui/src/tests/thread-shell.test.tsx b/webui/src/tests/thread-shell.test.tsx
index 42a4f5c2..b30268a4 100644
--- a/webui/src/tests/thread-shell.test.tsx
+++ b/webui/src/tests/thread-shell.test.tsx
@@ -931,6 +931,75 @@ describe("ThreadShell", () => {
     await waitFor(() => expect(screen.getByText("live assistant reply")).toBeInTheDocument());
   });
 
+  it("keeps live fork replies when a canonical refresh is missing an earlier assistant answer", async () => {
+    const client = makeClient();
+    let historyCalls = 0;
+    vi.stubGlobal(
+      "fetch",
+      vi.fn(async (input: RequestInfo | URL) => {
+        const url = String(input);
+        if (url.includes("websocket%3Achat-fork/webui-thread")) {
+          historyCalls += 1;
+          return httpJson(
+            transcriptFromSimpleMessages(
+              historyCalls === 1
+                ? [{ role: "user", content: "first fork question" }]
+                : [
+                    { role: "user", content: "first fork question" },
+                    { role: "user", content: "second fork question" },
+                  ],
+            ),
+          );
+        }
+        return {
+          ok: false,
+          status: 404,
+          json: async () => ({}),
+        };
+      }),
+    );
+
+    render(
+      wrap(
+        client,
+        <ThreadShell
+          session={session("chat-fork")}
+          title="Chat chat-fork"
+          onToggleSidebar={() => {}}
+          onNewChat={() => {}}
+        />,
+      ),
+    );
+
+    await waitFor(() => expect(screen.getByText("first fork question")).toBeInTheDocument());
+    await act(async () => {
+      client._emitChat("chat-fork", {
+        event: "message",
+        chat_id: "chat-fork",
+        text: "first fork answer",
+      });
+    });
+    expect(screen.getByText("first fork answer")).toBeInTheDocument();
+
+    fireEvent.change(screen.getByLabelText("Message input"), {
+      target: { value: "second fork question" },
+    });
+    fireEvent.click(screen.getByRole("button", { name: "Send message" }));
+    await waitFor(() =>
+      expectSendMessageWithTurn(client, "chat-fork", "second fork question"),
+    );
+    expect(screen.getByText("second fork question")).toBeInTheDocument();
+
+    await act(async () => {
+      client._emitSessionUpdate("chat-fork");
+    });
+
+    await waitFor(() => expect(historyCalls).toBe(2));
+    expect(screen.getByText("first fork question")).toBeInTheDocument();
+    expect(screen.getByText("first fork answer")).toBeInTheDocument();
+    expect(screen.getByText("second fork question")).toBeInTheDocument();
+  });
+
   it("does not refetch thread history on turn_end", async () => {
     const client = makeClient();
     let historyCalls = 0;