Merge pull request #306968 from microsoft/connor4312/ah-undo-redo

agentHost: support checkpointing and forking

Author: connor4312

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

@@ -102,6 +102,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. */
@@ -365,6 +367,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>;
+
 	/**
 	 * Receives client-provided customization refs and syncs them (e.g. copies
 	 * plugin files to local storage). Returns per-customization status with

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

@@ -184,17 +184,40 @@ 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) {
+				sourceTurns = sourceState.turns.slice(0, config.fork.turnIndex + 1)
+					.map(t => ({ ...t, id: generateUuid() }));
+			}
+
+			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(),
+			};
+			const state = this._stateManager.createSession(summary);
+			state.turns = 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

@@ -316,6 +316,23 @@ export class AgentSideEffects extends Disposable {
 				this._syncPendingMessages(action.session);
 				break;
 			}
+			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 => {
+					this._logService.error('[AgentSideEffects] truncateSession failed', err);
+				});
+				break;
+			}
 			case ActionType.SessionActiveClientChanged: {
 				const agent = this._options.getAgent(action.session);
 				const refs = action.activeClient?.customizations;