fix(composer): restore the prompt when a run fails after send

src/assistant/sessionStore.test.ts

@@ -26,7 +26,7 @@ const ASK_REQUEST = {
 beforeEach(() => {
   // Zustand exposes setState on the hook itself; reset the slice that
   // every test below mutates. Avoids cross-test bleed.
-  useAssistantStore.setState({ sessions: {}, activeSessionByTab: {} });
+  useAssistantStore.setState({ sessions: {}, activeSessionByTab: {}, recoverablePrompts: {} });
 });
 
 describe('initSession', () => {
@@ -236,3 +236,44 @@ describe('setRunStatus', () => {
     expect(useAssistantStore.getState().sessions[SESSION.id]!.isStreaming).toBe(true);
   });
 });
+
+describe('removeMessage — recoverable prompt capture', () => {
+  const userMsg = (id: string, text: string): AssistantMessage =>
+    ({ id, role: 'user', content: [{ type: 'text', text }] }) as unknown as AssistantMessage;
+
+  it('stashes a retracted user message text so the composer can restore it', () => {
+    const store = useAssistantStore.getState();
+    store.initSession(SESSION);
+    store.addMessage(SESSION.id, userMsg('u-1', '  hello there  '));
+    store.removeMessage(SESSION.id, 'u-1');
+    // Trimmed text is recoverable, keyed by session.
+    expect(useAssistantStore.getState().recoverablePrompts[SESSION.id]).toBe('hello there');
+  });
+
+  it('does not stash when an assistant message is removed', () => {
+    const store = useAssistantStore.getState();
+    store.initSession(SESSION);
+    const assistant = { id: 'a-1', role: 'assistant', content: [{ type: 'text', text: 'hi' }] } as unknown as AssistantMessage;
+    store.addMessage(SESSION.id, assistant);
+    store.removeMessage(SESSION.id, 'a-1');
+    expect(useAssistantStore.getState().recoverablePrompts[SESSION.id]).toBeUndefined();
+  });
+
+  it('does not stash an empty/whitespace-only user message', () => {
+    const store = useAssistantStore.getState();
+    store.initSession(SESSION);
+    store.addMessage(SESSION.id, userMsg('u-2', '   '));
+    store.removeMessage(SESSION.id, 'u-2');
+    expect(useAssistantStore.getState().recoverablePrompts[SESSION.id]).toBeUndefined();
+  });
+
+  it('clearRecoverablePrompt removes the stashed text', () => {
+    const store = useAssistantStore.getState();
+    store.initSession(SESSION);
+    store.addMessage(SESSION.id, userMsg('u-3', 'keep me'));
+    store.removeMessage(SESSION.id, 'u-3');
+    expect(useAssistantStore.getState().recoverablePrompts[SESSION.id]).toBe('keep me');
+    store.clearRecoverablePrompt(SESSION.id);
+    expect(useAssistantStore.getState().recoverablePrompts[SESSION.id]).toBeUndefined();
+  });
+});

src/assistant/sessionStore.ts

@@ -14,6 +14,7 @@ import type {
   AssistantMessageCursor,
   AssistantRun,
   AssistantSession,
+  ContentPart,
   ToolInvocation,
 } from '../generated/bindings';
 
@@ -68,6 +69,12 @@ export interface AssistantStoreState {
   initSession: (session: AssistantSession & { tabId?: string | null }) => void;
   addMessage: (sessionId: string, message: AssistantMessage) => void;
   removeMessage: (sessionId: string, messageId: string) => void;
+  /** Text of a user message whose run failed and was retracted by the
+   *  backend (429/400/token-limit/spawn error), keyed by session. The
+   *  composer reads this back into the input box so the typed prompt
+   *  isn't lost. Cleared once the composer consumes it. */
+  recoverablePrompts: Record<string, string>;
+  clearRecoverablePrompt: (sessionId: string) => void;
   markMessageQueued: (sessionId: string, messageId: string) => void;
   markQueuedMessagesDelivered: (sessionId: string, messageIds: string[]) => void;
   prependMessagePage: (
@@ -124,6 +131,7 @@ const useAssistantStore = create<AssistantStoreState>()(
     immer((set, get) => ({
       sessions: {},
       activeSessionByTab: {},
+      recoverablePrompts: {},
 
       setActiveSessionForTab: (tabId, sessionId) =>
         set((state) => {
@@ -167,6 +175,20 @@ const useAssistantStore = create<AssistantStoreState>()(
         set((state) => {
           const s = state.sessions[sessionId];
           if (!s) return;
+          // A retracted *user* message means the run failed before producing
+          // anything; stash its typed text so the composer can restore it
+          // instead of forcing the user to retype the prompt.
+          const removed = s.messages.find((m) => m.id === messageId);
+          if (removed && removed.role === 'user') {
+            const text = removed.content
+              .filter((p): p is Extract<ContentPart, { type: 'text' }> => p.type === 'text')
+              .map((p) => p.text)
+              .join('')
+              .trim();
+            if (text) {
+              state.recoverablePrompts[sessionId] = text;
+            }
+          }
           const before = s.messages.length;
           s.messages = s.messages.filter((m) => m.id !== messageId);
           if (s.messages.length < before && s.totalMessageCount !== null) {
@@ -176,6 +198,11 @@ const useAssistantStore = create<AssistantStoreState>()(
           delete s.streamingTextByMessageId[messageId];
         }),
 
+      clearRecoverablePrompt: (sessionId) =>
+        set((state) => {
+          delete state.recoverablePrompts[sessionId];
+        }),
+
       markMessageQueued: (sessionId, messageId) =>
         set((state) => {
           const s = state.sessions[sessionId];

src/components/TerminalEmulator/TerminalEmulator.tsx

@@ -41,6 +41,11 @@ interface TerminalEmulatorProps {
   onPickImage?: () => Promise<AttachImageResult>;
   onReadClipboardImage?: () => Promise<File | null>;
   agentWorking?: boolean;
+  /** Text of a prompt whose send failed asynchronously (run error after the
+   *  composer already cleared); restored into the input box when non-empty
+   *  and the box is empty. Cleared via onRecoverablePromptConsumed. */
+  recoverablePrompt?: string;
+  onRecoverablePromptConsumed?: () => void;
 }
 
 const TerminalEmulator = ({
@@ -50,6 +55,8 @@ const TerminalEmulator = ({
   onPickImage,
   onReadClipboardImage,
   agentWorking = false,
+  recoverablePrompt = '',
+  onRecoverablePromptConsumed,
 }: TerminalEmulatorProps) => {
   const location = useLocation();
   const [inputValue, setInputValue] = useState('');
@@ -279,6 +286,30 @@ const TerminalEmulator = ({
     }
   }
 
+  // Restore a prompt whose send failed asynchronously (the run errored after
+  // the composer already cleared — 429/400/token-limit). The BE retracts the
+  // user message and the store surfaces its text as `recoverablePrompt`. We
+  // apply it during render (same-component setState, converges — mirrors the
+  // draft swap above) so we never miss it, guarded so we don't clobber text
+  // the user has since typed, and only once per delivery.
+  const [appliedRecoverablePrompt, setAppliedRecoverablePrompt] = useState('');
+  if (recoverablePrompt && recoverablePrompt !== appliedRecoverablePrompt) {
+    setInputValue((current) => restoreFailedPrompt(current, recoverablePrompt));
+    setAppliedRecoverablePrompt(recoverablePrompt);
+  } else if (!recoverablePrompt && appliedRecoverablePrompt) {
+    // Slot cleared (consumed) — reset so a later, even identical, failure
+    // restores again.
+    setAppliedRecoverablePrompt('');
+  }
+  // Clear the store slot once we've surfaced it. `onRecoverablePromptConsumed`
+  // calls a zustand action (not a React state setter), so this is exempt from
+  // the set-state-in-effect rule.
+  useEffect(() => {
+    if (appliedRecoverablePrompt && onRecoverablePromptConsumed) {
+      onRecoverablePromptConsumed();
+    }
+  }, [appliedRecoverablePrompt, onRecoverablePromptConsumed]);
+
   // Ctrl+\ (or Cmd+\) toggles terminal mode. Matches the backslash key
   // (the 'Backslash' code or the '\\' character) and also the physical key
   // in the US-backtick position ('Backquote'): on a Spanish/ISO layout that

src/components/TerminalEmulator/TerminalEmulatorWrapper.tsx

@@ -100,6 +100,18 @@ const TerminalEmulatorWrapper = () => {
     (state) => (workspaceSessionId ? !!state.sessions[workspaceSessionId]?.isStreaming : false)
   );
 
+  // When a run fails before producing anything (429/400/token-limit), the BE
+  // retracts the user message and the store stashes its text here so the
+  // composer can restore the lost prompt instead of forcing a retype.
+  const recoverablePrompt = useAssistantStore(
+    (state) => (workspaceSessionId ? state.recoverablePrompts[workspaceSessionId] ?? '' : '')
+  );
+  const handleRecoverablePromptConsumed = useCallback(() => {
+    if (workspaceSessionId) {
+      useAssistantStore.getState().clearRecoverablePrompt(workspaceSessionId);
+    }
+  }, [workspaceSessionId]);
+
   const inputDisabled = isWorkspaceRoute && workspaceIsStreaming;
 
   /**
@@ -298,6 +310,8 @@ const TerminalEmulatorWrapper = () => {
       onPickImage={handlePickImage}
       onReadClipboardImage={readClipboardImageAsFile}
       agentWorking={inputDisabled}
+      recoverablePrompt={recoverablePrompt}
+      onRecoverablePromptConsumed={handleRecoverablePromptConsumed}
     />
   );
 };