feat(mobile): show prompt in chat thread immediately on submit (#2407)

Gilbert09 · web-flow · commit 32ebc986af2a · 2026-05-28T15:47:39.000+01:00
diff --git a/apps/mobile/src/app/task/[id].tsx b/apps/mobile/src/app/task/[id].tsx
@@ -32,6 +32,10 @@ import {
 } from "@/features/tasks/composer/options";
 import { TaskChatComposer } from "@/features/tasks/composer/TaskChatComposer";
 import { taskKeys } from "@/features/tasks/hooks/useTasks";
+import {
+  pendingTaskPromptStoreApi,
+  usePendingTaskPrompt,
+} from "@/features/tasks/stores/pendingTaskPromptStore";
 import { useTaskSessionStore } from "@/features/tasks/stores/taskSessionStore";
 import { useTaskStore } from "@/features/tasks/stores/taskStore";
 import type { Task } from "@/features/tasks/types";
@@ -92,6 +96,30 @@ export default function TaskDetailScreen() {
 
   const session = taskId ? getSessionForTask(taskId) : undefined;
 
+  // Optimistic echo set by the new-task screen (or the terminal-resume path
+  // below) so the user's prompt appears in the thread immediately, before
+  // the live session catches up.
+  const optimisticPrompt = usePendingTaskPrompt(taskId);
+
+  // Clear the echo once the canonical user_message_chunk with matching text
+  // arrives via SSE — `TaskSessionView` also dedups visually, but clearing
+  // the store frees it for the next submit. Only events with `ts >= setAt`
+  // qualify so a text-identical historical turn (e.g. resubmitting
+  // "Continue") doesn't drop the echo before the real copy lands.
+  useEffect(() => {
+    if (!taskId || !optimisticPrompt) return;
+    const matched = session?.events.some(
+      (e) =>
+        e.type === "session_update" &&
+        e.notification?.update?.sessionUpdate === "user_message_chunk" &&
+        e.notification.update.content?.text === optimisticPrompt.promptText &&
+        (e.ts ?? 0) >= optimisticPrompt.setAt,
+    );
+    if (matched) {
+      pendingTaskPromptStoreApi.clear(taskId);
+    }
+  }, [taskId, optimisticPrompt, session?.events]);
+
   // Per-task composer pill values. Persisted in taskStore so reopening the
   // task keeps the user's choices; defaults fall back to the same constants
   // the new-task composer uses.
@@ -217,6 +245,19 @@ export default function TaskDetailScreen() {
   const handleSendAfterTerminal = useCallback(
     async (text: string, attachments: PendingAttachment[]) => {
       if (!taskId || !task) return;
+      // Optimistically echo into the chat before tearing down the old session
+      // and waiting for the resume run's SSE stream to come up.
+      const echoAttachments = attachments.map((a) => ({
+        kind: a.kind,
+        uri: a.uri,
+        fileName: a.fileName,
+        mimeType: a.mimeType,
+      }));
+      pendingTaskPromptStoreApi.set(taskId, {
+        promptText: text,
+        attachments: echoAttachments.length > 0 ? echoAttachments : undefined,
+        setAt: Date.now(),
+      });
       try {
         setRetrying(true);
         disconnectFromTask(taskId);
@@ -242,6 +283,7 @@ export default function TaskDetailScreen() {
         updateTaskInCache(updatedTask);
       } catch (err) {
         log.error("Failed to send after terminal", err);
+        pendingTaskPromptStoreApi.clear(taskId);
         setRetrying(false);
         Alert.alert(
           "Failed to send",
@@ -390,7 +432,10 @@ export default function TaskDetailScreen() {
     !!session &&
     session.status === "connecting" &&
     session.events.length === 0;
-  const showLoading = loading || isHistoryLoading;
+  // Suppress the full-screen overlay when we have an optimistic prompt to
+  // show — the user just submitted and seeing their own text + a connecting
+  // indicator is friendlier than a blank spinner.
+  const showLoading = (loading || isHistoryLoading) && !optimisticPrompt;
   const showAutomationContext =
     fromAutomation === "1" || task?.origin_product === "automation";
   const automationContextLabel =
@@ -473,6 +518,15 @@ export default function TaskDetailScreen() {
           }
           onOpenTask={handleOpenTask}
           onSendPermissionResponse={handleSendPermissionResponse}
+          optimisticUserMessage={
+            optimisticPrompt
+              ? {
+                  text: optimisticPrompt.promptText,
+                  attachments: optimisticPrompt.attachments,
+                  setAt: optimisticPrompt.setAt,
+                }
+              : undefined
+          }
           contentContainerStyle={{
             paddingTop: 8,
             paddingBottom:
diff --git a/apps/mobile/src/app/task/index.tsx b/apps/mobile/src/app/task/index.tsx
@@ -62,6 +62,10 @@ import { Pill } from "@/features/tasks/composer/Pill";
 import { RepositoryPickerInline } from "@/features/tasks/composer/RepositoryPickerInline";
 import { SelectSheet } from "@/features/tasks/composer/SelectSheet";
 import { useUserIntegrations } from "@/features/tasks/hooks/useUserIntegrations";
+import {
+  generatePendingTaskKey,
+  pendingTaskPromptStoreApi,
+} from "@/features/tasks/stores/pendingTaskPromptStore";
 import { useTaskStore } from "@/features/tasks/stores/taskStore";
 import type {
   CreateTaskOptions,
@@ -263,8 +267,28 @@ export default function NewTaskScreen() {
 
     setCreating(true);
 
+    // Echo the prompt into the chat thread the moment the user taps send.
+    // The key is transient until `createTask` returns the real task id, at
+    // which point we `move` it so the detail screen can pick it up.
+    const pendingKey = generatePendingTaskKey();
+    const trimmedPrompt = prompt.trim();
+    const echoAttachments = attachments.map((a) => ({
+      kind: a.kind,
+      uri: a.uri,
+      fileName: a.fileName,
+      mimeType: a.mimeType,
+    }));
+    pendingTaskPromptStoreApi.set(pendingKey, {
+      promptText: trimmedPrompt,
+      attachments: echoAttachments.length > 0 ? echoAttachments : undefined,
+      setAt: Date.now(),
+    });
+
+    // Tracks where the optimistic echo currently lives so the catch block
+    // can clear the correct key regardless of how far the flow got.
+    let currentPendingKey = pendingKey;
+
     try {
-      const trimmedPrompt = prompt.trim();
       // The task description is plain text (it shows up as the task title and
       // in metadata). Attachments only enter the agent prompt via the cloud
       // payload below.
@@ -292,6 +316,9 @@ export default function NewTaskScreen() {
           : {}),
       } as CreateTaskOptions);
 
+      pendingTaskPromptStoreApi.move(pendingKey, task.id);
+      currentPendingKey = task.id;
+
       const pendingUserMessage =
         attachments.length > 0
           ? serializeCloudPrompt(
@@ -318,6 +345,7 @@ export default function NewTaskScreen() {
       router.replace(`/task/${task.id}`);
     } catch (creationError) {
       log.error("Failed to create task", creationError);
+      pendingTaskPromptStoreApi.clear(currentPendingKey);
     } finally {
       setCreating(false);
     }
diff --git a/apps/mobile/src/features/tasks/components/TaskSessionView.test.tsx b/apps/mobile/src/features/tasks/components/TaskSessionView.test.tsx
@@ -51,7 +51,84 @@ vi.mock("./PlanApprovalCard", () => ({
     createElement("PlanApprovalCard", props),
 }));
 
+function renderTaskSessionView(
+  props: Parameters<typeof TaskSessionView>[0],
+): ReturnType<typeof create> {
+  let renderer!: ReturnType<typeof create>;
+  act(() => {
+    renderer = create(createElement(TaskSessionView, props));
+  });
+  return renderer;
+}
+
+function findHumanMessages(renderer: ReturnType<typeof create>) {
+  // vi.mock'd `HumanMessage` is rendered as the literal string `"HumanMessage"`
+  // (an intrinsic), so node.type is a string at runtime even though the type
+  // says ElementType.
+  return renderer.root.findAll(
+    (node) => (node.type as unknown as string) === "HumanMessage",
+  );
+}
+
 describe("TaskSessionView", () => {
+  function userMessageEvent(text: string, ts: number) {
+    return {
+      type: "session_update" as const,
+      ts,
+      notification: {
+        update: {
+          sessionUpdate: "user_message_chunk",
+          content: { type: "text", text },
+        },
+      },
+    };
+  }
+
+  const SUBMIT_TS = 1000;
+
+  it.each([
+    {
+      name: "no SSE echo yet → optimistic renders",
+      events: [],
+      expectedCount: 1,
+    },
+    {
+      name: "matching SSE chunk after submit → optimistic suppressed",
+      events: [userMessageEvent("Ship it", SUBMIT_TS + 5)],
+      expectedCount: 1,
+    },
+    {
+      name: "text-identical historical turn → optimistic still renders",
+      // Same text but ts predates submit — a prior "Ship it" message shouldn't
+      // cause the new optimistic echo to be deduped.
+      events: [userMessageEvent("Ship it", SUBMIT_TS - 1000)],
+      expectedCount: 2,
+    },
+    {
+      name: "non-matching SSE text → optimistic still renders",
+      events: [userMessageEvent("Different text", SUBMIT_TS + 5)],
+      expectedCount: 2,
+    },
+  ])("optimistic echo: $name", ({ events, expectedCount }) => {
+    const renderer = renderTaskSessionView({
+      events,
+      optimisticUserMessage: { text: "Ship it", setAt: SUBMIT_TS },
+    });
+
+    expect(findHumanMessages(renderer)).toHaveLength(expectedCount);
+  });
+
+  it("optimistic echo carries the submitted text into the rendered bubble", () => {
+    const renderer = renderTaskSessionView({
+      events: [],
+      optimisticUserMessage: { text: "Ship it", setAt: SUBMIT_TS },
+    });
+
+    const humans = findHumanMessages(renderer);
+    expect(humans).toHaveLength(1);
+    expect(humans[0].props.content).toBe("Ship it");
+  });
+
   it("keeps question tools pending after the run goes idle", () => {
     const events = [
       {
diff --git a/apps/mobile/src/features/tasks/components/TaskSessionView.tsx b/apps/mobile/src/features/tasks/components/TaskSessionView.tsx
@@ -41,6 +41,15 @@ interface PermissionResponseArgs {
   displayText: string;
 }
 
+interface OptimisticUserMessage {
+  text: string;
+  attachments?: SessionNotificationAttachment[];
+  // Submit-time epoch ms. Dedup only fires against user messages whose `ts`
+  // is at or after this — protects against a text-identical historical turn
+  // suppressing the new optimistic echo.
+  setAt: number;
+}
+
 interface TaskSessionViewProps {
   events: SessionEvent[];
   pendingPermissions?: Record<string, CloudPendingPermissionRequest>;
@@ -52,6 +61,11 @@ interface TaskSessionViewProps {
   onOpenTask?: (taskId: string) => void;
   onSendPermissionResponse?: (args: PermissionResponseArgs) => void;
   contentContainerStyle?: object;
+  // Renders a user message at the bottom of the thread before the SSE echo
+  // arrives — for the gap between submit and the live session catching up.
+  // Suppressed automatically once a real user_message_chunk with matching
+  // text appears in `events`.
+  optimisticUserMessage?: OptimisticUserMessage;
 }
 
 interface ToolData {
@@ -792,6 +806,7 @@ export function TaskSessionView({
   onOpenTask,
   onSendPermissionResponse,
   contentContainerStyle,
+  optimisticUserMessage,
 }: TaskSessionViewProps) {
   const processorRef = useRef(createProcessorState());
   const prevEventsRef = useRef(events);
@@ -838,9 +853,34 @@ export function TaskSessionView({
   }
   prevAgentActive.current = agentActive;
 
+  // Append the optimistic user echo (if any) as the newest message, unless a
+  // real `user` message with matching text *and a ts at or after submit time*
+  // has already arrived via SSE. Gating on `ts` prevents a text-identical
+  // historical turn from suppressing a freshly-submitted echo.
+  const messagesWithOptimistic = useMemo(() => {
+    if (!optimisticUserMessage) return messages;
+    const alreadyEchoed = messages.some(
+      (m) =>
+        m.type === "user" &&
+        m.content === optimisticUserMessage.text &&
+        (m.ts ?? 0) >= optimisticUserMessage.setAt,
+    );
+    if (alreadyEchoed) return messages;
+    const optimistic: ParsedMessage = {
+      id: "optimistic-user",
+      type: "user",
+      content: optimisticUserMessage.text,
+      attachments: optimisticUserMessage.attachments,
+    };
+    return [...messages, optimistic];
+  }, [messages, optimisticUserMessage]);
+
   // Inverted FlatList renders data[0] at the visual bottom.
   // Reverse so newest messages are at index 0 = bottom.
-  const reversedMessages = useMemo(() => [...messages].reverse(), [messages]);
+  const reversedMessages = useMemo(
+    () => [...messagesWithOptimistic].reverse(),
+    [messagesWithOptimistic],
+  );
   const themeColors = useThemeColors();
   const flatListRef = useRef<FlatList>(null);
   const hasPendingQuestion = useMemo(
diff --git a/apps/mobile/src/features/tasks/stores/pendingTaskPromptStore.test.ts b/apps/mobile/src/features/tasks/stores/pendingTaskPromptStore.test.ts
@@ -0,0 +1,71 @@
+import { beforeEach, describe, expect, it } from "vitest";
+import {
+  pendingTaskPromptStoreApi,
+  usePendingTaskPromptStore,
+} from "./pendingTaskPromptStore";
+
+describe("pendingTaskPromptStore", () => {
+  beforeEach(() => {
+    usePendingTaskPromptStore.setState({ byKey: {} });
+  });
+
+  it("stores prompts keyed by an arbitrary id", () => {
+    pendingTaskPromptStoreApi.set("uuid-1", {
+      promptText: "Fix the login bug",
+      setAt: 1000,
+    });
+
+    expect(pendingTaskPromptStoreApi.get("uuid-1")).toEqual({
+      promptText: "Fix the login bug",
+      setAt: 1000,
+    });
+  });
+
+  it("moves a prompt from a transient key to the real task id", () => {
+    pendingTaskPromptStoreApi.set("uuid-1", {
+      promptText: "Do the thing",
+      setAt: 1000,
+    });
+    pendingTaskPromptStoreApi.move("uuid-1", "task-123");
+
+    expect(pendingTaskPromptStoreApi.get("uuid-1")).toBeUndefined();
+    expect(pendingTaskPromptStoreApi.get("task-123")).toEqual({
+      promptText: "Do the thing",
+      setAt: 1000,
+    });
+  });
+
+  it("ignores move when the source key has no prompt", () => {
+    pendingTaskPromptStoreApi.move("missing", "task-999");
+    expect(pendingTaskPromptStoreApi.get("task-999")).toBeUndefined();
+  });
+
+  it("clears prompts", () => {
+    pendingTaskPromptStoreApi.set("task-42", {
+      promptText: "Hi",
+      setAt: 1000,
+    });
+    pendingTaskPromptStoreApi.clear("task-42");
+    expect(pendingTaskPromptStoreApi.get("task-42")).toBeUndefined();
+  });
+
+  it("preserves attachments through move", () => {
+    pendingTaskPromptStoreApi.set("uuid-1", {
+      promptText: "Look at this",
+      setAt: 1000,
+      attachments: [
+        {
+          kind: "image",
+          uri: "file://x.png",
+          fileName: "x.png",
+          mimeType: "image/png",
+        },
+      ],
+    });
+    pendingTaskPromptStoreApi.move("uuid-1", "task-7");
+
+    expect(pendingTaskPromptStoreApi.get("task-7")?.attachments).toHaveLength(
+      1,
+    );
+  });
+});
diff --git a/apps/mobile/src/features/tasks/stores/pendingTaskPromptStore.ts b/apps/mobile/src/features/tasks/stores/pendingTaskPromptStore.ts
