Merge pull request #309669 from microsoft/connor4312/proper-forking

agentHost: use proper sdk-provided forking/truncation methods

Author: connor4312

src/vs/platform/agentHost/common/agentService.ts

@@ -101,8 +101,19 @@ export interface IAgentCreateSessionConfig {
 	readonly session?: URI;
 	readonly workingDirectory?: URI;
 	readonly config?: Record<string, string>;
-	/** Fork from an existing session at a specific turn index. */
-	readonly fork?: { readonly session: URI; readonly turnIndex: number };
+	/** Fork from an existing session at a specific turn. */
+	readonly fork?: {
+		readonly session: URI;
+		readonly turnIndex: number;
+		readonly turnId: string;
+		/**
+		 * Maps old protocol turn IDs to new protocol turn IDs.
+		 * Populated by the service layer after generating fresh UUIDs
+		 * for the forked session's turns. Used by the agent to remap
+		 * per-turn data (e.g. SDK event ID mappings) in the session database.
+		 */
+		readonly turnIdMapping?: ReadonlyMap<string, string>;
+	};
 }
 
 export const AgentHostSessionConfigBranchNameHintKey = 'branchNameHint';
@@ -406,21 +417,11 @@ export interface IAgent {
 	authenticate(resource: string, token: string): Promise<boolean>;
 
 	/**
-	 * Truncate a session's history. If `turnIndex` is provided (0-based), keeps
-	 * turns up to and including that turn. If omitted, all turns are removed.
+	 * Truncate a session's history. If `turnId` is provided, keeps turns up to
+	 * and including that turn. If omitted, all turns are removed.
 	 * Optional — not all providers support truncation.
 	 */
-	truncateSession?(session: URI, turnIndex?: number): Promise<void>;
-
-	/**
-	 * Fork a session at a specific turn, creating a new session on disk
-	 * with the source session's history up to and including the specified turn.
-	 * Optional — not all providers support forking.
-	 *
-	 * @param turnIndex 0-based turn index to fork at.
-	 * @returns The new session's raw ID.
-	 */
-	forkSession?(sourceSession: URI, newSessionId: string, turnIndex: number): Promise<void>;
+	truncateSession?(session: URI, turnId?: string): Promise<void>;
 
 	/**
 	 * Receives client-provided customization refs and syncs them (e.g. copies

src/vs/platform/agentHost/common/sessionDataService.ts

@@ -10,6 +10,9 @@ import type { FileEditKind } from './state/sessionState.js';
 
 export const ISessionDataService = createDecorator<ISessionDataService>('sessionDataService');
 
+/** Filename of the per-session SQLite database. */
+export const SESSION_DB_FILENAME = 'session.db';
+
 // ---- File-edit types ----------------------------------------------------
 
 /**
@@ -68,6 +71,48 @@ export interface ISessionDatabase extends IDisposable {
 	 */
 	deleteTurn(turnId: string): Promise<void>;
 
+	/**
+	 * Associates a Copilot SDK event ID with a turn. The event ID corresponds
+	 * to the `user.message` event in the SDK event stream and is used by
+	 * the SDK's `history.truncate` and `sessions.fork` RPCs.
+	 */
+	setTurnEventId(turnId: string, eventId: string): Promise<void>;
+
+	/**
+	 * Retrieves the SDK event ID previously stored for a turn.
+	 * Returns `undefined` if no event ID has been set.
+	 */
+	getTurnEventId(turnId: string): Promise<string | undefined>;
+
+	/**
+	 * Returns the SDK event ID of the turn inserted immediately after the
+	 * given turn, or `undefined` if the given turn is the last one.
+	 */
+	getNextTurnEventId(turnId: string): Promise<string | undefined>;
+
+	/**
+	 * Returns the SDK event ID of the earliest turn in insertion order,
+	 * or `undefined` if there are no turns.
+	 */
+	getFirstTurnEventId(): Promise<string | undefined>;
+
+	/**
+	 * Deletes the given turn and all turns inserted after it, along
+	 * with their associated file edits (cascade).
+	 */
+	truncateFromTurn(turnId: string): Promise<void>;
+
+	/**
+	 * Deletes all turns inserted after the given turn (but keeps the
+	 * given turn itself). Associated file edits cascade-delete.
+	 */
+	deleteTurnsAfter(turnId: string): Promise<void>;
+
+	/**
+	 * Deletes all turns and their associated file edits.
+	 */
+	deleteAllTurns(): Promise<void>;
+
 	/**
 	 * Store a file-edit snapshot (metadata + content) for a tool invocation
 	 * within a turn.
@@ -122,6 +167,12 @@ export interface ISessionDatabase extends IDisposable {
 	 */
 	setMetadata(key: string, value: string): Promise<void>;
 
+	/**
+	 * Bulk-remaps turn IDs using the provided old→new mapping.
+	 * Used after copying a database file for a forked session.
+	 */
+	remapTurnIds(mapping: ReadonlyMap<string, string>): Promise<void>;
+
 	/**
 	 * Close the database connection. After calling this method, the object is
 	 * considered disposed and all other methods will reject with an error.

src/vs/platform/agentHost/node/agentService.ts

@@ -196,6 +196,23 @@ export class AgentService extends Disposable implements IAgentService {
 			throw new Error(`No agent provider registered for: ${providerId ?? '(none)'}`);
 		}
 
+		// When forking, build the old→new turn ID mapping before creating the
+		// session so the agent can use it to remap per-turn data.
+		if (config?.fork) {
+			const sourceState = this._stateManager.getSessionState(config.fork.session.toString());
+			if (sourceState) {
+				const sourceTurns = sourceState.turns.slice(0, config.fork.turnIndex + 1);
+				const turnIdMapping = new Map<string, string>();
+				for (const t of sourceTurns) {
+					turnIdMapping.set(t.id, generateUuid());
+				}
+				config = {
+					...config,
+					fork: { ...config.fork, turnIdMapping },
+				};
+			}
+		}
+
 		// Ensure the command auto-approver is ready before any session events
 		// can arrive. This makes shell command auto-approval fully synchronous.
 		// Safe to run in parallel with createSession since no events flow until
@@ -219,9 +236,9 @@ export class AgentService extends Disposable implements IAgentService {
 		if (config?.fork) {
 			const sourceState = this._stateManager.getSessionState(config.fork.session.toString());
 			let sourceTurns: ITurn[] = [];
-			if (sourceState) {
+			if (sourceState && config.fork.turnIdMapping) {
 				sourceTurns = sourceState.turns.slice(0, config.fork.turnIndex + 1)
-					.map(t => ({ ...t, id: generateUuid() }));
+					.map(t => ({ ...t, id: config!.fork!.turnIdMapping!.get(t.id) ?? generateUuid() }));
 			}
 
 			const summary: ISessionSummary = {

src/vs/platform/agentHost/node/agentSideEffects.ts

@@ -626,17 +626,7 @@ export class AgentSideEffects extends Disposable {
 			}
 			case ActionType.SessionTruncated: {
 				const agent = this._options.getAgent(action.session);
-				let turnIndex: number | undefined;
-				if (action.turnId !== undefined) {
-					const state = this._stateManager.getSessionState(action.session);
-					if (state) {
-						const idx = state.turns.findIndex(t => t.id === action.turnId);
-						if (idx >= 0) {
-							turnIndex = idx;
-						}
-					}
-				}
-				agent?.truncateSession?.(URI.parse(action.session), turnIndex).catch(err => {
+				agent?.truncateSession?.(URI.parse(action.session), action.turnId).catch(err => {
 					this._logService.error('[AgentSideEffects] truncateSession failed', err);
 				});
 				// Turns were removed — recompute diffs from scratch (no changedTurnId)

src/vs/platform/agentHost/node/copilot/copilotAgent.ts

@@ -21,12 +21,11 @@ import { localize } from '../../../../nls.js';
 import { IAgentPluginManager, ISyncedCustomization } from '../../common/agentPluginManager.js';
 import { AgentHostSessionConfigBranchNameHintKey, AgentSession, IAgent, IAgentAttachment, IAgentCreateSessionConfig, IAgentCreateSessionResult, IAgentDescriptor, IAgentMessageEvent, IAgentModelInfo, IAgentProgressEvent, IAgentResolveSessionConfigParams, IAgentSessionConfigCompletionsParams, IAgentSessionMetadata, IAgentSessionProjectInfo, IAgentSubagentStartedEvent, IAgentToolCompleteEvent, IAgentToolStartEvent } from '../../common/agentService.js';
 import type { IResolveSessionConfigResult, ISessionConfigCompletionsResult } from '../../common/state/protocol/commands.js';
-import { ISessionDataService } from '../../common/sessionDataService.js';
+import { ISessionDataService, SESSION_DB_FILENAME } from '../../common/sessionDataService.js';
 import { CustomizationStatus, ICustomizationRef, SessionInputResponseKind, type ISessionInputAnswer, type IPendingMessage, type PolicyState } from '../../common/state/sessionState.js';
 import { CopilotAgentSession, SessionWrapperFactory } from './copilotAgentSession.js';
 import { parsedPluginsEqual, toSdkCustomAgents, toSdkHooks, toSdkMcpServers, toSdkSkillDirectories } from './copilotPluginConverters.js';
 import { CopilotSessionWrapper } from './copilotSessionWrapper.js';
-import { forkCopilotSessionOnDisk, getCopilotDataDir, truncateCopilotSessionOnDisk } from './copilotAgentForking.js';
 import { IProtectedResourceMetadata } from '../../common/state/protocol/state.js';
 import { IAgentHostGitService } from '../agentHostGitService.js';
 import { IAgentHostTerminalManager } from '../agentHostTerminalManager.js';
@@ -236,27 +235,50 @@ export class CopilotAgent extends Disposable implements IAgent {
 		const client = await this._ensureClient();
 		const parsedPlugins = await this._plugins.getAppliedPlugins();
 
-		// When forking, we manipulate the CLI's on-disk data and then resume
-		// instead of creating a fresh session via the SDK.
+		// When forking, use the SDK's sessions.fork RPC.
 		if (config?.fork) {
 			const sourceSessionId = AgentSession.id(config.fork.session);
-			const newSessionId = config.session ? AgentSession.id(config.session) : generateUuid();
 
 			// Serialize against the source session to prevent concurrent
-			// modifications while we read its on-disk data.
+			// modifications while we read its state.
 			return this._sessionSequencer.queue(sourceSessionId, async () => {
-				this._logService.info(`[Copilot] Forking session ${sourceSessionId} at index ${config.fork!.turnIndex} → ${newSessionId}`);
+				this._logService.info(`[Copilot] Forking session ${sourceSessionId} at turnId=${config.fork!.turnId}`);
 
-				// Ensure the source session is loaded so on-disk data is available
-				if (!this._sessions.has(sourceSessionId)) {
-					await this._resumeSession(sourceSessionId);
-				}
+				// Ensure the source session is loaded so we can read its event IDs
+				const sourceEntry = this._sessions.get(sourceSessionId) ?? await this._resumeSession(sourceSessionId);
 
-				const copilotDataDir = getCopilotDataDir();
-				await forkCopilotSessionOnDisk(copilotDataDir, sourceSessionId, newSessionId, config.fork!.turnIndex);
+				// Look up the SDK event ID for the turn *after* the fork point.
+				// toEventId is exclusive — events before it are included.
+				// If there's no next turn, omit toEventId to include all events.
+				const toEventId = await sourceEntry.getNextTurnEventId(config.fork!.turnId);
+
+				const forkResult = await client.rpc.sessions.fork({
+					sessionId: sourceSessionId,
+					...(toEventId ? { toEventId } : {}),
+				});
+				const newSessionId = forkResult.sessionId;
+
+				// Copy the source session's database file so the forked session
+				// inherits turn event IDs and file-edit snapshots.
+				const sourceDbDir = this._sessionDataService.getSessionDataDir(config.fork!.session);
+				const targetDbDir = this._sessionDataService.getSessionDataDirById(newSessionId);
+				const sourceDbPath = URI.joinPath(sourceDbDir, SESSION_DB_FILENAME);
+				const targetDbPath = URI.joinPath(targetDbDir, SESSION_DB_FILENAME);
+				try {
+					await fs.mkdir(targetDbDir.fsPath, { recursive: true });
+					await fs.copyFile(sourceDbPath.fsPath, targetDbPath.fsPath);
+				} catch (err) {
+					this._logService.warn(`[Copilot] Failed to copy session database for fork: ${err instanceof Error ? err.message : String(err)}`);
+				}
 
 				// Resume the forked session so the SDK loads the forked history
 				const agentSession = await this._resumeSession(newSessionId);
+
+				// Remap turn IDs to match the new protocol turn IDs
+				if (config.fork!.turnIdMapping) {
+					await agentSession.remapTurnIds(config.fork!.turnIdMapping);
+				}
+
 				const session = agentSession.sessionUri;
 				this._logService.info(`[Copilot] Forked session created: ${session.toString()}`);
 				const project = await projectFromCopilotContext({ cwd: config.workingDirectory?.fsPath }, this._gitService);
@@ -448,39 +470,33 @@ export class CopilotAgent extends Disposable implements IAgent {
 		});
 	}
 
-	async truncateSession(session: URI, turnIndex?: number): Promise<void> {
+	async truncateSession(session: URI, turnId?: string): Promise<void> {
 		const sessionId = AgentSession.id(session);
 		await this._sessionSequencer.queue(sessionId, async () => {
-			this._logService.info(`[Copilot:${sessionId}] Truncating session${turnIndex !== undefined ? ` at index ${turnIndex}` : ' (all turns)'}`);
-
-			const keepUpToTurnIndex = turnIndex ?? -1;
-
-			// Destroy the SDK session first and wait for cleanup to complete,
-			// ensuring on-disk data (events.jsonl, locks) is released before
-			// we modify it. Then dispose the wrapper.
-			const entry = this._sessions.get(sessionId);
-			if (entry) {
-				await entry.destroySession();
+			this._logService.info(`[Copilot:${sessionId}] Truncating session${turnId !== undefined ? ` at turnId=${turnId}` : ' (all turns)'}`);
+
+			// Ensure the session is loaded so we can use the SDK RPC
+			const entry = this._sessions.get(sessionId) ?? await this._resumeSession(sessionId);
+
+			// Look up the SDK event ID for the truncation boundary.
+			// The protocol semantics: turnId is the last turn to KEEP.
+			// The SDK semantics: eventId and all events after it are removed.
+			// So we need the event ID of the *next* turn after turnId.
+			// For "remove all", we need the first turn's event ID.
+			let eventId: string | undefined;
+			if (turnId) {
+				eventId = await entry.getNextTurnEventId(turnId);
+			} else {
+				eventId = await entry.getFirstTurnEventId();
 			}
-			this._sessions.deleteAndDispose(sessionId);
-
-			const copilotDataDir = getCopilotDataDir();
-			await truncateCopilotSessionOnDisk(copilotDataDir, sessionId, keepUpToTurnIndex);
 
-			// Resume the session from the modified on-disk data
-			await this._resumeSession(sessionId);
-			this._logService.info(`[Copilot:${sessionId}] Session truncated and resumed`);
-		});
-	}
-
-	async forkSession(sourceSession: URI, newSessionId: string, turnIndex: number): Promise<void> {
-		const sourceSessionId = AgentSession.id(sourceSession);
-		await this._sessionSequencer.queue(sourceSessionId, async () => {
-			this._logService.info(`[Copilot] Forking session ${sourceSessionId} at index ${turnIndex} → ${newSessionId}`);
+			if (eventId) {
+				await entry.truncateAtEventId(eventId, turnId);
+			} else {
+				this._logService.info(`[Copilot:${sessionId}] No event ID found for truncation, nothing to truncate`);
+			}
 
-			const copilotDataDir = getCopilotDataDir();
-			await forkCopilotSessionOnDisk(copilotDataDir, sourceSessionId, newSessionId, turnIndex);
-			this._logService.info(`[Copilot] Forked session ${newSessionId} created on disk`);
+			this._logService.info(`[Copilot:${sessionId}] Session truncated`);
 		});
 	}