From bccfa671dfd3af77a7cf779d24228844a14a7fcb Mon Sep 17 00:00:00 2001 From: juan <2930882+juacker@users.noreply.github.com> Date: Tue, 30 Jun 2026 22:50:18 +0200 Subject: [PATCH] fix(composer): restore the prompt when a run fails after send PR #78 restored the typed prompt only on *synchronous* send errors (no connection, off-workspace). But the common failures the user hits (429/400/token-exhausted) happen *asynchronously* during run execution, after assistant_send_message already returned success and the composer already cleared. On those, the backend retracts the unanswered user message (discard_unanswered_run_input -> MessageDeleted), so the prompt vanished from both the composer AND the conversation, forcing a retype. The discard's own doc even claimed the text "stays recoverable via the input history" -- a recovery path that never existed. Fix: capture the retracted text where the deletion lands. The store's removeMessage now stashes a removed *user* message's text in recoverablePrompts[sessionId]; the wrapper surfaces it to the composer, which restores it into the empty input box (during-render setState, converges -- mirrors the draft-swap pattern; lint-safe vs set-state-in-effect) and clears the store slot once consumed. Guarded by restoreFailedPrompt so it never clobbers text the user typed since. Image-bearing messages are unaffected: the backend preserves and auto-retries them, so no MessageDeleted fires for them. Tests: 4 store tests (capture trims; ignores assistant/empty; clear). tsc / eslint / vitest (132) / vite build clean. --- src/assistant/sessionStore.test.ts | 43 ++++++++++++++++++- src/assistant/sessionStore.ts | 27 ++++++++++++ .../TerminalEmulator/TerminalEmulator.tsx | 31 +++++++++++++ .../TerminalEmulatorWrapper.tsx | 14 ++++++ 4 files changed, 114 insertions(+), 1 deletion(-) diff --git a/src/assistant/sessionStore.test.ts b/src/assistant/sessionStore.test.ts index 4323360..721cfba 100644 --- a/src/assistant/sessionStore.test.ts +++ b/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(); + }); +}); diff --git a/src/assistant/sessionStore.ts b/src/assistant/sessionStore.ts index 79ee045..e2eca31 100644 --- a/src/assistant/sessionStore.ts +++ b/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; + clearRecoverablePrompt: (sessionId: string) => void; markMessageQueued: (sessionId: string, messageId: string) => void; markQueuedMessagesDelivered: (sessionId: string, messageIds: string[]) => void; prependMessagePage: ( @@ -124,6 +131,7 @@ const useAssistantStore = create()( immer((set, get) => ({ sessions: {}, activeSessionByTab: {}, + recoverablePrompts: {}, setActiveSessionForTab: (tabId, sessionId) => set((state) => { @@ -167,6 +175,20 @@ const useAssistantStore = create()( 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 => 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()( delete s.streamingTextByMessageId[messageId]; }), + clearRecoverablePrompt: (sessionId) => + set((state) => { + delete state.recoverablePrompts[sessionId]; + }), + markMessageQueued: (sessionId, messageId) => set((state) => { const s = state.sessions[sessionId]; diff --git a/src/components/TerminalEmulator/TerminalEmulator.tsx b/src/components/TerminalEmulator/TerminalEmulator.tsx index 2747a2d..f658ea9 100644 --- a/src/components/TerminalEmulator/TerminalEmulator.tsx +++ b/src/components/TerminalEmulator/TerminalEmulator.tsx @@ -41,6 +41,11 @@ interface TerminalEmulatorProps { onPickImage?: () => Promise; onReadClipboardImage?: () => Promise; 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 diff --git a/src/components/TerminalEmulator/TerminalEmulatorWrapper.tsx b/src/components/TerminalEmulator/TerminalEmulatorWrapper.tsx index 801ded1..41bea73 100644 --- a/src/components/TerminalEmulator/TerminalEmulatorWrapper.tsx +++ b/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} /> ); };