Harden remote agent host session lifecycle

Written by Copilot

Author: roblourens

src/vs/platform/agentHost/electron-browser/remoteAgentHostProtocolClient.ts

@@ -21,7 +21,7 @@ import { AgentSubscriptionManager, type IAgentSubscription } from '../common/sta
 import { agentHostAuthority, fromAgentHostUri, toAgentHostUri } from '../common/agentHostUri.js';
 import type { IClientNotificationMap, ICommandMap } from '../common/state/protocol/messages.js';
 import type { IActionEnvelope, INotification, ISessionAction, ITerminalAction } from '../common/state/sessionActions.js';
-import { ISessionSummary, ROOT_STATE_URI, StateComponents, type IRootState } from '../common/state/sessionState.js';
+import { ISessionSummary, ROOT_STATE_URI, StateComponents, type IRootState, type ISessionState } from '../common/state/sessionState.js';
 import { PROTOCOL_VERSION } from '../common/state/sessionCapabilities.js';
 import { isJsonRpcNotification, isJsonRpcRequest, isJsonRpcResponse, type IJsonRpcResponse, type IProtocolMessage, type IStateSnapshot } from '../common/state/sessionProtocol.js';
 import { isClientTransport, type IProtocolTransport } from '../common/state/sessionTransport.js';
@@ -182,16 +182,30 @@ export class RemoteAgentHostProtocolClient extends Disposable implements IAgentC
 	async createSession(config?: IAgentCreateSessionConfig): Promise<URI> {
 		const provider = config?.provider ?? 'copilot';
 		const session = AgentSession.uri(provider, generateUuid());
+		const fork = this._toProtocolFork(config?.fork);
 		await this._sendRequest('createSession', {
 			session: session.toString(),
 			provider,
 			model: config?.model,
 			workingDirectory: config?.workingDirectory ? fromAgentHostUri(config.workingDirectory).toString() : undefined,
+			fork,
 			config: config?.config,
 		});
 		return session;
 	}
 
+	private _toProtocolFork(fork: IAgentCreateSessionConfig['fork']): { readonly session: string; readonly turnId: string } | undefined {
+		if (!fork) {
+			return undefined;
+		}
+		const sourceState = this.getSubscriptionUnmanaged<ISessionState>(StateComponents.Session, fork.session)?.value;
+		const turnId = sourceState?.turns[fork.turnIndex]?.id;
+		if (!turnId) {
+			throw new Error(`Cannot fork: turn index ${fork.turnIndex} not found in protocol state for ${fork.session.toString()}`);
+		}
+		return { session: fork.session.toString(), turnId };
+	}
+
 	async resolveSessionConfig(params: IAgentResolveSessionConfigParams): Promise<IResolveSessionConfigResult> {
 		return this._sendRequest('resolveSessionConfig', {
 			provider: params.provider,

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

@@ -651,9 +651,10 @@ export class CopilotAgent extends Disposable implements IAgent {
 		}
 		try {
 			await this._gitService.removeWorktree(worktree.repositoryRoot, worktree.worktree);
-			this._createdWorktrees.delete(sessionId);
 		} catch (error) {
 			this._logService.warn(`[Copilot:${sessionId}] Failed to remove worktree '${worktree.worktree.fsPath}': ${error instanceof Error ? error.message : String(error)}`);
+		} finally {
+			this._createdWorktrees.delete(sessionId);
 		}
 	}
 

src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts

@@ -395,6 +395,12 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 		this._onDidDisconnect.fire();
 		this._connection = undefined;
 		this._defaultDirectory = undefined;
+		if (this._currentNewSession) {
+			this._clearNewSessionConfig(this._currentNewSession.id);
+			this._currentNewSession = undefined;
+		}
+		this._currentNewSessionStatus = undefined;
+		this._selectedModelId = undefined;
 
 		if (this._sessionTypes.length > 0) {
 			this._sessionTypes = [];

src/vs/sessions/contrib/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts

@@ -439,6 +439,27 @@ suite('RemoteAgentHostSessionsProvider', () => {
 		});
 	});
 
+	test('clearConnection clears pending new session config', () => {
+		const provider = createProvider(disposables, connection);
+		const workspace = {
+			label: 'my-project',
+			icon: { id: 'remote' },
+			repositories: [{ uri: URI.parse('vscode-agent-host://auth/home/user/project'), workingDirectory: undefined, detail: undefined, baseBranchName: undefined, baseBranchProtected: undefined }],
+			requiresWorkspaceTrust: false,
+		};
+
+		const session = provider.createNewSession(workspace);
+		provider.clearConnection();
+
+		assert.deepStrictEqual({
+			resolved: provider.getSessionByResource(session.resource),
+			config: provider.getSessionConfig(session.sessionId),
+		}, {
+			resolved: undefined,
+			config: undefined,
+		});
+	});
+
 	test('createNewSession throws when no repository URI', () => {
 		const provider = createProvider(disposables, connection);
 		const workspace = { label: 'empty', icon: { id: 'remote' }, repositories: [], requiresWorkspaceTrust: false };