fix(chat): restore fork prompts in the destination composer (#3387)

packages/ui/src/components/chat/ChatInput.tsx

@@ -952,7 +952,10 @@ const ChatInputComponent: React.FC<ChatInputProps> = ({
             identity: initialDraftIdentityRef.current,
         },
         onIdentityChange: () => setInputMode('normal'),
-        onDraftRestored: () => composerRef.current?.selectAll(),
+        onDraftRestored: (source) => {
+            if (source === 'fork') composerRef.current?.focus();
+            composerRef.current?.selectAll();
+        },
     });
 
     // Focus textarea when new session draft is opened

packages/ui/src/components/chat/composer/DOCUMENTATION.md

@@ -183,6 +183,11 @@ and the send path reading the same grammar.
   draft. Two orderings are load-bearing: the debounced write is skipped once
   while a draft is being restored, and a deleted draft's empty signature is
   recorded before a queued write could resurrect it.
+  Fork replay text and files arrive in `input-store.pendingComposerRestore`,
+  addressed to the fork's runtime, directory, and session. The hook consumes
+  them after loading that identity's draft. Selection alone is not enough:
+  the deferred chat column can still show the source composer. Ordinary
+  pending text insertions keep their existing path in `ChatInput`.
 - `state/useDraftTarget.ts` — the draft can target a directory that does not
   exist yet (a worktree being created). It must survive not appearing in the
   branch list, or the selector snaps back to the project root mid-creation. It

packages/ui/src/components/chat/composer/state/__tests__/useComposerDraft.fork.test.ts

@@ -0,0 +1,164 @@
+import { beforeEach, describe, expect, test } from 'bun:test';
+import React, { act } from 'react';
+import { createRoot } from 'react-dom/client';
+
+import { installHookTestDom } from '@/components/session/sidebar/test-utils/testDom';
+import { readChatDraft, writeChatDraft, type ChatDraftIdentity } from '@/lib/chatDraftPersistence';
+import { getDeferredSafeStorage } from '@/stores/utils/safeStorage';
+import { useInputStore } from '@/sync/input-store';
+import { useComposerDraft } from '../useComposerDraft';
+
+const source: ChatDraftIdentity = { runtimeKey: 'runtime-a', directory: '/repo', sessionId: 'source' };
+const fork: ChatDraftIdentity = { ...source, sessionId: 'fork' };
+const replayFile = { url: 'data:text/plain;base64,aGVsbG8=', mimeType: 'text/plain', filename: 'replay.txt' };
+
+function renderComposer(persistEnabled: boolean) {
+    const dom = installHookTestDom();
+    const originalRaf = globalThis.requestAnimationFrame;
+    const originalCancelRaf = globalThis.cancelAnimationFrame;
+    const frames = new Map<number, FrameRequestCallback>();
+    let frameId = 0;
+    globalThis.requestAnimationFrame = (callback) => {
+        frames.set(++frameId, callback);
+        return frameId;
+    };
+    globalThis.cancelAnimationFrame = (id) => { frames.delete(id); };
+    const root = createRoot(dom.container);
+    const restored: string[] = [];
+    const result = { text: '', mentions: new Set<string>(), restored };
+
+    function Probe({ identity }: { identity: ChatDraftIdentity }) {
+        const [message, setMessage] = React.useState('source draft @source.ts');
+        const messageRef = React.useRef(message);
+        const confirmedMentionsRef = React.useRef(new Set(['source.ts']));
+        React.useEffect(() => { messageRef.current = message; }, [message]);
+        useComposerDraft({
+            message, messageRef, setMessage, confirmedMentionsRef, identity, persistEnabled,
+            initialDraft: { text: '', identity: source },
+            onDraftRestored: (reason) => { result.restored.push(reason); },
+        });
+        result.text = message;
+        result.mentions = confirmedMentionsRef.current;
+        return null;
+    }
+
+    const render = (identity: ChatDraftIdentity) => {
+        act(() => { root.render(React.createElement(Probe, { identity })); });
+    };
+    render(source);
+    return {
+        result,
+        render,
+        flushFrames: () => {
+            const pending = [...frames.values()];
+            frames.clear();
+            act(() => { for (const callback of pending) callback(0); });
+        },
+        teardown: () => {
+            act(() => { root.unmount(); });
+            globalThis.requestAnimationFrame = originalRaf;
+            globalThis.cancelAnimationFrame = originalCancelRaf;
+            dom.restore();
+        },
+    };
+}
+
+beforeEach(() => {
+    getDeferredSafeStorage().removeItem('openchamber.chatDrafts.v2');
+    useInputStore.setState({ pendingComposerRestore: null });
+    useInputStore.getState().clearAttachedFiles();
+});
+
+describe('fork composer restoration', () => {
+    for (const persistEnabled of [true, false]) {
+        test(`waits for the rendered fork and preserves the source, persistence=${persistEnabled}`, () => {
+            writeChatDraft(fork, 'previous fork draft @old.ts', ['old.ts']);
+            useInputStore.getState().addRestoredAttachment({ ...replayFile, filename: 'source.txt' });
+            const sourceFiles = useInputStore.getState().attachedFiles;
+            const composer = renderComposer(persistEnabled);
+            try {
+                // Selection already changed, but the deferred chat column still renders source.
+                act(() => {
+                    useInputStore.setState({ pendingComposerRestore: { target: fork, text: 'replay prompt', files: [replayFile] } });
+                });
+                expect(composer.result.text).toBe('source draft @source.ts');
+                expect(useInputStore.getState().attachedFiles).toBe(sourceFiles);
+                expect(useInputStore.getState().pendingComposerRestore).not.toBeNull();
+
+                composer.render(fork);
+                expect(composer.result.text).toBe('replay prompt');
+                expect(composer.result.mentions.size).toBe(0);
+                expect(useInputStore.getState().attachedFiles.map((file) => file.filename)).toEqual(['replay.txt']);
+                expect(useInputStore.getState().pendingComposerRestore).toBeNull();
+                composer.flushFrames();
+                expect(composer.result.restored).toContain('fork');
+                expect(readChatDraft(source).text).toBe(persistEnabled ? 'source draft @source.ts' : '');
+
+                composer.render(source);
+                expect(composer.result.text).toBe(persistEnabled ? 'source draft @source.ts' : '');
+                expect(readChatDraft(fork).text).toBe(persistEnabled ? 'replay prompt' : '');
+            } finally {
+                composer.teardown();
+            }
+        });
+    }
+
+    test('waits through unrelated session, directory, and runtime renders', () => {
+        const composer = renderComposer(true);
+        try {
+            act(() => {
+                useInputStore.setState({ pendingComposerRestore: { target: fork, text: 'replay', files: [] } });
+            });
+            for (const identity of [
+                { ...fork, sessionId: 'other' },
+                { ...fork, directory: '/other' },
+                { ...fork, runtimeKey: 'runtime-b' },
+            ]) {
+                composer.render(identity);
+                expect(composer.result.text).toBe('');
+                expect(useInputStore.getState().pendingComposerRestore).not.toBeNull();
+            }
+            composer.render(fork);
+            expect(composer.result.text).toBe('replay');
+            expect(useInputStore.getState().pendingComposerRestore).toBeNull();
+        } finally {
+            composer.teardown();
+        }
+    });
+
+    test('persists a replay even when its text equals the outgoing source draft', async () => {
+        const composer = renderComposer(true);
+        try {
+            const text = composer.result.text;
+            act(() => {
+                useInputStore.setState({ pendingComposerRestore: { target: fork, text, files: [] } });
+            });
+            composer.render(fork);
+            await act(async () => { await new Promise((resolve) => setTimeout(resolve, 550)); });
+            expect(readChatDraft(fork)).toEqual({ text, confirmedMentions: new Set() });
+            expect(readChatDraft(source)).toEqual({ text, confirmedMentions: new Set(['source.ts']) });
+        } finally {
+            composer.teardown();
+        }
+    });
+
+    test('restores file-only and empty prompts without keeping destination text or files', () => {
+        const composer = renderComposer(true);
+        try {
+            writeChatDraft(fork, 'stale destination text', []);
+            act(() => {
+                useInputStore.setState({ pendingComposerRestore: { target: fork, text: '', files: [replayFile] } });
+            });
+            composer.render(fork);
+            expect(composer.result.text).toBe('');
+            expect(useInputStore.getState().attachedFiles.map((file) => file.filename)).toEqual(['replay.txt']);
+            act(() => {
+                useInputStore.setState({ pendingComposerRestore: { target: fork, text: '', files: [] } });
+            });
+            expect(composer.result.text).toBe('');
+            expect(useInputStore.getState().attachedFiles).toEqual([]);
+        } finally {
+            composer.teardown();
+        }
+    });
+});

packages/ui/src/components/chat/composer/state/useComposerDraft.ts

@@ -12,6 +12,7 @@
  */
 
 import React from 'react';
+import { useInputStore } from '@/sync/input-store';
 
 import {
     getChatDraftIdentityKey,
@@ -52,8 +53,8 @@ export interface ComposerDraftOptions {
     initialDraft: { text: string; identity: ChatDraftIdentity | null };
     /** Called when the composer switches to a different draft identity. */
     onIdentityChange?: () => void;
-    /** Called after a non-empty draft is restored, to select its text. */
-    onDraftRestored?: () => void;
+    /** Called after restoring a saved draft or fork replay, to select its text. */
+    onDraftRestored?: (source: 'saved' | 'fork') => void;
 }
 
 export interface ComposerDraftControls {
@@ -81,6 +82,7 @@ export function useComposerDraft(options: ComposerDraftOptions): ComposerDraftCo
     const skipNextPersistRef = React.useRef(false);
     const lastPersistedRef = React.useRef<Map<string, string>>(new Map());
     const currentIdentityRef = React.useRef<ChatDraftIdentity | null>(initialDraft.identity);
+    const pendingComposerRestore = useInputStore((state) => state.pendingComposerRestore);
 
     // Callbacks reach the effects through a ref so a caller passing inline
     // functions does not re-run the persistence effects on every render.
@@ -129,7 +131,7 @@ export function useComposerDraft(options: ComposerDraftOptions): ComposerDraftCo
             writeChatDraft(initialDraft.identity, '', []);
             return;
         }
-        requestAnimationFrame(() => callbacksRef.current.onDraftRestored?.());
+        requestAnimationFrame(() => callbacksRef.current.onDraftRestored?.('saved'));
         // Runs once; the initial draft is captured at mount by design.
         // eslint-disable-next-line react-hooks/exhaustive-deps
     }, [persistEnabled]);
@@ -160,10 +162,35 @@ export function useComposerDraft(options: ComposerDraftOptions): ComposerDraftCo
         setMessage(restored.text);
         confirmedMentionsRef.current = restored.confirmedMentions;
         if (restored.text) {
-            requestAnimationFrame(() => callbacksRef.current.onDraftRestored?.());
+            requestAnimationFrame(() => callbacksRef.current.onDraftRestored?.('saved'));
         }
     }, [clearPending, confirmedMentionsRef, identity, messageRef, persistEnabled, persistNow, setMessage]);
 
+    // The chat column can still show the source after navigation selects a fork.
+    // Apply its replay only after the destination's draft has been loaded above.
+    React.useEffect(() => {
+        if (!pendingComposerRestore) return;
+        const input = useInputStore.getState();
+        const pending = input.consumePendingComposerRestore(identity);
+        if (!pending) return;
+
+        clearPending();
+        skipNextPersistRef.current = true;
+        messageRef.current = pending.text;
+        confirmedMentionsRef.current = new Set();
+        setMessage(pending.text);
+        // Equal source/replay text need not trigger another render to persist.
+        if (persistEnabled) persistNow(pending.target, pending.text);
+        input.clearAttachedFiles();
+        for (const file of pending.files) input.addRestoredAttachment(file);
+        requestAnimationFrame(() => {
+            const current = currentIdentityRef.current;
+            if (current && getChatDraftIdentityKey(current) === getChatDraftIdentityKey(pending.target)) {
+                callbacksRef.current.onDraftRestored?.('fork');
+            }
+        });
+    }, [clearPending, confirmedMentionsRef, identity, messageRef, pendingComposerRestore, persistEnabled, persistNow, setMessage]);
+
     // A draft deleted elsewhere (session deleted, drafts cleared) clears the
     // composer if it is the one on screen.
     React.useEffect(() => subscribeChatDraftDeletion((deleted) => {