Wait for complete sandbox credentials before connecting

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

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { CancellationError } from '../../../../../base/common/errors.js';
+import { CancellationError, isCancellationError } from '../../../../../base/common/errors.js';
 import { CancellationToken } from '../../../../../base/common/cancellation.js';
 import { Disposable, DisposableMap, DisposableStore } from '../../../../../base/common/lifecycle.js';
 import { timeout } from '../../../../../base/common/async.js';
@@ -21,6 +21,7 @@ import {
 	ICloudSandboxConnectOptions,
 	ICloudSandboxApiService,
 	isCloudSandboxSealedToken,
+	type CloudSandboxConnectResult,
 	type ICloudSandboxClientToken,
 } from '../../../../../platform/agentHost/common/cloudSandboxAgentHost.js';
 import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js';
@@ -36,6 +37,15 @@ const LOG_PREFIX = '[CloudSandboxAgentHost]';
 /** Maximum number of `/connect` "waking" retries before giving up. */
 const MAX_WAKING_RETRIES = 20;
 
+/**
+ * Maximum number of `/connect` re-mints while the sealed token is missing, sized to cover the
+ * backend's own registration retry cycle.
+ */
+export const MAX_SEALED_TOKEN_RETRIES = 12;
+
+/** Delay between `/connect` re-mints while waiting for complete credentials. */
+const SEALED_TOKEN_RETRY_DELAY_MS = 5_000;
+
 /**
  * Renderer-side coordinator for Copilot cloud sandbox connections.
  *
@@ -54,6 +64,9 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa
 	/** Current Web PubSub credentials per connection address, including the sealed GitHub token. */
 	private readonly _creds = new Map<string, ICloudSandboxCreds>();
 
+	/** Overridable so tests can exercise the re-mint loop without waiting on real delays. */
+	protected readonly sealedTokenRetryDelayMs: number = SEALED_TOKEN_RETRY_DELAY_MS;
+
 	constructor(
 		@IRemoteAgentHostService private readonly _remoteAgentHostService: IRemoteAgentHostService,
 		@ICloudSandboxApiService private readonly _apiService: ICloudSandboxApiService,
@@ -112,7 +125,7 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa
 	 * Open the relay with an already-minted token, drive the AHP handshake, and register the
 	 * connection.
 	 */
-	private async _establish(options: ICloudSandboxConnectOptions, address: string, clientToken: ICloudSandboxClientToken, token: CancellationToken): Promise<string> {
+	protected async _establish(options: ICloudSandboxConnectOptions, address: string, clientToken: ICloudSandboxClientToken, token: CancellationToken): Promise<string> {
 		// Mutable holder read by the transport factory: the protocol client re-invokes the factory to
 		// soft-reconnect, picking up whatever credentials the refresh scheduler last wrote.
 		const creds: ICloudSandboxCreds = { token: clientToken };
@@ -172,6 +185,9 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa
 					this._logService.warn(`${LOG_PREFIX} Sealed-token authenticate failed for ${address}`, err);
 				}
 			}
+		} else if (!connectError) {
+			// Without an envelope every later request answers `-32007 AuthRequired`.
+			this._logService.error(`${LOG_PREFIX} Mission Control returned no sealed token for ${address}; this session will not be able to make authenticated requests.`);
 		}
 
 		try {
@@ -217,12 +233,47 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa
 			}
 			const result = await this._apiService.connect({ environmentId: options.environmentId, sessionId: options.sessionId }, token);
 			if (result.kind === 'token') {
-				return result.token;
+				return await this._awaitSealedToken(options, result.token, token);
 			}
 			const delayMs = Math.min(result.waking.retryAfterSeconds * 1000, MAX_WAKING_DELAY_MS);
 			this._logService.info(`${LOG_PREFIX} Environment ${options.environmentId} waking; retrying in ${delayMs}ms (attempt ${attempt + 1}/${MAX_WAKING_RETRIES})`);
 			await timeout(delayMs, token);
 		}
 		throw new Error(`Timed out waiting for sandbox environment ${options.environmentId} to wake.`);
 	}
+
+	/**
+	 * Re-mint credentials until they carry a sealed token, which a freshly provisioned environment
+	 * can omit for a short window after it comes up. Returns the last credentials either way, since
+	 * an environment may legitimately never seal one.
+	 */
+	private async _awaitSealedToken(options: ICloudSandboxConnectOptions, minted: ICloudSandboxClientToken, token: CancellationToken): Promise<ICloudSandboxClientToken> {
+		let clientToken = minted;
+		// Match what `_establish` accepts: an unsealed value would wrongly end the loop.
+		for (let attempt = 0; attempt < MAX_SEALED_TOKEN_RETRIES && !isCloudSandboxSealedToken(clientToken.encrypted_github_token); attempt++) {
+			if (token.isCancellationRequested) {
+				throw new CancellationError();
+			}
+			this._logService.info(`${LOG_PREFIX} Environment ${options.environmentId} has no sealed GitHub token yet; re-minting in ${this.sealedTokenRetryDelayMs}ms (attempt ${attempt + 1}/${MAX_SEALED_TOKEN_RETRIES})`);
+			await timeout(this.sealedTokenRetryDelayMs, token);
+
+			let result: CloudSandboxConnectResult;
+			try {
+				result = await this._apiService.connect({ environmentId: options.environmentId, sessionId: options.sessionId }, token);
+			} catch (err) {
+				if (isCancellationError(err) || token.isCancellationRequested) {
+					throw err;
+				}
+				// The initial mint still works, so degrade rather than discard it.
+				this._logService.warn(`${LOG_PREFIX} Re-mint for ${options.environmentId} failed; continuing without a sealed token`, err);
+				break;
+			}
+			if (result.kind !== 'token') {
+				// Went back to waking mid-wait; the handshake watchdog covers a host that is gone.
+				break;
+			}
+			clientToken = result.token;
+		}
+		return clientToken;
+	}
 }

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

@@ -0,0 +1,180 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import assert from 'assert';
+import { CancellationToken } from '../../../../../../base/common/cancellation.js';
+import { Event } from '../../../../../../base/common/event.js';
+import { URI } from '../../../../../../base/common/uri.js';
+import { mock } from '../../../../../../base/test/common/mock.js';
+import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js';
+import {
+	CloudSandboxEnabledSettingId,
+	cloudSandboxAddress,
+	ICloudSandboxApiService,
+	type CloudSandboxConnectResult,
+	type ICloudSandboxClientToken,
+} from '../../../../../../platform/agentHost/common/cloudSandboxAgentHost.js';
+import { IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js';
+import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js';
+import { TestConfigurationService } from '../../../../../../platform/configuration/test/common/testConfigurationService.js';
+import { IEnvironmentService } from '../../../../../../platform/environment/common/environment.js';
+import { TestInstantiationService } from '../../../../../../platform/instantiation/test/common/instantiationServiceMock.js';
+import { ILogService, NullLogService } from '../../../../../../platform/log/common/log.js';
+import { CloudSandboxAgentHostService, MAX_SEALED_TOKEN_RETRIES } from '../../browser/cloudSandboxAgentHostService.js';
+
+function clientToken(sealed: string | undefined): ICloudSandboxClientToken {
+	return {
+		access_token: 'wps-token',
+		expires_at: new Date(Date.now() + 60 * 60 * 1000).toISOString(),
+		wps_endpoint: 'wss://relay.example.com',
+		hub: 'hub',
+		subprotocol: 'json.reliable.webpubsub.azure.v1',
+		client_id: 'client-1',
+		groups: { to_host: 'to_host', to_client: 'to_client', broadcast: 'broadcast' },
+		...(sealed ? { encrypted_github_token: sealed } : {}),
+	} as ICloudSandboxClientToken;
+}
+
+/** Exposes the re-mint delay and skips the relay, so the mint loop runs in isolation. */
+class TestCloudSandboxAgentHostService extends CloudSandboxAgentHostService {
+	protected override readonly sealedTokenRetryDelayMs = 0;
+
+	/** The sealed token as it stood when minting finished. */
+	sealedTokenAtEstablish: string | undefined;
+
+	protected override async _establish(_options: never, address: string, clientToken: { encrypted_github_token?: string }): Promise<string> {
+		this.sealedTokenAtEstablish = clientToken.encrypted_github_token;
+		return address;
+	}
+}
+
+type ScriptedConnectResult = CloudSandboxConnectResult | Error;
+
+function createService(store: Pick<{ add<T extends { dispose(): void }>(t: T): T }, 'add'>, results: readonly ScriptedConnectResult[]): { service: TestCloudSandboxAgentHostService; connectCalls: () => number } {
+	let calls = 0;
+	const instantiationService = store.add(new TestInstantiationService());
+
+	const configurationService = new TestConfigurationService();
+	configurationService.setUserConfiguration(CloudSandboxEnabledSettingId, true);
+	configurationService.setUserConfiguration(RemoteAgentHostsEnabledSettingId, true);
+	instantiationService.stub(IConfigurationService, configurationService);
+
+	instantiationService.stub(ICloudSandboxApiService, new class extends mock<ICloudSandboxApiService>() {
+		override async connect(): Promise<CloudSandboxConnectResult> {
+			// Hold the last result so a caller can keep re-minting past the scripted responses.
+			const result = results[Math.min(calls, results.length - 1)];
+			calls++;
+			if (result instanceof Error) {
+				throw result;
+			}
+			return result;
+		}
+	}());
+	instantiationService.stub(IRemoteAgentHostService, new class extends mock<IRemoteAgentHostService>() {
+		override readonly onDidChangeConnections = Event.None;
+		override readonly connections = [];
+		override getConnection() { return undefined; }
+	}());
+	instantiationService.stub(IEnvironmentService, new class extends mock<IEnvironmentService>() {
+		override readonly logsHome = URI.file('/logs');
+	}());
+	instantiationService.stub(ILogService, new NullLogService());
+
+	return {
+		service: store.add(instantiationService.createInstance(TestCloudSandboxAgentHostService)),
+		connectCalls: () => calls,
+	};
+}
+
+suite('CloudSandboxAgentHostService sealed token', () => {
+
+	const store = ensureNoDisposablesAreLeakedInTestSuite();
+
+	test('re-mints credentials until the sealed GitHub token arrives', async () => {
+		// A fresh environment can answer `/connect` before its credentials are complete.
+		const { service, connectCalls } = createService(store, [
+			{ kind: 'token', token: clientToken(undefined) },
+			{ kind: 'token', token: clientToken(undefined) },
+			{ kind: 'token', token: clientToken('copilot-sealed.v1.key.payload') },
+		]);
+
+		await service.connect({ environmentId: 'env-1', name: 'Sandbox' }, CancellationToken.None);
+
+		assert.deepStrictEqual({ calls: connectCalls(), sealed: service.sealedTokenAtEstablish }, {
+			calls: 3,
+			sealed: 'copilot-sealed.v1.key.payload',
+		});
+	});
+
+	test('gives up re-minting and connects anyway, since a host may never seal one', async () => {
+		// Refusing to connect would be worse than a session that cannot reach GitHub APIs.
+		const { service, connectCalls } = createService(store, [
+			{ kind: 'token', token: clientToken(undefined) },
+		]);
+
+		await service.connect({ environmentId: 'env-1', name: 'Sandbox' }, CancellationToken.None);
+
+		// Bounded, and the connection still proceeds unsealed. Initial mint plus one per retry.
+		assert.deepStrictEqual({ calls: connectCalls(), sealed: service.sealedTokenAtEstablish }, {
+			calls: MAX_SEALED_TOKEN_RETRIES + 1,
+			sealed: undefined,
+		});
+	});
+
+	test('does not re-mint when the first credentials already carry a sealed token', async () => {
+		const { service, connectCalls } = createService(store, [
+			{ kind: 'token', token: clientToken('copilot-sealed.v1.key.payload') },
+		]);
+
+		await service.connect({ environmentId: 'env-1', name: 'Sandbox' }, CancellationToken.None);
+
+		assert.strictEqual(connectCalls(), 1);
+	});
+
+	test('keeps re-minting when the value is present but not a sealed envelope', async () => {
+		// A plaintext bearer is refused when forwarding, so accepting it here would skip re-minting.
+		const { service, connectCalls } = createService(store, [
+			{ kind: 'token', token: clientToken('ghu_plaintext') },
+			{ kind: 'token', token: clientToken('copilot-sealed.v1.key.payload') },
+		]);
+
+		await service.connect({ environmentId: 'env-1', name: 'Sandbox' }, CancellationToken.None);
+
+		assert.deepStrictEqual({ calls: connectCalls(), sealed: service.sealedTokenAtEstablish }, {
+			calls: 2,
+			sealed: 'copilot-sealed.v1.key.payload',
+		});
+	});
+
+	test('connects with the initial credentials when a re-mint fails', async () => {
+		// A transient failure while chasing the seal must not discard credentials that work.
+		const { service } = createService(store, [
+			{ kind: 'token', token: clientToken(undefined) },
+			new Error('network blip'),
+		]);
+
+		const address = await service.connect({ environmentId: 'env-1', name: 'Sandbox' }, CancellationToken.None);
+
+		assert.deepStrictEqual({ address, sealed: service.sealedTokenAtEstablish }, {
+			address: cloudSandboxAddress('env-1'),
+			sealed: undefined,
+		});
+	});
+
+	test('stops re-minting when the environment goes back to waking', async () => {
+		// Re-entering the wake loop would stack two waits; the handshake watchdog covers this.
+		const { service, connectCalls } = createService(store, [
+			{ kind: 'token', token: clientToken(undefined) },
+			{ kind: 'waking', waking: { retryAfterSeconds: 5 } as never },
+		]);
+
+		await service.connect({ environmentId: 'env-1', name: 'Sandbox' }, CancellationToken.None);
+
+		assert.deepStrictEqual({ calls: connectCalls(), sealed: service.sealedTokenAtEstablish }, {
+			calls: 2,
+			sealed: undefined,
+		});
+	});
+});