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. */
+    /** Called after restoring a saved draft or fork replay, to select its text. */
-    onDraftRestored?: () => void;
+    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) => {

packages/ui/src/sync/DOCUMENTATION.md

@@ -52,7 +52,7 @@ So:
 | `viewport-store.ts` | Scroll anchors, session memory, loading indicators | App UI state |
 | `attachment-files.ts` | Attachment picker allowlists, MIME/content validation, structured-text sanitization, and HEIC conversion | Local chat attachments across shared UI runtimes |
 | `document-attachments.ts` | Bounded Office/OpenDocument extraction, document text serialization, embedded-image extraction, and positional citations | DOCX, PPTX, XLSX, ODT, ODP, and ODS chat attachments |
-| `input-store.ts` | Draft input state, attached files, synthetic parts | App UI state |
+| `input-store.ts` | Draft input state, attached files, synthetic parts, destination-scoped fork replay handoff | App UI state; fork replay targets runtime + directory + session |
 | `selection-store.ts` | Model/agent/variant selections | App UI state |
 | `voice-store.ts` | Voice state | App UI state |
 

packages/ui/src/sync/input-store.test.ts

@@ -50,6 +50,38 @@ const waitForReaderCount = async (count: number) => {
 
 const pngBytes = new Uint8Array([0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A])
 
+describe("input-store composer restore", () => {
+  beforeEach(() => {
+    useInputStore.setState({ pendingComposerRestore: null })
+  })
+
+  test("only the destination can consume a restore, and only once", () => {
+    const target = { runtimeKey: "runtime", directory: "/repo", sessionId: "fork" }
+    const pending = { target, text: "replay", files: [] }
+    useInputStore.setState({ pendingComposerRestore: pending })
+    for (const identity of [
+      null,
+      { ...target, sessionId: "source" },
+      { ...target, directory: "/elsewhere" },
+      { ...target, runtimeKey: "other-runtime" },
+    ]) {
+      expect(useInputStore.getState().consumePendingComposerRestore(identity)).toBeNull()
+      expect(useInputStore.getState().pendingComposerRestore).toBe(pending)
+    }
+    expect(useInputStore.getState().consumePendingComposerRestore(target)).toBe(pending)
+    expect(useInputStore.getState().consumePendingComposerRestore(target)).toBeNull()
+  })
+
+  test("keeps ordinary pending text independent from fork restoration", () => {
+    const target = { runtimeKey: "runtime", directory: "/repo", sessionId: "fork" }
+    const pending = { target, text: "", files: [] }
+    useInputStore.setState({ pendingComposerRestore: pending })
+    useInputStore.getState().setPendingInputText("ordinary insertion", "append")
+    expect(useInputStore.getState().consumePendingInputText()).toEqual({ text: "ordinary insertion", mode: "append" })
+    expect(useInputStore.getState().consumePendingComposerRestore(target)).toBe(pending)
+  })
+})
+
 describe("input-store attachments", () => {
   beforeEach(() => {
     pendingReaders.length = 0

packages/ui/src/sync/input-store.ts

@@ -7,6 +7,7 @@ import { create } from "zustand"
 import type { ContextPartMetadata } from '@/lib/messages/contextParts'
 import type { AttachedFile } from "@/stores/types/sessionTypes"
 import { prepareAttachmentFiles } from "./attachment-files"
+import { getChatDraftIdentityKey, type ChatDraftIdentity } from "@/lib/chatDraftPersistence"
 
 const FILE_URI_PREFIX = "file://"
 const MAX_ATTACHMENT_PREPARATION_ATTEMPTS = 3
@@ -128,6 +129,12 @@ export type VSCodeActiveEditorFile = {
 }
 
 export type InputState = {
+  pendingComposerRestore: {
+    target: ChatDraftIdentity
+    text: string
+    files: Array<{ url: string; mimeType: string; filename: string }>
+  } | null
+  consumePendingComposerRestore: (target: ChatDraftIdentity | null) => InputState["pendingComposerRestore"]
   pendingInputText: string | null
   pendingInputMode: "replace" | "append" | "append-inline"
   pendingSyntheticParts: SyntheticContextPart[] | null
@@ -158,6 +165,13 @@ export type InputState = {
 }
 
 export const useInputStore = create<InputState>()((set, get) => ({
+  pendingComposerRestore: null,
+  consumePendingComposerRestore: (target) => {
+    const pending = get().pendingComposerRestore
+    if (!pending || !target || getChatDraftIdentityKey(pending.target) !== getChatDraftIdentityKey(target)) return null
+    set({ pendingComposerRestore: null })
+    return pending
+  },
   pendingInputText: null,
   pendingInputMode: "replace",
   pendingSyntheticParts: null,

packages/ui/src/sync/session-actions.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, test, beforeEach, mock } from "bun:test"
 import type { PermissionRequest } from "@/types/permission"
 import type { QuestionRequest } from "@/types/question"
+import type { InputState } from "./input-store"
 
 // Mock SDK client that records permission.reply / question.reply calls
 const replyCalls: Array<{ method: string; params: Record<string, unknown> }> = []
@@ -18,6 +19,10 @@ const failingRevertSessionIds = new Set<string>()
 const failingUnrevertSessionIds = new Set<string>()
 let afterUnrevertCall: ((sessionId: string) => void) | null = null
 let sessionDeleteError: unknown | null = null
+let sessionForkResult: Session | null = null
+let sessionForkError: Error | null = null
+let beforeSessionForkResolve: (() => void) | null = null
+const selectedSessions: Array<{ sessionId: string | null; directoryHint?: string | null }> = []
 let beforeSessionUpdateResolve: ((sessionId: string) => void) | null = null
 let beforeSessionDeleteResolve: ((sessionId: string) => void) | null = null
 let beforeControlPlaneMoveResolve: ((sessionId: string) => void) | null = null
@@ -179,6 +184,13 @@ mock.module("@/lib/opencode/client", () => ({
       replyCalls.push({ method: "session.messages", params: { sessionID: sessionId, directory } })
       return Promise.resolve(sessionMessageRecords.get(sessionId) ?? [])
     }),
+    forkSession: mock(async (sessionId: string, messageId?: string, directory?: string | null): Promise<Session> => {
+      replyCalls.push({ method: "session.fork", params: { sessionID: sessionId, messageID: messageId, directory } })
+      beforeSessionForkResolve?.()
+      if (sessionForkError) throw sessionForkError
+      if (!sessionForkResult) throw new Error("Missing fork session fixture")
+      return sessionForkResult
+    }),
     replyToPermission: mock((requestId: string, reply: string, options?: { directory?: string | null }) => {
       replyCalls.push({ method: "permission.reply", params: { requestID: requestId, reply, directory: options?.directory } })
       return Promise.resolve(true)
@@ -236,7 +248,9 @@ mock.module("./session-ui-store", () => ({
         return null
       },
       currentSessionId: null,
-      setCurrentSession: () => {},
+      setCurrentSession: (sessionId: string | null, directoryHint?: string | null) => {
+        selectedSessions.push({ sessionId, directoryHint })
+      },
       setWorktreeMetadata: () => {},
       setSessionDirectory: (sessionID: string, directory: string) => {
         movedSessionDirectories.push({ sessionID, directory })
@@ -246,15 +260,27 @@ mock.module("./session-ui-store", () => ({
 }))
 
 // Mock useInputStore
-const inputState = {
+const inputState: Pick<InputState,
+  "pendingComposerRestore" | "pendingInputText" | "pendingInputMode" | "attachedFiles"
+  | "clearAttachedFiles" | "addRestoredAttachment"
+> = {
+  pendingComposerRestore: null,
   pendingInputText: "",
-  pendingInputMode: "normal" as const,
+  pendingInputMode: "replace",
   attachedFiles: [],
   clearAttachedFiles: () => {
     inputState.attachedFiles = []
   },
-  addRestoredAttachment: (attachment: never) => {
+  addRestoredAttachment: (attachment) => {
-    inputState.attachedFiles = [...inputState.attachedFiles, attachment]
+    inputState.attachedFiles = [...inputState.attachedFiles, {
+      id: attachment.url,
+      file: new File([], attachment.filename, { type: attachment.mimeType }),
+      dataUrl: attachment.url,
+      mimeType: attachment.mimeType,
+      filename: attachment.filename,
+      size: 0,
+      source: "server",
+    }]
   },
 }
 
@@ -1935,6 +1961,197 @@ describe("respondToPermission passes directory", () => {
   })
 })
 
+describe("forkFromMessage composer restore", () => {
+  const sourceSession: Session = {
+    id: "session-a",
+    slug: "source-session",
+    projectID: "project-a",
+    directory: "/test/project",
+    title: "Source session",
+    version: "1",
+    time: { created: 1, updated: 1 },
+  }
+  const forkedSession: Session = { ...sourceSession, id: "session-fork", slug: "forked-session" }
+  const textPart: Part = {
+    id: "part-text",
+    sessionID: sourceSession.id,
+    messageID: "message-fork",
+    type: "text",
+    text: "Replay this prompt",
+  }
+  const filePart: Part = {
+    id: "part-file",
+    sessionID: sourceSession.id,
+    messageID: "message-fork",
+    type: "file",
+    url: "data:image/png;base64,aW1hZ2U=",
+    mime: "image/png",
+    filename: "screenshot.png",
+  }
+  const restoredFile = { url: filePart.url, mimeType: filePart.mime, filename: filePart.filename }
+
+  beforeEach(() => {
+    replyCalls.length = 0
+    selectedSessions.length = 0
+    runtimeKey = "fork-runtime"
+    sessionForkResult = forkedSession
+    sessionForkError = null
+    beforeSessionForkResolve = null
+    inputState.pendingComposerRestore = null
+    inputState.pendingInputText = "Keep the source draft"
+    inputState.pendingInputMode = "append"
+    inputState.attachedFiles = [{
+      id: "source-attachment",
+      file: new File(["source"], "source.txt", { type: "text/plain" }),
+      dataUrl: "data:text/plain;base64,c291cmNl",
+      mimeType: "text/plain",
+      filename: "source.txt",
+      size: 6,
+      source: "local",
+    }]
+  })
+
+  for (const directory of ["/test/project", "/canonical/project"]) {
+    test(`stages the replay for the returned session in ${directory} without changing the source composer`, async () => {
+      sessionForkResult = { ...forkedSession, directory }
+      const source = createStore({}, {
+        session: [sourceSession],
+        part: { "message-fork": [textPart, filePart] },
+      })
+      const sourceInput = { ...inputState }
+      const { forkFromMessage, setActionRefs } = await import("./session-actions")
+      setActionRefs(actionSdk, createChildStores([[sourceSession.directory, source]]), () => "/other/project")
+
+      await forkFromMessage(sourceSession.id, "message-fork")
+
+      expect(replyCalls).toEqual([{
+        method: "session.fork",
+        params: { sessionID: sourceSession.id, messageID: "message-fork", directory: sourceSession.directory },
+      }])
+      expect(inputState.pendingComposerRestore).toEqual({
+        target: { runtimeKey: "fork-runtime", directory, sessionId: forkedSession.id },
+        text: "Replay this prompt",
+        files: [restoredFile],
+      })
+      expect(inputState.pendingInputText).toBe(sourceInput.pendingInputText)
+      expect(inputState.pendingInputMode).toBe(sourceInput.pendingInputMode)
+      expect(inputState.attachedFiles).toBe(sourceInput.attachedFiles)
+      expect(inputState.attachedFiles).toHaveLength(1)
+      expect(selectedSessions).toEqual([{ sessionId: forkedSession.id, directoryHint: directory }])
+      expect(source.getState().session).toEqual([sourceSession, sessionForkResult])
+    })
+  }
+
+  test("uses the returned project worktree when the fork has no directory", async () => {
+    const forkWithProject: Session & { project: { worktree: string } } = {
+      ...forkedSession, directory: "", project: { worktree: "/canonical/worktree" },
+    }
+    sessionForkResult = forkWithProject
+    const source = createStore({}, { session: [sourceSession], part: { "message-fork": [textPart] } })
+    const { forkFromMessage, setActionRefs } = await import("./session-actions")
+    setActionRefs(actionSdk, createChildStores([[sourceSession.directory, source]]), () => sourceSession.directory)
+
+    await forkFromMessage(sourceSession.id, "message-fork")
+
+    expect(inputState.pendingComposerRestore?.target.directory).toBe("/canonical/worktree")
+    expect(selectedSessions).toEqual([{ sessionId: forkedSession.id, directoryHint: "/canonical/worktree" }])
+  })
+
+  test("stages a file-only prompt with empty text without replacing source attachments", async () => {
+    const source = createStore({}, {
+      session: [sourceSession],
+      part: { "message-fork": [filePart] },
+    })
+    const sourceInput = { ...inputState }
+    const { forkFromMessage, setActionRefs } = await import("./session-actions")
+    setActionRefs(actionSdk, createChildStores([[sourceSession.directory, source]]), () => sourceSession.directory)
+
+    await forkFromMessage(sourceSession.id, "message-fork")
+
+    expect(inputState.pendingComposerRestore).toEqual({
+      target: { runtimeKey: "fork-runtime", directory: sourceSession.directory, sessionId: forkedSession.id },
+      text: "",
+      files: [restoredFile],
+    })
+    expect(inputState.pendingInputText).toBe(sourceInput.pendingInputText)
+    expect(inputState.attachedFiles).toBe(sourceInput.attachedFiles)
+    expect(selectedSessions).toEqual([{ sessionId: forkedSession.id, directoryHint: sourceSession.directory }])
+  })
+
+  test("excludes synthetic text and files from the staged replay", async () => {
+    const syntheticFile: Part & { synthetic: boolean } = {
+      ...filePart,
+      id: "part-synthetic-file",
+      url: "file:///test/project/generated.txt",
+      mime: "text/plain",
+      filename: "generated.txt",
+      synthetic: true,
+    }
+    const source = createStore({}, {
+      session: [sourceSession],
+      part: { "message-fork": [
+        { ...textPart, id: "part-synthetic-text", text: "Generated file contents", synthetic: true },
+        textPart,
+        syntheticFile,
+        filePart,
+      ] },
+    })
+    const { forkFromMessage, setActionRefs } = await import("./session-actions")
+    setActionRefs(actionSdk, createChildStores([[sourceSession.directory, source]]), () => sourceSession.directory)
+
+    await forkFromMessage(sourceSession.id, "message-fork")
+
+    expect(inputState.pendingComposerRestore).toEqual({
+      target: { runtimeKey: "fork-runtime", directory: sourceSession.directory, sessionId: forkedSession.id },
+      text: "Replay this prompt",
+      files: [restoredFile],
+    })
+  })
+
+  test("leaves input, selection, and sessions unchanged when the fork fails", async () => {
+    sessionForkError = new Error("fork failed")
+    const source = createStore({}, {
+      session: [sourceSession],
+      part: { "message-fork": [textPart, filePart] },
+    })
+    const sourceState = source.getState()
+    const sourceInput = { ...inputState }
+    const { forkFromMessage, setActionRefs } = await import("./session-actions")
+    setActionRefs(actionSdk, createChildStores([[sourceSession.directory, source]]), () => sourceSession.directory)
+
+    await expect(forkFromMessage(sourceSession.id, "message-fork")).rejects.toThrow("fork failed")
+
+    expect(inputState).toEqual(sourceInput)
+    expect(inputState.attachedFiles).toBe(sourceInput.attachedFiles)
+    expect(selectedSessions).toEqual([])
+    expect(source.getState()).toBe(sourceState)
+  })
+
+  test("does not select, mutate, or stage a fork resolved after the runtime changes", async () => {
+    beforeSessionForkResolve = () => { runtimeKey = "other-runtime" }
+    const source = createStore({}, {
+      session: [sourceSession],
+      part: { "message-fork": [textPart, filePart] },
+    })
+    const sourceState = source.getState()
+    const sourceInput = { ...inputState }
+    const { forkFromMessage, setActionRefs } = await import("./session-actions")
+    setActionRefs(actionSdk, createChildStores([[sourceSession.directory, source]]), () => sourceSession.directory)
+
+    await forkFromMessage(sourceSession.id, "message-fork")
+
+    expect(replyCalls).toEqual([{
+      method: "session.fork",
+      params: { sessionID: sourceSession.id, messageID: "message-fork", directory: sourceSession.directory },
+    }])
+    expect(runtimeKey).toBe("other-runtime")
+    expect(inputState).toEqual(sourceInput)
+    expect(inputState.attachedFiles).toBe(sourceInput.attachedFiles)
+    expect(selectedSessions).toEqual([])
+    expect(source.getState()).toBe(sourceState)
+  })
+})
+
 describe("revertToMessage passes session directory", () => {
   beforeEach(() => {
     replyCalls.length = 0
@@ -1944,7 +2161,7 @@ describe("revertToMessage passes session directory", () => {
     failingRevertSessionIds.clear()
     Object.assign(inputState, {
       pendingInputText: "previous draft",
-      pendingInputMode: "normal" as const,
+      pendingInputMode: "replace",
       attachedFiles: [],
     })
   })

packages/ui/src/sync/session-actions.ts

@@ -3,7 +3,7 @@
  * Replaces the action methods from the old useSessionStore.
  */
 
-import type { OpencodeClient, Session, Message, Part } from "@opencode-ai/sdk/v2/client"
+import type { FilePart, OpencodeClient, Session, Message, Part } from "@opencode-ai/sdk/v2/client"
 import { Binary } from "./binary"
 import { useSessionUIStore } from "./session-ui-store"
 import { useInputStore } from "./input-store"
@@ -43,6 +43,7 @@ import { normalizePath } from "@/lib/pathNormalization"
 import { mergeMessages } from "./optimistic"
 import { messagesBefore, messagesFrom } from "./message-ordering"
 import { deleteChatDirectory } from "@/lib/chatDirectories"
+import { createChatDraftIdentity } from "@/lib/chatDraftPersistence"
 
 const MESSAGE_REFETCH_LIMIT = 100
 const SEND_CONFIRMATION_REFETCH_LIMIT = 30
@@ -2425,9 +2426,10 @@ export async function unrevertSession(sessionId: string): Promise<void> {
  * 1. Extract text from the message for input restoration
  * 2. Call the runtime fork endpoint
  * 3. Insert the new session into the child store (so sidebar updates immediately)
- * 4. Switch to new session and set pending input text
+ * 4. Switch to the new session and stage its composer replay
  */
 export async function forkFromMessage(sessionId: string, messageId: string): Promise<void> {
+  const expectedRuntimeKey = getRuntimeKey()
   const { store, directory } = dirStoreForSession(sessionId)
   const state = store.getState()
 
@@ -2442,9 +2444,12 @@ export async function forkFromMessage(sessionId: string, messageId: string): Pro
     .map((p: Part) => ((p as Record<string, unknown>).text as string) || ((p as Record<string, unknown>).content as string) || "")
     .join("\n")
     .trim()
-  const fileParts = parts.filter((p) => p.type === "file" && !isSyntheticPart(p)) as Array<Record<string, unknown>>
+  const fileParts = parts.filter((part): part is FilePart => part.type === "file" && !isSyntheticPart(part))
 
   const forkedSession = await opencodeClient.forkSession(sessionId, messageId, directory)
+  if (isStaleRuntime(expectedRuntimeKey)) return
+  const target = createChatDraftIdentity(expectedRuntimeKey, resolveSessionOwnedDirectory(forkedSession) ?? directory, forkedSession.id)
+  if (!target) throw new Error("Forked session has no composer directory")
 
   // Insert new session into child store so sidebar updates immediately
   const current = store.getState()
@@ -2456,22 +2461,24 @@ export async function forkFromMessage(sessionId: string, messageId: string): Pro
   }
 
   // Switch to new session
-  useSessionUIStore.getState().setCurrentSession(forkedSession.id)
+  useSessionUIStore.getState().setCurrentSession(forkedSession.id, target.directory)
 
-  // Restore forked message text and file attachments to input
+  // Navigation is deferred in the chat column. Leave the source composer alone
-  if (messageText) {
+  // until the rendered draft identity matches the fork, including for file-only prompts.
   useInputStore.setState({
-      pendingInputText: messageText,
+    pendingComposerRestore: {
-      pendingInputMode: "replace" as const,
+      target,
+      text: messageText,
+      files: fileParts.filter((part) => part.url).map((part) => ({
+        url: part.url,
+        mimeType: part.mime,
+        filename: part.filename ?? "attachment",
+      })),
+    },
   })
-  }
-  // Clear existing attachments and restore file parts from the forked message.
-  restoreFilePartsToInput(fileParts)
   // The forked session is a fresh draft target, so the attached context of the
   // forked message follows the text into its composer.
-  if (directory) {
+  restoreContextPartsToInput(parts, { directory: target.directory, sessionKey: forkedSession.id })
-    restoreContextPartsToInput(parts, { directory, sessionKey: forkedSession.id })
-  }
 }
 
 export async function fetchMessagesForSession(sessionID: string, directory?: string | null): Promise<void> {