Refill composer when cancelling before the agent starts

packages/core/src/sessions/sessionService.ts

@@ -2176,6 +2176,31 @@ export class SessionService {
     }
   }
 
+  /**
+   * Whether the agent has begun working on the turn currently in flight. Used to
+   * decide whether cancelling can refill the composer with the just-sent message:
+   * before the prompt echo lands the message is still optimistic and no output is
+   * possible (refill is safe); once any output has streamed the turn was real
+   * work, so the message stays put. Works for local and cloud — both populate
+   * optimistic items and events from their respective transports.
+   */
+  hasAgentStartedCurrentTurn(taskId: string): boolean {
+    const session = this.d.store.getSessionByTaskId(taskId);
+    if (!session) return false;
+    // Echo not yet processed: the just-sent prompt is still an optimistic item
+    // (cleared atomically when the echo arrives), so no output can exist for it.
+    if (session.optimisticItems.length > 0) return false;
+    // Echo landed: the most recent session/prompt is this turn — any agent event
+    // after it (text, thinking, or a tool call) means work has started.
+    for (let i = session.events.length - 1; i >= 0; i--) {
+      const msg = session.events[i].message;
+      if (isJsonRpcRequest(msg) && msg.method === "session/prompt")
+        return false;
+      if (classifyTurnEventKind(msg) !== "other") return true;
+    }
+    return false;
+  }
+
   /**
    * Cancel the current prompt.
    */

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

@@ -1,3 +1,4 @@
+import { xmlToContent } from "@posthog/core/message-editor/content";
 import {
   combineQueuedCloudPrompts,
   promptToQueuedEditorContent,
@@ -48,8 +49,17 @@ export function useSessionCallbacks({
   const sessionRef = useRef(session);
   sessionRef.current = session;
 
+  // Serialized text of the most recent non-steer send, used to refill the
+  // composer if that turn is cancelled before the agent starts working.
+  const lastSentTextRef = useRef<string | null>(null);
+
   const messagingMode = useMessagingMode(taskId);
 
+  const isViewingTask = useCallback(() => {
+    const view = getAppViewSnapshot();
+    return view?.type === "task-detail" && view?.taskId === taskId;
+  }, [taskId]);
+
   const handleSendPrompt = useCallback(
     async (text: string) => {
       const currentSession = sessionRef.current;
@@ -71,14 +81,17 @@ export function useSessionCallbacks({
       try {
         markAsViewed(taskId);
         markActivity(taskId);
-        await sessionService.sendPrompt(taskId, text, {
-          steer: messagingMode === "steer",
-        });
 
-        const view = getAppViewSnapshot();
-        const isViewingTask =
-          view?.type === "task-detail" && view?.taskId === taskId;
-        if (isViewingTask) {
+        const steer = messagingMode === "steer";
+        // A steer folds into the running turn rather than starting its own, so it
+        // isn't a candidate for refill-on-cancel. Whether a non-steer send starts
+        // a turn or is queued, refill stays correct: a queued message is restored
+        // from the queue on cancel (below), which takes priority over this text.
+        lastSentTextRef.current = steer ? null : text;
+
+        await sessionService.sendPrompt(taskId, text, { steer });
+
+        if (isViewingTask()) {
           markAsViewed(taskId);
         }
       } catch (error) {
@@ -96,10 +109,18 @@ export function useSessionCallbacks({
       task.latest_run,
       sessionService,
       messagingMode,
+      isViewingTask,
     ],
   );
 
   const handleCancelPrompt = useCallback(async () => {
+    // Consume the stash up front so a second cancel can't refill twice. The
+    // cancelled message stays in history (it's already in the agent's context);
+    // we only refill the composer when the agent hadn't started working yet.
+    const justSent = lastSentTextRef.current;
+    lastSentTextRef.current = null;
+    const agentStarted = sessionService.hasAgentStartedCurrentTurn(taskId);
+
     const queuedMessages = sessionStoreSetters.dequeueMessages(taskId);
     const result = await sessionService.cancelPrompt(taskId);
     log.info("Prompt cancelled", { success: result });
@@ -109,6 +130,7 @@ export function useSessionCallbacks({
       : queuedMessages.map((message) => message.content).join("\n\n");
 
     if (queuedPrompt) {
+      // Queued messages are the more recent intent, so they win over justSent.
       const pendingContent = sessionRef.current?.isCloud
         ? promptToQueuedEditorContent(queuedPrompt)
         : {
@@ -121,9 +143,13 @@ export function useSessionCallbacks({
           };
 
       setPendingContent(taskId, pendingContent);
+    } else if (justSent && !agentStarted && isViewingTask()) {
+      // Refill the just-sent message so it can be edited and re-sent, but only
+      // while focused on this chat. xmlToContent restores attachment chips.
+      setPendingContent(taskId, xmlToContent(justSent));
     }
     requestFocus(taskId);
-  }, [taskId, setPendingContent, requestFocus, sessionService]);
+  }, [taskId, setPendingContent, requestFocus, sessionService, isViewingTask]);
 
   const handleRetry = useCallback(async () => {
     try {