diff --git a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts index 6566e7253d0f1a..6998e3f3d59352 100644 --- a/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts +++ b/src/vs/platform/agentHost/browser/agentHostProtocolClient.ts @@ -1800,6 +1800,11 @@ export class AgentHostProtocolClient extends Disposable implements IAgentConnect return this._dispatchRequest(method, {}, { allowIncompatibleUpgrade: true }); } + /** Low-level transport for typed host-extension adapters, not part of the shared agent connection. */ + requestHostExtension(method: string, params: Record): Promise { + return this._dispatchRequest(method, params); + } + private _handleMessage(msg: ProtocolMessage): void { if (this._state.kind === AgentHostClientState.Closed) { // After close, the transport may still emit late messages (e.g. diff --git a/src/vs/platform/agentHost/common/cloudSandboxAgentHost.ts b/src/vs/platform/agentHost/common/cloudSandboxAgentHost.ts index 833bd604044aeb..0f9a3fef9f1cf0 100644 --- a/src/vs/platform/agentHost/common/cloudSandboxAgentHost.ts +++ b/src/vs/platform/agentHost/common/cloudSandboxAgentHost.ts @@ -11,10 +11,12 @@ // to reach an agent host over one transport; it does not define a new kind of agent host. import { CancellationToken } from '../../../base/common/cancellation.js'; +import { URI } from '../../../base/common/uri.js'; import { IConfigurationService } from '../../configuration/common/configuration.js'; import { createDecorator } from '../../instantiation/common/instantiation.js'; import { RemoteAgentHostsEnabledSettingId } from './remoteAgentHostService.js'; import { IReplayedTaskHistory } from './taskEventReplay.js'; +import type { IAgentConnection } from './agentService.js'; /** Configuration key gating the cloud-sandbox connection path. Disabled by default. */ export const CloudSandboxEnabledSettingId = 'chat.agentHost.cloudSandbox.enabled'; @@ -360,4 +362,7 @@ export interface ICloudSandboxAgentHostService { * `/connect` and refreshed by `/reconnect`, or `undefined` when there is no connection. */ getSealedGitHubToken(environmentId: string): string | undefined; + + /** Prepare a repository through the typed project client owned by this sandbox connection. */ + prepareWorkingDirectory(connection: IAgentConnection, workingDirectory: URI | undefined, token: CancellationToken): Promise; } diff --git a/src/vs/platform/agentHost/common/meta/agentHostProjectMeta.ts b/src/vs/platform/agentHost/common/meta/agentHostProjectMeta.ts new file mode 100644 index 00000000000000..544574bade81d8 --- /dev/null +++ b/src/vs/platform/agentHost/common/meta/agentHostProjectMeta.ts @@ -0,0 +1,13 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import type { RootState } from '../state/sessionState.js'; + +/** Whether the host exposes the Copilot project-management extension. */ +export function supportsAgentHostProjectManagement(state: RootState): boolean { + const capability = state._meta?.['copilot.projectManagement']; + return typeof capability === 'object' && capability !== null + && (capability as { available?: unknown }).available === true; +} diff --git a/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts b/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts index ffdd813ae47fcc..82092e68f5ccd7 100644 --- a/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/agentHostProtocolClient.test.ts @@ -409,6 +409,29 @@ suite('AgentHostProtocolClient', () => { await connectPromise; } + test('host extension requests preserve parameters and return the host response', async () => { + const { client, transport } = createClient(); + await connectClient(client, transport); + const params = { channel: ROOT_STATE_URI, repository: 'https://example.com/owner/repo' }; + const response = client.requestHostExtension('x-test/prepareRepository', params); + await timeout(0); + const request = transport.sentMessages.find((message): message is JsonRpcRequest => 'id' in message && 'method' in message && message.method === 'x-test/prepareRepository'); + assert.ok(request); + transport.fireMessage({ jsonrpc: '2.0', id: request.id, result: { directory: '/checkout/repo' } }); + assert.deepStrictEqual({ params: request.params, result: await response }, { params, result: { directory: '/checkout/repo' } }); + }); + + test('host extension requests propagate protocol errors', async () => { + const { client, transport } = createClient(); + await connectClient(client, transport); + const response = client.requestHostExtension('x-test/prepareRepository', {}); + await timeout(0); + const request = transport.sentMessages.find((message): message is JsonRpcRequest => 'id' in message && 'method' in message && message.method === 'x-test/prepareRepository'); + assert.ok(request); + transport.fireMessage({ jsonrpc: '2.0', id: request.id, error: { code: JsonRpcErrorCodes.MethodNotFound, message: 'Project management unavailable' } }); + await assert.rejects(response, /Project management unavailable/); + }); + for (const identity of [LOCAL_AGENT_HOST_RESOURCE_IDENTITY, 'test.example:1234', 'vscode-remote://ssh-remote+test'] as const) { test(`workspace trust forwards only the target host's trusted roots (${String(identity)})`, async () => { const transport = disposables.add(new TestProtocolTransport()); diff --git a/src/vs/sessions/common/gitHubRepository.ts b/src/vs/sessions/common/gitHubRepository.ts new file mode 100644 index 00000000000000..39f9f5b7f1c744 --- /dev/null +++ b/src/vs/sessions/common/gitHubRepository.ts @@ -0,0 +1,9 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +export function getGitHubRepositoryId(repository: string): string | undefined { + const match = /^(?:(?:https?|ssh|git):\/\/(?:git@)?github\.com\/|git@github\.com:)?(?[^/:\s]+)\/(?[^/\s]+?)(?:\.git)?\/?$/i.exec(repository); + return match?.groups ? `${match.groups.owner}/${match.groups.repo}` : undefined; +} diff --git a/src/vs/sessions/contrib/providers/copilotChatSessions/browser/copilotChatSessionsProvider.ts b/src/vs/sessions/contrib/providers/copilotChatSessions/browser/copilotChatSessionsProvider.ts index 9f41b1551cc31e..687daaac1144f1 100644 --- a/src/vs/sessions/contrib/providers/copilotChatSessions/browser/copilotChatSessionsProvider.ts +++ b/src/vs/sessions/contrib/providers/copilotChatSessions/browser/copilotChatSessionsProvider.ts @@ -65,6 +65,7 @@ import { CloudSandboxAgentHostContribution, type ICloudSandboxProvisionedSession import { IFileService } from '../../../../../platform/files/common/files.js'; import { resolveGitRepositoryFromGitConfig } from '../../../../services/sessions/browser/gitHubRepositoryResolver.js'; import { IPathService } from '../../../../../workbench/services/path/common/pathService.js'; +import { getGitHubRepositoryId } from '../../../../common/gitHubRepository.js'; /** Copilot Cloud session type - cloud-hosted agent. */ export const CopilotCloudSessionType: ISessionType = { @@ -79,11 +80,6 @@ const STORAGE_KEY_ISOLATION_MODE = 'sessions.isolationPicker.selectedMode'; /** Remembers the cloud sandbox choice across new sessions, like the isolation picker above. */ const STORAGE_KEY_USE_SANDBOX = 'sessions.cloudSandboxPicker.useSandbox'; -function getGitHubRepositoryId(repository: string): string | undefined { - const match = /^(?:(?:https?|ssh|git):\/\/(?:git@)?github\.com\/|git@github\.com:)?(?[^/:\s]+)\/(?[^/\s]+?)(?:\.git)?\/?$/i.exec(repository); - return match?.groups ? `${match.groups.owner}/${match.groups.repo}` : undefined; -} - export type IsolationMode = 'worktree' | 'workspace'; export interface ICopilotChatSession { diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md b/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md index 5c08c520f9562f..e70e34fce0c645 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md @@ -61,6 +61,16 @@ Drafts expose the shared untitled `ISession` contract and use remote workspace m Remote session and chat resources preserve connection-specific routing identity through creation, hydration, and replacement. Backend session identifiers are translated only inside the provider. +### Preparing a remote workspace + +A connection customization may prepare a new session's working directory after authentication and before backend session creation. The shared session handler owns this ordering; the customization owns host-specific capability checks, requests, progress, and validation. The resolved directory is retained by the connection's working-directory resolver and passed through the normal resource URI mapping. + +Cloud sandbox repository selections use this boundary to resolve a checkout on hosts that advertise project management. They reuse a ready checkout or await the host's clone result before creating the session. The client does not clone locally or acquire a second workload credential. Hosts without the capability retain the existing host-selected directory behavior. + +The sandbox connection factory creates a typed project adapter for each connection. It owns capability checks, wire requests, and validation of both responses and catalogue entries; the resolver consumes only typed project operations. The sandbox service looks up the adapter by connection identity, not address, and connection teardown cancels outstanding preparation and removes the adapter without affecting a replacement connection. Arbitrary RPC dispatch is not exposed through `IAgentConnection`; only the concrete transport client supplies the low-level host-extension sender. + +Preparation failures and cancellation stop session creation and the first turn; they do not fall back to an unrelated directory. Existing backend sessions keep their established directories and do not run preparation again. Host-specific requests stay outside the standard protocol command map and are used only after checking the advertised capability. + ## Authentication and recovery Authentication challenges, credential refresh, and transport retries remain connection policy. The request that encountered a challenge observes its actual success, cancellation, or failure; provider operations do not silently convert authentication failures into availability results. diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts index bb0cfc14542ffa..58777e2dc75875 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostContribution.ts @@ -178,9 +178,6 @@ export class CloudSandboxAgentHostContribution extends Disposable implements IWo ) { super(); - // Supply the generic remote-agent-host contribution with the sandbox host's per-connection - // deviations (sealed-token auth + `ahp-session` backend scheme) without leaking sandbox - // specifics into that shared code path. this._register(this._connectionCustomizations.register( isCloudSandboxConnectionAddress, address => createCloudSandboxConnectionCustomization(address, this._cloudSandboxService)!, diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts index 45ce6c557ec03c..ce0fdb72032cbe 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxAgentHostService.ts @@ -5,15 +5,17 @@ import { CancellationError, isCancellationError } from '../../../../../base/common/errors.js'; import { CancellationToken } from '../../../../../base/common/cancellation.js'; -import { Disposable, DisposableStore, MutableDisposable } from '../../../../../base/common/lifecycle.js'; +import { Disposable, DisposableStore, MutableDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; import { IObservable, observableValue } from '../../../../../base/common/observable.js'; +import { URI } from '../../../../../base/common/uri.js'; import { raceCancellationError, timeout } from '../../../../../base/common/async.js'; +import { localize } from '../../../../../nls.js'; import { IProtocolTransport } from '../../../../../platform/agentHost/common/state/sessionTransport.js'; import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; import { editorWindowAgentHostClientInfo } from '../../../../../platform/agentHost/common/agentHostClientInfo.js'; import { WebPubSubRelayTransport } from '../../../../../platform/agentHost/browser/webPubSubRelayTransport.js'; import { AhpJsonlLogger } from '../../../../../platform/agentHost/common/ahpJsonlLogger.js'; -import { GITHUB_COPILOT_PROTECTED_RESOURCE, AgentHostAhpJsonlLoggingSettingId } from '../../../../../platform/agentHost/common/agentService.js'; +import { GITHUB_COPILOT_PROTECTED_RESOURCE, AgentHostAhpJsonlLoggingSettingId, IAgentConnection } from '../../../../../platform/agentHost/common/agentService.js'; import { buildWpsUrl, cloudSandboxAddress, @@ -31,6 +33,8 @@ import { IEnvironmentService } from '../../../../../platform/environment/common/ import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; import { ILogService } from '../../../../../platform/log/common/log.js'; import { CloudSandboxCredentialRefresher, MAX_WAKING_DELAY_MS, type ICloudSandboxCreds } from './cloudSandboxCredentialRefresh.js'; +import { CloudSandboxProjectResolver } from './cloudSandboxProjectResolver.js'; +import { CloudSandboxProjectsClient, ICloudSandboxProjectsClient } from './cloudSandboxProjectsClient.js'; const LOG_PREFIX = '[CloudSandboxAgentHost]'; @@ -60,6 +64,8 @@ class CloudSandboxConnectionFactory extends Disposable implements IRemoteAgentHo private readonly _stagedConnections = new Map(); private readonly _entries = observableValue(this, []); + private readonly _projectClients = new WeakMap(); + private readonly _projectResolver: CloudSandboxProjectResolver; constructor( private readonly _instantiationService: IInstantiationService, @@ -68,6 +74,7 @@ class CloudSandboxConnectionFactory extends Disposable implements IRemoteAgentHo ) { super(); this.entries = this._entries; + this._projectResolver = this._instantiationService.createInstance(CloudSandboxProjectResolver); // Staging is cleared only by an explicit `unstageConfiguration`, never by // observing the connection disappear. The service withdraws an entry // before arming a retry, so treating that as removal would delete the @@ -107,6 +114,17 @@ class CloudSandboxConnectionFactory extends Disposable implements IRemoteAgentHo return this._stagedConnections.get(cloudSandboxAddress(environmentId))?.creds.token.encrypted_github_token; } + async prepareWorkingDirectory(connection: IAgentConnection, workingDirectory: URI | undefined, token: CancellationToken): Promise { + if (token.isCancellationRequested) { + throw new CancellationError(); + } + const projects = this._projectClients.get(connection); + if (!projects) { + throw new Error(localize('cloudSandbox.projectConnectionUnavailable', "The sandbox connection is no longer available for repository preparation.")); + } + return this._projectResolver.resolve(projects, workingDirectory, token); + } + async createConnection(entry: IRemoteAgentHostEntry, _options: IRemoteAgentHostConnectOptions): Promise { if (entry.connection.type !== RemoteAgentHostEntryType.CloudSandbox) { throw new Error(`Cloud sandbox factory cannot create a ${entry.connection.type} connection.`); @@ -142,6 +160,9 @@ class CloudSandboxConnectionFactory extends Disposable implements IRemoteAgentHo }, ); const store = new DisposableStore(); + const projects = store.add(new CloudSandboxProjectsClient(client)); + this._projectClients.set(client, projects); + store.add(toDisposable(() => this._projectClients.delete(client))); const refresher = store.add(new MutableDisposable()); store.add(client.onDidChangeConnectionState(state => { if (state === 'connected' && !refresher.value) { @@ -207,6 +228,10 @@ export class CloudSandboxAgentHostService extends Disposable implements ICloudSa return this._connectionFactory.getSealedGitHubToken(environmentId); } + prepareWorkingDirectory(connection: IAgentConnection, workingDirectory: URI | undefined, token: CancellationToken): Promise { + return this._connectionFactory.prepareWorkingDirectory(connection, workingDirectory, token); + } + async connect(options: ICloudSandboxConnectOptions, token: CancellationToken): Promise { if (!this._configurationService.getValue(CloudSandboxEnabledSettingId)) { throw new Error('Copilot cloud sandbox connections are not enabled.'); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxConnectionCustomization.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxConnectionCustomization.ts index ffe5d03b79f81d..2f3ef7728add65 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxConnectionCustomization.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxConnectionCustomization.ts @@ -29,16 +29,7 @@ function isGitHubResource(resource: string): boolean { || host.endsWith('.ghe.com'); } -/** - * The {@link IRemoteAgentHostConnectionCustomization} for a cloud sandbox address, supplying the two - * ways the sandbox host deviates from the generic path: - * - * - **Auth**: the host only accepts a sealed envelope, so the connection's `encrypted_github_token` - * is presented instead of the resolved bearer. Fails closed if no sealed token is available. - * - **Scheme**: the host advertises provider `copilot` but addresses sessions as `ahp-session`. - * - * Returns `undefined` for non-sandbox addresses. - */ +/** Adapts authentication, session identity and repository preparation for a cloud sandbox. */ export function createCloudSandboxConnectionCustomization( address: string, sandboxService: ICloudSandboxAgentHostService, @@ -66,6 +57,7 @@ export function createCloudSandboxConnectionCustomization( }, backendSessionScheme: (provider: string): string | undefined => provider === CLOUD_SANDBOX_AGENT_PROVIDER ? CLOUD_SANDBOX_SESSION_SCHEME : undefined, + prepareWorkingDirectory: (connection, workingDirectory, token) => sandboxService.prepareWorkingDirectory(connection, workingDirectory, token), }; } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxProjectResolver.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxProjectResolver.ts new file mode 100644 index 00000000000000..a01e2fde3e550e --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxProjectResolver.ts @@ -0,0 +1,167 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { DeferredPromise, disposableTimeout } from '../../../../../base/common/async.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../../base/common/cancellation.js'; +import { CancellationError } from '../../../../../base/common/errors.js'; +import { DisposableStore } from '../../../../../base/common/lifecycle.js'; +import { Schemas } from '../../../../../base/common/network.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { localize } from '../../../../../nls.js'; +import { IProgress, IProgressService, IProgressStep, ProgressLocation } from '../../../../../platform/progress/common/progress.js'; +import { getGitHubRepositoryId } from '../../../../common/gitHubRepository.js'; +import { ICloudSandboxProject, ICloudSandboxProjectsClient, invalidCloudSandboxProject } from './cloudSandboxProjectsClient.js'; + +const CLONE_TIMEOUT_MS = 180_000; + +function findProject(client: ICloudSandboxProjectsClient, repository: string): ICloudSandboxProject | undefined { + return client.getProjects().find(project => + project.remoteUrl && getGitHubRepositoryId(project.remoteUrl)?.toLowerCase() === repository.toLowerCase()); +} + +function projectDirectory(client: ICloudSandboxProjectsClient, project: ICloudSandboxProject): URI | undefined { + if (project.status === 'failed') { + throw new Error(localize('cloudSandbox.cloneFailed', "The remote repository could not be prepared: {0}", project.error ?? localize('cloudSandbox.cloneFailedUnknown', "Cloning failed."))); + } + if (project.status !== 'ready') { + return undefined; + } + if (!project.git) { + throw new Error(localize('cloudSandbox.projectNotRepository', "The remote project is not a Git repository.")); + } + return client.toResourceUri(project.path); +} + +/** Resolves a sandbox repository through the host's project-management extension. */ +export class CloudSandboxProjectResolver { + constructor( + @IProgressService private readonly progressService: IProgressService, + ) { } + + async resolve(client: ICloudSandboxProjectsClient, workingDirectory: URI | undefined, token: CancellationToken): Promise { + if (token.isCancellationRequested) { + throw new CancellationError(); + } + if (!workingDirectory || workingDirectory.scheme !== Schemas.https || workingDirectory.authority.toLowerCase() !== 'github.com') { + return workingDirectory; + } + if (!client.isAvailable()) { + return workingDirectory; + } + const repository = getGitHubRepositoryId(workingDirectory.toString()); + if (!repository || workingDirectory.query || workingDirectory.fragment) { + throw new Error(localize('cloudSandbox.invalidRepository', "Choose a GitHub repository before starting this session.")); + } + const existing = findProject(client, repository); + if (existing?.status === 'ready') { + return projectDirectory(client, existing); + } + const store = new DisposableStore(); + const cts = store.add(new CancellationTokenSource(token)); + try { + return await this.progressService.withProgress({ + location: ProgressLocation.Notification, + title: localize('cloudSandbox.prepareRepository', "Preparing Repository"), + cancellable: true, + total: 100, + delay: 500, + }, progress => this.waitForProject(client, repository, existing, cts.token, progress), () => cts.cancel()); + } finally { + store.dispose(); + } + } + + private async waitForProject(client: ICloudSandboxProjectsClient, repository: string, existing: ICloudSandboxProject | undefined, token: CancellationToken, progress: IProgress): Promise { + const store = new DisposableStore(); + const cts = store.add(new CancellationTokenSource(token)); + const completion = new DeferredPromise(); + const deadline = new DeferredPromise(); + let projectId = existing?.id; + let awaitingRetryCatalogue = existing?.status === 'failed'; + let lastProgress: number | undefined; + const isPreviousFailure = (project: ICloudSandboxProject): boolean => + awaitingRetryCatalogue && project.id === existing?.id && project.status === 'failed'; + const update = (project: ICloudSandboxProject): void => { + if (completion.isSettled) { + return; + } + try { + if (projectId && project.id !== projectId) { + throw invalidCloudSandboxProject(); + } + const directory = projectDirectory(client, project); + if (directory) { + completion.complete(directory); + return; + } + const percent = Math.max(lastProgress ?? 0, project.progress ?? 0); + if (percent !== lastProgress) { + progress.report({ + message: localize('cloudSandbox.cloningRepository', "Cloning {0} ({1}%)...", repository, percent), + increment: percent - (lastProgress ?? 0), + }); + lastProgress = percent; + } + } catch (error) { + completion.error(error); + } + }; + const observe = (): void => { + try { + const project = findProject(client, repository); + if (project) { + if (isPreviousFailure(project)) { + return; + } + awaitingRetryCatalogue = false; + update(project); + } else if (projectId) { + completion.error(new Error(localize('cloudSandbox.projectRemoved', "The remote repository was removed while it was being prepared."))); + } + } catch (error) { + completion.error(error); + } + }; + store.add(client.onDidChange(observe)); + store.add(client.onDidError(error => completion.error(error))); + store.add(cts.token.onCancellationRequested(() => completion.cancel())); + store.add(disposableTimeout(() => deadline.error(new Error(localize('cloudSandbox.cloneTimeout', "Timed out waiting for the remote repository to be prepared."))), CLONE_TIMEOUT_MS)); + + const start = async (): Promise => { + if (cts.token.isCancellationRequested) { + throw new CancellationError(); + } + if (existing?.status === 'cloning') { + update(existing); + observe(); + return; + } + progress.report({ message: localize('cloudSandbox.startClone', "Cloning {0}...", repository) }); + const project = await client.cloneProject({ + url: `https://github.com/${repository}`, + depth: 1, + }, cts.token); + projectId = project.id; + const current = findProject(client, repository); + if (current && !isPreviousFailure(current)) { + if (current.id !== projectId) { + throw invalidCloudSandboxProject(); + } + awaitingRetryCatalogue = false; + update(current); + } else { + update(project); + } + }; + + try { + const [, directory] = await Promise.race([Promise.all([start(), completion.p]), deadline.p]); + return directory; + } finally { + cts.cancel(); + store.dispose(); + } + } +} diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxProjectsClient.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxProjectsClient.ts new file mode 100644 index 00000000000000..b1cf085ca4e2b3 --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/cloudSandboxProjectsClient.ts @@ -0,0 +1,158 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { raceCancellationError } from '../../../../../base/common/async.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../../base/common/cancellation.js'; +import { CancellationError } from '../../../../../base/common/errors.js'; +import { Emitter, Event } from '../../../../../base/common/event.js'; +import { Disposable, DisposableStore } from '../../../../../base/common/lifecycle.js'; +import { posix } from '../../../../../base/common/path.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { localize } from '../../../../../nls.js'; +import { AgentHostProtocolClient } from '../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; +import { supportsAgentHostProjectManagement } from '../../../../../platform/agentHost/common/meta/agentHostProjectMeta.js'; +import { RootState } from '../../../../../platform/agentHost/common/state/sessionState.js'; +import { getGitHubRepositoryId } from '../../../../common/gitHubRepository.js'; + +export interface ICloudSandboxProject { + readonly id: string; + readonly path: string; + readonly git: boolean; + readonly status: 'cloning' | 'ready' | 'failed'; + readonly remoteUrl?: string; + readonly progress?: number; + readonly error?: string; +} + +export interface ICloudSandboxCloneProjectOptions { + readonly url: string; + readonly depth: 1; +} + +/** Validated project operations for one sandbox connection. */ +export interface ICloudSandboxProjectsClient { + readonly onDidChange: Event; + readonly onDidError: Event; + isAvailable(): boolean; + getProjects(): readonly ICloudSandboxProject[]; + cloneProject(options: ICloudSandboxCloneProjectOptions, token: CancellationToken): Promise; + toResourceUri(path: string): URI; +} + +export function invalidCloudSandboxProject(): Error { + return new Error(localize('cloudSandbox.invalidProject', "The remote host returned invalid repository information.")); +} + +function readProject(value: unknown): ICloudSandboxProject { + if (typeof value !== 'object' || value === null) { + throw invalidCloudSandboxProject(); + } + const candidate = value as Partial; + const { id, path, git, status, remoteUrl, progress, error } = candidate; + if (typeof id !== 'string' || !id + || typeof path !== 'string' || !posix.isAbsolute(path) + || typeof git !== 'boolean' + || (status !== 'cloning' && status !== 'ready' && status !== 'failed') + || (remoteUrl !== undefined && typeof remoteUrl !== 'string') + || (progress !== undefined && (typeof progress !== 'number' || !Number.isInteger(progress) || progress < 0 || progress > 100)) + || (error !== undefined && typeof error !== 'string')) { + throw invalidCloudSandboxProject(); + } + return { id, path, git, status, remoteUrl, progress, error }; +} + +/** Owns the sandbox project's wire contract and validation, scoped to the transport lifetime. */ +export class CloudSandboxProjectsClient extends Disposable implements ICloudSandboxProjectsClient { + private readonly _lifetime = this._register(new CancellationTokenSource()); + private readonly _onDidClose = this._register(new Emitter()); + readonly onDidChange: Event; + readonly onDidError: Event; + + constructor(private readonly _connection: AgentHostProtocolClient) { + super(); + this.onDidChange = Event.map(_connection.rootState.onDidChange, () => undefined); + this.onDidError = Event.any(_connection.rootState.onDidError ?? Event.None, this._onDidClose.event); + } + + isAvailable(): boolean { + return supportsAgentHostProjectManagement(this._readRoot()); + } + + getProjects(): readonly ICloudSandboxProject[] { + const namespace = this._readRoot().config?.values.copilot; + const projects = typeof namespace === 'object' && namespace !== null + ? (namespace as { projects?: unknown }).projects + : undefined; + if (!Array.isArray(projects)) { + throw invalidCloudSandboxProject(); + } + return projects.map(readProject); + } + + async cloneProject(options: ICloudSandboxCloneProjectOptions, token: CancellationToken): Promise { + this._throwIfClosed(); + if (token.isCancellationRequested) { + throw new CancellationError(); + } + if (!this.isAvailable()) { + throw new Error(localize('cloudSandbox.projectRequestsUnsupported', "This connection cannot prepare repositories on the remote host.")); + } + const repository = getGitHubRepositoryId(options.url); + if (!repository) { + throw new Error(localize('cloudSandbox.invalidRepository', "Choose a GitHub repository before starting this session.")); + } + const store = new DisposableStore(); + const cts = store.add(new CancellationTokenSource(token)); + store.add(this._lifetime.token.onCancellationRequested(() => cts.cancel())); + try { + const response = await raceCancellationError(this._connection.requestHostExtension('extensions/cloneProject', { + url: options.url, + depth: options.depth, + }), cts.token); + this._throwIfClosed(); + if (typeof response !== 'object' || response === null) { + throw invalidCloudSandboxProject(); + } + const project = readProject((response as { project?: unknown }).project); + if (!project.remoteUrl || getGitHubRepositoryId(project.remoteUrl)?.toLowerCase() !== repository.toLowerCase()) { + throw invalidCloudSandboxProject(); + } + return project; + } finally { + store.dispose(); + } + } + + toResourceUri(path: string): URI { + this._throwIfClosed(); + return this._connection.resourceUris.fromAgentHost(URI.file(path)); + } + + private _readRoot(): RootState { + this._throwIfClosed(); + const state = this._connection.rootState.value; + if (state instanceof Error) { + throw state; + } + if (!state) { + throw new Error(localize('cloudSandbox.noProjectCatalogue', "Repository information is not available from the remote host.")); + } + return state; + } + + private _throwIfClosed(): void { + if (this._lifetime.token.isCancellationRequested) { + throw new CancellationError(); + } + } + + override dispose(): void { + if (!this._store.isDisposed) { + this._lifetime.cancel(); + this._onDidClose.fire(new CancellationError()); + super.dispose(); + } + } +} diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts index ef6db93b0a3104..6fe1aa288b044b 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHost.contribution.ts @@ -343,6 +343,8 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc // Per-agent working directory cache, scoped to the agent store lifetime const sessionWorkingDirs = new Map(); agentStore.add(toDisposable(() => sessionWorkingDirs.clear())); + const connectionCustomization = this._connectionCustomizations.get(address); + const prepareWorkingDirectory = connectionCustomization?.prepareWorkingDirectory; // Capture the working directory from the session that is being created. const resolveWorkingDirectory = (sessionResource: URI): URI | undefined => { @@ -423,7 +425,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc const sessionHandler = agentStore.add(this._instantiationService.createInstance( AgentHostSessionHandler, { provider: agent.provider, - backendSessionScheme: this._connectionCustomizations.get(address)?.backendSessionScheme?.(agent.provider), + backendSessionScheme: connectionCustomization?.backendSessionScheme?.(agent.provider), agentId, sessionType, fullName: displayName, @@ -433,6 +435,13 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc extensionId: 'vscode.remote-agent-host', extensionDisplayName: 'Remote Agent Host', resolveWorkingDirectory, + prepareWorkingDirectory: prepareWorkingDirectory ? async (sessionResource, workingDirectory, token) => { + const prepared = await prepareWorkingDirectory(connection, workingDirectory, token); + if (prepared) { + sessionWorkingDirs.set(sessionResource.toString(), prepared); + } + return prepared; + } : undefined, isNewSession, resolveAuthentication: (resources) => this._resolveAuthenticationInteractively(address, connection, resources), })); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostConnectionCustomization.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostConnectionCustomization.ts index df09f85fcbf3a1..19ce04bea9a1e7 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostConnectionCustomization.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostConnectionCustomization.ts @@ -4,6 +4,9 @@ *--------------------------------------------------------------------------------------------*/ import { IDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; +import { CancellationToken } from '../../../../../base/common/cancellation.js'; +import { URI } from '../../../../../base/common/uri.js'; +import { IAgentConnection } from '../../../../../platform/agentHost/common/agentService.js'; import { createDecorator } from '../../../../../platform/instantiation/common/instantiation.js'; import { IAgentHostAuthenticateRequest } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostAuth.js'; @@ -23,6 +26,9 @@ export interface IRemoteAgentHostConnectionCustomization { * Return `undefined` to keep scheme == provider. */ readonly backendSessionScheme?: (provider: string) => string | undefined; + + /** Prepare a host-addressable directory after authentication and before creating a session. */ + readonly prepareWorkingDirectory?: (connection: IAgentConnection, workingDirectory: URI | undefined, token: CancellationToken) => Promise; } /** Builds a {@link IRemoteAgentHostConnectionCustomization} for a concrete connection address. */ 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 index 5f421a4609f1ff..3f134c8af31c8b 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxAgentHostService.test.ts @@ -4,11 +4,16 @@ *--------------------------------------------------------------------------------------------*/ 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 { toDisposable } from '../../../../../../base/common/lifecycle.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 { AgentHostProtocolClient } from '../../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; +import { toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; import { CloudSandboxEnabledSettingId, cloudSandboxAddress, @@ -16,13 +21,16 @@ import { type CloudSandboxConnectResult, type ICloudSandboxClientToken, } from '../../../../../../platform/agentHost/common/cloudSandboxAgentHost.js'; -import { IRemoteAgentHostConnectionFactory, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IRemoteAgentHostConnectionFactory, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, 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 { IProgressService } from '../../../../../../platform/progress/common/progress.js'; import { CloudSandboxAgentHostService, MAX_SEALED_TOKEN_RETRIES } from '../../browser/cloudSandboxAgentHostService.js'; +import { createCloudSandboxConnectionCustomization } from '../../browser/cloudSandboxConnectionCustomization.js'; +import { createCloudSandboxProject as project, createCloudSandboxProjectsTestConnection } from './cloudSandboxProjectsTestUtils.js'; function clientToken(sealed: string | undefined): ICloudSandboxClientToken { return { @@ -82,6 +90,7 @@ function createService(store: Pick<{ add(t: T): T override readonly logsHome = URI.file('/logs'); }()); instantiationService.stub(ILogService, new NullLogService()); + instantiationService.stub(IProgressService, { withProgress: (_options, task) => task({ report: () => { } }) }); return { service: store.add(instantiationService.createInstance(TestCloudSandboxAgentHostService)), @@ -89,7 +98,7 @@ function createService(store: Pick<{ add(t: T): T }; } -suite('CloudSandboxAgentHostService sealed token', () => { +suite('CloudSandboxAgentHostService', () => { const store = ensureNoDisposablesAreLeakedInTestSuite(); @@ -109,6 +118,102 @@ suite('CloudSandboxAgentHostService sealed token', () => { }); }); + suite('project connections', () => { + const repository = URI.parse('https://github.com/owner/repo'); + + async function createHarness() { + const instantiationService = store.add(new TestInstantiationService()); + const remote = new class extends mock() { + factory: IRemoteAgentHostConnectionFactory | undefined; + override readonly connections = []; + override getConnection() { return undefined; } + override registerConnectionFactory(factory: IRemoteAgentHostConnectionFactory) { + this.factory = factory; + return toDisposable(() => { this.factory = undefined; }); + } + override reconnect(): void { } + override async waitForConnection(address: string) { + return { address, name: 'Sandbox', clientId: 'client-1', status: RemoteAgentHostConnectionStatus.connected }; + } + }(); + instantiationService.stub(IRemoteAgentHostService, remote); + instantiationService.stub(IConfigurationService, new TestConfigurationService({ + [CloudSandboxEnabledSettingId]: true, + [RemoteAgentHostsEnabledSettingId]: true, + })); + instantiationService.stub(ICloudSandboxApiService, { + connect: async () => ({ kind: 'token', token: clientToken('copilot-sealed.v1.key.payload') }), + }); + instantiationService.stub(IEnvironmentService, { logsHome: URI.file('/logs') }); + instantiationService.stub(ILogService, new NullLogService()); + instantiationService.stub(IProgressService, { withProgress: (_options, task) => task({ report: () => { } }) }); + const service = store.add(instantiationService.createInstance(CloudSandboxAgentHostService)); + await service.connect({ environmentId: 'env-1', name: 'Sandbox' }, CancellationToken.None); + const factory = remote.factory; + assert.ok(factory); + const entry = factory.entries.get()[0]; + assert.ok(entry); + return { + service, + createConnection: async (client: AgentHostProtocolClient) => { + instantiationService.stubInstance(AgentHostProtocolClient, client); + const created = await factory.createConnection(entry, { userInitiated: true }); + assert.ok(created.transportDisposable); + return { connection: created.connection, lifetime: store.add(created.transportDisposable) }; + }, + }; + } + + test('the sandbox customization uses the typed adapter created with its connection', async () => { + const h = await createHarness(); + const raw = createCloudSandboxProjectsTestConnection(store, { projects: [project()] }); + const created = await h.createConnection(raw.connection); + const prepare = createCloudSandboxConnectionCustomization(cloudSandboxAddress('env-1'), h.service)?.prepareWorkingDirectory; + assert.ok(prepare); + const result = await prepare(created.connection, repository, CancellationToken.None); + assert.deepStrictEqual({ directory: result?.toString(), requests: raw.requests }, { + directory: toAgentHostUri(URI.file('/checkout/owner/repo'), 'sandbox').toString(), + requests: [], + }); + }); + + test('disposing an old connection cannot remove the replacement adapter at the same address', async () => { + const h = await createHarness(); + const firstRaw = createCloudSandboxProjectsTestConnection(store, { projects: [project()] }); + const secondRaw = createCloudSandboxProjectsTestConnection(store, { projects: [project({ path: '/replacement/owner/repo' })] }); + const first = await h.createConnection(firstRaw.connection); + const second = await h.createConnection(secondRaw.connection); + first.lifetime.dispose(); + await assert.rejects(h.service.prepareWorkingDirectory(first.connection, repository, CancellationToken.None), /connection is no longer available/); + const result = await h.service.prepareWorkingDirectory(second.connection, repository, CancellationToken.None); + assert.strictEqual(result?.toString(), toAgentHostUri(URI.file('/replacement/owner/repo'), 'sandbox').toString()); + }); + + for (const awaitingResponse of [false, true]) { + test(`disposing the connection cancels preparation (${awaitingResponse ? 'clone response' : 'catalogue readiness'})`, async () => { + const response = new DeferredPromise(); + const h = await createHarness(); + const raw = createCloudSandboxProjectsTestConnection(store, { + projects: awaitingResponse ? [] : [project({ status: 'cloning', git: false })], + request: () => response.p, + }); + const created = await h.createConnection(raw.connection); + const result = h.service.prepareWorkingDirectory(created.connection, repository, CancellationToken.None); + created.lifetime.dispose(); + await assert.rejects(result, CancellationError); + assert.strictEqual(raw.hasListeners(), false); + response.complete({ project: project() }); + }); + } + + test('a connection not created by the sandbox factory cannot prepare a project', async () => { + const h = await createHarness(); + const raw = createCloudSandboxProjectsTestConnection(store); + await assert.rejects(h.service.prepareWorkingDirectory(raw.connection, repository, CancellationToken.None), /connection is no longer available/); + assert.deepStrictEqual(raw.requests, []); + }); + }); + 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, [ diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectResolver.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectResolver.test.ts new file mode 100644 index 00000000000000..5e11c68b16e533 --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectResolver.test.ts @@ -0,0 +1,298 @@ +/*--------------------------------------------------------------------------------------------- + * 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 { DeferredPromise, timeout } from '../../../../../../base/common/async.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; +import { CancellationError } from '../../../../../../base/common/errors.js'; +import { URI } from '../../../../../../base/common/uri.js'; +import { runWithFakedTimers } from '../../../../../../base/test/common/timeTravelScheduler.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; +import { IProgress, IProgressOptions, IProgressService, IProgressStep } from '../../../../../../platform/progress/common/progress.js'; +import { CloudSandboxProjectResolver } from '../../browser/cloudSandboxProjectResolver.js'; +import { CloudSandboxProjectsClient } from '../../browser/cloudSandboxProjectsClient.js'; +import { createCloudSandboxProject as project, createCloudSandboxProjectsTestConnection } from './cloudSandboxProjectsTestUtils.js'; + +class TestProgressService implements IProgressService { + declare readonly _serviceBrand: undefined; + readonly reports: IProgressStep[] = []; + shown = 0; + cancel: () => void = () => { }; + + withProgress(_options: IProgressOptions, task: (progress: IProgress) => Promise, onDidCancel?: () => void): Promise { + this.shown++; + this.cancel = () => onDidCancel?.(); + return task({ report: step => this.reports.push(step) }); + } +} + +suite('CloudSandboxProjectResolver', () => { + const store = ensureNoDisposablesAreLeakedInTestSuite(); + const repository = URI.parse('https://github.com/owner/repo'); + const checkout = toAgentHostUri(URI.file('/checkout/owner/repo'), 'sandbox'); + + function createHarness(options: { + projects?: readonly unknown[]; + capability?: unknown; + legacy?: boolean; + request?: () => Promise; + } = {}) { + const connection = createCloudSandboxProjectsTestConnection(store, options); + const client = store.add(new CloudSandboxProjectsClient(connection.connection)); + const progress = new TestProgressService(); + const resolver = new CloudSandboxProjectResolver(progress); + return { ...connection, client, resolver, progress }; + } + + for (const options of [{ legacy: true }, { capability: { available: false } }, { capability: { available: 'true' } }]) { + test(`preserves the pre-cloned path without an advertised capability (${JSON.stringify(options)})`, async () => { + const h = createHarness(options); + const result = await h.resolver.resolve(h.client, repository, CancellationToken.None); + assert.deepStrictEqual({ directory: result?.toString(), requests: h.requests, progress: h.progress.shown }, { + directory: repository.toString(), requests: [], progress: 0, + }); + }); + } + + test('leaves an existing filesystem working directory unchanged', async () => { + const h = createHarness(); + const result = await h.resolver.resolve(h.client, checkout, CancellationToken.None); + assert.deepStrictEqual({ directory: result?.toString(), requests: h.requests }, { directory: checkout.toString(), requests: [] }); + }); + + test('reuses a ready checkout matched across SSH, case and .git spellings', async () => { + const h = createHarness({ projects: [project({ remoteUrl: 'git@github.com:OWNER/REPO.git' })] }); + const result = await h.resolver.resolve(h.client, repository, CancellationToken.None); + assert.deepStrictEqual({ directory: result?.toString(), requests: h.requests, progress: h.progress.shown }, { + directory: checkout.toString(), requests: [], progress: 0, + }); + }); + + test('requests a shallow clone and waits for the catalogue to report ready', async () => { + const h = createHarness(); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + await h.requested.p; + h.setProjects([project({ status: 'cloning', git: false, progress: 40 })]); + let completed = false; + void result.then(() => completed = true); + await Promise.resolve(); + const completedWhileCloning = completed; + h.setProjects([project()]); + assert.deepStrictEqual({ + directory: (await result)?.toString(), + completedWhileCloning, + requests: h.requests, + reportedProgress: h.progress.reports.some(report => report.message?.includes('40%')), + hasListeners: h.hasListeners(), + }, { + directory: checkout.toString(), + completedWhileCloning: false, + requests: [{ method: 'extensions/cloneProject', params: { url: repository.toString(), depth: 1 } }], + reportedProgress: true, + hasListeners: false, + }); + }); + + test('does not replace a ready catalogue entry with a late cloning response', async () => { + const response = new DeferredPromise(); + const h = createHarness({ request: () => response.p }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + await h.requested.p; + h.setProjects([project()]); + response.complete({ project: project({ status: 'cloning', git: false }) }); + assert.strictEqual((await result)?.toString(), checkout.toString()); + }); + + test('joins a clone started by another client without starting another one', async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.setProjects([project()]); + assert.deepStrictEqual({ directory: (await result)?.toString(), requests: h.requests }, { directory: checkout.toString(), requests: [] }); + }); + + test('keeps newer catalogue progress without repeating or regressing announcements', async () => { + const response = new DeferredPromise(); + const h = createHarness({ request: () => response.p }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + await h.requested.p; + h.setProjects([project({ status: 'cloning', progress: 40 })]); + response.complete({ project: project({ status: 'cloning', progress: 0 }) }); + await timeout(0); + h.setProjects([project({ status: 'cloning', progress: 40 })]); + h.setProjects([project()]); + await result; + assert.deepStrictEqual(h.progress.reports, [ + { message: 'Cloning owner/repo...' }, + { message: 'Cloning owner/repo (40%)...', increment: 40 }, + ]); + }); + + test('retries an existing failed clone on a new user attempt', async () => { + const failed = project({ status: 'failed', git: false, error: 'Temporary network failure' }); + const response = new DeferredPromise(); + const h = createHarness({ projects: [failed], request: () => response.p }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.setProjects([failed]); + response.complete({ project: project({ status: 'cloning', git: false }) }); + h.setProjects([project()]); + assert.deepStrictEqual({ directory: (await result)?.toString(), requests: h.requests, hasListeners: h.hasListeners() }, { + directory: checkout.toString(), + requests: [{ method: 'extensions/cloneProject', params: { url: repository.toString(), depth: 1 } }], + hasListeners: false, + }); + }); + + test('a clone failure ends the attempt but a user resend can recover', async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const first = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.setProjects([project({ status: 'failed', git: false, error: 'Temporary network failure' })]); + await assert.rejects(first, /Temporary network failure/); + const requestsAfterFailure = h.requests.length; + const retry = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.setProjects([project({ status: 'cloning', git: false })]); + h.setProjects([project()]); + assert.deepStrictEqual({ directory: (await retry)?.toString(), requestsAfterFailure, requestsAfterRetry: h.requests.length }, { + directory: checkout.toString(), requestsAfterFailure: 0, requestsAfterRetry: 1, + }); + }); + + test('a failed retry is surfaced without an automatic retry loop', async () => { + const failed = project({ status: 'failed', git: false, error: 'Repository access denied' }); + const h = createHarness({ projects: [failed], request: async () => ({ project: failed }) }); + await assert.rejects(h.resolver.resolve(h.client, repository, CancellationToken.None), /Repository access denied/); + assert.deepStrictEqual({ requestCount: h.requests.length, hasListeners: h.hasListeners() }, { requestCount: 1, hasListeners: false }); + }); + + test('ignores the previous failure until the retry publishes its catalogue state', async () => { + const failed = project({ status: 'failed', git: false, error: 'Temporary network failure' }); + const h = createHarness({ projects: [failed] }); + let finished = false; + const result = assert.rejects(h.resolver.resolve(h.client, repository, CancellationToken.None), /Temporary network failure/).then(() => { + finished = true; + }); + await timeout(0); + h.setProjects([failed]); + await timeout(0); + const finishedBeforeRearm = finished; + h.setProjects([project({ status: 'cloning', git: false })]); + h.setProjects([failed]); + await result; + assert.deepStrictEqual({ finishedBeforeRearm, requestCount: h.requests.length, hasListeners: h.hasListeners() }, { + finishedBeforeRearm: false, requestCount: 1, hasListeners: false, + }); + }); + + test('surfaces failure while a clone is running and releases its listeners', async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.setProjects([project({ status: 'failed', error: 'Clone failed on the host' })]); + await assert.rejects(result, /Clone failed on the host/); + assert.strictEqual(h.hasListeners(), false); + }); + + for (const response of [{}, { project: project({ remoteUrl: 'https://github.com/another/repository' }) }, { project: project({ path: 'relative/path' }) }]) { + test(`rejects an invalid clone response (${JSON.stringify(response)})`, async () => { + const h = createHarness({ request: async () => response }); + await assert.rejects(h.resolver.resolve(h.client, repository, CancellationToken.None), /invalid repository information/); + assert.strictEqual(h.hasListeners(), false); + }); + } + + test('rejects a malformed catalogue without invoking a clone', async () => { + const h = createHarness({ projects: [{}] }); + await assert.rejects(h.resolver.resolve(h.client, repository, CancellationToken.None), /invalid repository information/); + assert.deepStrictEqual(h.requests, []); + }); + + test('reports a project removed while cloning', async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.setProjects([]); + await assert.rejects(result, /removed while it was being prepared/); + }); + + test('propagates subscription failures', async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.failSubscription(new Error('Connection closed')); + await assert.rejects(result, /Connection closed/); + }); + + test('cancels even while the clone request has not returned', async () => { + const cts = store.add(new CancellationTokenSource()); + const response = new DeferredPromise(); + const h = createHarness({ request: () => response.p }); + const result = h.resolver.resolve(h.client, repository, cts.token); + await h.requested.p; + cts.cancel(); + await assert.rejects(result, CancellationError); + assert.strictEqual(h.hasListeners(), false); + response.complete({ project: project() }); + }); + + test('supports cancellation from the progress notification', async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.progress.cancel(); + await assert.rejects(result, CancellationError); + assert.strictEqual(h.hasListeners(), false); + }); + + test('times out instead of leaving session creation waiting forever', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const start = Date.now(); + await assert.rejects(h.resolver.resolve(h.client, repository, CancellationToken.None), /Timed out/); + assert.deepStrictEqual({ elapsed: Date.now() - start, hasListeners: h.hasListeners() }, { elapsed: 180_000, hasListeners: false }); + })); + + for (const readyBeforeResponse of [false, true]) { + test(`times out an unacknowledged clone even when its catalogue is ready (${readyBeforeResponse})`, () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const response = new DeferredPromise(); + const h = createHarness({ request: () => response.p }); + const cts = store.add(new CancellationTokenSource()); + const start = Date.now(); + let elapsed: number | undefined; + let message: string | undefined; + const result = h.resolver.resolve(h.client, repository, cts.token).then( + () => { message = 'Unexpected success'; }, + (error: Error) => { + message = error.message; + elapsed = Date.now() - start; + }, + ); + try { + await h.requested.p; + if (readyBeforeResponse) { + h.setProjects([project()]); + } + await timeout(180_001); + assert.deepStrictEqual({ elapsed, message, hasListeners: h.hasListeners() }, { + elapsed: 180_000, + message: 'Timed out waiting for the remote repository to be prepared.', + hasListeners: false, + }); + } finally { + cts.cancel(); + await result; + response.complete({ project: project() }); + } + })); + } + + test('does not start a clone after cancellation', async () => { + const h = createHarness(); + await assert.rejects(h.resolver.resolve(h.client, repository, CancellationToken.Cancelled), CancellationError); + assert.deepStrictEqual(h.requests, []); + }); + + test('closing the connection cancels a readiness wait and releases its subscriptions', async () => { + const h = createHarness({ projects: [project({ status: 'cloning', git: false })] }); + const result = h.resolver.resolve(h.client, repository, CancellationToken.None); + h.client.dispose(); + await assert.rejects(result, CancellationError); + assert.strictEqual(h.hasListeners(), false); + }); +}); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectsClient.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectsClient.test.ts new file mode 100644 index 00000000000000..a7268a74ea3c9f --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectsClient.test.ts @@ -0,0 +1,115 @@ +/*--------------------------------------------------------------------------------------------- + * 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 { DeferredPromise } from '../../../../../../base/common/async.js'; +import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; +import { CancellationError } from '../../../../../../base/common/errors.js'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; +import { CloudSandboxProjectsClient, ICloudSandboxCloneProjectOptions } from '../../browser/cloudSandboxProjectsClient.js'; +import { createCloudSandboxProject as project, createCloudSandboxProjectsTestConnection } from './cloudSandboxProjectsTestUtils.js'; + +suite('CloudSandboxProjectsClient', () => { + const store = ensureNoDisposablesAreLeakedInTestSuite(); + const options: ICloudSandboxCloneProjectOptions = { url: 'https://github.com/owner/repo', depth: 1 }; + + test('sends the existing clone RPC and returns a validated project', async () => { + const expected = { ...project({ status: 'cloning', git: false, progress: 10 }), error: undefined }; + const h = createCloudSandboxProjectsTestConnection(store, { request: async () => ({ project: expected }) }); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + const result = await client.cloneProject(options, CancellationToken.None); + assert.deepStrictEqual({ requests: h.requests, result }, { + requests: [{ method: 'extensions/cloneProject', params: options }], + result: expected, + }); + }); + + for (const capability of [null, false, {}, { available: false }, { available: 'true' }]) { + test(`does not send a clone without the advertised capability (${JSON.stringify(capability)})`, async () => { + const h = createCloudSandboxProjectsTestConnection(store, { capability }); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + await assert.rejects(client.cloneProject(options, CancellationToken.None), /cannot prepare repositories/); + assert.deepStrictEqual({ available: client.isAvailable(), requests: h.requests }, { available: false, requests: [] }); + }); + } + + test('missing root state is an error, not an unsupported-host fallback', async () => { + const h = createCloudSandboxProjectsTestConnection(store); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + h.setRoot(undefined); + await assert.rejects(client.cloneProject(options, CancellationToken.None), /not available from the remote host/); + assert.deepStrictEqual(h.requests, []); + }); + + test('propagates the transport error without returning a project', async () => { + const error = new Error('Connection lost'); + const h = createCloudSandboxProjectsTestConnection(store, { request: async () => { throw error; } }); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + await assert.rejects(client.cloneProject(options, CancellationToken.None), candidate => candidate === error); + }); + + for (const response of [ + undefined, + null, + {}, + { project: { ...project(), id: '' } }, + { project: { ...project(), path: 'relative/path' } }, + { project: { ...project(), git: 'true' } }, + { project: { ...project(), status: 'unknown' } }, + { project: { ...project(), progress: 101 } }, + { project: { ...project(), progress: -1 } }, + { project: { ...project(), progress: 10.5 } }, + { project: { ...project(), error: 123 } }, + { project: project({ remoteUrl: undefined }) }, + { project: project({ remoteUrl: 'https://github.com/another/repo' }) }, + ]) { + test(`rejects malformed or mismatched clone results (${JSON.stringify(response)})`, async () => { + const h = createCloudSandboxProjectsTestConnection(store, { request: async () => response }); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + await assert.rejects(client.cloneProject(options, CancellationToken.None), /invalid repository information/); + }); + } + + test('validates catalogue entries before exposing them', () => { + const h = createCloudSandboxProjectsTestConnection(store, { projects: [{}] }); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + assert.throws(() => client.getProjects(), /invalid repository information/); + }); + + test('does not send a clone for a cancelled caller', async () => { + const h = createCloudSandboxProjectsTestConnection(store); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + await assert.rejects(client.cloneProject(options, CancellationToken.Cancelled), CancellationError); + assert.deepStrictEqual(h.requests, []); + }); + + for (const closeConnection of [false, true]) { + test(`cancels an outstanding request without needing its response (${closeConnection ? 'connection closed' : 'caller cancelled'})`, async () => { + const response = new DeferredPromise(); + const h = createCloudSandboxProjectsTestConnection(store, { request: () => response.p }); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + const cts = store.add(new CancellationTokenSource()); + const result = client.cloneProject(options, cts.token); + await h.requested.p; + if (closeConnection) { + client.dispose(); + } else { + cts.cancel(); + } + await assert.rejects(result, CancellationError); + response.complete({ project: project() }); + }); + } + + test('a disposed adapter cannot read state or send another clone', async () => { + const h = createCloudSandboxProjectsTestConnection(store); + const client = store.add(new CloudSandboxProjectsClient(h.connection)); + client.dispose(); + await assert.rejects(client.cloneProject(options, CancellationToken.None), CancellationError); + assert.throws(() => client.getProjects(), CancellationError); + assert.throws(() => client.isAvailable(), CancellationError); + assert.deepStrictEqual(h.requests, []); + }); +}); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectsTestUtils.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectsTestUtils.ts new file mode 100644 index 00000000000000..d92518a8c5dfab --- /dev/null +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/cloudSandboxProjectsTestUtils.ts @@ -0,0 +1,84 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { DeferredPromise } from '../../../../../../base/common/async.js'; +import { Emitter, Event } from '../../../../../../base/common/event.js'; +import { DisposableStore } from '../../../../../../base/common/lifecycle.js'; +import { upcastPartial } from '../../../../../../base/test/common/mock.js'; +import { AgentHostProtocolClient } from '../../../../../../platform/agentHost/browser/agentHostProtocolClient.js'; +import { toAgentHostUri } from '../../../../../../platform/agentHost/common/agentHostUri.js'; +import { IAgentSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js'; +import { RootState } from '../../../../../../platform/agentHost/common/state/sessionState.js'; +import { ICloudSandboxProject } from '../../browser/cloudSandboxProjectsClient.js'; + +export function createCloudSandboxProject(overrides: Partial = {}): ICloudSandboxProject { + return { + id: 'project-1', + path: '/checkout/owner/repo', + git: true, + status: 'ready', + remoteUrl: 'https://github.com/owner/repo', + ...overrides, + }; +} + +export function createCloudSandboxProjectsTestConnection(store: Pick, options: { + projects?: readonly unknown[]; + capability?: unknown; + legacy?: boolean; + request?: () => Promise; +} = {}) { + const capability = options.legacy ? undefined : options.capability === undefined ? { available: true } : options.capability; + const root = (projects: readonly unknown[]): RootState => ({ + agents: [], + _meta: capability === undefined ? {} : { 'copilot.projectManagement': capability }, + config: { schema: { type: 'object', properties: {} }, values: { copilot: { projects } } }, + }); + let state: RootState | Error | undefined = root(options.projects ?? []); + const changes = store.add(new Emitter()); + const errors = store.add(new Emitter()); + const requested = new DeferredPromise(); + const requests: { method: string; params: Record }[] = []; + const subscription: IAgentSubscription = { + get value() { return state; }, + get verifiedValue() { return state instanceof Error ? undefined : state; }, + onDidChange: changes.event, + onDidError: errors.event, + onWillApplyAction: Event.None, + onDidApplyAction: Event.None, + }; + const connection = store.add(upcastPartial({ + rootState: subscription, + onDidChangeConnectionState: Event.None, + resourceUris: upcastPartial({ + fromAgentHost: uri => toAgentHostUri(uri, 'sandbox'), + }), + async requestHostExtension(method: string, params: Record): Promise { + requests.push({ method, params }); + requested.complete(); + return options.request ? options.request() : { project: createCloudSandboxProject({ status: 'cloning', git: false, progress: 0 }) }; + }, + dispose: () => { }, + })); + return { + connection, requests, requested, + setProjects: (projects: readonly unknown[]) => { + const next = root(projects); + state = next; + changes.fire(next); + }, + setRoot: (next: RootState | undefined) => { + state = next; + if (next) { + changes.fire(next); + } + }, + failSubscription: (error: Error) => { + state = error; + errors.fire(error); + }, + hasListeners: () => changes.hasListeners() || errors.hasListeners(), + }; +} diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index a73e3aa391e869..3ae2fe4b994c81 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts @@ -4,7 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import { status } from '../../../../../../base/browser/ui/aria/aria.js'; -import { Delayer, disposableTimeout, raceCancellation } from '../../../../../../base/common/async.js'; +import { Delayer, disposableTimeout, raceCancellation, raceCancellationError } from '../../../../../../base/common/async.js'; import { decodeBase64, encodeBase64, VSBuffer } from '../../../../../../base/common/buffer.js'; import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; import { CancellationError, getErrorCode, isCancellationError } from '../../../../../../base/common/errors.js'; @@ -882,6 +882,8 @@ export interface IAgentHostSessionHandlerConfig { * falling back to the first workspace folder. */ readonly resolveWorkingDirectory?: (sessionResource: URI) => URI | undefined; + /** Prepare a new session's directory after authentication, before any session is created. */ + readonly prepareWorkingDirectory?: (sessionResource: URI, workingDirectory: URI | undefined, token: CancellationToken) => Promise; /** Whether a final-looking chat resource is still a client-side draft. */ readonly isNewSession?: (sessionResource: URI) => boolean; /** Called after a locally-created session has been accepted by the backend. */ @@ -932,6 +934,7 @@ class ActiveClientEntry extends Disposable { constructor( private readonly _scope: IAgentCustomizationScope, + readonly scopeRoots: readonly URI[], clientId: string, debounceDelay: number, private readonly _getSessionState: (backendSession: URI) => SessionState | undefined, @@ -1901,6 +1904,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC Object.keys(initialConfig).length > 0 ? initialConfig : undefined, imported ? { turns: imported.turns, model: imported.model } : undefined, stage => failureStage = stage, + cancellationToken, ); } else { failureStage = 'authentication'; @@ -2323,15 +2327,18 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC await entry.claim(backendSession, cancellationToken); } - private _ensureActiveClientEntry(sessionResource: URI): ActiveClientEntry { + private _ensureActiveClientEntry(sessionResource: URI, scopeRoots?: readonly URI[]): ActiveClientEntry { const existing = this._activeClientEntries.get(sessionResource); - if (existing) { + if (existing && (!scopeRoots || this._activeClientService.areScopeRootsEqual(existing.scopeRoots, scopeRoots))) { return existing; } + this._disposeActiveClientEntry(sessionResource); - const scope = this._activeClientService.acquireScope(this._config.sessionType, this._resolveCustomizationScopeRoots(sessionResource)); + const roots = scopeRoots ?? this._resolveCustomizationScopeRoots(sessionResource); + const scope = this._activeClientService.acquireScope(this._config.sessionType, roots); const entry = new ActiveClientEntry( scope, + roots, this._config.connection.clientId, AgentHostSessionHandler.ACTIVE_CLIENT_RECONCILIATION_DEBOUNCE_MS, backendSession => this._getSessionState(backendSession.toString()), @@ -5583,8 +5590,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } /** Creates a new backend session and subscribes to its state. */ - private async _createAndSubscribe(sessionResource: URI, model: ModelSelection | undefined, config?: Record, importConversation?: { readonly turns: readonly Turn[]; readonly model?: ModelSelection }, onFailureStage?: (stage: AgentHostInvocationFailureStage) => void): Promise { - const workingDirectories = this._resolveRequestedWorkingDirectories(sessionResource); + private async _createAndSubscribe(sessionResource: URI, model: ModelSelection | undefined, config?: Record, importConversation?: { readonly turns: readonly Turn[]; readonly model?: ModelSelection }, onFailureStage?: (stage: AgentHostInvocationFailureStage) => void, cancellationToken: CancellationToken = CancellationToken.None): Promise { const requestedSession = this._resolveSessionUri(sessionResource); const meta = this._provisionalService.getInitialSessionMetadata(sessionResource); @@ -5593,8 +5599,18 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC onFailureStage?.('authentication'); const protectedResources = await this._ensureRequiredAuthentication(model); - const activeClientEntry = this._ensureActiveClientEntry(sessionResource); - await activeClientEntry.whenSettled(); + onFailureStage?.('createSession'); + const requestedDirectory = this._resolveRequestedWorkingDirectory(sessionResource); + const preparedDirectory = this._config.prepareWorkingDirectory + ? await this._config.prepareWorkingDirectory(sessionResource, requestedDirectory, cancellationToken) + : requestedDirectory; + if (cancellationToken.isCancellationRequested) { + throw new CancellationError(); + } + const workingDirectories = this._resolveRequestedWorkingDirectories(sessionResource, preparedDirectory); + + const activeClientEntry = this._ensureActiveClientEntry(sessionResource, workingDirectories ?? []); + await raceCancellationError(activeClientEntry.whenSettled(), cancellationToken); const activeClient = this._getCurrentActiveClient(sessionResource); // Opt in to bring-up progress (chiefly the lazy first-use SDK download) @@ -6073,8 +6089,7 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC } /** `undefined` is preserved for createSession to let the host choose its working directories. */ - private _resolveRequestedWorkingDirectories(sessionResource: URI): readonly URI[] | undefined { - const primary = this._resolveRequestedWorkingDirectory(sessionResource); + private _resolveRequestedWorkingDirectories(sessionResource: URI, primary = this._resolveRequestedWorkingDirectory(sessionResource)): readonly URI[] | undefined { return this._hostAddressableWorkingDirectories( computeWorkingDirectories(primary, this._workspaceContextService.getWorkspace().folders.map(folder => folder.uri), this._getRootState(), this._config.provider) ); diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts index e96c92f5e55b01..20cfdb0d686d88 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts @@ -8,6 +8,7 @@ import * as dom from '../../../../../../base/browser/dom.js'; import { setARIAContainer } from '../../../../../../base/browser/ui/aria/aria.js'; import { encodeBase64, VSBuffer } from '../../../../../../base/common/buffer.js'; import { CancellationToken, CancellationTokenSource } from '../../../../../../base/common/cancellation.js'; +import { CancellationError } from '../../../../../../base/common/errors.js'; import { Codicon } from '../../../../../../base/common/codicons.js'; import { Emitter, Event } from '../../../../../../base/common/event.js'; import { DisposableStore, IDisposable, IReference, toDisposable } from '../../../../../../base/common/lifecycle.js'; @@ -10946,6 +10947,161 @@ suite('AgentHostChatContribution', () => { assert.strictEqual(agentHostService.createSessionCalls[0].workingDirectories?.[0]?.toString(), URI.file('/custom/working/dir').toString()); })); + test('handler prepares the working directory after authentication and before creation', () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { instantiationService, agentHostService, chatAgentService } = createTestServices(disposables); + const calls: string[] = []; + const repository = URI.parse('https://example.com/owner/repo'); + const directory = URI.file('/host/checkout'); + agentHostService.setRootState({ + agents: [{ + provider: 'copilot', + displayName: 'Test', + description: 'test', + models: [], + protectedResources: [{ + resource: 'https://example.com', + authorization_servers: ['https://example.com/login'], + required: true, + }], + }], + activeSessions: 0, + }); + const handler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { + provider: 'copilot', + agentId: 'prepare-directory-test', + sessionType: 'prepare-directory-test', + fullName: 'Test', + description: 'test', + connection: agentHostService, + connectionAuthority: 'local', + resolveWorkingDirectory: () => repository, + resolveAuthentication: async () => { + calls.push('authenticate'); + return true; + }, + prepareWorkingDirectory: async (_sessionResource, requestedDirectory) => { + calls.push(`prepare:${requestedDirectory?.toString()}`); + calls.push(`created:${agentHostService.createSessionCalls.length}`); + return directory; + }, + })); + const { turnPromise, session, turnId, fire } = await startTurn(handler, agentHostService, chatAgentService, disposables, { agentId: 'prepare-directory-test' }); + fire({ type: 'chat/turnComplete', endedAt: '2025-01-01T00:00:00.000Z', session, turnId } as ChatAction); + await turnPromise; + assert.deepStrictEqual({ + calls, + directories: agentHostService.createSessionCalls.map(call => call.workingDirectories?.map(uri => uri.toString())), + }, { + calls: ['authenticate', `prepare:${repository.toString()}`, 'created:0'], + directories: [[directory.toString()]], + }); + })); + + test('handler does not create a session or send a turn when directory preparation fails', async () => { + const { instantiationService, agentHostService, chatAgentService } = createTestServices(disposables); + disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { + provider: 'copilot', + agentId: 'prepare-failure-test', + sessionType: 'prepare-failure-test', + fullName: 'Test', + description: 'test', + connection: agentHostService, + connectionAuthority: 'local', + prepareWorkingDirectory: async () => { throw new Error('Repository could not be prepared'); }, + })); + const registered = chatAgentService.registeredAgents.get('prepare-failure-test'); + assert.ok(registered); + await assert.rejects(registered.impl.invoke(makeRequest({ + agentId: 'prepare-failure-test', + sessionResource: URI.from({ scheme: 'prepare-failure-test', path: '/new-failure' }), + }), () => { }, [], CancellationToken.None), /Repository could not be prepared/); + assert.deepStrictEqual({ sessions: agentHostService.createSessionCalls.length, turns: agentHostService.turnActions.length }, { sessions: 0, turns: 0 }); + }); + + for (const changesDirectory of [true, false]) { + test(`handler rebinds the prepared customization scope only when its roots change (${changesDirectory})`, () => runWithFakedTimers({ useFakeTimers: true }, async () => { + const { instantiationService, agentHostService, chatAgentService, activeClientService, seedActiveClient } = createTestServices(disposables); + const directory = URI.file('/host/checkout'); + const requestedDirectory = changesDirectory ? URI.parse('https://example.com/owner/repo') : directory; + const initialRoots = changesDirectory ? [] : [directory.toString()]; + const scopes: { roots: string[]; disposed: boolean }[] = []; + const acquireScope = activeClientService.acquireScope; + activeClientService.acquireScope = (sessionType, roots) => { + const scope = acquireScope(sessionType, roots); + const record = { roots: roots.map(root => root.toString()), disposed: false }; + scopes.push(record); + return { + ...scope, + dispose: () => { + record.disposed = true; + scope.dispose(); + }, + }; + }; + const customizations: ClientPluginCustomization[] = [{ + type: CustomizationType.Plugin, id: 'checkout-mcp', uri: 'file:///checkout-mcp', name: 'Checkout MCP', + }]; + disposables.add(seedActiveClient('prepare-scope-test', { customizations: constObservable(customizations) }, [directory])); + agentHostService.setInitializeResult({ defaultDirectory: URI.file('/host').toString() }); + const handler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { + provider: 'copilot', + agentId: 'prepare-scope-test', + sessionType: 'prepare-scope-test', + fullName: 'Test', + description: 'test', + connection: agentHostService, + connectionAuthority: 'local', + resolveWorkingDirectory: () => requestedDirectory, + prepareWorkingDirectory: async () => directory, + })); + let scopesBeforeSend: string[][] = []; + const { turnPromise, session, turnId, fire } = await startTurn(handler, agentHostService, chatAgentService, disposables, { + agentId: 'prepare-scope-test', + beforeInvoke: () => { scopesBeforeSend = scopes.map(scope => scope.roots); }, + }); + fire({ type: 'chat/turnComplete', endedAt: '2025-01-01T00:00:00.000Z', session, turnId } as ChatAction); + await turnPromise; + assert.deepStrictEqual({ + scopesBeforeSend, + scopes, + directories: agentHostService.createSessionCalls.map(call => call.workingDirectories?.map(uri => uri.toString())), + customizations: agentHostService.createSessionCalls.map(call => call.activeClient?.customizations), + }, { + scopesBeforeSend: [initialRoots], + scopes: changesDirectory + ? [{ roots: [], disposed: true }, { roots: [directory.toString()], disposed: false }] + : [{ roots: initialRoots, disposed: false }], + directories: [[directory.toString()]], + customizations: [customizations], + }); + })); + } + + test('handler does not create a session after directory preparation is cancelled', async () => { + const { instantiationService, agentHostService, chatAgentService } = createTestServices(disposables); + const cts = disposables.add(new CancellationTokenSource()); + disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { + provider: 'copilot', + agentId: 'prepare-cancellation-test', + sessionType: 'prepare-cancellation-test', + fullName: 'Test', + description: 'test', + connection: agentHostService, + connectionAuthority: 'local', + prepareWorkingDirectory: async () => { + cts.cancel(); + return URI.file('/host/checkout'); + }, + })); + const registered = chatAgentService.registeredAgents.get('prepare-cancellation-test'); + assert.ok(registered); + await assert.rejects(registered.impl.invoke(makeRequest({ + agentId: 'prepare-cancellation-test', + sessionResource: URI.from({ scheme: 'prepare-cancellation-test', path: '/new-cancellation' }), + }), () => { }, [], cts.token), CancellationError); + assert.deepStrictEqual({ sessions: agentHostService.createSessionCalls.length, turns: agentHostService.turnActions.length }, { sessions: 0, turns: 0 }); + }); + test('handler forwards request session config to createSession', () => runWithFakedTimers({ useFakeTimers: true }, async () => { const { instantiationService, agentHostService, chatAgentService } = createTestServices( disposables, @@ -11159,6 +11315,7 @@ suite('AgentHostChatContribution', () => { test('handler resolves authentication before sending to an eager-created session', async () => { const authenticationRequests: ProtectedResourceMetadata[][] = []; + let preparations = 0; const { instantiationService, agentHostService, chatAgentService } = createTestServices(disposables); const sessionHandler = disposables.add(instantiationService.createInstance(AgentHostSessionHandler, { provider: 'copilot', @@ -11172,6 +11329,10 @@ suite('AgentHostChatContribution', () => { authenticationRequests.push(protectedResources); return true; }, + prepareWorkingDirectory: async (_sessionResource, workingDirectory) => { + preparations++; + return workingDirectory; + }, })); const protectedResource: ProtectedResourceMetadata = { resource: 'https://api.github.com', @@ -11215,9 +11376,11 @@ suite('AgentHostChatContribution', () => { assert.deepStrictEqual({ authenticationRequests, turnActionCount: agentHostService.turnActions.length, + preparations, }, { authenticationRequests: [[protectedResource]], turnActionCount: 1, + preparations: 0, }); });