fix: show sent message in chat without waiting for ws broadcast (#851)

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<string, unknown> = {};
 
@@ -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", () => {

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 | undefined>(
             "message.add",
             buildMessagePayload(sessionId, taskId, content, planModeEnabled, comment),
             10000,
           );
+          if (created && created.id && created.session_id) {
+            storeApi.getState().addMessage(created);
+          }
         }
 
         markCommentsSent([comment.id]);

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<void> {
+async function sendMessageRequest(payload: SendMessagePayload): Promise<Message | undefined> {
   const client = getWebSocketClient();
-  if (!client) return;
+  if (!client) return undefined;
 
   const {
     taskId,
@@ -108,7 +109,7 @@ async function sendMessageRequest(payload: SendMessagePayload): Promise<void> {
   } = payload;
   const hasAttachments = attachments && attachments.length > 0;
 
-  await client.request(
+  return client.request<Message | undefined>(
     "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,
     ],
   );