agentHost: support checkpointing and forking

This enables restore/undo checkpoint for agent host sessions. It also
does some initial work for forking although this is not yet fully
implemented.

The copilot SDK does not actually support these yet, so to do these we
shut down the session, rewrite copilot SDK's disk storage, and start it
back up again. It actually works. We'll need to make sure it works when
we upgrade the SDK, but I don't expect it to break terribly, as the
Copilot SDK folks must already be backwards-compatible to arbitrary old
SDK versions that exist on the user's device, and we'd essentially just
be an 'old SDK' with some dependency on internals. Of course that should
all be swapped out when they eventually add proper support for it.

I just flagged the specific scheme of agent host sessions thus far while
developing, but will clean up prior to merging.

Author: connor4312

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

@@ -101,6 +101,8 @@ export interface IAgentCreateSessionConfig {
 	readonly model?: string;
 	readonly session?: URI;
 	readonly workingDirectory?: URI;
+	/** Fork from an existing session at a specific turn index. */
+	readonly fork?: { readonly session: URI; readonly turnIndex: number };
 }
 
 /** Serializable attachment passed alongside a message to the agent host. */
@@ -364,6 +366,23 @@ 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.
+	 * 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>;
+
 	/** Gracefully shut down all sessions. */
 	shutdown(): Promise<void>;
 

src/vs/platform/agentHost/common/state/protocol/action-origin.generated.ts

@@ -9,7 +9,7 @@
 // Generated from types/actions.ts — do not edit
 // Run `npm run generate` to regenerate.
 
-import { ActionType, type IStateAction, type IRootAgentsChangedAction, type IRootActiveSessionsChangedAction, type ISessionReadyAction, type ISessionCreationFailedAction, type ISessionTurnStartedAction, type ISessionDeltaAction, type ISessionResponsePartAction, type ISessionToolCallStartAction, type ISessionToolCallDeltaAction, type ISessionToolCallReadyAction, type ISessionToolCallConfirmedAction, type ISessionToolCallCompleteAction, type ISessionToolCallResultConfirmedAction, type ISessionTurnCompleteAction, type ISessionTurnCancelledAction, type ISessionErrorAction, type ISessionTitleChangedAction, type ISessionUsageAction, type ISessionReasoningAction, type ISessionModelChangedAction, type ISessionServerToolsChangedAction, type ISessionActiveClientChangedAction, type ISessionActiveClientToolsChangedAction, type ISessionPendingMessageSetAction, type ISessionPendingMessageRemovedAction, type ISessionQueuedMessagesReorderedAction, type ISessionCustomizationsChangedAction, type ISessionCustomizationToggledAction } from './actions.js';
+import { ActionType, type IStateAction, type IRootAgentsChangedAction, type IRootActiveSessionsChangedAction, type ISessionReadyAction, type ISessionCreationFailedAction, type ISessionTurnStartedAction, type ISessionDeltaAction, type ISessionResponsePartAction, type ISessionToolCallStartAction, type ISessionToolCallDeltaAction, type ISessionToolCallReadyAction, type ISessionToolCallConfirmedAction, type ISessionToolCallCompleteAction, type ISessionToolCallResultConfirmedAction, type ISessionTurnCompleteAction, type ISessionTurnCancelledAction, type ISessionErrorAction, type ISessionTitleChangedAction, type ISessionUsageAction, type ISessionReasoningAction, type ISessionModelChangedAction, type ISessionServerToolsChangedAction, type ISessionActiveClientChangedAction, type ISessionActiveClientToolsChangedAction, type ISessionPendingMessageSetAction, type ISessionPendingMessageRemovedAction, type ISessionQueuedMessagesReorderedAction, type ISessionCustomizationsChangedAction, type ISessionCustomizationToggledAction, type ISessionTruncatedAction } from './actions.js';
 
 
 // ─── Root vs Session Action Unions ───────────────────────────────────────────
@@ -48,6 +48,7 @@ export type ISessionAction =
 	| ISessionQueuedMessagesReorderedAction
 	| ISessionCustomizationsChangedAction
 	| ISessionCustomizationToggledAction
+	| ISessionTruncatedAction
 	;
 
 /** Union of session actions that clients may dispatch. */
@@ -65,6 +66,7 @@ export type IClientSessionAction =
 	| ISessionPendingMessageRemovedAction
 	| ISessionQueuedMessagesReorderedAction
 	| ISessionCustomizationToggledAction
+	| ISessionTruncatedAction
 	;
 
 /** Union of session actions that only the server may produce. */
@@ -119,4 +121,5 @@ export const IS_CLIENT_DISPATCHABLE: { readonly [K in IStateAction['type']]: boo
 	[ActionType.SessionQueuedMessagesReordered]: true,
 	[ActionType.SessionCustomizationsChanged]: false,
 	[ActionType.SessionCustomizationToggled]: true,
+	[ActionType.SessionTruncated]: true,
 };

src/vs/platform/agentHost/common/state/protocol/actions.ts

@@ -45,6 +45,7 @@ export const enum ActionType {
 	SessionQueuedMessagesReordered = 'session/queuedMessagesReordered',
 	SessionCustomizationsChanged = 'session/customizationsChanged',
 	SessionCustomizationToggled = 'session/customizationToggled',
+	SessionTruncated = 'session/truncated',
 }
 
 // ─── Action Envelope ─────────────────────────────────────────────────────────
@@ -562,6 +563,31 @@ export interface ISessionCustomizationToggledAction {
 	enabled: boolean;
 }
 
+// ─── Truncation ──────────────────────────────────────────────────────────────
+
+/**
+ * Truncates a session's history. If `turnId` is provided, all turns after that
+ * turn are removed and the specified turn is kept. If `turnId` is omitted, all
+ * turns are removed.
+ *
+ * If there is an active turn it is silently dropped and the session status
+ * returns to `idle`.
+ *
+ * Common use-case: truncate old data then dispatch a new
+ * `session/turnStarted` with an edited message.
+ *
+ * @category Session Actions
+ * @version 1
+ * @clientDispatchable
+ */
+export interface ISessionTruncatedAction {
+	type: ActionType.SessionTruncated;
+	/** Session URI */
+	session: URI;
+	/** Keep turns up to and including this turn. Omit to clear all turns. */
+	turnId?: string;
+}
+
 // ─── Pending Message Actions ─────────────────────────────────────────────────
 
 /**
@@ -664,4 +690,5 @@ export type IStateAction =
 	| ISessionPendingMessageRemovedAction
 	| ISessionQueuedMessagesReorderedAction
 	| ISessionCustomizationsChangedAction
-	| ISessionCustomizationToggledAction;
+	| ISessionCustomizationToggledAction
+	| ISessionTruncatedAction;

src/vs/platform/agentHost/common/state/protocol/commands.ts

@@ -163,6 +163,20 @@ export interface ISubscribeResult {
  * { "jsonrpc": "2.0", "id": 2, "error": { "code": -32003, "message": "Session already exists" } }
  * ```
  */
+/**
+ * Identifies a source session and turn to fork from.
+ *
+ * When provided in `createSession`, the server populates the new session with
+ * content from the source session up to and including the response of the
+ * specified turn.
+ */
+export interface ISessionForkSource {
+	/** URI of the existing session to fork from */
+	session: URI;
+	/** Turn ID in the source session; content up to and including this turn's response is copied */
+	turnId: string;
+}
+
 export interface ICreateSessionParams {
 	/** Session URI (client-chosen, e.g. `copilot:/<uuid>`) */
 	session: URI;
@@ -172,6 +186,11 @@ export interface ICreateSessionParams {
 	model?: string;
 	/** Working directory for the session */
 	workingDirectory?: URI;
+	/**
+	 * Fork from an existing session. The new session is populated with content
+	 * from the source session up to and including the specified turn's response.
+	 */
+	fork?: ISessionForkSource;
 }
 
 // ─── disposeSession ──────────────────────────────────────────────────────────

src/vs/platform/agentHost/common/state/protocol/reducers.ts

@@ -485,6 +485,27 @@ export function sessionReducer(state: ISessionState, action: ISessionAction, log
 			return { ...state, customizations: updated };
 		}
 
+		// ── Truncation ────────────────────────────────────────────────────────
+
+		case ActionType.SessionTruncated: {
+			let turns: typeof state.turns;
+			if (action.turnId === undefined) {
+				turns = [];
+			} else {
+				const idx = state.turns.findIndex(t => t.id === action.turnId);
+				if (idx < 0) {
+					return state;
+				}
+				turns = state.turns.slice(0, idx + 1);
+			}
+			return {
+				...state,
+				turns,
+				activeTurn: undefined,
+				summary: { ...state.summary, status: SessionStatus.Idle, modifiedAt: Date.now() },
+			};
+		}
+
 		// ── Pending Messages ──────────────────────────────────────────────────
 
 		case ActionType.SessionPendingMessageSet: {

src/vs/platform/agentHost/common/state/protocol/version/registry.ts

@@ -52,6 +52,7 @@ export const ACTION_INTRODUCED_IN: { readonly [K in IStateAction['type']]: numbe
 	[ActionType.SessionQueuedMessagesReordered]: 1,
 	[ActionType.SessionCustomizationsChanged]: 1,
 	[ActionType.SessionCustomizationToggled]: 1,
+	[ActionType.SessionTruncated]: 1,
 };
 
 /**

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

@@ -173,17 +173,41 @@ export class AgentService extends Disposable implements IAgentService {
 		this._sessionToProvider.set(session.toString(), provider.id);
 		this._logService.trace(`[AgentService] createSession returned: ${session.toString()}`);
 
-		// Create state in the state manager
-		const summary: ISessionSummary = {
-			resource: session.toString(),
-			provider: provider.id,
-			title: 'New Session',
-			status: SessionStatus.Idle,
-			createdAt: Date.now(),
-			modifiedAt: Date.now(),
-			workingDirectory: config?.workingDirectory?.toString(),
-		};
-		this._stateManager.createSession(summary);
+		// When forking, populate the new session's protocol state with
+		// the source session's turns so the client sees the forked history.
+		if (config?.fork) {
+			const sourceState = this._stateManager.getSessionState(config.fork.session.toString());
+			let sourceTurns: ITurn[] = [];
+			if (sourceState) {
+				const forkIdx = sourceState.turns.findIndex(t => t.id === config.fork!.turnId);
+				if (forkIdx >= 0) {
+					sourceTurns = sourceState.turns.slice(0, forkIdx + 1);
+				}
+			}
+
+			const summary: ISessionSummary = {
+				resource: session.toString(),
+				provider: provider.id,
+				title: sourceState?.summary.title ?? 'Forked Session',
+				status: SessionStatus.Idle,
+				createdAt: Date.now(),
+				modifiedAt: Date.now(),
+				workingDirectory: config.workingDirectory?.toString(),
+			};
+			this._stateManager.restoreSession(summary, sourceTurns);
+		} else {
+			// Create empty state for new sessions
+			const summary: ISessionSummary = {
+				resource: session.toString(),
+				provider: provider.id,
+				title: 'New Session',
+				status: SessionStatus.Idle,
+				createdAt: Date.now(),
+				modifiedAt: Date.now(),
+				workingDirectory: config?.workingDirectory?.toString(),
+			};
+			this._stateManager.createSession(summary);
+		}
 		this._stateManager.dispatchServerAction({ type: ActionType.SessionReady, session: session.toString() });
 
 		return session;

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

@@ -252,6 +252,24 @@ export class AgentSideEffects extends Disposable {
 				this._syncPendingMessages(action.session);
 				break;
 			}
+			case ActionType.SessionTruncated: {
+				const agent = this._options.getAgent(action.session);
+				// Resolve the protocol turnId to a 0-based index using the
+				// state manager's turn list. The reducer has already applied
+				// the truncation, so we look at the pre-truncation state via
+				// the turnId position.
+				let turnIndex: number | undefined;
+				if (action.turnId !== undefined) {
+					const state = this._stateManager.getSessionState(action.session);
+					// After the reducer, the turns array is already truncated.
+					// The kept turns include the target, so its index = length - 1.
+					turnIndex = state ? state.turns.length - 1 : undefined;
+				}
+				agent?.truncateSession?.(URI.parse(action.session), turnIndex).catch(err => {
+					this._logService.error('[AgentSideEffects] truncateSession failed', err);
+				});
+				break;
+			}
 		}
 	}