diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index 86bae87a..10613f32 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -110,7 +110,10 @@ export class AgentSessionRuntime { return { cancelled: result?.cancel === true }; } - private async emitBeforeFork(entryId: string): Promise<{ cancelled: boolean }> { + private async emitBeforeFork( + entryId: string, + options: { position: "before" | "at" }, + ): Promise<{ cancelled: boolean }> { const runner = this.session.extensionRunner; if (!runner?.hasHandlers("session_before_fork")) { return { cancelled: false }; @@ -119,6 +122,7 @@ export class AgentSessionRuntime { const result = await runner.emit({ type: "session_before_fork", entryId, + ...options, }); return { cancelled: result?.cancel === true }; } @@ -191,26 +195,41 @@ export class AgentSessionRuntime { return { cancelled: false }; } - async fork(entryId: string): Promise<{ cancelled: boolean; selectedText?: string }> { - const beforeResult = await this.emitBeforeFork(entryId); + async fork( + entryId: string, + options?: { position?: "before" | "at" }, + ): Promise<{ cancelled: boolean; selectedText?: string }> { + const position = options?.position ?? "before"; + const beforeResult = await this.emitBeforeFork(entryId, { position }); if (beforeResult.cancelled) { return { cancelled: true }; } + let targetLeafId: string | null; + let selectedText: string | undefined; const selectedEntry = this.session.sessionManager.getEntry(entryId); - if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { + if (!selectedEntry) { throw new Error("Invalid entry ID for forking"); } + if (position === "at") { + targetLeafId = selectedEntry.id; + } else { + if (selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { + throw new Error("Invalid entry ID for forking"); + } + targetLeafId = selectedEntry.parentId; + selectedText = extractUserMessageText(selectedEntry.message.content); + } + const previousSessionFile = this.session.sessionFile; - const selectedText = extractUserMessageText(selectedEntry.message.content); if (this.session.sessionManager.isPersisted()) { const currentSessionFile = this.session.sessionFile; if (!currentSessionFile) { throw new Error("Persisted session is missing a session file"); } const sessionDir = this.session.sessionManager.getSessionDir(); - if (!selectedEntry.parentId) { + if (!targetLeafId) { const sessionManager = SessionManager.create(this.cwd, sessionDir); sessionManager.newSession({ parentSession: currentSessionFile }); await this.teardownCurrent(); @@ -226,7 +245,7 @@ export class AgentSessionRuntime { } const sourceManager = SessionManager.open(currentSessionFile, sessionDir); - const forkedSessionPath = sourceManager.createBranchedSession(selectedEntry.parentId); + const forkedSessionPath = sourceManager.createBranchedSession(targetLeafId); if (!forkedSessionPath) { throw new Error("Failed to create forked session"); } @@ -244,10 +263,10 @@ export class AgentSessionRuntime { } const sessionManager = this.session.sessionManager; - if (!selectedEntry.parentId) { + if (!targetLeafId) { sessionManager.newSession({ parentSession: this.session.sessionFile }); } else { - sessionManager.createBranchedSession(selectedEntry.parentId); + sessionManager.createBranchedSession(targetLeafId); } await this.teardownCurrent(); this.apply( diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index f2257faa..7fd63c53 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -143,7 +143,10 @@ export type NewSessionHandler = (options?: { setup?: (sessionManager: SessionManager) => Promise; }) => Promise<{ cancelled: boolean }>; -export type ForkHandler = (entryId: string) => Promise<{ cancelled: boolean }>; +export type ForkHandler = ( + entryId: string, + options?: { position?: "before" | "at" }, +) => Promise<{ cancelled: boolean }>; export type NavigateTreeHandler = ( targetId: string, @@ -559,7 +562,7 @@ export class ExtensionRunner { ...this.createContext(), waitForIdle: () => this.waitForIdleFn(), newSession: (options) => this.newSessionHandler(options), - fork: (entryId) => this.forkHandler(entryId), + fork: (entryId, options) => this.forkHandler(entryId, options), navigateTree: (targetId, options) => this.navigateTreeHandler(targetId, options), switchSession: (sessionPath) => this.switchSessionHandler(sessionPath), reload: () => this.reloadHandler(), diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index b135b5a2..d4e540f4 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -309,7 +309,7 @@ export interface ExtensionCommandContext extends ExtensionContext { }): Promise<{ cancelled: boolean }>; /** Fork from a specific entry, creating a new session file. */ - fork(entryId: string): Promise<{ cancelled: boolean }>; + fork(entryId: string, options?: { position?: "before" | "at" }): Promise<{ cancelled: boolean }>; /** Navigate to a different point in the session tree. */ navigateTree( @@ -473,6 +473,7 @@ export interface SessionBeforeSwitchEvent { export interface SessionBeforeForkEvent { type: "session_before_fork"; entryId: string; + position: "before" | "at"; } /** Fired before context compaction (can be cancelled or customized) */ @@ -1423,7 +1424,7 @@ export interface ExtensionCommandContextActions { parentSession?: string; setup?: (sessionManager: SessionManager) => Promise; }) => Promise<{ cancelled: boolean }>; - fork: (entryId: string) => Promise<{ cancelled: boolean }>; + fork: (entryId: string, options?: { position?: "before" | "at" }) => Promise<{ cancelled: boolean }>; navigateTree: ( targetId: string, options?: { summarize?: boolean; customInstructions?: string; replaceInstructions?: boolean; label?: string }, diff --git a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts index c9f28f74..d02c27fd 100644 --- a/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/user-message-selector.ts @@ -18,11 +18,12 @@ class UserMessageList implements Component { public onCancel?: () => void; private maxVisible: number = 10; // Max messages visible - constructor(messages: UserMessageItem[]) { + constructor(messages: UserMessageItem[], initialSelectedId?: string) { // Store messages in chronological order (oldest to newest) this.messages = messages; - // Start with the last (most recent) message selected - this.selectedIndex = Math.max(0, messages.length - 1); + const initialIndex = initialSelectedId ? messages.findIndex((message) => message.id === initialSelectedId) : -1; + // Start with selected message if provided, else default to the most recent + this.selectedIndex = initialIndex >= 0 ? initialIndex : Math.max(0, messages.length - 1); } invalidate(): void { @@ -109,19 +110,24 @@ class UserMessageList implements Component { export class UserMessageSelectorComponent extends Container { private messageList: UserMessageList; - constructor(messages: UserMessageItem[], onSelect: (entryId: string) => void, onCancel: () => void) { + constructor( + messages: UserMessageItem[], + onSelect: (entryId: string) => void, + onCancel: () => void, + initialSelectedId?: string, + ) { super(); // Add header this.addChild(new Spacer(1)); this.addChild(new Text(theme.bold("Branch from Message"), 1, 0)); - this.addChild(new Text(theme.fg("muted", "Select a message to create a new branch from that point"), 1, 0)); + this.addChild(new Text(theme.fg("muted", "Select a message or duplicate the current session"), 1, 0)); this.addChild(new Spacer(1)); this.addChild(new DynamicBorder()); this.addChild(new Spacer(1)); // Create message list - this.messageList = new UserMessageList(messages); + this.messageList = new UserMessageList(messages, initialSelectedId); this.messageList.onSelect = onSelect; this.messageList.onCancel = onCancel; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index b8df40f2..3bf3fbba 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -1451,9 +1451,9 @@ export class InteractiveMode { return this.handleFatalRuntimeError("Failed to create session", error); } }, - fork: async (entryId) => { + fork: async (entryId, options) => { try { - const result = await this.runtimeHost.fork(entryId); + const result = await this.runtimeHost.fork(entryId, options); if (!result.cancelled) { await this.handleRuntimeSessionChange(); this.renderCurrentSessionState(); @@ -3970,6 +3970,7 @@ export class InteractiveMode { } private showUserMessageSelector(): void { + const duplicateSessionOptionId = "__duplicate_session__"; const userMessages = this.session.getUserMessagesForForking(); if (userMessages.length === 0) { @@ -3977,26 +3978,49 @@ export class InteractiveMode { return; } + const selectableEntries = [ + ...userMessages.map((m) => ({ id: m.entryId, text: m.text })), + { id: duplicateSessionOptionId, text: "Duplicate session" }, + ]; + const initialSelectedId = userMessages[userMessages.length - 1]?.entryId; + this.showSelector((done) => { const selector = new UserMessageSelectorComponent( - userMessages.map((m) => ({ id: m.entryId, text: m.text })), + selectableEntries, async (entryId) => { - const result = await this.runtimeHost.fork(entryId); - if (result.cancelled) { + try { + let result: { cancelled: boolean; selectedText?: string }; + if (entryId === duplicateSessionOptionId) { + const leafId = this.sessionManager.getLeafId(); + if (!leafId) { + throw new Error("Cannot duplicate session: no current entry selected"); + } + result = await this.runtimeHost.fork(leafId, { position: "at" }); + } else { + result = await this.runtimeHost.fork(entryId); + } + + if (result.cancelled) { + done(); + this.ui.requestRender(); + return; + } + + await this.handleRuntimeSessionChange(); + this.renderCurrentSessionState(); + this.editor.setText(result.selectedText ?? ""); done(); - this.ui.requestRender(); - return; + this.showStatus("Branched to new session"); + } catch (error: unknown) { + done(); + this.showError(error instanceof Error ? error.message : String(error)); } - await this.handleRuntimeSessionChange(); - this.renderCurrentSessionState(); - this.editor.setText(result.selectedText ?? ""); - done(); - this.showStatus("Branched to new session"); }, () => { done(); this.ui.requestRender(); }, + initialSelectedId, ); return { component: selector, focus: selector.getMessageList() }; }); diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index 267d6846..0c79956c 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -76,8 +76,8 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr } return result; }, - fork: async (entryId) => { - const result = await runtimeHost.fork(entryId); + fork: async (entryId, forkOptions) => { + const result = await runtimeHost.fork(entryId, forkOptions); if (!result.cancelled) { await rebindSession(); } diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index 32e61b00..b775adc0 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -298,8 +298,8 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise { - const result = await runtimeHost.fork(entryId); + fork: async (entryId, forkOptions) => { + const result = await runtimeHost.fork(entryId, forkOptions); if (!result.cancelled) { await rebindSession(); } diff --git a/packages/coding-agent/test/agent-session-runtime-events.test.ts b/packages/coding-agent/test/agent-session-runtime-events.test.ts index 3413332d..d47dce97 100644 --- a/packages/coding-agent/test/agent-session-runtime-events.test.ts +++ b/packages/coding-agent/test/agent-session-runtime-events.test.ts @@ -176,7 +176,7 @@ describe("AgentSessionRuntime session lifecycle events", () => { expect(successResult.selectedText).toBe("hello"); await runtimeHost.session.bindExtensions({}); expect(events).toEqual([ - { type: "session_before_fork", entryId: userMessage.entryId }, + { type: "session_before_fork", entryId: userMessage.entryId, position: "before" }, { type: "session_start", reason: "fork", previousSessionFile }, ]); @@ -184,6 +184,12 @@ describe("AgentSessionRuntime session lifecycle events", () => { cancelNextFork = true; const cancelResult = await runtimeHost.fork(userMessage.entryId); expect(cancelResult).toEqual({ cancelled: true }); - expect(events).toEqual([{ type: "session_before_fork", entryId: userMessage.entryId }]); + expect(events).toEqual([{ type: "session_before_fork", entryId: userMessage.entryId, position: "before" }]); + + events.length = 0; + cancelNextFork = true; + const cancelAtResult = await runtimeHost.fork("missing-entry", { position: "at" }); + expect(cancelAtResult).toEqual({ cancelled: true }); + expect(events).toEqual([{ type: "session_before_fork", entryId: "missing-entry", position: "at" }]); }); }); diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index e60572d6..21d519af 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -695,6 +695,30 @@ describe("ExtensionRunner", () => { }); }); + describe("command context", () => { + it("passes fork options through to the bound handler", async () => { + const runtime = createExtensionRuntime(); + const runner = new ExtensionRunner([], runtime, tempDir, sessionManager, modelRegistry); + const fork = vi.fn(async () => ({ cancelled: false })); + + runner.bindCommandContext({ + waitForIdle: async () => {}, + newSession: async () => ({ cancelled: false }), + fork, + navigateTree: async () => ({ cancelled: false }), + switchSession: async () => ({ cancelled: false }), + reload: async () => {}, + }); + + const commandContext = runner.createCommandContext(); + await commandContext.fork("entry-1"); + expect(fork).toHaveBeenCalledWith("entry-1", undefined); + + await commandContext.fork("entry-2", { position: "at" }); + expect(fork).toHaveBeenLastCalledWith("entry-2", { position: "at" }); + }); + }); + describe("hasHandlers", () => { it("returns true when handlers exist for event type", async () => { const extCode = ` diff --git a/packages/coding-agent/test/suite/agent-session-runtime.test.ts b/packages/coding-agent/test/suite/agent-session-runtime.test.ts index e6d4b84e..99152fc1 100644 --- a/packages/coding-agent/test/suite/agent-session-runtime.test.ts +++ b/packages/coding-agent/test/suite/agent-session-runtime.test.ts @@ -217,7 +217,7 @@ describe("AgentSessionRuntime characterization", () => { expect(successResult.selectedText).toBe("hello"); await runtime.session.bindExtensions({}); expect(events).toEqual([ - { type: "session_before_fork", entryId: userMessage.entryId }, + { type: "session_before_fork", entryId: userMessage.entryId, position: "before" }, { type: "session_start", reason: "fork", previousSessionFile }, ]); @@ -225,7 +225,13 @@ describe("AgentSessionRuntime characterization", () => { cancelNextFork = true; const cancelResult = await runtime.fork(userMessage.entryId); expect(cancelResult).toEqual({ cancelled: true }); - expect(events).toEqual([{ type: "session_before_fork", entryId: userMessage.entryId }]); + expect(events).toEqual([{ type: "session_before_fork", entryId: userMessage.entryId, position: "before" }]); + + events.length = 0; + cancelNextFork = true; + const cancelAtResult = await runtime.fork("missing-entry", { position: "at" }); + expect(cancelAtResult).toEqual({ cancelled: true }); + expect(events).toEqual([{ type: "session_before_fork", entryId: "missing-entry", position: "at" }]); }); it("throws when forking with an invalid entry id", async () => {