diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts index 3577d685cc51c..9fb0a8c27bd40 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts +++ b/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(); + /** 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 { + protected async _establish(options: ICloudSandboxConnectOptions, address: string, clientToken: ICloudSandboxClientToken, token: CancellationToken): Promise { // 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,7 +233,7 @@ 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})`); @@ -225,4 +241,39 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa } 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 { + 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; + } } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts new file mode 100644 index 0000000000000..04d223ea19bbd --- /dev/null +++ b/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 { + this.sealedTokenAtEstablish = clientToken.encrypted_github_token; + return address; + } +} + +type ScriptedConnectResult = CloudSandboxConnectResult | Error; + +function createService(store: Pick<{ add(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() { + override async connect(): Promise { + // 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() { + override readonly onDidChangeConnections = Event.None; + override readonly connections = []; + override getConnection() { return undefined; } + }()); + instantiationService.stub(IEnvironmentService, new class extends mock() { + 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, + }); + }); +});