diff --git a/packages/ui/src/components/chat/ChatInput.tsx b/packages/ui/src/components/chat/ChatInput.tsx index d05b8b6c..fba574ef 100644 --- a/packages/ui/src/components/chat/ChatInput.tsx +++ b/packages/ui/src/components/chat/ChatInput.tsx @@ -952,7 +952,10 @@ const ChatInputComponent: React.FC = ({ 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 diff --git a/packages/ui/src/components/chat/composer/DOCUMENTATION.md b/packages/ui/src/components/chat/composer/DOCUMENTATION.md index 0e060e0e..3f6ee746 100644 --- a/packages/ui/src/components/chat/composer/DOCUMENTATION.md +++ b/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 diff --git a/packages/ui/src/components/chat/composer/state/__tests__/useComposerDraft.fork.test.ts b/packages/ui/src/components/chat/composer/state/__tests__/useComposerDraft.fork.test.ts new file mode 100644 index 00000000..46173940 --- /dev/null +++ b/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(); + 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(), 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(); + } + }); +}); diff --git a/packages/ui/src/components/chat/composer/state/useComposerDraft.ts b/packages/ui/src/components/chat/composer/state/useComposerDraft.ts index 513a3e7a..471178a6 100644 --- a/packages/ui/src/components/chat/composer/state/useComposerDraft.ts +++ b/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>(new Map()); const currentIdentityRef = React.useRef(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) => { diff --git a/packages/ui/src/sync/DOCUMENTATION.md b/packages/ui/src/sync/DOCUMENTATION.md index 9fe4f434..18cd4a90 100644 --- a/packages/ui/src/sync/DOCUMENTATION.md +++ b/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 | diff --git a/packages/ui/src/sync/input-store.test.ts b/packages/ui/src/sync/input-store.test.ts index 5303c2d5..ba7a9dc5 100644 --- a/packages/ui/src/sync/input-store.test.ts +++ b/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 diff --git a/packages/ui/src/sync/input-store.ts b/packages/ui/src/sync/input-store.ts index be2d25e2..87b0d45e 100644 --- a/packages/ui/src/sync/input-store.ts +++ b/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()((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, diff --git a/packages/ui/src/sync/session-actions.test.ts b/packages/ui/src/sync/session-actions.test.ts index c3080348..9d7d7ba0 100644 --- a/packages/ui/src/sync/session-actions.test.ts +++ b/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 }> = [] @@ -18,6 +19,10 @@ const failingRevertSessionIds = new Set() const failingUnrevertSessionIds = new Set() 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 => { + 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 = { + pendingComposerRestore: null, pendingInputText: "", - pendingInputMode: "normal" as const, + pendingInputMode: "replace", attachedFiles: [], clearAttachedFiles: () => { inputState.attachedFiles = [] }, - addRestoredAttachment: (attachment: never) => { - inputState.attachedFiles = [...inputState.attachedFiles, attachment] + addRestoredAttachment: (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: [], }) }) diff --git a/packages/ui/src/sync/session-actions.ts b/packages/ui/src/sync/session-actions.ts index c29e7ab1..837b7571 100644 --- a/packages/ui/src/sync/session-actions.ts +++ b/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 { * 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 { + 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).text as string) || ((p as Record).content as string) || "") .join("\n") .trim() - const fileParts = parts.filter((p) => p.type === "file" && !isSyntheticPart(p)) as Array> + 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 - if (messageText) { - useInputStore.setState({ - pendingInputText: messageText, - pendingInputMode: "replace" as const, - }) - } - // Clear existing attachments and restore file parts from the forked message. - restoreFilePartsToInput(fileParts) + // Navigation is deferred in the chat column. Leave the source composer alone + // until the rendered draft identity matches the fork, including for file-only prompts. + useInputStore.setState({ + pendingComposerRestore: { + target, + text: messageText, + files: fileParts.filter((part) => part.url).map((part) => ({ + url: part.url, + mimeType: part.mime, + filename: part.filename ?? "attachment", + })), + }, + }) // 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, sessionKey: forkedSession.id }) - } + restoreContextPartsToInput(parts, { directory: target.directory, sessionKey: forkedSession.id }) } export async function fetchMessagesForSession(sessionID: string, directory?: string | null): Promise {