fix: keep cloud user message visible while log echo is pending

When sending a follow-up to a cloud run, drain queued messages into
optimistic user bubbles so the conversation stays visible during the
gap between submit and the cloud's session/prompt echo. Mirror the
local optimistic flow: clear the bubbles when the echo lands, on send
failure, on retry exhaustion, and on terminal-status cleanup.

Generated-By: PostHog Code
Task-Id: 8aeaf6f9-8b18-426a-9453-c668ca17d227

Author: VojtechBartos

apps/code/src/renderer/features/sessions/service/service.test.ts

@@ -1379,6 +1379,60 @@ describe("SessionService", () => {
       );
     });
 
+    it("appends an optimistic user_message before sending a direct cloud prompt", async () => {
+      const service = getSessionService();
+      mockSessionStoreSetters.getSessionByTaskId.mockReturnValue(
+        createMockSession({
+          isCloud: true,
+          cloudStatus: "in_progress",
+        }),
+      );
+      mockTrpcCloudTask.sendCommand.mutate.mockResolvedValue({
+        success: true,
+        result: { stopReason: "end_turn" },
+      });
+
+      await service.sendPrompt("task-123", "Hello cloud");
+
+      expect(mockSessionStoreSetters.appendOptimisticItem).toHaveBeenCalledWith(
+        "run-123",
+        expect.objectContaining({
+          type: "user_message",
+          content: "Hello cloud",
+        }),
+      );
+      const optimisticCallOrder =
+        mockSessionStoreSetters.appendOptimisticItem.mock
+          .invocationCallOrder[0];
+      const sendCallOrder =
+        mockTrpcCloudTask.sendCommand.mutate.mock.invocationCallOrder[0];
+      expect(optimisticCallOrder).toBeLessThan(sendCallOrder);
+    });
+
+    it("clears optimistic items if a direct cloud send fails", async () => {
+      const service = getSessionService();
+      mockSessionStoreSetters.getSessionByTaskId.mockReturnValue(
+        createMockSession({
+          isCloud: true,
+          cloudStatus: "in_progress",
+        }),
+      );
+      mockTrpcCloudTask.sendCommand.mutate.mockRejectedValue(
+        new Error("network down"),
+      );
+
+      await expect(
+        service.sendPrompt("task-123", "Hello cloud"),
+      ).rejects.toThrow("network down");
+
+      expect(
+        mockSessionStoreSetters.appendOptimisticItem,
+      ).toHaveBeenCalledTimes(1);
+      expect(mockSessionStoreSetters.clearOptimisticItems).toHaveBeenCalledWith(
+        "run-123",
+      );
+    });
+
     it("attempts automatic recovery on fatal error", async () => {
       const service = getSessionService();
       const mockSession = createMockSession({

apps/code/src/renderer/features/sessions/service/service.ts

@@ -92,6 +92,13 @@ const LOCAL_SESSION_RECOVERY_MESSAGE =
 const LOCAL_SESSION_RECOVERY_FAILED_MESSAGE =
   "Connecting to to the agent has been lost. Retry, or start a new session.";
 
+function isUserPromptEcho(acpMsg: AcpMessage): boolean {
+  return (
+    isJsonRpcRequest(acpMsg.message) &&
+    acpMsg.message.method === "session/prompt"
+  );
+}
+
 /**
  * Build default configOptions for cloud sessions so the mode switcher
  * is available in the UI even without a local agent connection.
@@ -993,19 +1000,17 @@ export class SessionService {
     const session = sessionStoreSetters.getSessions()[taskRunId];
     if (!session) return;
 
-    const isUserPromptEcho =
-      isJsonRpcRequest(acpMsg.message) &&
-      acpMsg.message.method === "session/prompt";
+    const isPromptEcho = isUserPromptEcho(acpMsg);
 
     // Once the agent starts responding, clear initialPrompt so that
     // retry reconnects to this session instead of creating a new one.
-    if (!isUserPromptEcho && session.initialPrompt?.length) {
+    if (!isPromptEcho && session.initialPrompt?.length) {
       sessionStoreSetters.updateSession(taskRunId, {
         initialPrompt: undefined,
       });
     }
 
-    if (isUserPromptEcho) {
+    if (isPromptEcho) {
       sessionStoreSetters.replaceOptimisticWithEvent(taskRunId, acpMsg);
     } else {
       sessionStoreSetters.appendEvents(taskRunId, [acpMsg]);
@@ -1518,6 +1523,17 @@ export class SessionService {
       isPromptPending: true,
     });
 
+    // Show the bubble immediately while we wait for the cloud log stream to
+    // echo the user_message back. Without this, the user sees a gap between
+    // submit and the agent's response on cloud runs.
+    if (!options?.skipQueueGuard) {
+      sessionStoreSetters.appendOptimisticItem(session.taskRunId, {
+        type: "user_message",
+        content: transport.promptText,
+        timestamp: Date.now(),
+      });
+    }
+
     track(ANALYTICS_EVENTS.PROMPT_SENT, {
       task_id: session.taskId,
       is_initial: session.events.length === 0,
@@ -1565,6 +1581,12 @@ export class SessionService {
       sessionStoreSetters.updateSession(session.taskRunId, {
         isPromptPending: false,
       });
+      // Drop optimistic items so a failed send doesn't leave a ghost bubble.
+      // The combined-prompt path (skipQueueGuard) clears its own optimistic
+      // items in sendQueuedCloudMessages on retry exhaustion.
+      if (!options?.skipQueueGuard) {
+        sessionStoreSetters.clearOptimisticItems(session.taskRunId);
+      }
       throw error;
     }
   }
@@ -1574,10 +1596,29 @@ export class SessionService {
     attempt = 0,
     pendingPrompt?: string | ContentBlock[],
   ): Promise<{ stopReason: string }> {
-    // First attempt: atomically dequeue. Retries reuse the already-dequeued prompt.
-    const combinedPrompt =
-      pendingPrompt ??
-      combineQueuedCloudPrompts(sessionStoreSetters.dequeueMessages(taskId));
+    // First attempt: atomically dequeue and convert each entry into an
+    // optimistic bubble. Retries reuse the already-dequeued prompt and must
+    // not stack additional bubbles.
+    let combinedPrompt: string | ContentBlock[] | null;
+    if (pendingPrompt) {
+      combinedPrompt = pendingPrompt;
+    } else {
+      const dequeued = sessionStoreSetters.dequeueMessages(taskId);
+      combinedPrompt = combineQueuedCloudPrompts(dequeued);
+      if (combinedPrompt) {
+        const taskRunId =
+          sessionStoreSetters.getSessionByTaskId(taskId)?.taskRunId;
+        if (taskRunId) {
+          for (const msg of dequeued) {
+            sessionStoreSetters.appendOptimisticItem(taskRunId, {
+              type: "user_message",
+              content: msg.content,
+              timestamp: msg.queuedAt,
+            });
+          }
+        }
+      }
+    }
     if (!combinedPrompt) return { stopReason: "skipped" };
 
     const session = sessionStoreSetters.getSessionByTaskId(taskId);
@@ -1632,6 +1673,10 @@ export class SessionService {
         taskId,
         attempts: attempt + 1,
       });
+      const failedSession = sessionStoreSetters.getSessionByTaskId(taskId);
+      if (failedSession) {
+        sessionStoreSetters.clearOptimisticItems(failedSession.taskRunId);
+      }
       toast.error("Failed to send follow-up message. Please try again.");
       return { stopReason: "error" };
     }
@@ -2888,6 +2933,8 @@ export class SessionService {
       const expectedCount = update.totalEntryCount;
       const delta = expectedCount - currentCount;
 
+      let appendedEvents: AcpMessage[] = [];
+
       if (delta <= 0) {
         // Already caught up — skip duplicate entries
       } else if (delta <= update.newEntries.length) {
@@ -2901,6 +2948,7 @@ export class SessionService {
         );
         sessionStoreSetters.appendEvents(taskRunId, newEvents, expectedCount);
         this.updatePromptStateFromEvents(taskRunId, newEvents);
+        appendedEvents = newEvents;
       } else {
         // Gap in data — append everything we have but don't jump processedLineCount
         log.warn("Cloud task log count inconsistency", {
@@ -2921,6 +2969,16 @@ export class SessionService {
           currentCount + update.newEntries.length,
         );
         this.updatePromptStateFromEvents(taskRunId, newEvents);
+        appendedEvents = newEvents;
+      }
+
+      // The cloud log stream has caught up to a user_message we already
+      // showed optimistically — drop the optimistic bubbles so they don't
+      // duplicate the now-real event. Mirrors replaceOptimisticWithEvent
+      // on the local path; we clear in bulk because cloud entries arrive
+      // batched alongside other event kinds.
+      if (appendedEvents.some(isUserPromptEcho)) {
+        sessionStoreSetters.clearOptimisticItems(taskRunId);
       }
     }
 
@@ -2975,6 +3033,7 @@ export class SessionService {
           (session.messageQueue.length > 0 || session.isPromptPending)
         ) {
           sessionStoreSetters.clearMessageQueue(session.taskId);
+          sessionStoreSetters.clearOptimisticItems(taskRunId);
           sessionStoreSetters.updateSession(taskRunId, {
             isPromptPending: false,
           });