fix(agent): preserve messages until send succeeds

Author: tatoalo

packages/ui/src/features/sessions/components/SessionView.tsx

@@ -1,4 +1,5 @@
 import { Pause, Spinner, Warning } from "@phosphor-icons/react";
+import { contentToXml } from "@posthog/core/message-editor/content";
 import {
   createLatestPlanTracker,
   SESSION_SERVICE,
@@ -68,7 +69,7 @@ interface SessionViewProps {
   isPromptPending?: boolean | null;
   promptStartedAt?: number | null;
   onBeforeSubmit?: (text: string, clearEditor: () => void) => boolean;
-  onSendPrompt: (text: string) => void;
+  onSendPrompt: (text: string) => Promise<boolean>;
   onBashCommand?: (command: string) => void;
   onCancelPrompt: () => void;
   repoPath?: string | null;
@@ -299,6 +300,9 @@ export function SessionView({
   ]);
 
   const isCloudRun = useIsWorkspaceCloudRun(taskId);
+  const editorRef = useRef<PromptInputHandle>(null);
+  const sendInFlightRef = useRef(false);
+  const [isSendingPrompt, setIsSendingPrompt] = useState(false);
 
   const latestPlanTrackerRef = useRef<ReturnType<
     typeof createLatestPlanTracker
@@ -309,11 +313,22 @@ export function SessionView({
     (): Plan | null => latestPlanTracker.update(events) as Plan | null,
     [events, latestPlanTracker],
   );
-
   const handleSubmit = useCallback(
-    (text: string) => {
-      if (text.trim()) {
-        onSendPrompt(text);
+    async (text: string): Promise<void> => {
+      if (!text.trim() || sendInFlightRef.current) return;
+
+      sendInFlightRef.current = true;
+      setIsSendingPrompt(true);
+      try {
+        if (await onSendPrompt(text)) {
+          const editor = editorRef.current;
+          if (editor && contentToXml(editor.getContent()) === text) {
+            editor.clear();
+          }
+        }
+      } finally {
+        sendInFlightRef.current = false;
+        setIsSendingPrompt(false);
       }
     },
     [onSendPrompt],
@@ -337,7 +352,6 @@ export function SessionView({
   const cancelQueuedEdit = useCancelQueuedMessageEdit(taskId);
 
   const [isDraggingFile, setIsDraggingFile] = useState(false);
-  const editorRef = useRef<PromptInputHandle>(null);
   const promptRecallRef = useRef<PromptRecallHandler | null>(null);
   const handlePromptRecall = useCallback<PromptRecallHandler>(
     (direction) => promptRecallRef.current?.(direction) ?? null,
@@ -678,8 +692,9 @@ export function SessionView({
                           placeholder="Type a message... @ to mention files, ! for bash mode, / for skills"
                           disabled={!isRunning && !handoffInProgress}
                           submitDisabledExternal={
-                            handoffInProgress || !isOnline
+                            handoffInProgress || !isOnline || isSendingPrompt
                           }
+                          clearOnSubmit={false}
                           submitTooltipOverride={
                             !isOnline ? "No internet connection" : undefined
                           }

packages/ui/src/features/sessions/hooks/useSessionCallbacks.test.tsx

@@ -163,6 +163,24 @@ describe("useSessionCallbacks.handleSendPrompt while editing a queued message",
   });
 });
 
+describe("useSessionCallbacks.handleSendPrompt", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    sessionState.editingQueuedId = undefined;
+    sessionState.messageQueue = [];
+  });
+
+  it("reports a failed send so the composer keeps its content", async () => {
+    sessionService.sendPrompt.mockRejectedValue(new Error("fetch failed"));
+
+    const { result } = renderCallbacks();
+    const sent = await result.current.handleSendPrompt("keep this message");
+
+    expect(sent).toBe(false);
+    expect(toastError).toHaveBeenCalledWith("fetch failed");
+  });
+});
+
 describe("useSessionCallbacks.handleCancelPrompt", () => {
   beforeEach(() => {
     vi.clearAllMocks();

packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts

@@ -62,7 +62,7 @@ export function useSessionCallbacks({
   const messagingMode = useMessagingMode(taskId);
 
   const handleSendPrompt = useCallback(
-    async (text: string) => {
+    async (text: string): Promise<boolean> => {
       const currentSession = sessionRef.current;
       const currentEvents = currentSession?.events ?? [];
       const handled = await tryExecuteCodeCommand(text, {
@@ -77,7 +77,7 @@ export function useSessionCallbacks({
           : null,
         taskRun: task.latest_run ?? null,
       });
-      if (handled) return;
+      if (handled) return true;
 
       let promptText =
         rewriteLocalSkillCommandPrompt(
@@ -109,7 +109,7 @@ export function useSessionCallbacks({
           );
           if (updated) {
             markAsViewed(taskId);
-            return;
+            return true;
           }
           // Target no longer queued — drop the stale hold and send as new.
           sessionService.clearEditingQueuedMessage(taskId);
@@ -127,7 +127,7 @@ export function useSessionCallbacks({
             setPendingContent(taskId, xmlToContent(promptText ?? text));
             requestFocus(taskId);
           }
-          return;
+          return false;
         }
       }
 
@@ -144,18 +144,13 @@ export function useSessionCallbacks({
         if (isViewingTask) {
           markAsViewed(taskId);
         }
+        return true;
       } catch (error) {
         const message =
           error instanceof Error ? error.message : "Failed to send message";
-        // The composer clears optimistically on submit, so a failed send
-        // would otherwise lose the typed prompt. Restore it — unless the
-        // user has already started typing a new message.
-        if (isContentEmpty(useDraftStore.getState().drafts[taskId] ?? null)) {
-          setPendingContent(taskId, xmlToContent(text));
-          requestFocus(taskId);
-        }
         toast.error(message);
         log.error("Failed to send prompt", error);
+        return false;
       }
     },
     [