feat(sessions): refill composer when cancelling before the agent starts

packages/core/src/sessions/hasAgentStartedTurn.test.ts

@@ -0,0 +1,87 @@
+import type { AcpMessage } from "@posthog/shared";
+import { describe, expect, it } from "vitest";
+import { createUserMessageEvent } from "./sessionEvents";
+import { hasAgentStartedTurn } from "./sessionService";
+
+function agentUpdate(
+  sessionUpdate: string,
+  ts: number,
+  content?: { type: string },
+): AcpMessage {
+  return {
+    type: "acp_message",
+    ts,
+    message: {
+      jsonrpc: "2.0",
+      method: "session/update",
+      params: { update: { sessionUpdate, ...(content ? { content } : {}) } },
+    },
+  } as AcpMessage;
+}
+
+describe("hasAgentStartedTurn", () => {
+  it.each<[string, AcpMessage[], boolean]>([
+    // A prior completed turn is present, but the just-sent prompt is still an
+    // optimistic item — its echo has not landed, so the agent has not started.
+    [
+      "still optimistic (echo not landed)",
+      [
+        createUserMessageEvent("first", 1),
+        agentUpdate("agent_message_chunk", 2, { type: "text" }),
+      ],
+      true,
+    ],
+    ["there are no events", [], false],
+    [
+      "only the prompt echo has landed",
+      [createUserMessageEvent("hi", 1)],
+      false,
+    ],
+  ])("is false when %s", (_label, events, hasOptimisticItems) => {
+    expect(hasAgentStartedTurn({ events, hasOptimisticItems })).toBe(false);
+  });
+
+  it.each([
+    [
+      "agent text chunk",
+      agentUpdate("agent_message_chunk", 2, { type: "text" }),
+    ],
+    ["agent thought chunk", agentUpdate("agent_thought_chunk", 2)],
+    ["tool call", agentUpdate("tool_call", 2)],
+    ["tool call update", agentUpdate("tool_call_update", 2)],
+  ])("is true after %s", (_label, output) => {
+    expect(
+      hasAgentStartedTurn({
+        events: [createUserMessageEvent("hi", 1), output],
+        hasOptimisticItems: false,
+      }),
+    ).toBe(true);
+  });
+
+  it.each<[string, AcpMessage[], boolean]>([
+    // Only looks after the latest prompt: the prior turn's output is ignored.
+    [
+      "prior turn has output but the latest prompt does not",
+      [
+        createUserMessageEvent("first", 1),
+        agentUpdate("agent_message_chunk", 2, { type: "text" }),
+        createUserMessageEvent("second", 3),
+      ],
+      false,
+    ],
+    [
+      "the latest turn has output after an earlier turn",
+      [
+        createUserMessageEvent("first", 1),
+        agentUpdate("agent_message_chunk", 2, { type: "text" }),
+        createUserMessageEvent("second", 3),
+        agentUpdate("tool_call", 4),
+      ],
+      true,
+    ],
+  ])("turn isolation: %s", (_label, events, expected) => {
+    expect(hasAgentStartedTurn({ events, hasOptimisticItems: false })).toBe(
+      expected,
+    );
+  });
+});

packages/core/src/sessions/sessionService.ts

@@ -601,6 +601,42 @@ function classifyTurnEventKind(
   return "other";
 }
 
+/**
+ * Whether the agent has begun working on the turn currently in flight.
+ *
+ * While the just-sent prompt is still optimistic (`hasOptimisticItems`, i.e. the
+ * `session/prompt` echo has not landed) no output is possible — refill is safe.
+ * Skipping this guard would let a prior completed turn's output look like the
+ * current turn's. Once the echo has landed, returns true if any agent
+ * text/output event follows the latest prompt. Reads only persisted state, so
+ * it is the same decision for local and cloud.
+ */
+export function hasAgentStartedTurn({
+  events,
+  hasOptimisticItems,
+}: {
+  events: AcpMessage[];
+  hasOptimisticItems: boolean;
+}): boolean {
+  if (hasOptimisticItems) return false;
+
+  let lastPromptIndex = -1;
+  for (let i = events.length - 1; i >= 0; i--) {
+    const msg = events[i].message;
+    if (isJsonRpcRequest(msg) && msg.method === "session/prompt") {
+      lastPromptIndex = i;
+      break;
+    }
+  }
+  if (lastPromptIndex === -1) return false;
+
+  for (let i = lastPromptIndex + 1; i < events.length; i++) {
+    const kind = classifyTurnEventKind(events[i].message);
+    if (kind === "text" || kind === "output") return true;
+  }
+  return false;
+}
+
 export class SessionService {
   private connectingTasks = new Map<string, Promise<void>>();
   private reconcilingTasks = new Set<string>();
@@ -2866,6 +2902,28 @@ export class SessionService {
     }
   }
 
+  /**
+   * Whether the agent has produced any output for the latest turn.
+   *
+   * Used to decide if a just-cancelled message can be safely refilled into the
+   * composer: if the agent never started, the message is effectively lost from
+   * the user's point of view, so we put it back.
+   *
+   * While the prompt is still optimistic (the `session/prompt` echo has not
+   * landed) no output is possible, so this returns false — refill is safe.
+   * Once the echo has landed it looks for any agent text/output event after the
+   * latest prompt. Reads only persisted store events, so it works for both
+   * local and cloud (avoids the local-only `liveTurnContent`).
+   */
+  hasAgentStartedCurrentTurn(taskId: string): boolean {
+    const session = this.d.store.getSessionByTaskId(taskId);
+    if (!session) return false;
+    return hasAgentStartedTurn({
+      events: session.events,
+      hasOptimisticItems: session.optimisticItems.length > 0,
+    });
+  }
+
   /**
    * Begin an in-place edit of a queued message: mark it as the edit target so
    * that, until the edit is saved or cancelled, it and everything queued after

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

@@ -37,6 +37,11 @@ import { useCallback, useRef } from "react";
 
 const log = logger.scope("session-callbacks");
 
+function isViewingTask(taskId: string): boolean {
+  const view = getAppViewSnapshot();
+  return view?.type === "task-detail" && view?.taskId === taskId;
+}
+
 interface UseSessionCallbacksOptions {
   taskId: string;
   task: Task;
@@ -59,6 +64,10 @@ export function useSessionCallbacks({
   const sessionRef = useRef(session);
   sessionRef.current = session;
 
+  // Serialized text of the in-flight message, captured on send so a cancel
+  // before the agent produces output can refill it into the composer.
+  const lastSentTextRef = useRef<string | null>(null);
+
   const messagingMode = useMessagingMode(taskId);
 
   const handleSendPrompt = useCallback(
@@ -134,14 +143,15 @@ export function useSessionCallbacks({
       try {
         markAsViewed(taskId);
         markActivity(taskId);
-        await sessionService.sendPrompt(taskId, promptText ?? 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.
+        lastSentTextRef.current = steer ? null : text;
+
+        await sessionService.sendPrompt(taskId, promptText ?? text, { steer });
+
+        if (isViewingTask(taskId)) {
           markAsViewed(taskId);
         }
       } catch (error) {
@@ -183,12 +193,20 @@ export function useSessionCallbacks({
       editingId &&
       currentSession?.messageQueue.some((m) => m.id === editingId)
     ) {
+      lastSentTextRef.current = null;
       const result = await sessionService.cancelPrompt(taskId);
       log.info("Prompt cancelled during queued edit", { success: result });
       requestFocus(taskId);
       return;
     }
 
+    // Snapshot the refill inputs at cancel time: whether the agent had started,
+    // and the just-sent text. Null the stash so a repeated cancel can't refill
+    // the same message twice.
+    const agentStarted = sessionService.hasAgentStartedCurrentTurn(taskId);
+    const lastSentText = lastSentTextRef.current;
+    lastSentTextRef.current = null;
+
     const queuedMessages = sessionStoreSetters.dequeueMessages(taskId);
     const result = await sessionService.cancelPrompt(taskId);
     log.info("Prompt cancelled", { success: result });
@@ -203,6 +221,9 @@ export function useSessionCallbacks({
         : textToContent(typeof queuedPrompt === "string" ? queuedPrompt : "");
 
       setPendingContent(taskId, pendingContent);
+    } else if (!agentStarted && lastSentText && isViewingTask(taskId)) {
+      // xmlToContent restores attachments and mentions as chips, not raw XML.
+      setPendingContent(taskId, xmlToContent(lastSentText));
     }
     requestFocus(taskId);
   }, [taskId, setPendingContent, requestFocus, sessionService]);