From 4915b5e4cfa39a2e1cbec566b4151329e3645e94 Mon Sep 17 00:00:00 2001
From: Huixin615 <enjoyyourlife615@gmail.com>
Date: Tue, 7 Jul 2026 07:10:48 +0800
Subject: [PATCH] fix(frontend): enable regenerate in custom agent chats
 (#3967)

---
 .../[agent_name]/chats/[thread_id]/page.tsx   |  14 +++
 frontend/tests/e2e/agent-chat.spec.ts         | 114 +++++++++++++++++-
 2 files changed, 127 insertions(+), 1 deletion(-)

diff --git a/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx b/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx
index 50573ba27..d22c3be5d 100644
--- a/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx
+++ b/frontend/src/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx
@@ -91,6 +91,7 @@ export default function AgentChatPage() {
     thread,
     pendingUsageMessages,
     sendMessage,
+    regenerateMessage,
     isUploading,
     isHistoryLoading,
     hasMoreHistory,
@@ -203,6 +204,11 @@ export default function AgentChatPage() {
   const handleStop = useCallback(async () => {
     await thread.stop();
   }, [thread]);
+  const handleRegenerate = useCallback(
+    (messageId: string, supersededMessageIds: string[]) =>
+      regenerateMessage(threadId, messageId, supersededMessageIds),
+    [regenerateMessage, threadId],
+  );
 
   const tokenUsageInlineMode = tokenUsageEnabled
     ? localSettings.tokenUsage.inlineMode
@@ -293,6 +299,14 @@ export default function AgentChatPage() {
                   loadMoreHistory={loadMoreHistory}
                   isHistoryLoading={isHistoryLoading}
                   tokenUsageInlineMode={tokenUsageInlineMode}
+                  canRegenerate={
+                    !isNewThread &&
+                    !isMock &&
+                    env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY !== "true" &&
+                    !isUploading &&
+                    !thread.isLoading
+                  }
+                  onRegenerateMessage={handleRegenerate}
                   onSubmitHumanInput={
                     isMock || env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true"
                       ? undefined
diff --git a/frontend/tests/e2e/agent-chat.spec.ts b/frontend/tests/e2e/agent-chat.spec.ts
index 35616705b..e9fc29cb4 100644
--- a/frontend/tests/e2e/agent-chat.spec.ts
+++ b/frontend/tests/e2e/agent-chat.spec.ts
@@ -1,6 +1,10 @@
 import { expect, test } from "@playwright/test";
 
-import { mockLangGraphAPI } from "./utils/mock-api";
+import {
+  handleRunStream,
+  mockLangGraphAPI,
+  MOCK_THREAD_ID,
+} from "./utils/mock-api";
 
 const MOCK_AGENTS = [
   {
@@ -43,4 +47,112 @@ test.describe("Agent chat", () => {
       page.locator("header span", { hasText: "test-agent" }),
     ).toBeVisible({ timeout: 15_000 });
   });
+
+  test("agent chat can regenerate its latest response", async ({ page }) => {
+    const humanMessage = {
+      type: "human",
+      id: "msg-human-agent",
+      content: [{ type: "text", text: "Original agent question" }],
+    };
+    const aiMessage = {
+      type: "ai",
+      id: "msg-ai-agent",
+      content: "Custom agent response",
+    };
+    mockLangGraphAPI(page, {
+      agents: MOCK_AGENTS,
+      threads: [
+        {
+          thread_id: MOCK_THREAD_ID,
+          title: "Agent conversation",
+          agent_name: "test-agent",
+          messages: [humanMessage, aiMessage],
+        },
+      ],
+    });
+
+    let prepareMessageId: string | undefined;
+    let streamBody: Record<string, unknown> | undefined;
+    await page.route(
+      `**/api/threads/${MOCK_THREAD_ID}/runs/regenerate/prepare`,
+      (route) => {
+        prepareMessageId = (
+          route.request().postDataJSON() as { message_id?: string }
+        ).message_id;
+        return route.fulfill({
+          status: 200,
+          contentType: "application/json",
+          body: JSON.stringify({
+            input: { messages: [humanMessage] },
+            checkpoint: {
+              checkpoint_id: "checkpoint-before-human",
+              checkpoint_ns: "",
+              checkpoint_map: null,
+            },
+            metadata: {
+              regenerate_from_message_id: aiMessage.id,
+              regenerate_from_run_id: `run-${MOCK_THREAD_ID}`,
+              regenerate_checkpoint_id: "checkpoint-before-human",
+            },
+            target_run_id: `run-${MOCK_THREAD_ID}`,
+          }),
+        });
+      },
+    );
+    await page.route(
+      `**/api/langgraph/threads/${MOCK_THREAD_ID}/runs/stream`,
+      (route) => {
+        streamBody = route.request().postDataJSON() as Record<string, unknown>;
+        return handleRunStream(route);
+      },
+    );
+
+    await page.goto(`/workspace/agents/test-agent/chats/${MOCK_THREAD_ID}`);
+    await expect(page.getByText(aiMessage.content)).toBeVisible({
+      timeout: 15_000,
+    });
+
+    await page.evaluate((selectedText) => {
+      const element = Array.from(document.querySelectorAll("p")).find(
+        (candidate) => candidate.textContent?.includes(selectedText),
+      );
+      const textNode = element?.firstChild;
+      if (!element || !textNode) {
+        throw new Error("Unable to find the custom agent response");
+      }
+      const range = document.createRange();
+      range.selectNodeContents(textNode);
+      const selection = window.getSelection();
+      selection?.removeAllRanges();
+      selection?.addRange(range);
+      element.dispatchEvent(new MouseEvent("mouseup", { bubbles: true }));
+    }, aiMessage.content);
+    await expect(
+      page.getByRole("button", { name: "Ask in side chat" }),
+    ).toBeVisible();
+    await page.keyboard.press("Escape");
+
+    const assistantTurn = page.locator("[data-assistant-turn]").last();
+    await assistantTurn.hover();
+    await page.getByRole("button", { name: "Regenerate" }).click();
+
+    await expect.poll(() => prepareMessageId).toBe(aiMessage.id);
+    await expect.poll(() => streamBody).toBeDefined();
+    expect(streamBody).toMatchObject({
+      checkpoint: {
+        checkpoint_id: "checkpoint-before-human",
+        checkpoint_ns: "",
+        checkpoint_map: null,
+      },
+      metadata: {
+        regenerate_from_message_id: aiMessage.id,
+        regenerate_from_run_id: `run-${MOCK_THREAD_ID}`,
+        regenerate_checkpoint_id: "checkpoint-before-human",
+      },
+      context: {
+        agent_name: "test-agent",
+        thread_id: MOCK_THREAD_ID,
+      },
+    });
+  });
 });