diff --git a/apps/web/src/components/ChatView.logic.test.ts b/apps/web/src/components/ChatView.logic.test.ts index 426818dba7a2..320fdecb9d75 100644 --- a/apps/web/src/components/ChatView.logic.test.ts +++ b/apps/web/src/components/ChatView.logic.test.ts @@ -3,6 +3,7 @@ import { findRecordedWorktreeSetup, mergeDirectPreviewAnnotation, planComposerTurn, + restorePlanFollowUpComposer, snapshotComposerSend, withMessageContext, deriveAvailableRevertTurnCounts, @@ -1136,6 +1137,59 @@ describe("planComposerTurn", () => { }); }); +describe("restorePlanFollowUpComposer", () => { + // The draft is the failed follow-up's thread; the refs and cursor are whatever thread is on + // screen. + it.each([ + { + name: "puts the follow-up back while its thread is open", + before: { draft: "", refs: "", cursor: null }, + routeIsCurrent: true, + after: { draft: "Tighten step 2", refs: "Tighten step 2", cursor: "Tighten step 2" }, + }, + { + name: "keeps a draft typed while the follow-up was in flight", + before: { draft: "Next idea", refs: "Next idea", cursor: null }, + routeIsCurrent: true, + after: { draft: "Next idea", refs: "Next idea", cursor: null }, + }, + { + name: "leaves another thread's composer alone after navigating away", + before: { draft: "", refs: "Other thread", cursor: null }, + routeIsCurrent: false, + after: { draft: "Tighten step 2", refs: "Other thread", cursor: null }, + }, + ])("$name", ({ before, routeIsCurrent, after }) => { + const composer: { draft: string; refs: string; cursor: string | null } = { ...before }; + restorePlanFollowUpComposer({ + snapshot: { + prompt: "Tighten step 2", + terminalContexts: [], + reviewComments: [], + previewAnnotations: [], + threadContexts: [], + }, + draftHasUserContent: composer.draft.trim().length > 0, + routeIsCurrent, + writeComposerRefs: (snapshot) => { + composer.refs = snapshot.prompt; + }, + writePrompt: (prompt) => { + composer.draft = prompt; + }, + writeTerminalContexts: () => undefined, + writeReviewComments: () => undefined, + writePreviewAnnotations: () => undefined, + writeThreadContexts: () => undefined, + resetCursor: ({ prompt }) => { + composer.cursor = prompt; + }, + }); + + expect(composer).toEqual(after); + }); +}); + describe("withMessageContext", () => { const context = { records: [] } as unknown as NonNullable< Parameters[0]["context"] diff --git a/apps/web/src/components/ChatView.logic.ts b/apps/web/src/components/ChatView.logic.ts index 231135d9c607..9c51eb0e9805 100644 --- a/apps/web/src/components/ChatView.logic.ts +++ b/apps/web/src/components/ChatView.logic.ts @@ -1284,8 +1284,13 @@ export interface PlanFollowUpComposerSnapshot { readonly threadContexts: ReadonlyArray; } +// A draft typed while the send was in flight is newer and stays. Composer refs and the cursor +// belong to the thread on screen, so they follow only while the follow-up's route is current. export function restorePlanFollowUpComposer(input: { readonly snapshot: PlanFollowUpComposerSnapshot; + readonly draftHasUserContent: boolean; + readonly routeIsCurrent: boolean; + readonly writeComposerRefs: (snapshot: PlanFollowUpComposerSnapshot) => void; readonly writePrompt: (prompt: string) => void; readonly writeTerminalContexts: (contexts: ReadonlyArray) => void; readonly writeReviewComments: (comments: ReadonlyArray) => void; @@ -1297,11 +1302,14 @@ export function restorePlanFollowUpComposer(input: { detectTrigger: boolean; }) => void; }): void { + if (input.draftHasUserContent) return; + if (input.routeIsCurrent) input.writeComposerRefs(input.snapshot); input.writePrompt(input.snapshot.prompt); input.writeTerminalContexts(input.snapshot.terminalContexts); input.writeReviewComments(input.snapshot.reviewComments); input.writePreviewAnnotations(input.snapshot.previewAnnotations); input.writeThreadContexts(input.snapshot.threadContexts); + if (!input.routeIsCurrent) return; input.resetCursor({ cursor: collapseExpandedComposerCursor(input.snapshot.prompt, input.snapshot.prompt.length), prompt: input.snapshot.prompt, diff --git a/apps/web/src/components/chat/useComposerSend.ts b/apps/web/src/components/chat/useComposerSend.ts index 597cafb35e92..5b31e78d62fa 100644 --- a/apps/web/src/components/chat/useComposerSend.ts +++ b/apps/web/src/components/chat/useComposerSend.ts @@ -684,8 +684,6 @@ export function useComposerSend(input: ComposerSendInput) { interactionMode: followUp.interactionMode, }); if (!followUpSent) { - promptRef.current = followUpPromptSnapshot; - composerTerminalContextsRef.current = [...followUpTerminalContexts]; restorePlanFollowUpComposer({ snapshot: { prompt: followUpPromptSnapshot, @@ -694,6 +692,14 @@ export function useComposerSend(input: ComposerSendInput) { previewAnnotations: followUpPreviewAnnotations, threadContexts: followUpThreadContexts, }, + draftHasUserContent: composerDraftHasUserContent( + useComposerDraftStore.getState().getComposerDraft(composerDraftTarget), + ), + routeIsCurrent: currentRouteThreadKeyRef.current === routeThreadKey, + writeComposerRefs: (snapshot) => { + promptRef.current = snapshot.prompt; + composerTerminalContextsRef.current = [...snapshot.terminalContexts]; + }, writePrompt: (prompt) => setComposerDraftPrompt(composerDraftTarget, prompt), writeTerminalContexts: (contexts) => setComposerDraftTerminalContexts(composerDraftTarget, [...contexts]),