|
| 1 | +/*--------------------------------------------------------------------------------------------- |
| 2 | + * Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | + * Licensed under the MIT License. See License.txt in the project root for license information. |
| 4 | + *--------------------------------------------------------------------------------------------*/ |
| 5 | + |
| 6 | +import { raceCancellationError } from '../../../../../base/common/async.js'; |
| 7 | +import { CancellationToken, CancellationTokenSource } from '../../../../../base/common/cancellation.js'; |
| 8 | +import { CancellationError, isCancellationError } from '../../../../../base/common/errors.js'; |
| 9 | +import { Event } from '../../../../../base/common/event.js'; |
| 10 | +import { Disposable, DisposableStore, IDisposable, MutableDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; |
| 11 | +import { autorun, IObservable, waitForState } from '../../../../../base/common/observable.js'; |
| 12 | +import { URI } from '../../../../../base/common/uri.js'; |
| 13 | +import { IAgentHostConnectionsService } from '../../../../../platform/agentHost/common/agentHostConnectionsService.js'; |
| 14 | +import { isRemoteAgentHostSessionType } from '../../../../../platform/agentHost/common/agentHostSessionType.js'; |
| 15 | +import { IEvaluationSessionAttachment, IEvaluationSessionAttachmentService, IEvaluationSessionIdentity } from '../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/evaluationSessionAttachmentService.js'; |
| 16 | +import { ISessionsService } from '../../../../services/sessions/browser/sessionsService.js'; |
| 17 | +import { ISession } from '../../../../services/sessions/common/session.js'; |
| 18 | +import { ISessionsManagementService } from '../../../../services/sessions/common/sessionsManagement.js'; |
| 19 | + |
| 20 | +export interface IEvaluationSessionAttachmentStartupServices { |
| 21 | + readonly sessionsManagementService: Pick<ISessionsManagementService, 'getSession' | 'onDidChangeSessions'>; |
| 22 | + readonly sessionsService: Pick<ISessionsService, 'canOpenSession' | 'openSession'> & { readonly activeSession: IObservable<ISession | undefined>; readonly initialRestoreComplete: IObservable<boolean> }; |
| 23 | + readonly connectionsService: Pick<IAgentHostConnectionsService, 'connections' | 'resolveSessionResource'>; |
| 24 | + readonly attachmentService: IEvaluationSessionAttachmentService; |
| 25 | + readonly whenWorkbenchRestored: Promise<void>; |
| 26 | + readonly reconcileClientToolSets: () => void; |
| 27 | +} |
| 28 | +export function parseEvaluationSessionResource(value: string): URI { |
| 29 | + const resource = URI.parse(value, true); |
| 30 | + const id = resource.path.substring(1); |
| 31 | + if (!isRemoteAgentHostSessionType(resource.scheme) || resource.authority || !id |
| 32 | + || resource.path !== `/${id}` || id.includes('/') || resource.query || resource.fragment |
| 33 | + || resource.toString() !== value) { |
| 34 | + throw new Error('The evaluation session URI must be a canonical remote session URI.'); |
| 35 | + } |
| 36 | + return resource; |
| 37 | +} |
| 38 | +export async function startEvaluationSessionAttachment(value: string | undefined, getServices: () => IEvaluationSessionAttachmentStartupServices, token: CancellationToken, onFailure: (error: Error) => void = () => { }): Promise<IDisposable | undefined> { |
| 39 | + if (value === undefined) { |
| 40 | + return undefined; |
| 41 | + } |
| 42 | + const resource = parseEvaluationSessionResource(value); |
| 43 | + const services = getServices(); |
| 44 | + let attachment: IEvaluationSessionAttachment | undefined; |
| 45 | + try { |
| 46 | + await waitForState(services.sessionsService.initialRestoreComplete, complete => complete, undefined, token); |
| 47 | + const session = await waitForExactSession(services.sessionsManagementService, resource, token); |
| 48 | + const identity = resolveEvaluationSessionIdentity(resource, session, services.connectionsService); |
| 49 | + if (!await raceCancellationError(services.sessionsService.canOpenSession(session), token)) { |
| 50 | + throw new Error('The evaluation session workspace is not trusted.'); |
| 51 | + } |
| 52 | + attachment = services.attachmentService.attach(identity); |
| 53 | + await raceCancellationError(services.sessionsService.openSession(resource), token); |
| 54 | + await waitForState( |
| 55 | + services.sessionsService.activeSession, |
| 56 | + active => active?.resource.toString() === resource.toString(), |
| 57 | + active => active && active.resource.toString() !== resource.toString() |
| 58 | + ? new Error('Opening the evaluation session activated a different session.') |
| 59 | + : false, |
| 60 | + token, |
| 61 | + ); |
| 62 | + await raceCancellationError(services.whenWorkbenchRestored, token); |
| 63 | + if (services.sessionsService.activeSession.get()?.resource.toString() !== resource.toString()) { |
| 64 | + throw new Error('The active evaluation session changed before publication was ready.'); |
| 65 | + } |
| 66 | + services.reconcileClientToolSets(); |
| 67 | + attachment.markActiveClientPublicationReady(); |
| 68 | + const retained = new DisposableStore(); |
| 69 | + retained.add(attachment); |
| 70 | + attachment = undefined; |
| 71 | + retained.add(autorun(reader => { |
| 72 | + if (services.sessionsService.activeSession.read(reader)?.resource.toString() !== resource.toString()) { |
| 73 | + retained.dispose(); |
| 74 | + onFailure(new Error('The active evaluation session changed.')); |
| 75 | + } |
| 76 | + })); |
| 77 | + // The SDK supplies the workspace working directory; the driver waits for active-client inventory before sending a turn. |
| 78 | + return retained; |
| 79 | + } catch (error) { |
| 80 | + attachment?.dispose(); |
| 81 | + throw error; |
| 82 | + } |
| 83 | +} |
| 84 | +async function waitForExactSession(service: Pick<ISessionsManagementService, 'getSession' | 'onDidChangeSessions'>, resource: URI, token: CancellationToken): Promise<ISession> { |
| 85 | + for (;;) { |
| 86 | + if (token.isCancellationRequested) { |
| 87 | + throw new CancellationError(); |
| 88 | + } |
| 89 | + const session = service.getSession(resource); |
| 90 | + if (session?.resource.toString() === resource.toString()) { |
| 91 | + return session; |
| 92 | + } |
| 93 | + const change = Event.toPromise(service.onDidChangeSessions); |
| 94 | + try { |
| 95 | + await raceCancellationError(change, token); |
| 96 | + } finally { |
| 97 | + change.cancel(); |
| 98 | + } |
| 99 | + } |
| 100 | +} |
| 101 | +export function resolveEvaluationSessionIdentity(resource: URI, session: ISession, connectionsService: Pick<IAgentHostConnectionsService, 'connections' | 'resolveSessionResource'>): IEvaluationSessionIdentity { |
| 102 | + const resolution = connectionsService.resolveSessionResource(resource); |
| 103 | + const connection = connectionsService.connections.find(candidate => !candidate.isAmbient && candidate.connection === resolution?.connection); |
| 104 | + const backendSession = (session as ISession & { readonly backendUri?: URI }).backendUri; |
| 105 | + if (!resolution || !connection || session.resource.toString() !== resource.toString() |
| 106 | + || !URI.isUri(backendSession) || backendSession.authority || backendSession.path !== resource.path |
| 107 | + || backendSession.query || backendSession.fragment) { |
| 108 | + throw new Error('The evaluation session is not backed by the exact connected remote agent host.'); |
| 109 | + } |
| 110 | + return { connectionAuthority: connection.authority, backendSession }; |
| 111 | +} |
| 112 | +export class EvaluationSessionAttachmentLifecycle extends Disposable { |
| 113 | + private readonly _attachment = this._register(new MutableDisposable<IDisposable>()); |
| 114 | + constructor(value: string | undefined, getServices: () => IEvaluationSessionAttachmentStartupServices, onFailure: (error: Error) => void) { |
| 115 | + super(); |
| 116 | + if (value === undefined) { |
| 117 | + return; |
| 118 | + } |
| 119 | + const cancellation = new CancellationTokenSource(); |
| 120 | + this._register(toDisposable(() => cancellation.dispose(true))); |
| 121 | + void startEvaluationSessionAttachment(value, getServices, cancellation.token, error => { |
| 122 | + this._attachment.clear(); |
| 123 | + onFailure(error); |
| 124 | + }).then(attachment => { |
| 125 | + if (cancellation.token.isCancellationRequested) { |
| 126 | + attachment?.dispose(); |
| 127 | + } else { |
| 128 | + this._attachment.value = attachment; |
| 129 | + } |
| 130 | + }).catch(error => { |
| 131 | + if (!isCancellationError(error) && !cancellation.token.isCancellationRequested) { |
| 132 | + onFailure(error); |
| 133 | + } |
| 134 | + }); |
| 135 | + } |
| 136 | +} |
0 commit comments