Fix sandbox activation races and initial connection interactivity

Reject stale activation results after cancellation, feature teardown, or provider replacement. Keep initial sandbox connections read-only while preserving input during self-healing reconnects.

Add regression coverage for review feedback on #334644, including actual environment lookup failures.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: osortega

src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts

@@ -463,44 +463,45 @@ export class CloudSandboxAgentHostContribution extends Disposable implements IWo
 		return authority ? byAuthority.get(authority) : undefined;
 	}
 
-	/**
-	 * Async-activation hook for a sandbox session type: make the session openable, either by
-	 * connecting to an environment that is already online — resolving once the host advertises the
-	 * agent backing this session type, so the chat can load — or by serving its persisted history
-	 * read-only. Returns false only when the environment is unknown, or nothing can be shown at
-	 * all: no live host and no history to fall back on.
-	 */
+	/** Opens an online environment through its host, or an offline session from persisted history. */
 	protected async _waitForActivation(sessionType: string): Promise<boolean> {
 		const address = this._findAddressForSessionType(sessionType);
 		const env = address ? this._environments.get(address) : undefined;
-		if (!address || !env) {
+		const provider = address ? this._providerInstances.get(address) : undefined;
+		if (!address || !env || !provider) {
 			return false;
 		}
-		// Resuming is a side effect the user has to ask for. Mission Control wakes the environment
-		// behind `/connect`, and it cannot say in advance whether a dormant one will come back, so
-		// opening a session must not gamble minutes of wake on the guess. Read the environment's
-		// state first and dial only what is already online; anything else opens from history with
-		// the connection banner offering the connect.
-		//
-		// Without a task there is no history to serve, and refusing to open would leave the
-		// session unreachable entirely — so fall through to the connect, which is the only way
-		// such a session can show anything at all.
-		if (env.taskId && !await this._isEnvironmentOnline(env)) {
+		const token = this._enabledCts.token;
+		const isCurrentActivation = () => {
+			const current = !token.isCancellationRequested
+				&& this._isEnabled()
+				&& this._environments.has(address)
+				&& this._providerInstances.get(address) === provider;
+			if (!current) {
+				this._logService.trace(`${LOG_PREFIX} Abandoning activation for ${address} after teardown.`);
+			}
+			return current;
+		};
+
+		// Without a task there is no history fallback, so connecting is the only way to open it.
+		const shouldConnect = !env.taskId || await this._isEnvironmentOnline(env, token);
+		if (!isCurrentActivation()) {
+			return false;
+		}
+		if (!shouldConnect) {
 			this._logService.info(`${LOG_PREFIX} Environment for ${address} is not online; serving history and leaving the connect to the user.`);
-			return this._activateReadOnly(sessionType, address, env, this._fetchTaskHistory(env));
+			return this._activateReadOnly(sessionType, address, env, this._fetchTaskHistory(env, token));
 		}
 
 		const connectError = await this
 			.connect({ environmentId: env.environmentId, sessionId: env.sessionId, name: env.name })
 			.then(() => undefined, (error: unknown) => error ?? new Error('connect failed'));
+		if (!isCurrentActivation()) {
+			return false;
+		}
 		if (connectError !== undefined) {
 			this._logService.warn(`${LOG_PREFIX} connect-on-open failed for ${address}: ${connectError instanceof Error ? connectError.message : String(connectError)}`);
-			// Serve history whatever the reason: `/connect` fails in several ways for a deleted
-			// sandbox, so gating on any one of them would leave the rest with no history.
-			if (this._isEnabled() && !this._enabledCts.token.isCancellationRequested) {
-				return this._activateReadOnly(sessionType, address, env, this._fetchTaskHistory(env));
-			}
-			return false;
+			return this._activateReadOnly(sessionType, address, env, this._fetchTaskHistory(env, token));
 		}
 		const authority = agentHostAuthority(address);
 		while (true) {
@@ -519,34 +520,23 @@ export class CloudSandboxAgentHostContribution extends Disposable implements IWo
 		}
 	}
 
-	/**
-	 * Whether the environment currently has a daemon listening on the relay. Only `online` does, so
-	 * only `online` can be dialled without triggering a resume.
-	 *
-	 * An unreadable record answers `false`: this gates a side effect the user has not asked for, so
-	 * a Mission Control blip must cost a click rather than an unrequested wake.
-	 */
-	private async _isEnvironmentOnline(env: ICloudSandboxEnvironment): Promise<boolean> {
+	/** An unreadable record must not trigger an automatic resume. */
+	private async _isEnvironmentOnline(env: ICloudSandboxEnvironment, token: CancellationToken): Promise<boolean> {
 		try {
-			const record = await this._apiService.getEnvironment(env.environmentId, this._enabledCts.token);
+			const record = await this._apiService.getEnvironment(env.environmentId, token);
 			return record.status === 'online';
 		} catch (error) {
 			this._logService.trace(`${LOG_PREFIX} Could not read the state of ${env.environmentId}; treating it as not online: ${error instanceof Error ? error.message : String(error)}`);
 			return false;
 		}
 	}
 
-	/**
-	 * A task's persisted history, or `undefined` when there is no task or the read failed. Served
-	 * by Mission Control rather than the sandbox, so it stays readable while the environment is
-	 * asleep or gone. Never rejects.
-	 */
-	private _fetchTaskHistory(env: ICloudSandboxEnvironment): Promise<IReplayedTaskHistory | undefined> | undefined {
+	/** Reads history from Mission Control without connecting to the sandbox. */
+	private _fetchTaskHistory(env: ICloudSandboxEnvironment, token: CancellationToken): Promise<IReplayedTaskHistory | undefined> | undefined {
 		const taskId = env.taskId;
 		if (!taskId) {
 			return undefined;
 		}
-		const token = this._enabledCts.token;
 		return this._apiService.getSessionHistory(taskId, token).catch((error: unknown) => {
 			this._logService.trace(`${LOG_PREFIX} History read for ${env.environmentId} did not complete: ${error instanceof Error ? error.message : String(error)}`);
 			return undefined;

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

@@ -91,12 +91,7 @@ export interface IRemoteAgentHostSessionsProviderConfig {
 	readonly omitHostFromWorkspaceLabel?: boolean;
 	/** Type icon for this host's workspaces. See {@link ISessionWorkspace.typeIcon}. */
 	readonly workspaceTypeIcon?: ThemeIcon;
-	/**
-	 * Forces this host's sessions read-only whenever it is not connected. Set by hosts that cannot
-	 * accept work offline — a cloud sandbox has to be resumed before it can be sent to, so a
-	 * composer would take input that nothing will ever deliver. Hosts left with the default keep
-	 * their sessions writable while disconnected, queuing the input for the reconnect.
-	 */
+	/** Keeps unavailable and initially connecting sessions read-only, but permits self-healing reconnects. */
 	readonly readOnlyWhenDisconnected?: boolean;
 	/** See {@link IAgentHostAdapterOptions.defaultChangesetKind}. */
 	readonly defaultChangesetKind?: ChangesetKind.Branch | ChangesetKind.Uncommitted | ChangesetKind.Session;
@@ -148,14 +143,6 @@ export class RemoteAgentHostSessionsProvider extends BaseAgentHostSessionsProvid
 	private readonly _automationStore: ReconnectableAgentHostAutomationStore;
 
 	private readonly _connectionStatus = observableValue<RemoteAgentHostConnectionStatus>('connectionStatus', RemoteAgentHostConnectionStatus.disconnected);
-	/**
-	 * Whether every session on this host is read-only, which hides the composer.
-	 *
-	 * Only ever true for hosts configured with
-	 * {@link IRemoteAgentHostSessionsProviderConfig.readOnlyWhenDisconnected}, and only while they
-	 * are disconnected: elsewhere a dropped host keeps its sessions writable so the input queues
-	 * for the reconnect.
-	 */
 	private readonly _readOnly: IObservable<boolean>;
 	readonly connectionStatus: IObservable<RemoteAgentHostConnectionStatus> = this._connectionStatus;
 
@@ -256,13 +243,10 @@ export class RemoteAgentHostSessionsProvider extends BaseAgentHostSessionsProvid
 		this.connectionLabels = config.connectionLabels;
 		this.canConnectOnDemand = !!config.connectOnDemand;
 		this._readOnly = config.readOnlyWhenDisconnected
-			// Deliberately not `!isConnected`: that would include `reconnecting`, which is the
-			// protocol client restoring a dropped transport while the host itself is up. Input
-			// sent then is delivered once the transport returns, so hiding the composer would
-			// interrupt a conversation mid-sentence over a blip the user should not notice.
 			? derived(this, reader => {
 				const status = this._connectionStatus.read(reader);
 				return RemoteAgentHostConnectionStatus.isDisconnected(status)
+					|| RemoteAgentHostConnectionStatus.isConnecting(status)
 					|| RemoteAgentHostConnectionStatus.isIncompatible(status);
 			})
 			: constObservable(false);

src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostContribution.test.ts

@@ -4,7 +4,9 @@
  *--------------------------------------------------------------------------------------------*/
 
 import assert from 'assert';
+import { DeferredPromise } from '../../../../../../base/common/async.js';
 import { CancellationToken } from '../../../../../../base/common/cancellation.js';
+import { CancellationError } from '../../../../../../base/common/errors.js';
 import { Event } from '../../../../../../base/common/event.js';
 import { Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js';
 import { mock, upcastPartial } from '../../../../../../base/test/common/mock.js';
@@ -31,7 +33,7 @@ import {
 } from '../../../../../../platform/agentHost/common/cloudSandboxAgentHost.js';
 import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js';
 import { IObservable, observableValue } from '../../../../../../base/common/observable.js';
-import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js';
+import { ConfigurationTarget, IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js';
 import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js';
 import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js';
 import { ILogService, NullLogService } from '../../../../../../platform/log/common/log.js';
@@ -156,6 +158,7 @@ const GITHUB_SANDBOX_GROUP: IAgentHostGroup = {
 interface ITestHarness {
 	readonly contribution: TestCloudSandboxContribution;
 	readonly configurationService: TestConfigurationService;
+	setEnabled(enabled: boolean): Promise<void>;
 	/** Discovery's answer, mutable so a test can change what a later pass reports. */
 	discovered: readonly ICloudSandboxDiscoveredSession[];
 	/** Runs a discovery pass and waits for it to reconcile. */
@@ -170,6 +173,7 @@ interface ITestHarness {
 	activate(environmentId: string): Promise<boolean>;
 	readonly created: ICloudSandboxCreateSessionRequest[];
 	readonly connectedTo: string[];
+	readonly historyRequests: string[];
 	/** Host groups currently declared to the filter service. */
 	readonly hostGroups: IAgentHostGroup[];
 }
@@ -182,6 +186,7 @@ interface ITestHarness {
 async function createContribution(store: Pick<DisposableStore, 'add'>, sessions: readonly ICloudSandboxDiscoveredSession[], options?: {
 	/** Task Mission Control returns from `createSession`, or a rejection. */
 	readonly createSession?: () => Promise<ICloudSandboxCreatedSession>;
+	readonly getEnvironment?: (id: string, token: CancellationToken) => Promise<ICloudSandboxEnvironmentRecord>;
 	/** Whether the sandbox feature settings start on. Defaults to `true`. */
 	readonly enabled?: boolean;
 }): Promise<ITestHarness> {
@@ -191,13 +196,24 @@ async function createContribution(store: Pick<DisposableStore, 'add'>, sessions:
 	const instantiationService = store.add(new TestInstantiationService());
 	const created: ICloudSandboxCreateSessionRequest[] = [];
 	const connectedTo: string[] = [];
+	const historyRequests: string[] = [];
 	const harness: ITestHarness = {
 		discovered: sessions,
 		environmentStatus: 'offline',
 		readOnlySessionTypes,
 		created,
 		connectedTo,
+		historyRequests,
 		hostGroups,
+		setEnabled: async (enabled: boolean) => {
+			await configurationService.setUserConfiguration(CloudSandboxEnabledSettingId, enabled);
+			configurationService.onDidChangeConfigurationEmitter.fire({
+				affectsConfiguration: key => key === CloudSandboxEnabledSettingId,
+				affectedKeys: new Set([CloudSandboxEnabledSettingId]),
+				change: { keys: [CloudSandboxEnabledSettingId], overrides: [] },
+				source: ConfigurationTarget.USER,
+			});
+		},
 		runDiscovery: async () => { await Promise.all(discoveryHandlers.map(handler => handler())); },
 		activate: async (environmentId: string) => {
 			const sessionType = remoteAgentHostSessionTypeId(agentHostAuthority(cloudSandboxAddress(environmentId)), CLOUD_SANDBOX_AGENT_PROVIDER);
@@ -209,10 +225,14 @@ async function createContribution(store: Pick<DisposableStore, 'add'>, sessions:
 		override async listSessions(_token: CancellationToken): Promise<ICloudSandboxDiscoveryResult> {
 			return { kind: 'complete', sessions: harness.discovered };
 		}
-		override async getEnvironment(id: string): Promise<ICloudSandboxEnvironmentRecord> {
+		override async getEnvironment(id: string, token: CancellationToken): Promise<ICloudSandboxEnvironmentRecord> {
+			if (options?.getEnvironment) {
+				return options.getEnvironment(id, token);
+			}
 			return { id, status: harness.environmentStatus };
 		}
-		override async getSessionHistory(): Promise<IReplayedTaskHistory> {
+		override async getSessionHistory(taskId: string): Promise<IReplayedTaskHistory> {
+			historyRequests.push(taskId);
 			return { sessions: [], truncated: false };
 		}
 		override async createSession(request: ICloudSandboxCreateSessionRequest): Promise<ICloudSandboxCreatedSession> {
@@ -414,13 +434,120 @@ suite('CloudSandboxAgentHostContribution', () => {
 	});
 
 	test('does not wake an environment whose state could not be read', async () => {
-		// A Mission Control blip must cost a click, not an unrequested resume.
-		const harness = await createContribution(store, [discoveredSession()]);
-		harness.environmentStatus = 'degraded';
+		const harness = await createContribution(store, [discoveredSession()], {
+			getEnvironment: async () => { throw new Error('Expected environment lookup failure'); },
+		});
 
 		const opened = await harness.activate('env-1');
 
-		assert.deepStrictEqual({ opened, connectedTo: harness.connectedTo }, { opened: true, connectedTo: [] });
+		assert.deepStrictEqual({
+			opened,
+			connectedTo: harness.connectedTo,
+			historyRequests: harness.historyRequests,
+			servedFromHistory: harness.readOnlySessionTypes.length,
+		}, { opened: true, connectedTo: [], historyRequests: ['task-1'], servedFromHistory: 1 });
+	});
+
+	for (const reenable of [false, true]) {
+		test(`abandons a cancelled environment lookup when the feature is ${reenable ? 're-enabled' : 'disabled'}`, async () => {
+			const environment = new DeferredPromise<ICloudSandboxEnvironmentRecord>();
+			const requestedToken = new DeferredPromise<CancellationToken>();
+			const harness = await createContribution(store, [discoveredSession()], {
+				getEnvironment: (_id, token) => {
+					void requestedToken.complete(token);
+					return environment.p;
+				},
+			});
+			const activation = harness.activate('env-1');
+			const token = await requestedToken.p;
+
+			await harness.setEnabled(false);
+			if (reenable) {
+				await harness.setEnabled(true);
+				await harness.runDiscovery();
+			}
+			await environment.error(new CancellationError());
+
+			assert.deepStrictEqual({
+				opened: await activation,
+				cancelled: token.isCancellationRequested,
+				connectedTo: harness.connectedTo,
+				historyRequests: harness.historyRequests,
+				readOnlySessionTypes: harness.readOnlySessionTypes,
+			}, { opened: false, cancelled: true, connectedTo: [], historyRequests: [], readOnlySessionTypes: [] });
+		});
+	}
+
+	for (const status of ['online', 'offline'] as const) {
+		test(`does not reactivate a removed environment after a late ${status} record`, async () => {
+			const environment = new DeferredPromise<ICloudSandboxEnvironmentRecord>();
+			const harness = await createContribution(store, [discoveredSession()], {
+				getEnvironment: () => environment.p,
+			});
+			const activation = harness.activate('env-1');
+			harness.discovered = [];
+			await harness.runDiscovery();
+			await environment.complete({ id: 'env-1', status });
+
+			assert.deepStrictEqual({
+				opened: await activation,
+				connectedTo: harness.connectedTo,
+				historyRequests: harness.historyRequests,
+				readOnlySessionTypes: harness.readOnlySessionTypes,
+			}, { opened: false, connectedTo: [], historyRequests: [], readOnlySessionTypes: [] });
+		});
+	}
+
+	test('does not register old history against a replacement provider at the same address', async () => {
+		const environment = new DeferredPromise<ICloudSandboxEnvironmentRecord>();
+		const harness = await createContribution(store, [discoveredSession()], {
+			getEnvironment: () => environment.p,
+		});
+		const activation = harness.activate('env-1');
+		harness.discovered = [];
+		await harness.runDiscovery();
+		harness.discovered = [discoveredSession({ taskId: 'task-2' })];
+		await harness.runDiscovery();
+		await environment.complete({ id: 'env-1', status: 'offline' });
+
+		assert.deepStrictEqual({
+			opened: await activation,
+			historyRequests: harness.historyRequests,
+			readOnlySessionTypes: harness.readOnlySessionTypes,
+		}, { opened: false, historyRequests: [], readOnlySessionTypes: [] });
+	});
+
+	test('keeps activation valid across a discovery refresh of the same provider', async () => {
+		const environment = new DeferredPromise<ICloudSandboxEnvironmentRecord>();
+		const harness = await createContribution(store, [discoveredSession()], {
+			getEnvironment: () => environment.p,
+		});
+		const activation = harness.activate('env-1');
+		await harness.runDiscovery();
+		await environment.complete({ id: 'env-1', status: 'offline' });
+
+		assert.deepStrictEqual({
+			opened: await activation,
+			connectedTo: harness.connectedTo,
+			historyRequests: harness.historyRequests,
+		}, { opened: true, connectedTo: [], historyRequests: ['task-1'] });
+	});
+
+	test('does not restore history after an old connect fails across disable and re-enable', async () => {
+		const harness = await createContribution(store, [discoveredSession()]);
+		harness.environmentStatus = 'online';
+		harness.onConnect = async () => {
+			await harness.setEnabled(false);
+			await harness.setEnabled(true);
+			await harness.runDiscovery();
+			throw new Error('Expected connection failure after teardown');
+		};
+
+		assert.deepStrictEqual({
+			opened: await harness.activate('env-1'),
+			historyRequests: harness.historyRequests,
+			readOnlySessionTypes: harness.readOnlySessionTypes,
+		}, { opened: false, historyRequests: [], readOnlySessionTypes: [] });
 	});
 
 	test('connects a dormant environment that has no history to fall back on', async () => {