From 3b26601b77d761e203ca4c641c7f6676382a6fcc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Flor=C3=AAncio?= Date: Sat, 9 May 2026 11:39:14 +0100 Subject: [PATCH] fix: show sent message in chat without waiting for ws broadcast (#851) --- .../domains/comments/use-run-comment.test.ts | 40 +++++++++++++++++++ .../hooks/domains/comments/use-run-comment.ts | 9 ++++- apps/web/hooks/use-message-handler.ts | 19 ++++++--- 3 files changed, 62 insertions(+), 6 deletions(-) diff --git a/apps/web/hooks/domains/comments/use-run-comment.test.ts b/apps/web/hooks/domains/comments/use-run-comment.test.ts index 027a17a008..b868df8f09 100644 --- a/apps/web/hooks/domains/comments/use-run-comment.test.ts +++ b/apps/web/hooks/domains/comments/use-run-comment.test.ts @@ -9,6 +9,7 @@ import type { DiffComment, PlanComment } from "@/lib/state/slices/comments"; const mockRequest = vi.fn(); const mockAppendToQueue = vi.fn(); const mockMarkCommentsSent = vi.fn(); +const mockAddMessage = vi.fn(); const mockGetWebSocketClient = vi.fn(() => ({ request: mockRequest })); let mockStoreState: Record = {}; @@ -49,6 +50,7 @@ function makeStoreState(sessionState: string, planMode = false) { chatInput: { planModeBySessionId: { "sess-1": planMode }, }, + addMessage: mockAddMessage, }; } @@ -148,6 +150,44 @@ describe("useRunComment — idle agent sends directly", () => { expect(mockMarkCommentsSent).not.toHaveBeenCalled(); }); + + // Regression: comments sent via "Run" sometimes did not appear in the chat + // until a page refresh, because the hook depended entirely on the + // session.message.added broadcast — which can be missed if the client's + // session subscription is briefly absent or its send buffer drops. + // The message returned in the message.add response must be added to the + // store optimistically; addMessage is idempotent so a later broadcast for + // the same id is a no-op. + it("adds returned message to the store so chat updates without waiting for broadcast", async () => { + const returnedMessage = { + id: "msg-42", + session_id: "sess-1", + task_id: "task-1", + author_type: "user", + content: "[diff] fix this", + type: "message", + created_at: "2026-05-08T00:00:00Z", + }; + mockRequest.mockResolvedValueOnce(returnedMessage); + const { result } = renderCommentHook(); + + await act(async () => { + await result.current.runComment(makeDiffComment()); + }); + + expect(mockAddMessage).toHaveBeenCalledWith(returnedMessage); + }); + + it("does not call addMessage when message.add returns no message", async () => { + mockRequest.mockResolvedValueOnce(undefined); + const { result } = renderCommentHook(); + + await act(async () => { + await result.current.runComment(makeDiffComment()); + }); + + expect(mockAddMessage).not.toHaveBeenCalled(); + }); }); describe("useRunComment — busy agent queues", () => { diff --git a/apps/web/hooks/domains/comments/use-run-comment.ts b/apps/web/hooks/domains/comments/use-run-comment.ts index 235eb21c1c..4946983214 100644 --- a/apps/web/hooks/domains/comments/use-run-comment.ts +++ b/apps/web/hooks/domains/comments/use-run-comment.ts @@ -15,6 +15,7 @@ import type { FileEditorComment, PRFeedbackComment, } from "@/lib/state/slices/comments"; +import type { Message } from "@/lib/types/http"; /** * Format a single comment into markdown suitable for sending to the agent. @@ -125,11 +126,17 @@ export function useRunComment({ sessionId, taskId }: UseRunCommentParams) { } else { const client = getWebSocketClient(); if (!client) throw new Error("WebSocket client unavailable"); - await client.request( + // Add the returned message to the store directly so the chat updates + // even if the session.message.added broadcast is missed (subscription + // gap, dropped frame, etc.). addMessage is idempotent on id. + const created = await client.request( "message.add", buildMessagePayload(sessionId, taskId, content, planModeEnabled, comment), 10000, ); + if (created && created.id && created.session_id) { + storeApi.getState().addMessage(created); + } } markCommentsSent([comment.id]); diff --git a/apps/web/hooks/use-message-handler.ts b/apps/web/hooks/use-message-handler.ts index c3762bad9a..cda6532208 100644 --- a/apps/web/hooks/use-message-handler.ts +++ b/apps/web/hooks/use-message-handler.ts @@ -1,12 +1,13 @@ import { useCallback } from "react"; import { getWebSocketClient } from "@/lib/ws/connection"; +import { useAppStoreApi } from "@/components/state-provider"; import { useQueue } from "./domains/session/use-queue"; import type { MessageAttachment } from "@/components/task/chat/chat-input-container"; import type { ActiveDocument } from "@/lib/state/slices/ui/types"; import type { PlanComment } from "@/lib/state/slices/comments"; import { toBlockquote } from "@/lib/state/slices/comments/format"; import type { ContextFile } from "@/lib/state/context-files-store"; -import type { CustomPrompt } from "@/lib/types/http"; +import type { CustomPrompt, Message } from "@/lib/types/http"; function buildDocumentContext( activeDocument: ActiveDocument | null, @@ -92,9 +93,9 @@ type SendMessagePayload = { contextFilesMeta?: Array<{ path: string; name: string }>; }; -async function sendMessageRequest(payload: SendMessagePayload): Promise { +async function sendMessageRequest(payload: SendMessagePayload): Promise { const client = getWebSocketClient(); - if (!client) return; + if (!client) return undefined; const { taskId, @@ -108,7 +109,7 @@ async function sendMessageRequest(payload: SendMessagePayload): Promise { } = payload; const hasAttachments = attachments && attachments.length > 0; - await client.request( + return client.request( "message.add", { task_id: taskId, @@ -137,6 +138,7 @@ export function useMessageHandler({ prompts = [], }: UseMessageHandlerParams) { const { queue } = useQueue(resolvedSessionId); + const storeApi = useAppStoreApi(); const buildFinalMessage = useCallback( (message: string, inlineMentions?: ContextFile[]) => { @@ -182,7 +184,10 @@ export function useMessageHandler({ return; } - await sendMessageRequest({ + // Add the returned message to the store directly so the chat updates + // even if the session.message.added broadcast is missed (subscription + // gap, dropped frame, etc.). addMessage is idempotent on id. + const created = await sendMessageRequest({ taskId, resolvedSessionId, finalMessage, @@ -192,6 +197,9 @@ export function useMessageHandler({ attachments, contextFilesMeta, }); + if (created && created.id && created.session_id) { + storeApi.getState().addMessage(created); + } }, [ resolvedSessionId, @@ -202,6 +210,7 @@ export function useMessageHandler({ isAgentBusy, queue, buildFinalMessage, + storeApi, ], );