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 },