feat(coding-agent): add fork position and duplicate session option (#3431)

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(

packages/coding-agent/src/core/extensions/runner.ts

@@ -143,7 +143,10 @@ export type NewSessionHandler = (options?: {
 	setup?: (sessionManager: SessionManager) => Promise<void>;
 }) => 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(),

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<void>;
 	}) => 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 },

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;
 

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() };
 		});

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();
 					}

packages/coding-agent/src/modes/rpc/rpc-mode.ts

@@ -298,8 +298,8 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
 					}
 					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();
 					}