From 9788797ee06c9e058b978c1f5a3ec75b22ff7388 Mon Sep 17 00:00:00 2001 From: Alex Knight Date: Sat, 6 Jun 2026 21:29:22 -0700 Subject: [PATCH] Add fork-from-message for branching conversations Adds a Codex-style fork affordance on assistant messages. Hovering an assistant response reveals a subtle Fork control below it; clicking opens a dialog to branch the conversation into a new sidebar thread containing history up to and including that response, with an empty composer. The dialog offers a worktree choice (same worktree or a new worktree), mirroring the new-thread environment toggle and building on the existing tree/branch infrastructure. - session-driver: ForkPosition ("before"|"at"|"after"), ForkSessionOptions/Result - pi-sdk-driver: forkSession using native getBranch + createBranchedSession - electron: forkThread IPC wiring (additive; no config/userData changes) - renderer: fork button (timeline), fork-modal, App wiring, styles - tests: core fork-from-message spec; on-demand demo spec (ignored by default) --- apps/desktop/electron/app-store-worktree.ts | 60 +++++++- apps/desktop/electron/app-store.ts | 5 + apps/desktop/electron/main.ts | 2 + apps/desktop/electron/preload.ts | 3 + apps/desktop/playwright.config.ts | 2 + apps/desktop/src/App.tsx | 91 ++++++++++++ apps/desktop/src/conversation-timeline.tsx | 39 ++++- apps/desktop/src/desktop-state.ts | 14 ++ apps/desktop/src/fork-modal.tsx | 139 ++++++++++++++++++ apps/desktop/src/icons.tsx | 12 ++ apps/desktop/src/ipc.ts | 3 + apps/desktop/src/styles/main.css | 70 +++++++++ apps/desktop/src/timeline-item.tsx | 40 ++++- .../tests/core/fork-from-message.spec.ts | 79 ++++++++++ apps/desktop/tests/demo/fork-demo.spec.ts | 106 +++++++++++++ apps/desktop/tests/helpers/electron-app.ts | 61 ++++++++ packages/pi-sdk-driver/src/pi-sdk-driver.ts | 6 + .../pi-sdk-driver/src/session-supervisor.ts | 117 +++++++++++++++ .../src/vendor/session-driver.d.ts | 15 ++ packages/session-driver/src/index.ts | 3 + packages/session-driver/src/types.ts | 23 +++ 21 files changed, 885 insertions(+), 5 deletions(-) create mode 100644 apps/desktop/src/fork-modal.tsx create mode 100644 apps/desktop/tests/core/fork-from-message.spec.ts create mode 100644 apps/desktop/tests/demo/fork-demo.spec.ts diff --git a/apps/desktop/electron/app-store-worktree.ts b/apps/desktop/electron/app-store-worktree.ts index 1490e290..69bbe20b 100644 --- a/apps/desktop/electron/app-store-worktree.ts +++ b/apps/desktop/electron/app-store-worktree.ts @@ -4,7 +4,13 @@ import { homedir } from "node:os"; import { sessionKey } from "@pi-gui/pi-sdk-driver"; import type { WorktreeCatalogEntry } from "@pi-gui/catalogs"; import type { WorkspaceRef } from "@pi-gui/session-driver"; -import type { CreateWorktreeInput, DesktopAppState, RemoveWorktreeInput, StartThreadInput } from "../src/desktop-state"; +import type { + CreateWorktreeInput, + DesktopAppState, + ForkThreadInput, + RemoveWorktreeInput, + StartThreadInput, +} from "../src/desktop-state"; import { sendMessageToSession } from "./app-store-composer"; import type { CreateWorktreeOptions } from "./worktree-manager"; import type { AppStoreInternals } from "./app-store-internals"; @@ -157,6 +163,58 @@ export async function startThread(store: AppStoreInternals, input: StartThreadIn }); } +export async function forkThread(store: AppStoreInternals, input: ForkThreadInput): Promise { + await store.initialize(); + const sourceWorkspace = store.workspaceRefFromState(input.sourceWorkspaceId); + if (!sourceWorkspace) { + return store.withError(`Unknown workspace: ${input.sourceWorkspaceId}`); + } + + return store.withErrorHandling(async () => { + let targetWorkspace = sourceWorkspace; + if (input.environment === "worktree") { + const rootWorkspace = store.workspaceRefFromState(input.rootWorkspaceId) ?? sourceWorkspace; + const worktreeOptions = buildWorktreeOptions( + store, + rootWorkspace, + input.sourceWorkspaceId, + input.sourceSessionId, + ); + const created = await store.worktreeManager.createWorktree(rootWorkspace, worktreeOptions); + const synced = await store.driver.syncWorkspace(created.path, created.displayName); + targetWorkspace = synced.workspace; + } + + const sourceRef = { workspaceId: input.sourceWorkspaceId, sessionId: input.sourceSessionId }; + const { snapshot: session, selectedText } = await store.driver.forkSession(sourceRef, { + targetWorkspace, + userMessageIndex: input.userMessageIndex, + ...(input.position ? { position: input.position } : {}), + }); + store.updateSessionConfig(session.ref, session.config); + + // Set selection eagerly so subscription replay events read the new session ID. + store.state = { + ...store.state, + selectedWorkspaceId: session.ref.workspaceId, + selectedSessionId: session.ref.sessionId, + }; + + // Load the branched history transcript from the driver before publishing state. + await store.reloadTranscriptFromDriver(session.ref); + + return store.refreshState({ + selectedWorkspaceId: session.ref.workspaceId, + selectedSessionId: session.ref.sessionId, + composerDraft: selectedText ?? "", + composerDraftSyncSource: "selection", + clearLastError: true, + refreshWorktrees: input.environment === "worktree", + activeView: "threads", + }); + }); +} + export async function syncAndListWorktrees( store: AppStoreInternals, workspaces: readonly { diff --git a/apps/desktop/electron/app-store.ts b/apps/desktop/electron/app-store.ts index 16818f0f..b50a4432 100644 --- a/apps/desktop/electron/app-store.ts +++ b/apps/desktop/electron/app-store.ts @@ -43,6 +43,7 @@ import { type CreateSessionInput, type CreateWorktreeInput, type DesktopAppState, + type ForkThreadInput, type NotificationPreferences, type QueuedComposerMessage, type RemoveWorktreeInput, @@ -335,6 +336,10 @@ export class DesktopAppStore implements AppStoreInternals { return worktree.createWorktree(this, input); } + async forkThread(input: ForkThreadInput): Promise { + return worktree.forkThread(this, input); + } + async removeWorktree(input: RemoveWorktreeInput): Promise { return worktree.removeWorktree(this, input); } diff --git a/apps/desktop/electron/main.ts b/apps/desktop/electron/main.ts index dbbcf24b..92c34524 100644 --- a/apps/desktop/electron/main.ts +++ b/apps/desktop/electron/main.ts @@ -34,6 +34,7 @@ import type { ComposerImageAttachment, CreateSessionInput, CreateWorktreeInput, + ForkThreadInput, RemoveWorktreeInput, StartThreadInput, WorkspaceSessionTarget, @@ -631,6 +632,7 @@ app.whenReady().then(async () => { store.createSession(input), ); ipcMain.handle(desktopIpc.startThread, (_event, input: StartThreadInput) => store.startThread(input)); + ipcMain.handle(desktopIpc.forkThread, (_event, input: ForkThreadInput) => store.forkThread(input)); ipcMain.handle(desktopIpc.openSkillInFinder, async (_event, workspaceId: string, filePath: string) => { const resolved = store.getSkillFilePath(workspaceId, filePath); if (!resolved) { diff --git a/apps/desktop/electron/preload.ts b/apps/desktop/electron/preload.ts index 7af681bc..048372b4 100644 --- a/apps/desktop/electron/preload.ts +++ b/apps/desktop/electron/preload.ts @@ -26,6 +26,7 @@ import type { CreateSessionInput, CreateWorktreeInput, DesktopAppState, + ForkThreadInput, NotificationPreferences, RemoveWorktreeInput, SelectedTranscriptRecord, @@ -145,6 +146,8 @@ contextBridge.exposeInMainWorld("piApp", { ipcRenderer.invoke(desktopIpc.createSession, input) as Promise, startThread: (input: StartThreadInput) => ipcRenderer.invoke(desktopIpc.startThread, input) as Promise, + forkThread: (input: ForkThreadInput) => + ipcRenderer.invoke(desktopIpc.forkThread, input) as Promise, cancelCurrentRun: () => ipcRenderer.invoke(desktopIpc.cancelCurrentRun) as Promise, setActiveView: (view: AppView) => ipcRenderer.invoke(desktopIpc.setActiveView, view) as Promise, diff --git a/apps/desktop/playwright.config.ts b/apps/desktop/playwright.config.ts index 9993e06a..24b1b1ca 100644 --- a/apps/desktop/playwright.config.ts +++ b/apps/desktop/playwright.config.ts @@ -2,6 +2,8 @@ import { defineConfig } from "@playwright/test"; export default defineConfig({ testDir: "./tests", + // Demo specs record marketing videos on demand; keep them out of default/CI discovery. + testIgnore: "**/demo/**", timeout: 60_000, // Electron user-surface tests are materially more reliable when one app owns the input loop at a time. workers: 1, diff --git a/apps/desktop/src/App.tsx b/apps/desktop/src/App.tsx index 7343456f..262546a2 100644 --- a/apps/desktop/src/App.tsx +++ b/apps/desktop/src/App.tsx @@ -8,6 +8,7 @@ import { type ComposerAttachment, type ComposerImageAttachment, type DesktopAppState, + type ForkThreadInput, type NewThreadEnvironment, type SelectedTranscriptRecord, type StartThreadInput, @@ -44,6 +45,7 @@ import { useThreadSearch } from "./hooks/use-thread-search"; import { useWorkspaceMenu } from "./hooks/use-workspace-menu"; import { buildExtensionDockModel, ExtensionDialog, hasExtensionDockContent } from "./extension-session-ui"; import { TreeModal } from "./tree-modal"; +import { ForkModal } from "./fork-modal"; import { getEffectiveModelRuntime } from "./model-settings"; import { resolveRepoWorkspaceId } from "./workspace-roots"; import { @@ -181,6 +183,17 @@ export default function App() { loading: false, submitting: false, }); + const [forkModalState, setForkModalState] = useState<{ + readonly open: boolean; + readonly submitting: boolean; + readonly userMessageIndex: number; + readonly messagePreview?: string; + readonly error?: string; + }>({ + open: false, + submitting: false, + userMessageIndex: 0, + }); const composerRef = useRef(null); const newThreadComposerRef = useRef(null); const timelinePaneRef = useRef(null); @@ -756,6 +769,73 @@ export default function App() { [api, selectedSession, selectedWorkspace], ); + const closeForkModal = useCallback(() => { + setForkModalState((current) => + current.submitting + ? current + : { + open: false, + submitting: false, + userMessageIndex: 0, + }, + ); + }, []); + + const openForkModal = useCallback( + (turnIndex: number, preview?: string) => { + if (!api || !selectedWorkspace || !selectedSession) { + return; + } + const trimmed = preview?.trim(); + setForkModalState({ + open: true, + submitting: false, + userMessageIndex: turnIndex, + messagePreview: trimmed ? trimmed.slice(0, 280) : undefined, + }); + }, + [api, selectedSession, selectedWorkspace], + ); + + const handleForkSubmit = useCallback( + (environment: NewThreadEnvironment) => { + if (!api || !selectedWorkspace || !selectedSession) { + return; + } + const rootWorkspaceId = + (snapshot ? resolveRepoWorkspaceId(snapshot.workspaces, selectedWorkspace.id) : undefined) ?? + selectedWorkspace.id; + const input: ForkThreadInput = { + sourceWorkspaceId: selectedWorkspace.id, + sourceSessionId: selectedSession.id, + rootWorkspaceId, + environment, + userMessageIndex: forkModalState.userMessageIndex, + position: "after", + }; + setForkModalState((current) => ({ ...current, submitting: true, error: undefined })); + void api + .forkThread(input) + .then((state) => { + setSnapshot(state); + setForkModalState({ + open: false, + submitting: false, + userMessageIndex: 0, + }); + focusComposer(); + }) + .catch((error) => { + setForkModalState((current) => ({ + ...current, + submitting: false, + error: error instanceof Error ? error.message : String(error), + })); + }); + }, + [api, forkModalState.userMessageIndex, selectedSession, selectedWorkspace, snapshot], + ); + const slashMenu = useSlashMenu({ composerDraft, setComposerDraft, @@ -2142,6 +2222,7 @@ export default function App() { onJumpToLatest={jumpToLatest} onContentHeightChange={handleTimelineContentHeightChange} onViewFileInDiff={handleViewFileInDiff} + onForkFromMessage={openForkModal} /> @@ -2213,6 +2294,16 @@ export default function App() { onNavigate={navigateTreeSelection} /> ) : null} + {forkModalState.open ? ( + + ) : null} ) : selectedWorkspace ? (
diff --git a/apps/desktop/src/conversation-timeline.tsx b/apps/desktop/src/conversation-timeline.tsx index 450ae7aa..93b53428 100644 --- a/apps/desktop/src/conversation-timeline.tsx +++ b/apps/desktop/src/conversation-timeline.tsx @@ -1,4 +1,4 @@ -import { useCallback, useLayoutEffect, useRef, useState, type MutableRefObject, type RefCallback, type RefObject } from "react"; +import { useCallback, useLayoutEffect, useMemo, useRef, useState, type MutableRefObject, type RefCallback, type RefObject } from "react"; import type { TranscriptMessage } from "./desktop-state"; import { ThreadSearchBar } from "./thread-search"; import { TimelineItem } from "./timeline-item"; @@ -31,6 +31,7 @@ interface ConversationTimelineProps { readonly onJumpToLatest: () => void; readonly onContentHeightChange: () => void; readonly onViewFileInDiff?: (path: string) => void; + readonly onForkFromMessage?: (turnIndex: number, preview?: string) => void; } export function ConversationTimeline({ @@ -46,7 +47,27 @@ export function ConversationTimeline({ onJumpToLatest, onContentHeightChange, onViewFileInDiff, + onForkFromMessage, }: ConversationTimelineProps) { + // Map each assistant message id to the 0-based turn index it belongs to (the ordinal + // of the user message that started the turn). The fork affordance lives on assistant + // messages (codex-style); forking branches the conversation after that turn so the new + // thread keeps the full history up to and including the assistant response. + const forkTurnIndexByAssistantId = useMemo(() => { + const map = new Map(); + let turnIndex = -1; + for (const item of transcript) { + if (item.kind !== "message") { + continue; + } + if (item.role === "user") { + turnIndex += 1; + } else if (item.role === "assistant" && turnIndex >= 0) { + map.set(item.id, turnIndex); + } + } + return map; + }, [transcript]); // Giant prose blocks and attachment-heavy rows routinely blow past the estimator, // so keep those transcripts on the exact DOM path instead of restoring to a fake bottom. const hasUnreliableVirtualizedHeights = transcript.some( @@ -173,6 +194,8 @@ export function ConversationTimeline({ onHeightChange={updateMeasuredHeight} onToggleToolCall={toggleToolCall} onViewFileInDiff={onViewFileInDiff} + forkTurnIndexByAssistantId={forkTurnIndexByAssistantId} + onForkFromMessage={onForkFromMessage} /> ) : (
@@ -184,6 +207,8 @@ export function ConversationTimeline({ expandedToolCallIds={expandedToolCallIds} onToggleToolCall={toggleToolCall} onViewFileInDiff={onViewFileInDiff} + forkTurnIndex={forkTurnIndexByAssistantId.get(item.id)} + onForkFromMessage={onForkFromMessage} /> ))}
@@ -207,6 +232,8 @@ function VirtualizedTranscriptList({ onHeightChange, onToggleToolCall, onViewFileInDiff, + forkTurnIndexByAssistantId, + onForkFromMessage, }: { readonly transcript: readonly TranscriptMessage[]; readonly timelinePaneRef: MutableRefObject; @@ -217,6 +244,8 @@ function VirtualizedTranscriptList({ readonly onHeightChange: (id: string, height: number) => void; readonly onToggleToolCall: (callId: string) => void; readonly onViewFileInDiff?: (path: string) => void; + readonly forkTurnIndexByAssistantId: ReadonlyMap; + readonly onForkFromMessage?: (turnIndex: number, preview?: string) => void; }) { const [viewport, setViewport] = useState({ scrollTop: 0, height: 0 }); const previousTotalHeightRef = useRef(0); @@ -289,6 +318,8 @@ function VirtualizedTranscriptList({ expandedToolCallIds={expandedToolCallIds} onToggleToolCall={onToggleToolCall} onViewFileInDiff={onViewFileInDiff} + forkTurnIndex={forkTurnIndexByAssistantId.get(item.id)} + onForkFromMessage={onForkFromMessage} /> ); })} @@ -304,6 +335,8 @@ function MeasuredTimelineItem({ expandedToolCallIds, onToggleToolCall, onViewFileInDiff, + forkTurnIndex, + onForkFromMessage, }: { readonly item: TranscriptMessage; readonly className?: string; @@ -312,6 +345,8 @@ function MeasuredTimelineItem({ readonly expandedToolCallIds: ReadonlySet; readonly onToggleToolCall: (callId: string) => void; readonly onViewFileInDiff?: (path: string) => void; + readonly forkTurnIndex?: number; + readonly onForkFromMessage?: (turnIndex: number, preview?: string) => void; }) { const rowRef = useRef(null); @@ -347,6 +382,8 @@ function MeasuredTimelineItem({ expandedToolCallIds={expandedToolCallIds} onToggleToolCall={onToggleToolCall} onViewFileInDiff={onViewFileInDiff} + forkTurnIndex={forkTurnIndex} + onForkFromMessage={onForkFromMessage} /> ); diff --git a/apps/desktop/src/desktop-state.ts b/apps/desktop/src/desktop-state.ts index 45186685..83888a1a 100644 --- a/apps/desktop/src/desktop-state.ts +++ b/apps/desktop/src/desktop-state.ts @@ -146,6 +146,20 @@ export type StartThreadInput = { readonly thinkingLevel?: string; }; +export type ForkThreadPosition = "before" | "at" | "after"; + +export type ForkThreadInput = { + readonly sourceWorkspaceId: string; + readonly sourceSessionId: string; + /** Root workspace used as the base when forking into a new worktree. */ + readonly rootWorkspaceId: string; + /** "local" forks into the source workspace; "worktree" creates a new worktree. */ + readonly environment: NewThreadEnvironment; + /** 0-based index of the user message to fork from (counted in the rendered transcript). */ + readonly userMessageIndex: number; + readonly position?: ForkThreadPosition; +}; + export interface RemoveWorktreeInput { readonly workspaceId: string; readonly worktreeId: string; diff --git a/apps/desktop/src/fork-modal.tsx b/apps/desktop/src/fork-modal.tsx new file mode 100644 index 00000000..c6a3bbbd --- /dev/null +++ b/apps/desktop/src/fork-modal.tsx @@ -0,0 +1,139 @@ +import { useEffect, useRef, useState, type KeyboardEvent as ReactKeyboardEvent } from "react"; +import type { NewThreadEnvironment } from "./desktop-state"; + +interface ForkModalProps { + readonly submitting: boolean; + readonly error?: string; + /** Preview of the assistant response the fork will branch after. */ + readonly messagePreview?: string; + /** Whether forking into a new worktree is available for the source workspace. */ + readonly canUseWorktree: boolean; + readonly onClose: () => void; + readonly onSubmit: (environment: NewThreadEnvironment) => void; +} + +export function ForkModal({ + submitting, + error, + messagePreview, + canUseWorktree, + onClose, + onSubmit, +}: ForkModalProps) { + const [environment, setEnvironment] = useState("local"); + const dialogRef = useRef(null); + + useEffect(() => { + dialogRef.current?.querySelector("[data-fork-confirm='true']")?.focus(); + }, []); + + const handleKeyDown = (event: ReactKeyboardEvent) => { + if (event.key === "Escape" && !submitting) { + event.preventDefault(); + onClose(); + } + }; + + return ( +
{ + if (event.target !== event.currentTarget || submitting) { + return; + } + onClose(); + }} + > +
+
+
+
Fork conversation
+

Start a new thread

+
+ +
+ + {error ? ( +
+ {error} +
+ ) : null} + +
+
+ Forks the conversation up to and including this response into a new sidebar thread with an empty + composer, so you can continue it in a different direction. The original thread stays untouched. +
+ + {messagePreview ? ( +
+ {messagePreview} +
+ ) : null} + +
+ + +
+ +
+
+ {environment === "worktree" + ? "A fresh worktree is created and the forked thread opens there." + : "The forked thread opens in the same folder as the original."} +
+
+ + +
+
+
+
+
+ ); +} diff --git a/apps/desktop/src/icons.tsx b/apps/desktop/src/icons.tsx index 65e1baf2..c74118cc 100644 --- a/apps/desktop/src/icons.tsx +++ b/apps/desktop/src/icons.tsx @@ -322,6 +322,18 @@ export function WorktreeIcon() { ); } +export function ForkIcon() { + return ( + + + + + + + + ); +} + export function GripIcon() { return ( diff --git a/apps/desktop/src/ipc.ts b/apps/desktop/src/ipc.ts index 233b3f7e..c7afee79 100644 --- a/apps/desktop/src/ipc.ts +++ b/apps/desktop/src/ipc.ts @@ -11,6 +11,7 @@ import type { CreateSessionInput, CreateWorktreeInput, DesktopAppState, + ForkThreadInput, ModelSettingsScopeMode, NotificationPreferences, RemoveWorktreeInput, @@ -51,6 +52,7 @@ export const desktopIpc = { unarchiveSession: "pi-gui:unarchive-session", createSession: "pi-gui:create-session", startThread: "pi-gui:start-thread", + forkThread: "pi-gui:fork-thread", cancelCurrentRun: "pi-gui:cancel-current-run", setActiveView: "pi-gui:set-active-view", setSidebarCollapsed: "pi-gui:set-sidebar-collapsed", @@ -236,6 +238,7 @@ export interface PiDesktopApi { unarchiveSession(target: WorkspaceSessionTarget): Promise; createSession(input: CreateSessionInput): Promise; startThread(input: StartThreadInput): Promise; + forkThread(input: ForkThreadInput): Promise; cancelCurrentRun(): Promise; setActiveView(view: AppView): Promise; setSidebarCollapsed(collapsed: boolean): Promise; diff --git a/apps/desktop/src/styles/main.css b/apps/desktop/src/styles/main.css index 67c6ccba..f59e1662 100644 --- a/apps/desktop/src/styles/main.css +++ b/apps/desktop/src/styles/main.css @@ -364,6 +364,10 @@ .timeline-item--assistant { max-width: 760px; + display: flex; + flex-direction: column; + align-items: flex-start; + gap: 2px; } .timeline-item--summary-card { @@ -385,7 +389,54 @@ .timeline-item--user { display: flex; + align-items: center; justify-content: flex-end; + gap: 8px; +} + +.timeline-item__actions { + display: flex; + align-items: center; + gap: 4px; + margin-top: 2px; + opacity: 0; + transition: opacity 0.12s ease; +} + +.timeline-item--assistant:hover .timeline-item__actions, +.timeline-item--user:hover .timeline-item__actions, +.timeline-item__actions:focus-within { + opacity: 1; +} + +.timeline-item__action { + display: inline-flex; + align-items: center; + gap: 5px; + height: 24px; + padding: 0 8px; + border-radius: 7px; + color: var(--muted-strong); + border: 1px solid transparent; + background: transparent; + font-size: 12px; + line-height: 1; + cursor: pointer; +} + +.timeline-item__action svg { + width: 13px; + height: 13px; +} + +.timeline-item__action-label { + font-weight: 500; +} + +.timeline-item__action:hover { + background: var(--surface-muted); + border-color: var(--line); + color: var(--ink); } .timeline-item__bubble { @@ -2717,6 +2768,25 @@ overflow: hidden; } +.tree-modal--compact { + width: min(520px, 100%); + gap: 14px; +} + +.fork-modal__preview { + padding: 12px 14px; + border-radius: 14px; + border: 1px solid var(--line); + background: var(--surface-muted); + color: var(--text-strong); + font-size: 13px; + line-height: 1.5; + max-height: 140px; + overflow: auto; + white-space: pre-wrap; + word-break: break-word; +} + .tree-modal__header { display: flex; align-items: flex-start; diff --git a/apps/desktop/src/timeline-item.tsx b/apps/desktop/src/timeline-item.tsx index 954414cd..d75f2e46 100644 --- a/apps/desktop/src/timeline-item.tsx +++ b/apps/desktop/src/timeline-item.tsx @@ -2,7 +2,7 @@ import type { SessionTranscriptMessage } from "@pi-gui/pi-sdk-driver"; import type { TimelineActivity, TimelineToolCall, TimelineSummary, TranscriptMessage } from "./timeline-types"; import { MessageMarkdown } from "./message-markdown"; import { InlineDiff, extractDiffFromOutput } from "./diff-inline"; -import { ChevronRightIcon, CopyIcon, DiffIcon, FileIcon } from "./icons"; +import { ChevronRightIcon, CopyIcon, DiffIcon, FileIcon, ForkIcon } from "./icons"; import { extensionToLanguage } from "./syntax-highlight"; export function TimelineItem({ @@ -10,15 +10,25 @@ export function TimelineItem({ expandedToolCallIds, onToggleToolCall, onViewFileInDiff, + forkTurnIndex, + onForkFromMessage, }: { readonly item: TranscriptMessage; readonly expandedToolCallIds?: ReadonlySet; readonly onToggleToolCall?: (callId: string) => void; readonly onViewFileInDiff?: (path: string) => void; + readonly forkTurnIndex?: number; + readonly onForkFromMessage?: (turnIndex: number, preview?: string) => void; }) { switch (item.kind) { case "message": - return ; + return ( + + ); case "activity": return ; case "tool": @@ -37,7 +47,15 @@ export function TimelineItem({ } } -function TimelineMessage({ item }: { readonly item: SessionTranscriptMessage }) { +function TimelineMessage({ + item, + forkTurnIndex, + onForkFromMessage, +}: { + readonly item: SessionTranscriptMessage; + readonly forkTurnIndex?: number; + readonly onForkFromMessage?: (turnIndex: number, preview?: string) => void; +}) { if (item.role === "user") { return (
@@ -84,9 +102,25 @@ function TimelineMessage({ item }: { readonly item: SessionTranscriptMessage }) ); } + const canFork = onForkFromMessage != null && forkTurnIndex != null; return (
+ {canFork ? ( +
+ +
+ ) : null}
); } diff --git a/apps/desktop/tests/core/fork-from-message.spec.ts b/apps/desktop/tests/core/fork-from-message.spec.ts new file mode 100644 index 00000000..fef4db7c --- /dev/null +++ b/apps/desktop/tests/core/fork-from-message.spec.ts @@ -0,0 +1,79 @@ +import { expect, test } from "@playwright/test"; +import { join } from "node:path"; +import { + getDesktopState, + launchDesktop, + makeUserDataDir, + makeWorkspace, + seedAgentDir, + seedForkSessionFixture, + selectSession, + waitForWorkspaceByPath, +} from "../helpers/electron-app"; + +test("forks a thread from an assistant response into a new sidebar session", async () => { + test.setTimeout(90_000); + const userDataDir = await makeUserDataDir(); + const agentDir = join(userDataDir, "agent"); + const workspacePath = await makeWorkspace("fork-from-message-workspace"); + await seedAgentDir(agentDir); + await seedForkSessionFixture(agentDir, workspacePath); + + const harness = await launchDesktop(userDataDir, { + agentDir, + initialWorkspaces: [workspacePath], + testMode: "background", + }); + + try { + const window = await harness.firstWindow(); + await selectSession(window, "Fork fixture session"); + + const transcript = window.getByTestId("transcript"); + await expect(transcript).toContainText("Third fork question"); + + const workspace = await waitForWorkspaceByPath(window, workspacePath); + const before = await getDesktopState(window); + const beforeWorkspace = before.workspaces.find((entry) => entry.id === workspace.id); + const beforeSessionCount = beforeWorkspace?.sessions.length ?? 0; + const beforeSelectedSessionId = before.selectedSessionId; + + // Fork after the second assistant response (codex-style); history up to and including it branches off. + const secondAnswer = transcript.locator(".timeline-item--assistant", { hasText: "Second fork answer" }); + await secondAnswer.hover(); + await secondAnswer.getByTestId("fork-from-message").click(); + + const forkModal = window.getByTestId("fork-modal"); + await expect(forkModal).toBeVisible(); + await expect(window.getByTestId("fork-modal-preview")).toContainText("Second fork answer"); + await expect(window.getByTestId("fork-environment-local")).toHaveAttribute("aria-pressed", "true"); + + await window.getByTestId("fork-modal-confirm").click(); + await expect(forkModal).toHaveCount(0); + + // A new session is created and selected. + await expect + .poll(async () => { + const state = await getDesktopState(window); + const ws = state.workspaces.find((entry) => entry.id === workspace.id); + return ws?.sessions.length ?? 0; + }) + .toBe(beforeSessionCount + 1); + + const after = await getDesktopState(window); + expect(after.selectedSessionId).not.toBe(beforeSelectedSessionId); + expect(after.selectedWorkspaceId).toBe(workspace.id); + + // The composer starts empty so the forked thread continues from the existing history (codex-style). + await expect(window.getByTestId("composer")).toHaveValue(""); + + // The branched transcript keeps the full history up to and including the forked response. + await expect(transcript).toContainText("First fork answer"); + await expect(transcript).toContainText("Second fork question"); + await expect(transcript).toContainText("Second fork answer"); + // Everything after the fork point is dropped. + await expect(transcript).not.toContainText("Third fork question"); + } finally { + await harness.close(); + } +}); diff --git a/apps/desktop/tests/demo/fork-demo.spec.ts b/apps/desktop/tests/demo/fork-demo.spec.ts new file mode 100644 index 00000000..9bcf2d09 --- /dev/null +++ b/apps/desktop/tests/demo/fork-demo.spec.ts @@ -0,0 +1,106 @@ +import { expect, test } from "@playwright/test"; +import { copyFile, mkdir } from "node:fs/promises"; +import { join } from "node:path"; +import { + getDesktopState, + launchDesktop, + makeUserDataDir, + makeWorkspace, + seedAgentDir, + seedForkSessionFixture, + selectSession, + waitForWorkspaceByPath, +} from "../helpers/electron-app"; + +// Slow-paced screen recording of the fork-from-message feature for demo purposes. +// Run directly: PI_APP_TEST_MODE=background playwright test apps/desktop/tests/demo/fork-demo.spec.ts +test("records a fork-from-message walkthrough", async () => { + test.setTimeout(120_000); + const videoDir = join(process.cwd(), "videos"); + await mkdir(videoDir, { recursive: true }); + + const userDataDir = await makeUserDataDir(); + const agentDir = join(userDataDir, "agent"); + const workspacePath = await makeWorkspace("fork-demo-workspace"); + await seedAgentDir(agentDir); + await seedForkSessionFixture(agentDir, workspacePath); + + const harness = await launchDesktop(userDataDir, { + agentDir, + initialWorkspaces: [workspacePath], + testMode: "background", + recordVideoDir: videoDir, + recordVideoSize: { width: 1480, height: 940 }, + }); + + const window = await harness.firstWindow(); + const video = window.video(); + + try { + // Settle on the seeded conversation. + await window.waitForTimeout(1_200); + await selectSession(window, "Fork fixture session"); + + const transcript = window.getByTestId("transcript"); + await expect(transcript).toContainText("Third fork question"); + await window.waitForTimeout(2_000); + + const workspace = await waitForWorkspaceByPath(window, workspacePath); + const before = await getDesktopState(window); + const beforeWorkspace = before.workspaces.find((entry) => entry.id === workspace.id); + const beforeSessionCount = beforeWorkspace?.sessions.length ?? 0; + const beforeSelectedSessionId = before.selectedSessionId; + + // Reveal the fork affordance on the second assistant response. + const secondAnswer = transcript.locator(".timeline-item--assistant", { hasText: "Second fork answer" }); + await secondAnswer.scrollIntoViewIfNeeded(); + await window.waitForTimeout(800); + await secondAnswer.hover(); + await window.waitForTimeout(1_400); + + // Open the fork dialog. + await secondAnswer.getByTestId("fork-from-message").click(); + const forkModal = window.getByTestId("fork-modal"); + await expect(forkModal).toBeVisible(); + await expect(window.getByTestId("fork-modal-preview")).toContainText("Second fork answer"); + await window.waitForTimeout(2_200); + + // Show the worktree choice, then settle back on "Same worktree". + await window.getByTestId("fork-environment-worktree").hover(); + await window.waitForTimeout(1_400); + await window.getByTestId("fork-environment-local").hover(); + await window.waitForTimeout(1_200); + + // Confirm the fork. + await window.getByTestId("fork-modal-confirm").click(); + await expect(forkModal).toHaveCount(0); + + await expect + .poll(async () => { + const state = await getDesktopState(window); + const ws = state.workspaces.find((entry) => entry.id === workspace.id); + return ws?.sessions.length ?? 0; + }) + .toBe(beforeSessionCount + 1); + + const after = await getDesktopState(window); + expect(after.selectedSessionId).not.toBe(beforeSelectedSessionId); + + // Linger on the result: new session selected, history up to and including the fork shown, empty composer. + await expect(window.getByTestId("composer")).toHaveValue(""); + await expect(transcript).toContainText("First fork answer"); + await expect(transcript).toContainText("Second fork answer"); + await expect(transcript).not.toContainText("Third fork question"); + await window.waitForTimeout(3_500); + } finally { + await harness.close(); + } + + const recordedPath = await video?.path(); + if (recordedPath) { + const finalPath = join(videoDir, "fork-from-message-demo.webm"); + await copyFile(recordedPath, finalPath); + // eslint-disable-next-line no-console + console.log(`\nFORK_DEMO_VIDEO=${finalPath}\n`); + } +}); diff --git a/apps/desktop/tests/helpers/electron-app.ts b/apps/desktop/tests/helpers/electron-app.ts index cb6c8849..8e73a343 100644 --- a/apps/desktop/tests/helpers/electron-app.ts +++ b/apps/desktop/tests/helpers/electron-app.ts @@ -59,6 +59,8 @@ export interface LaunchDesktopOptions { readonly scrubProviderEnv?: boolean; readonly envOverrides?: Readonly>; readonly inheritParentEnv?: boolean; + readonly recordVideoDir?: string; + readonly recordVideoSize?: { readonly width: number; readonly height: number }; } export interface SeedAgentDirOptions { @@ -106,6 +108,14 @@ export async function launchDesktop( args: [desktopDir], cwd: desktopDir, env, + ...(normalized.recordVideoDir + ? { + recordVideo: { + dir: normalized.recordVideoDir, + ...(normalized.recordVideoSize ? { size: normalized.recordVideoSize } : {}), + }, + } + : {}), }); return createDesktopHarness(electronApp); @@ -634,6 +644,57 @@ export async function seedToolResultTreeSessionFixture( }); } +export async function seedForkSessionFixture( + agentDir: string, + workspacePath: string, +): Promise<{ + readonly sessionId: string; + readonly title: "Fork fixture session"; +}> { + const { SessionManager } = (await import( + "../../../../node_modules/@earendil-works/pi-coding-agent/dist/core/session-manager.js" + )) as { + SessionManager: { + create(cwd: string): { + appendMessage(message: { role: "user" | "assistant"; content: string; timestamp: number }): string; + appendModelChange(provider: string, modelId: string): string; + appendSessionInfo(name: string): string; + appendThinkingLevelChange(thinkingLevel: string): string; + getSessionId(): string; + }; + }; + }; + + return withAgentDirEnv(agentDir, async () => { + const sessionManager = SessionManager.create(workspacePath); + let timestamp = Date.now(); + const nextTimestamp = () => { + timestamp += 1_000; + return timestamp; + }; + const appendUser = (content: string) => + sessionManager.appendMessage({ role: "user", content, timestamp: nextTimestamp() }); + const appendAssistant = (content: string) => + sessionManager.appendMessage({ role: "assistant", content, timestamp: nextTimestamp() }); + + sessionManager.appendModelChange("openai", "gpt-5.4"); + sessionManager.appendThinkingLevelChange("high"); + appendUser("First fork question"); + appendAssistant("First fork answer"); + appendUser("Second fork question"); + appendAssistant("Second fork answer"); + appendUser("Third fork question"); + appendAssistant("Third fork answer"); + + sessionManager.appendSessionInfo("Fork fixture session"); + + return { + sessionId: sessionManager.getSessionId(), + title: "Fork fixture session", + }; + }); +} + async function withAgentDirEnv(agentDir: string, action: () => Promise): Promise { const previousAgentDir = process.env.PI_CODING_AGENT_DIR; process.env.PI_CODING_AGENT_DIR = agentDir; diff --git a/packages/pi-sdk-driver/src/pi-sdk-driver.ts b/packages/pi-sdk-driver/src/pi-sdk-driver.ts index 1e0e20b8..103a720f 100644 --- a/packages/pi-sdk-driver/src/pi-sdk-driver.ts +++ b/packages/pi-sdk-driver/src/pi-sdk-driver.ts @@ -8,6 +8,8 @@ import type { } from "@pi-gui/session-driver/types"; import type { CreateSessionOptions, + ForkSessionOptions, + ForkSessionResult, HostUiResponse, SessionDriver, SessionEventListener, @@ -54,6 +56,10 @@ export class PiSdkDriver implements SessionDriver { return this.supervisor.createSession(workspace, options); } + forkSession(sourceRef: SessionRef, options: ForkSessionOptions): Promise { + return this.supervisor.forkSession(sourceRef, options); + } + openSession(sessionRef: SessionRef): Promise { return this.supervisor.openSession(sessionRef); } diff --git a/packages/pi-sdk-driver/src/session-supervisor.ts b/packages/pi-sdk-driver/src/session-supervisor.ts index 9adb8144..b9b7b123 100644 --- a/packages/pi-sdk-driver/src/session-supervisor.ts +++ b/packages/pi-sdk-driver/src/session-supervisor.ts @@ -25,6 +25,8 @@ import type { } from "@pi-gui/session-driver/types"; import type { CreateSessionOptions, + ForkSessionOptions, + ForkSessionResult, HostUiRequest, HostUiResponse, SessionConfig, @@ -335,6 +337,121 @@ export class SessionSupervisor { return snapshot; } + async forkSession(sourceRef: SessionRef, options: ForkSessionOptions): Promise { + const sourceRecord = await this.ensureRecord(sourceRef); + const sourceSession = this.requireSession(sourceRecord); + const sourceManager = sourceSession.sessionManager; + const sourceFile = sourceRecord.sessionFile ?? sourceManager.getSessionFile(); + if (!sourceFile) { + throw new Error(`Session ${sessionKey(sourceRef)} cannot be forked because no session file is tracked.`); + } + + // Resolve the nth user message in the source branch (root -> leaf order). + const branch = sourceManager.getBranch(); + const userEntries = branch.filter( + (entry): entry is Extract => + entry.type === "message" && entry.message.role === "user", + ); + const selectedEntry = userEntries[options.userMessageIndex]; + if (!selectedEntry) { + throw new Error( + `Cannot fork session ${sessionKey(sourceRef)}: no user message at index ${options.userMessageIndex}.`, + ); + } + + const position = options.position ?? "before"; + let targetLeafId: string | undefined; + let selectedText: string | undefined; + if (position === "after") { + // Codex-style fork: branch at the END of the selected user message's turn so the + // new thread contains the full history up to and including that turn's assistant + // response, with an empty composer. The turn ends just before the next user message + // (or at the branch leaf when this is the final turn). + const nextUserEntry = userEntries[options.userMessageIndex + 1]; + if (nextUserEntry) { + const nextIndex = branch.findIndex((entry) => entry.id === nextUserEntry.id); + targetLeafId = nextIndex > 0 ? branch[nextIndex - 1]?.id : selectedEntry.id; + } else { + targetLeafId = branch[branch.length - 1]?.id ?? selectedEntry.id; + } + } else if (position === "at") { + targetLeafId = selectedEntry.id; + } else { + targetLeafId = selectedEntry.parentId ?? undefined; + selectedText = messageText(selectedEntry.message as unknown as Record) || undefined; + } + + const targetWorkspace = options.targetWorkspace; + await this.touchWorkspace(targetWorkspace); + const sameWorkspace = resolve(targetWorkspace.path) === resolve(sourceRecord.workspace.path); + + // Build a branched SessionManager containing only the history up to the fork point. + let branchedManager: SessionManager; + if (!targetLeafId) { + // Forking before the first user message: start a fresh empty session in the target. + branchedManager = SessionManager.create(targetWorkspace.path); + branchedManager.newSession({ parentSession: sourceFile }); + } else if (sameWorkspace) { + const opened = SessionManager.open(sourceFile); + const forkedPath = opened.createBranchedSession(targetLeafId); + if (!forkedPath) { + throw new Error(`Failed to create forked session from ${sessionKey(sourceRef)}.`); + } + branchedManager = SessionManager.open(forkedPath); + } else { + const forked = SessionManager.forkFrom(sourceFile, targetWorkspace.path); + const forkedPath = forked.createBranchedSession(targetLeafId); + if (!forkedPath) { + throw new Error(`Failed to create forked session from ${sessionKey(sourceRef)}.`); + } + branchedManager = SessionManager.open(forkedPath); + } + + const createOptions: CreateAgentSessionOptions = { + cwd: targetWorkspace.path, + sessionManager: branchedManager, + ...(this.modelRegistry ? { modelRegistry: this.modelRegistry } : {}), + }; + const sourceConfig = sourceRecord.config; + if (sourceConfig?.provider && sourceConfig?.modelId) { + try { + createOptions.model = this.resolveModel(sourceConfig.provider, sourceConfig.modelId); + } catch { + // Source model is no longer available; fall back to the runtime default. + } + } + if (sourceConfig?.thinkingLevel) { + createOptions.thinkingLevel = sourceConfig.thinkingLevel as NonNullable< + CreateAgentSessionOptions["thinkingLevel"] + >; + } + + const runtime = await this.createAgentSessionRuntimeImpl(createOptions); + const session = runtime.session; + + const title = options.title ?? sourceRecord.title; + const record = this.createRecord(targetWorkspace, runtime, title); + forcePersistSession(session.sessionManager); + record.config = deriveSessionConfig(session.sessionManager); + const sessionFile = record.sessionFile ?? session.sessionManager.getSessionFile(); + if (sessionFile) { + record.sessionFile = sessionFile; + await this.catalogs.setSessionFile(record.ref, sessionFile); + } + + this.records.set(sessionKey(record.ref), record); + await this.bindSessionRuntime(record); + await this.persistSnapshot(record); + const snapshot = buildSnapshot(record); + await this.emit(record, { + type: "sessionOpened", + sessionRef: record.ref, + timestamp: nowIso(), + snapshot, + }); + return selectedText === undefined ? { snapshot } : { snapshot, selectedText }; + } + async openSession(sessionRef: SessionRef): Promise { const record = await this.ensureRecord(sessionRef); await this.touchWorkspace(record.workspace); diff --git a/packages/pi-sdk-driver/src/vendor/session-driver.d.ts b/packages/pi-sdk-driver/src/vendor/session-driver.d.ts index f8a76a6e..34c93b4c 100644 --- a/packages/pi-sdk-driver/src/vendor/session-driver.d.ts +++ b/packages/pi-sdk-driver/src/vendor/session-driver.d.ts @@ -91,6 +91,20 @@ declare module "@pi-gui/session-driver" { readonly initialThinkingLevel?: string; } + export type ForkPosition = "before" | "at" | "after"; + + export interface ForkSessionOptions { + readonly targetWorkspace: WorkspaceRef; + readonly userMessageIndex: number; + readonly position?: ForkPosition; + readonly title?: string; + } + + export interface ForkSessionResult { + readonly snapshot: SessionSnapshot; + readonly selectedText?: string; + } + export interface SessionErrorInfo { readonly message: string; readonly code?: string; @@ -301,6 +315,7 @@ declare module "@pi-gui/session-driver" { export interface SessionDriver { createSession(workspace: WorkspaceRef, options?: CreateSessionOptions): Promise; + forkSession(sourceRef: SessionRef, options: ForkSessionOptions): Promise; openSession(sessionRef: SessionRef): Promise; archiveSession(sessionRef: SessionRef): Promise; unarchiveSession(sessionRef: SessionRef): Promise; diff --git a/packages/session-driver/src/index.ts b/packages/session-driver/src/index.ts index 7df77224..15141937 100644 --- a/packages/session-driver/src/index.ts +++ b/packages/session-driver/src/index.ts @@ -1,6 +1,9 @@ export type { AssistantDeltaEvent, CreateSessionOptions, + ForkPosition, + ForkSessionOptions, + ForkSessionResult, ExtensionCompatibilityIssue, ExtensionCompatibilityIssueEvent, HostUiResponse, diff --git a/packages/session-driver/src/types.ts b/packages/session-driver/src/types.ts index eb1b3638..8c6afe82 100644 --- a/packages/session-driver/src/types.ts +++ b/packages/session-driver/src/types.ts @@ -121,6 +121,28 @@ export interface CreateSessionOptions { readonly initialThinkingLevel?: string; } +export type ForkPosition = "before" | "at" | "after"; + +export interface ForkSessionOptions { + /** Target workspace for the forked session (the source workspace, or a new worktree). */ + readonly targetWorkspace: WorkspaceRef; + /** 0-based index of the user message to fork from, counted within the source branch. */ + readonly userMessageIndex: number; + /** + * "before" (default) forks before the selected user message so it can be edited and + * re-sent; "at" keeps the selected user message in the forked history. + */ + readonly position?: ForkPosition; + /** Optional title for the forked session. Defaults to the source session title. */ + readonly title?: string; +} + +export interface ForkSessionResult { + readonly snapshot: SessionSnapshot; + /** When forking "before", the text of the selected user message for composer prefill. */ + readonly selectedText?: string; +} + export interface SessionEventBase { readonly type: string; readonly sessionRef: SessionRef; @@ -306,6 +328,7 @@ export type Unsubscribe = () => void; export interface SessionDriver { createSession(workspace: WorkspaceRef, options?: CreateSessionOptions): Promise; + forkSession(sourceRef: SessionRef, options: ForkSessionOptions): Promise; openSession(sessionRef: SessionRef): Promise; archiveSession(sessionRef: SessionRef): Promise; unarchiveSession(sessionRef: SessionRef): Promise;