From a95b96100de612bd55a7287e628c4ce5ff382edc Mon Sep 17 00:00:00 2001 From: "Molecule AI Dev Engineer A (Kimi)" Date: Sun, 14 Jun 2026 02:29:08 +0000 Subject: [PATCH] fix(chat): render push-mode echo replies when WS completion races ahead (#2802 #2759) The token-per-send guard added in #2759 dropped push-mode HTTP replies when a WebSocket completion event (ACTIVITY_LOGGED status=ok / AGENT_MESSAGE) finished the token before the HTTP 200 landed. Fast echo runtimes hit this race, so the user message appeared but the echo reply never rendered. - In useChatSend.ts, process push-mode replies unconditionally and only gate the poll-mode queued short-circuit on token membership. Token cleanup stays idempotent. - Add unit regression test: finishSendByMessageId before HTTP reply resolves, assert onAgentMessage still receives the echo content. - Harden chat-separation E2E selectors to scope sub-tab/textarea lookups to the visible chat panel (), avoiding intermittent matches against a hidden ConciergeShell ChatTab copy. Refs #2802 / #2759 --- canvas/e2e/chat-separation.spec.ts | 14 +++--- .../useChatSend.concurrentReplies.test.tsx | 43 +++++++++++++++++++ .../components/tabs/chat/hooks/useChatSend.ts | 43 ++++++++----------- 3 files changed, 67 insertions(+), 33 deletions(-) diff --git a/canvas/e2e/chat-separation.spec.ts b/canvas/e2e/chat-separation.spec.ts index 1235cf738..21499569c 100644 --- a/canvas/e2e/chat-separation.spec.ts +++ b/canvas/e2e/chat-separation.spec.ts @@ -128,7 +128,7 @@ test.describe("Chat Sub-Tabs", () => { }); test("chat tab shows My Chat and Agent Comms sub-tabs", async ({ page }) => { - const panel = page.locator("#panel-chat"); + const panel = page.locator("#panel-chat [data-testid='chat-panel']:visible"); await expect(panel.getByRole("button", { name: "My Chat" })).toBeVisible(); await expect(panel.getByRole("button", { name: "Agent Comms" })).toBeVisible(); }); @@ -141,7 +141,7 @@ test.describe("Chat Sub-Tabs", () => { }); test("switching to Agent Comms shows different content", async ({ page }) => { - const panel = page.locator("#panel-chat"); + const panel = page.locator("#panel-chat [data-testid='chat-panel']:visible"); await panel.getByRole("button", { name: "Agent Comms" }).click(); // Agent Comms should be selected and My Chat's textarea should not be visible. @@ -152,7 +152,7 @@ test.describe("Chat Sub-Tabs", () => { }); test("My Chat has input box, Agent Comms does not", async ({ page }) => { - const panel = page.locator("#panel-chat"); + const panel = page.locator("#panel-chat [data-testid='chat-panel']:visible"); // My Chat has the textarea. await expect(panel.locator("textarea").first()).toBeVisible(); @@ -163,7 +163,7 @@ test.describe("Chat Sub-Tabs", () => { }); test("switching back to My Chat preserves messages", async ({ page }) => { - const panel = page.locator("#panel-chat"); + const panel = page.locator("#panel-chat [data-testid='chat-panel']:visible"); // Send a message so there is content to preserve. const textarea = panel.locator("textarea").first(); @@ -343,13 +343,13 @@ test.describe("Data Flow — Initial Prompt in Chat", () => { }); test("seeded chat history appears in My Chat", async ({ page }) => { - const panel = page.locator("#panel-chat"); + const panel = page.locator("#panel-chat [data-testid='chat-panel']:visible"); await expect(panel.getByText('Hello from seed with "quotes"')).toBeVisible({ timeout: 5_000 }); await expect(panel.getByText('Hello back from seed with "quotes"')).toBeVisible({ timeout: 5_000 }); }); test("My Chat empty state is not shown when history exists", async ({ page }) => { - const panel = page.locator("#panel-chat"); + const panel = page.locator("#panel-chat [data-testid='chat-panel']:visible"); await expect(panel.getByText("No messages yet")).not.toBeVisible(); }); }); @@ -384,7 +384,7 @@ test.describe("No JS Errors", () => { await openChatPanel(page, workspaceName); // Switch between tabs. - const panel = page.locator("#panel-chat"); + const panel = page.locator("#panel-chat [data-testid='chat-panel']:visible"); await panel.getByRole("button", { name: "Agent Comms" }).click(); await panel.getByRole("button", { name: "My Chat" }).click(); diff --git a/canvas/src/components/tabs/chat/hooks/__tests__/useChatSend.concurrentReplies.test.tsx b/canvas/src/components/tabs/chat/hooks/__tests__/useChatSend.concurrentReplies.test.tsx index 35c8b3ff4..5c4ba7f90 100644 --- a/canvas/src/components/tabs/chat/hooks/__tests__/useChatSend.concurrentReplies.test.tsx +++ b/canvas/src/components/tabs/chat/hooks/__tests__/useChatSend.concurrentReplies.test.tsx @@ -609,3 +609,46 @@ describe("useChatSend — legacy no-messageId fallback is exact-one-token conser expect(result.current.error).toBeNull(); }); }); + +describe("useChatSend — push-mode reply is not dropped by a racing WS completion (core#2802)", () => { + it("renders the push-mode echo reply even when onSendComplete finishes the token first", async () => { + const send = deferred(); + apiPostMock.mockImplementationOnce(() => send.promise); + + const onAgentMessage = vi.fn(); + const { result } = renderHook(() => + useChatSend("push-echo-ws-race", { + getHistoryMessages: () => [], + onAgentMessage, + }), + ); + + await act(async () => { + await result.current.sendMessage("echo me"); + await Promise.resolve(); + }); + expect(result.current.sending).toBe(true); + + const messageId = (apiPostMock.mock.calls[0][1] as any).params.message.messageId; + + // A WebSocket completion event (ACTIVITY_LOGGED status=ok with messageId, + // or AGENT_MESSAGE) can arrive before the HTTP 200 on a fast echo/reply. + // It finishes the token — the spinner must drop. + act(() => { + result.current.finishSendByMessageId?.(messageId); + }); + expect(result.current.sending).toBe(false); + + // The late HTTP push-mode reply still carries the actual content and MUST + // be rendered; otherwise the echo bubble never appears (core#2802). + await act(async () => { + send.resolve({ result: { parts: [{ kind: "text", text: "Echo: echo me" }] } }); + await Promise.resolve(); + }); + + expect(onAgentMessage).toHaveBeenCalledTimes(1); + expect((onAgentMessage.mock.calls[0][0] as { content: string }).content).toBe( + "Echo: echo me", + ); + }); +}); diff --git a/canvas/src/components/tabs/chat/hooks/useChatSend.ts b/canvas/src/components/tabs/chat/hooks/useChatSend.ts index e06c29151..37730fc11 100644 --- a/canvas/src/components/tabs/chat/hooks/useChatSend.ts +++ b/canvas/src/components/tabs/chat/hooks/useChatSend.ts @@ -366,33 +366,17 @@ export function useChatSend(workspaceId: string, options: UseChatSendOptions) { { timeoutMs: 30 * 60 * 1000 }, ) .then((resp) => { - // core#2725: only process the reply that belongs to this token. - // If the token is neither in-flight nor pending-WS, it has already - // been finished or superseded — drop it to avoid misrouted replies - // / duplicates. - if ( - !inFlightTokensRef.current.has(myToken) && - !pendingWSTokensRef.current.has(myToken) - ) return; - - // Task #227 — poll-mode (external/MCP workspace) queued-200 - // short-circuit. ws-server's `proxyA2ARequest` returns - // `{status:"queued", delivery_mode:"poll", ...}` immediately - // when the target has no URL (delivery_mode=poll), BEFORE the - // agent has produced any reply. There is no `result.parts` - // payload here — the actual reply will arrive separately via - // the AGENT_MESSAGE WebSocket event after the agent's next - // `wait_for_message` poll. - // - // Keep the spinner up by moving the token to the WS-pending set; - // releaseSendGuards will prune it when the AGENT_MESSAGE lands - // (handled by useChatSocket `onAgentMessage`/`onSendComplete`) - // or an explicit error fires (`onSendError` from an - // ACTIVITY_LOGGED status="error"). Don't synthesise an empty - // agent bubble. + // Poll-mode queued short-circuit (external/MCP workspace): the + // server returns `{status:"queued"}` immediately and the real reply + // arrives later via the AGENT_MESSAGE WebSocket event. if (resp?.status === "queued") { - // If WS already completed for this token via the legacy fallback - // path, finish immediately instead of re-pending forever. + if ( + !inFlightTokensRef.current.has(myToken) && + !pendingWSTokensRef.current.has(myToken) + ) { + // Already finished by a WS event (legacy fallback raced ahead). + return; + } if (wsCompletedTokensRef.current.has(myToken)) { finishSendToken(myToken); } else { @@ -401,6 +385,13 @@ export function useChatSend(workspaceId: string, options: UseChatSendOptions) { return; } + // Push-mode synchronous reply: process the agent message even if a + // WebSocket completion event (ACTIVITY_LOGGED or AGENT_MESSAGE) + // already finished this token. Without this, a fast echo/reply that + // triggers a WS completion before the HTTP 200 lands would have its + // token removed here and the reply bubble would never render + // (core#2802). Token cleanup is idempotent; ChatTab's message dedup + // handles the rare case where both paths carry the same content. const replyText = extractReplyText(resp); const replyFiles = extractFilesFromTask( (resp?.result ?? {}) as Record, -- 2.52.0