From f6cd176377997c998e2914b5196805a7991d0255 Mon Sep 17 00:00:00 2001 From: Rob Lourens Date: Tue, 15 Sep 2026 14:04:35 -0700 Subject: [PATCH] agentHost: enforce fresh remote transport lifecycle Separate ordinary ensure/open operations from explicit reconnects so initialized relays retain their protocol clients while replacement clients always receive fresh transports. Make SSH and WSL relay replacement explicit, preserve tunnel metadata and location-change behavior, and cover connection races across transports.\n\nFixes #335919.\n\n(Written by Copilot)\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../browser/remoteAgentHostServiceImpl.ts | 183 ++++++++++++++---- .../common/remoteAgentHostService.ts | 21 +- .../agentHost/common/sshRemoteAgentHost.ts | 24 ++- .../agentHost/common/tunnelAgentHost.ts | 6 +- .../agentHost/common/wslRemoteAgentHost.ts | 13 +- .../sshRemoteAgentHostServiceImpl.ts | 88 ++++----- .../wslRemoteAgentHostServiceImpl.ts | 105 +++++----- .../node/sshRemoteAgentHostService.ts | 44 +++-- .../node/wslRemoteAgentHostService.ts | 37 ++-- .../remoteAgentHostService.test.ts | 144 ++++++++++++-- .../sshRemoteAgentHostService.test.ts | 163 ++++++++++++---- .../wslRemoteAgentHostService.test.ts | 17 ++ .../node/sshRemoteAgentHostService.test.ts | 77 +++++--- .../node/wslRemoteAgentHostService.test.ts | 56 +++++- .../REMOTE_AGENT_HOST_SESSIONS_PROVIDER.md | 10 +- .../browser/browserTunnelAgentHostService.ts | 25 ++- .../entryDrivenProviderContribution.ts | 2 + .../remoteAgentHostSessionsProvider.ts | 12 +- .../browser/sshAgentHost.contribution.ts | 24 +-- .../browser/tunnelAgentHost.contribution.ts | 113 +++++++---- .../webTunnelAgentHostService.contribution.ts | 4 + .../browser/webTunnelAgentHostService.ts | 25 ++- .../browser/wslAgentHost.contribution.ts | 8 +- .../tunnelAgentHostServiceImpl.ts | 25 ++- .../browserTunnelAgentHostService.test.ts | 96 ++++++++- .../remoteAgentHostSessionsProvider.test.ts | 17 +- .../tunnelAgentHost.contribution.test.ts | 24 ++- .../tunnelAgentHostServiceImpl.test.ts | 69 ++++++- 28 files changed, 1073 insertions(+), 359 deletions(-) diff --git a/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts b/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts index 622809cdfdad0..c56b1d39b83bd 100644 --- a/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts +++ b/src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts @@ -52,19 +52,13 @@ import { agentsWindowAgentHostClientInfo, editorWindowAgentHostClientInfo } from interface IConnectionEntry { readonly store: DisposableStore; client?: IRemoteAgentHostProtocolClient; - /** - * Optional teardown for the shared-process tunnel that this entry's - * transport is using (SSH or dev-tunnels). Tracked separately from - * {@link store} because on reconnect the new entry takes ownership of - * the same underlying connectionId — running the old teardown would - * disconnect the freshly-established tunnel as a side effect. - */ + /** Optional teardown for resources established alongside the transport. */ readonly transportDisposable?: IDisposable; - /** Whether a replacement connection assumes transport teardown ownership. */ - readonly reconnectTransfersTransportOwnership: boolean; connected: boolean; /** Current connection status for UI display. */ status: RemoteAgentHostConnectionStatus; + /** Most recent failure while this entry has no usable connection. */ + connectionError?: Error; } function disposeEntry(entry: IConnectionEntry): void { @@ -72,6 +66,10 @@ function disposeEntry(entry: IConnectionEntry): void { entry.transportDisposable?.dispose(); } +function asError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); +} + /** * Whether a failed connection attempt must not be retried automatically. * @@ -161,6 +159,8 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo }); /** In-flight connection attempts, keyed by normalized address. */ private readonly _pendingConnects = new Map>(); + /** Addresses with an explicit replacement serialized behind the current attempt. */ + private readonly _queuedReconnects = new Set(); private readonly _names = new Map(); private readonly _tokens = new Map(); private readonly _pendingConnectionWaits = new Map>(); @@ -305,25 +305,73 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo return result; } - reconnect(address: string, userInitiated = true): void { + ensureConnection(address: string, userInitiated = true): void { if (this._store.isDisposed) { return; } const normalized = normalizeRemoteAgentHostAddress(address); - // A dial already in flight is itself a fresh attempt, so neither a - // retry nor a user request gains anything by tearing it down and - // starting a second one — that is what produced concurrent remote - // bootstraps. Join it instead. A user-initiated request still restores - // the retry budget, so pressing reconnect while a slow bootstrap runs - // is not silently useless if that bootstrap ultimately fails. - if (this._pendingConnects.has(normalized)) { - if (userInitiated) { - this._failedReconnects.delete(normalized); - this._cancelReconnect(normalized); - this._reconnectAttempts.delete(normalized); + this._failedReconnects.delete(normalized); + if (userInitiated) { + this._cancelReconnect(normalized); + this._reconnectAttempts.delete(normalized); + } + + const existing = this._entries.get(normalized); + if (existing?.client && ( + RemoteAgentHostConnectionStatus.isConnected(existing.status) + || RemoteAgentHostConnectionStatus.isConnecting(existing.status) + || RemoteAgentHostConnectionStatus.isReconnecting(existing.status) + || RemoteAgentHostConnectionStatus.isIncompatible(existing.status) + )) { + return; + } + const pendingConnect = this._pendingConnects.get(normalized); + if (pendingConnect) { + if ( + userInitiated + && existing + && RemoteAgentHostConnectionStatus.isDisconnected(existing.status) + && !this._queuedReconnects.has(normalized) + ) { + const configuredEntry = this._configuredEntries.get().find(entry => this._entryAddress(entry) === normalized); + if (configuredEntry && this._connectionFactories.has(configuredEntry.connection.type)) { + this._queueReconnect( + pendingConnect, + { + ...configuredEntry, + connectionToken: this._tokens.get(normalized) ?? configuredEntry.connectionToken, + }, + normalized, + { userInitiated }, + ); + } } return; } + + const configuredEntry = this._configuredEntries.get().find( + entry => this._entryAddress(entry) === normalized + ); + if (!configuredEntry) { + this._failedReconnects.set(normalized, new Error(`No remote agent host entry is staged for ${normalized}.`)); + return; + } + if (!this._connectionFactories.has(configuredEntry.connection.type)) { + this._failedReconnects.set(normalized, new Error(`No connection factory is registered for ${configuredEntry.connection.type}.`)); + return; + } + + void this._connectTo({ + ...configuredEntry, + connectionToken: this._tokens.get(normalized) ?? configuredEntry.connectionToken, + }, { userInitiated }); + } + + reconnect(address: string, userInitiated = true): void { + if (this._store.isDisposed) { + return; + } + const normalized = normalizeRemoteAgentHostAddress(address); this._failedReconnects.delete(normalized); const configuredEntry = this._configuredEntries.get().find( @@ -350,21 +398,50 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._reconnectAttempts.delete(normalized); } - // Tear down existing connection if present - const entry = this._entries.get(normalized); - if (entry) { - this._entries.delete(normalized); - entry.store.dispose(); - if (!entry.reconnectTransfersTransportOwnership) { - entry.transportDisposable?.dispose(); + const pendingConnect = this._pendingConnects.get(normalized); + if (pendingConnect) { + if (!userInitiated) { + return; } - this._onDidChangeConnections.fire(); + this._queueReconnect(pendingConnect, entryToReconnect, normalized, { userInitiated }); + return; } - // Start fresh connection attempt + this._disposeEntryForReconnect(normalized); void this._connectTo(entryToReconnect, { userInitiated }); } + private _queueReconnect(previousConnect: Promise, entry: IRemoteAgentHostEntry, address: string, options: IRemoteAgentHostConnectOptions): void { + const pendingConnect = new DeferredPromise(); + this._queuedReconnects.add(address); + this._pendingConnects.set(address, pendingConnect.p); + void (async () => { + try { + await previousConnect; + this._disposeEntryForReconnect(address); + await this._createAndConnect(entry, address, options); + } catch (error) { + this._logService.error(`[RemoteAgentHost] Unexpected error reconnecting to ${address}`, error); + } finally { + if (this._pendingConnects.get(address) === pendingConnect.p) { + this._pendingConnects.delete(address); + this._queuedReconnects.delete(address); + } + void pendingConnect.complete(); + } + })(); + } + + private _disposeEntryForReconnect(address: string): void { + const entry = this._entries.get(address); + if (!entry) { + return; + } + this._entries.delete(address); + disposeEntry(entry); + this._onDidChangeConnections.fire(); + } + /** * Skips a protocol client's pending backoff, or starts a fresh user-initiated dial. */ @@ -394,25 +471,40 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo throw reconnectFailure; } - const wait = this._getOrCreateConnectionWait(normalizedAddress); + let wait = this._getOrCreateConnectionWait(normalizedAddress); // Follow an in-flight dial rather than a wall clock. Establishment cost // varies by kind — an SSH host may install the remote CLI first, taking // far longer than a WebSocket dial — and a timeout here would report // failure while that attempt keeps running and later succeeds. The // timeout only guards the case where nothing is in flight to follow. - const pendingConnect = this._pendingConnects.get(normalizedAddress); - if (pendingConnect) { + let pendingConnect = this._pendingConnects.get(normalizedAddress); + while (pendingConnect) { await pendingConnect; const connected = this._getConnectionInfo(normalizedAddress); if (connected) { return connected; } + const successor = this._pendingConnects.get(normalizedAddress); + if (successor && successor !== pendingConnect) { + wait = this._getOrCreateConnectionWait(normalizedAddress); + pendingConnect = successor; + continue; + } + const connectionError = this._entries.get(normalizedAddress)?.connectionError; + if (connectionError) { + throw connectionError; + } // The dial finished without producing a usable connection: surface // the reason it recorded rather than waiting out the timeout. return wait.p; } + const connectionError = this._entries.get(normalizedAddress)?.connectionError; + if (connectionError) { + throw connectionError; + } + const connection = await raceTimeout(wait.p, RemoteAgentHostService.ConnectionWaitTimeout, () => { this._pendingConnectionWaits.delete(normalizedAddress); }); @@ -595,16 +687,12 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo store: new DisposableStore(), connected: false, status: RemoteAgentHostConnectionStatus.disconnectedBecause(disconnectReason), - reconnectTransfersTransportOwnership: false, + connectionError: asError(err), }); this._rejectPendingConnectionWait(address, err); - // Clear the in-flight marker before notifying. A consumer may dial again - // from this notification — an automatic start does — and `reconnect` - // joins a pending dial rather than racing it. Leaving the marker set - // would join *this* dial, which has already failed, so the new request - // would never connect and its `waitForConnection` would never settle. - // `_connectTo` clears by identity, so a fresh dial started here survives. - this._pendingConnects.delete(address); + // Keep the in-flight marker through notification. An explicit + // reconnect raised by a listener is serialized behind this attempt, + // and `_connectTo` clears its own marker by identity in `finally`. // Nothing else reports this failure: no entry was created, so consumers // only learn the address became unavailable from this notification. this._onDidChangeConnections.fire(); @@ -632,7 +720,6 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo store, client, transportDisposable: createdConnection.transportDisposable, - reconnectTransfersTransportOwnership: createdConnection.reconnectTransfersTransportOwnership ?? false, connected: false, status: RemoteAgentHostConnectionStatus.connecting, }; @@ -649,6 +736,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._logService.warn(`[RemoteAgentHost] Connection closed: ${address}`); entry.connected = false; entry.status = RemoteAgentHostConnectionStatus.disconnectedBecause(reason ?? AgentHostTransportFailureReason.Unknown); + entry.connectionError = new Error(`Connection closed: ${address}`); entry.client = undefined; disposeEntry(entry); this._onDidChangeConnections.fire(); @@ -683,9 +771,11 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo case 'connected': entry.connected = true; entry.status = RemoteAgentHostConnectionStatus.connected; + entry.connectionError = undefined; // A soft reconnect that restores the transport settles any // wait started before the drop, which would otherwise sit // until its timeout even though the host is reachable again. + this._cancelReconnect(address); this._reconnectAttempts.delete(address); this._resolvePendingConnectionWait(address); this._onDidChangeConnections.fire(); @@ -696,6 +786,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo case 'incompatible': entry.connected = false; entry.status = RemoteAgentHostConnectionStatus.incompatible('Authentication failed during connection initialization.', [PROTOCOL_VERSION]); + entry.connectionError = new Error('Authentication failed during connection initialization.'); this._reconnectAttempts.delete(address); this._onDidChangeConnections.fire(); break; @@ -715,6 +806,8 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._logService.info(`[RemoteAgentHost] Connected to ${address}`); entry.connected = true; entry.status = RemoteAgentHostConnectionStatus.connected; + entry.connectionError = undefined; + this._cancelReconnect(address); this._reconnectAttempts.delete(address); this._resolvePendingConnectionWait(address); this._onDidChangeConnections.fire(); @@ -739,6 +832,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo if (incompatible) { this._logService.warn(`[RemoteAgentHost] Incompatible with ${address}: ${incompatible.kind === 'incompatible' ? incompatible.message : ''}`); entry.status = incompatible; + entry.connectionError = asError(err); this._reconnectAttempts.delete(address); this._rejectPendingConnectionWait(address, err); this._onDidChangeConnections.fire(); @@ -752,6 +846,9 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo // then rejects this promise. Tearing the entry down here would // cancel that retry and lose the client's replay state, so let it // restore itself instead of rebuilding from scratch. + if (RemoteAgentHostConnectionStatus.isConnected(entry.status)) { + return; + } if (RemoteAgentHostConnectionStatus.isReconnecting(entry.status)) { this._logService.info(`[RemoteAgentHost] Handshake with ${address} was interrupted; the protocol client is restoring the connection`); return; @@ -759,6 +856,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo const disconnectReason = err instanceof NonReconnectableTransportError ? err.reason : AgentHostTransportFailureReason.Unknown; entry.status = RemoteAgentHostConnectionStatus.disconnectedBecause(disconnectReason); + entry.connectionError = asError(err); entry.client = undefined; // Clean up the failed client while retaining its entry and status. disposeEntry(entry); @@ -928,6 +1026,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo this._reconnectTimeouts.clear(); this._reconnectAttempts.clear(); this._pendingConnects.clear(); + this._queuedReconnects.clear(); for (const [address, wait] of this._pendingConnectionWaits) { void wait.error(new Error(`Remote agent host service disposed before connecting to ${address}`)); } diff --git a/src/vs/platform/agentHost/common/remoteAgentHostService.ts b/src/vs/platform/agentHost/common/remoteAgentHostService.ts index de4adc40bec18..705ebfdcc537b 100644 --- a/src/vs/platform/agentHost/common/remoteAgentHostService.ts +++ b/src/vs/platform/agentHost/common/remoteAgentHostService.ts @@ -290,18 +290,13 @@ export interface IRemoteAgentHostConnectOptions { /** A built, not-yet-handshaken connection and its owned resources. */ export interface IRemoteAgentHostCreatedConnection { - /** The client the service will handshake and own. */ + /** The client the service will handshake and own. Its transport must be fresh and uninitialized. */ readonly connection: IRemoteAgentHostProtocolClient; /** * Teardown for resources the factory established alongside the client * (e.g. a shared-process relay channel). Disposed with the connection entry. */ readonly transportDisposable?: IDisposable; - /** - * Whether a redial transfers transport teardown ownership to the new connection. - * Defaults to `false`. - */ - readonly reconnectTransfersTransportOwnership?: boolean; } /** Builds agent host connections of one {@link RemoteAgentHostEntryType}. */ @@ -311,10 +306,12 @@ export interface IRemoteAgentHostConnectionFactory { /** Entries owned by this factory. */ readonly entries: IObservable; /** - * Build a client bound to a transport for `entry`. + * Build a new client bound to a fresh, uninitialized transport for `entry`. * * Must NOT perform the protocol handshake — the service calls `connect()` * itself so handshake outcome classification lives in exactly one place. + * A factory backed by a lower process must replace any retained relay rather + * than bind the new client to a transport that another client initialized. */ createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise; } @@ -724,6 +721,15 @@ export interface IRemoteAgentHostService { /** Waits for a configured remote host to establish a connection. */ waitForConnection(address: string): Promise; + /** + * Ensures a configured remote host has a usable protocol client. + * + * Retains an existing client while it is connected, connecting, + * reconnecting, or incompatibly alive. Starts a fresh factory connection + * only when no live client exists. + */ + ensureConnection(address: string, userInitiated?: boolean): void; + /** * Disconnects an active remote host connection by address. */ @@ -803,6 +809,7 @@ export class NullRemoteAgentHostService implements IRemoteAgentHostService { async waitForConnection(): Promise { throw new Error('Remote agent host connections are not supported in this environment.'); } + ensureConnection(_address: string, _userInitiated?: boolean): void { } async removeRemoteAgentHost(_address: string): Promise { } reconnect(_address: string, _userInitiated?: boolean): void { } reconnectNow(_address: string): void { } diff --git a/src/vs/platform/agentHost/common/sshRemoteAgentHost.ts b/src/vs/platform/agentHost/common/sshRemoteAgentHost.ts index 5c4cecb57ba9c..9685b51f578d5 100644 --- a/src/vs/platform/agentHost/common/sshRemoteAgentHost.ts +++ b/src/vs/platform/agentHost/common/sshRemoteAgentHost.ts @@ -38,6 +38,19 @@ export const enum SSHAuthMethod { Password = 'password', } +/** + * How the shared process handles an SSH connection already retained for the + * same stable connection key. + */ +export const enum SSHConnectionMode { + /** Return the retained relay unchanged. Only its existing protocol client may keep using it. */ + Reuse = 'reuse', + /** Replace only the WebSocket relay while retaining the SSH client and selected endpoint. */ + ReplaceRelay = 'replaceRelay', + /** Replace the full SSH connection and select an endpoint again. */ + ReplaceConnection = 'replaceConnection', +} + export interface ISSHAgentHostConfig { /** Remote hostname or IP. */ readonly host: string; @@ -172,7 +185,8 @@ export interface ISSHRemoteAgentHostService { * 4. Creates a WebSocket relay over the SSH channel * 5. Waits for {@link IRemoteAgentHostService} to complete the protocol handshake * - * Resolves with the connection handle once the agent host is reachable. + * Resolves with the existing or newly established connection handle once + * the agent host is reachable. Does not replace a live protocol client. */ connect(config: ISSHAgentHostConfig): Promise; @@ -204,8 +218,8 @@ export interface ISSHRemoteAgentHostService { resolveSSHConfig(host: string): Promise; /** - * Re-establish an SSH tunnel on startup for a previously connected host. - * Resolves once the service-owned protocol connection is ready. + * Explicitly replace the connection for a previously connected host. + * Resolves once the new service-owned protocol connection is ready. * * @param userInitiated See {@link ISSHAgentHostConfig.userInitiated}. * Defaults to `true` (picker-eligible) when omitted; background/auto @@ -560,7 +574,7 @@ export interface ISSHRemoteAgentHostMainService { * Bootstrap a remote agent host over SSH. Returns serializable * connection info for the renderer to register. */ - connect(config: ISSHAgentHostConfig): Promise; + connect(config: ISSHAgentHostConfig, mode: SSHConnectionMode): Promise; /** * Send a message to a remote agent host through the SSH relay. @@ -598,5 +612,5 @@ export interface ISSHRemoteAgentHostMainService { * The renderer computes this from its stored preference for this host's * {@link computeSSHConnectionKey stable key} before calling reconnect. */ - reconnect(sshConfigHost: string, name: string, remoteAgentHostCommand?: string, agentForward?: boolean, userInitiated?: boolean, preferredAgentLocation?: RemoteAgentHostLocationPreference): Promise; + reconnect(sshConfigHost: string, name: string, mode: SSHConnectionMode, remoteAgentHostCommand?: string, agentForward?: boolean, userInitiated?: boolean, preferredAgentLocation?: RemoteAgentHostLocationPreference): Promise; } diff --git a/src/vs/platform/agentHost/common/tunnelAgentHost.ts b/src/vs/platform/agentHost/common/tunnelAgentHost.ts index 8a45b562bae06..b532a5033b232 100644 --- a/src/vs/platform/agentHost/common/tunnelAgentHost.ts +++ b/src/vs/platform/agentHost/common/tunnelAgentHost.ts @@ -484,8 +484,7 @@ export interface ITunnelAgentHostService { getAutoConnectMode(tunnel: ITunnelInfo): TunnelAutoConnectMode; /** - * Connect to a tunnel's agent host and register the connection - * with {@link IRemoteAgentHostService}. + * Ensure a tunnel's agent host is connected and refresh its cached metadata. * * @param tunnel The tunnel to connect to. * @param authProvider Optional auth provider to use. If omitted, uses cached/last known. @@ -497,6 +496,9 @@ export interface ITunnelAgentHostService { */ connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise; + /** Replace a tunnel's live connection after refreshing its cached metadata. */ + reconnect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise; + /** Whether {@link deleteTunnel} is supported by this implementation. */ readonly canDeleteTunnels: boolean; diff --git a/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts b/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts index abcda4cd5456b..6a0ac5116bff1 100644 --- a/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts +++ b/src/vs/platform/agentHost/common/wslRemoteAgentHost.ts @@ -60,6 +60,14 @@ export interface IWSLConnectResult { readonly connectionToken: string | undefined; } +/** How the shared process handles a live relay for the same WSL distro. */ +export const enum WSLConnectionMode { + /** Return the retained relay unchanged. Only its existing protocol client may keep using it. */ + Reuse = 'reuse', + /** Stop the retained relay and bootstrap a fresh one. */ + Replace = 'replace', +} + export interface IWSLAgentHostConnection extends IDisposable { readonly distro: string; readonly localAddress: string; @@ -98,9 +106,10 @@ export interface IWSLRemoteAgentHostService { isWSLAvailable(): Promise; listDistros(): Promise; listRunningDistros(): Promise; + /** Ensure the distro is connected without replacing a live protocol client. */ connect(config: IWSLAgentHostConfig): Promise; disconnect(distro: string): Promise; - /** Reconnect a cached distro, optionally as an automatic recovery attempt. */ + /** Explicitly replace a cached distro connection, optionally as an automatic recovery attempt. */ reconnect(distro: string, name: string, userInitiated?: boolean): Promise; /** * Distros the user has connected to, persisted across windows. Drives the @@ -131,7 +140,7 @@ export interface IWSLRemoteAgentHostMainService { isWSLAvailable(): Promise; listDistros(): Promise; listRunningDistros(): Promise; - connect(config: IWSLAgentHostConfig): Promise; + connect(config: IWSLAgentHostConfig, mode: WSLConnectionMode): Promise; disconnect(distro: string): Promise; reconnect(distro: string, name: string, remoteAgentHostCommand?: string, userInitiated?: boolean): Promise; } diff --git a/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts b/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts index d54829f107677..c4d4a02847a07 100644 --- a/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts +++ b/src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts @@ -40,6 +40,7 @@ import { isSSHHostKeyDeniedError, SSH_HOST_KEY_DENIED_ERROR_NAME, SSHAuthMethod, + SSHConnectionMode, type ISSHAgentHostConfig, type ISSHAgentHostConnection, type ISSHConnectResult, @@ -181,6 +182,10 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect return entry; } + clearStagedConfiguration(address: string): void { + this._stagedConfigurations.delete(address); + } + getEntryForSSHConfigHost(sshConfigHost: string): IRemoteAgentHostEntry | undefined { return readSSHRemoteAgentHostEntries(this._storageService).find(entry => entry.connection.type === RemoteAgentHostEntryType.SSH && entry.connection.sshConfigHost === sshConfigHost @@ -193,15 +198,15 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect } const stagedConfig = this._stagedConfigurations.get(entry.connection.address); - this._stagedConfigurations.delete(entry.connection.address); let result; try { result = stagedConfig - ? await this._mainService.connect(this._augmentConfig({ ...stagedConfig, userInitiated: stagedConfig.userInitiated ?? options.userInitiated })) + ? await this._mainService.connect(this._augmentConfig({ ...stagedConfig, userInitiated: stagedConfig.userInitiated ?? options.userInitiated }), SSHConnectionMode.ReplaceConnection) : entry.connection.sshConfigHost ? await this._mainService.reconnect( entry.connection.sshConfigHost, entry.name, + SSHConnectionMode.ReplaceConnection, this._getRemoteAgentHostCommand(), this._isSSHAgentForwardingEnabled(), options.userInitiated, @@ -214,7 +219,7 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect authMethod: SSHAuthMethod.Agent, name: entry.name, userInitiated: options.userInitiated, - })); + }), SSHConnectionMode.ReplaceConnection); } catch (error) { // A refused host key is the user's decision, not a transient fault. // Report it in the shared vocabulary for "do not retry" while keeping @@ -248,25 +253,12 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect }, }; if (existing) { - if (this._remoteAgentHostService.getConnection(result.address)) { - this._logService.trace('[SSHRemoteAgentHost] Returning existing connection handle'); - this.storeEntry(persistedEntry); - return { - connection: this._createRelayClient(result), - transportDisposable: this._createTransportDisposable(result.connectionId, existing, this._observeSuccessfulConnection(result, options.userInitiated)), - reconnectTransfersTransportOwnership: true, - }; - } - this._logService.info(`[SSHRemoteAgentHost] Replacing stale connection handle for ${result.address}`); - this._connections.delete(result.connectionId); - // The main service retained the SSH client while replacing its relay. - // Marking this handle closed keeps disposal from disconnecting it. - existing.fireClose(); - existing.dispose(); - this._onDidChangeConnections(); + await this._mainService.disconnect(result.connectionId); + throw new Error(`SSH shared process reused relay ${result.connectionId} while creating a fresh protocol client for ${result.address}.`); } const handle = new SSHAgentHostConnectionHandle( + result.connectionId, result.config, result.address, result.name, @@ -274,7 +266,7 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect result.instanceId, result.primary, result.lifecycle, - () => this._mainService.disconnect(result.connectionId), + connectionId => this._mainService.disconnect(connectionId), ); try { this._connections.set(result.connectionId, handle); @@ -282,9 +274,8 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect this.storeEntry(persistedEntry); const endpointSelectionObserver = this._observeSuccessfulConnection(result, options.userInitiated); return { - connection: this._createRelayClient(result), - transportDisposable: this._createTransportDisposable(result.connectionId, handle, endpointSelectionObserver), - reconnectTransfersTransportOwnership: true, + connection: this._createRelayClient(result, handle), + transportDisposable: this._createTransportDisposable(handle, endpointSelectionObserver), }; } catch (err) { this._logService.error('[SSHRemoteAgentHost] Connection setup failed', err); @@ -332,9 +323,10 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect } } - private _createTransportDisposable(connectionId: string, handle: SSHAgentHostConnectionHandle, endpointSelectionObserver?: IDisposable): IDisposable { + private _createTransportDisposable(handle: SSHAgentHostConnectionHandle, endpointSelectionObserver?: IDisposable): IDisposable { return toDisposable(() => { endpointSelectionObserver?.dispose(); + const connectionId = handle.connectionId; if (this._connections.get(connectionId) === handle) { this._connections.delete(connectionId); this._onDidChangeConnections(); @@ -345,13 +337,19 @@ class SSHConnectionFactory extends Disposable implements IRemoteAgentHostConnect }); } - private _createRelayClient(result: Pick): AgentHostProtocolClient { + private _createRelayClient(result: Pick, handle: SSHAgentHostConnectionHandle): AgentHostProtocolClient { const reestablish = async (): Promise => { if (!result.sshConfigHost) { throw new NonReconnectableTransportError('Cannot automatically reconnect an SSH connection without an SSH config host.'); } const preferredAgentLocation = this._locationPreferenceService.getPreference(computeSSHConnectionKey({ sshConfigHost: result.sshConfigHost })); - const reconnected = await this._mainService.reconnect(result.sshConfigHost, result.name, this._getRemoteAgentHostCommand(), this._isSSHAgentForwardingEnabled(), false, preferredAgentLocation); + const reconnected = await this._mainService.reconnect(result.sshConfigHost, result.name, SSHConnectionMode.ReplaceRelay, this._getRemoteAgentHostCommand(), this._isSSHAgentForwardingEnabled(), false, preferredAgentLocation); + if (this._connections.get(handle.connectionId) === handle) { + this._connections.delete(handle.connectionId); + handle.replaceConnectionId(reconnected.connectionId); + this._connections.set(handle.connectionId, handle); + this._onDidChangeConnections(); + } return { connectionId: reconnected.connectionId }; }; return this._relayClientFactory.createClient(this._mainService, result.connectionId, result.address, reestablish); @@ -473,21 +471,8 @@ export class SSHRemoteAgentHostService extends Disposable implements ISSHRemoteA this._logService.info(`[SSHRemoteAgentHost] onDidCloseConnection: connectionId=${connectionId}`); const handle = this._connections.get(connectionId); if (handle) { - this._logService.info(`[SSHRemoteAgentHost] onDidCloseConnection: found handle for ${connectionId}, cleaning up`); - this._connections.delete(connectionId); - handle.fireClose(); - handle.dispose(); - this._onDidChangeConnections.fire(); - - // Defense-in-depth: also signal the protocol client directly. The - // ReconnectingRelayTransport normally observes `onDidRelayClose` (fired from - // the same shared-process code path as this event) and calls back - // into the client. If that IPC delivery is missed for any reason, - // the renderer-side client would stay in `Connected` until its - // liveness watchdog fires — which can take hours when the - // renderer is backgrounded and Chromium throttles `setTimeout`. - // Use the handle's address (e.g., "ssh:macbook-air") since - // RemoteAgentHostService keys its clients by address, not connectionId. + // Keep the logical handle while its protocol client restores the + // relay. The reestablish callback rekeys it to the new relay id. this._logService.info(`[SSHRemoteAgentHost] onDidCloseConnection: notifying protocol client for ${handle.localAddress}`); this._remoteAgentHostService.notifyConnectionClosed(handle.localAddress); } else { @@ -536,8 +521,12 @@ export class SSHRemoteAgentHostService extends Disposable implements ISSHRemoteA const entry = this._connectionFactory.stageConfiguration({ ...config, userInitiated: config.userInitiated ?? true }); const address = getEntryAddress(entry); - this._remoteAgentHostService.reconnect(address, true); - await this._remoteAgentHostService.waitForConnection(address); + this._remoteAgentHostService.ensureConnection(address, true); + try { + await this._remoteAgentHostService.waitForConnection(address); + } finally { + this._connectionFactory.clearStagedConfiguration(address); + } return this._getConnectionHandle(address); } @@ -981,6 +970,7 @@ class SSHAgentHostConnectionHandle extends Disposable implements ISSHAgentHostCo private _closedByMain = false; constructor( + private _connectionId: string, readonly config: ISSHAgentHostConnection['config'], readonly localAddress: string, readonly name: string, @@ -988,7 +978,7 @@ class SSHAgentHostConnectionHandle extends Disposable implements ISSHAgentHostCo readonly instanceId: ISSHAgentHostConnection['instanceId'], readonly primary: ISSHAgentHostConnection['primary'], readonly lifecycle: ISSHAgentHostConnection['lifecycle'], - disconnectFn: () => Promise, + disconnectFn: (connectionId: string) => Promise, ) { super(); @@ -996,11 +986,19 @@ class SSHAgentHostConnectionHandle extends Disposable implements ISSHAgentHostCo // (skip if already closed from the main process side) this._register(toDisposable(() => { if (!this._closedByMain) { - disconnectFn().catch(() => { /* best effort */ }); + disconnectFn(this._connectionId).catch(() => { /* best effort */ }); } })); } + get connectionId(): string { + return this._connectionId; + } + + replaceConnectionId(connectionId: string): void { + this._connectionId = connectionId; + } + /** Called by the service when the main process signals connection closure. */ fireClose(): void { this._closedByMain = true; diff --git a/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts b/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts index 7162484488d9c..5c14b682cdfd1 100644 --- a/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts +++ b/src/vs/platform/agentHost/electron-browser/wslRemoteAgentHostServiceImpl.ts @@ -33,13 +33,14 @@ import { type IWSLDistro, type IWSLRemoteAgentHostMainService, WSL_ADDRESS_PREFIX, + WSLConnectionMode, } from '../common/wslRemoteAgentHost.js'; export const IWSLRelayClientFactory = createDecorator('wslRelayClientFactory'); export interface IWSLRelayClientFactory { readonly _serviceBrand: undefined; - createClient(mainService: IWSLRemoteAgentHostMainService, connectionId: string, address: string, connection: IWSLConnectResult, remoteAgentHostCommand: string | undefined): AgentHostProtocolClient; + createClient(mainService: IWSLRemoteAgentHostMainService, connectionId: string, address: string, connection: IWSLConnectResult, remoteAgentHostCommand: string | undefined, onDidReestablish: (connectionId: string) => void): AgentHostProtocolClient; } export class WSLRelayClientFactory implements IWSLRelayClientFactory { @@ -52,7 +53,7 @@ export class WSLRelayClientFactory implements IWSLRelayClientFactory { @ILogService private readonly _logService: ILogService, ) { } - createClient(mainService: IWSLRemoteAgentHostMainService, connectionId: string, address: string, connection: IWSLConnectResult, remoteAgentHostCommand: string | undefined): AgentHostProtocolClient { + createClient(mainService: IWSLRemoteAgentHostMainService, connectionId: string, address: string, connection: IWSLConnectResult, remoteAgentHostCommand: string | undefined, onDidReestablish: (connectionId: string) => void): AgentHostProtocolClient { const config: IWSLAgentHostConfig = { distro: connection.distro, name: connection.name, @@ -73,6 +74,7 @@ export class WSLRelayClientFactory implements IWSLRelayClientFactory { throw new NonReconnectableTransportError(`WSL distro '${config.distro}' is not running.`, AgentHostTransportFailureReason.HostNotRunning); } const result = await mainService.reconnect(config.distro, config.name, config.remoteAgentHostCommand, false); + onDidReestablish(result.connectionId); return { connectionId: result.connectionId, }; @@ -144,7 +146,7 @@ class WSLConnectionFactory extends Disposable implements IRemoteAgentHostConnect readonly kind = RemoteAgentHostEntryType.WSL; readonly entries: IObservable; - private readonly _stagedConfigurations = new Map(); + private readonly _stagedConfigurations = new Map(); constructor( private readonly _storageService: IStorageService, @@ -155,7 +157,7 @@ class WSLConnectionFactory extends Disposable implements IRemoteAgentHostConnect private readonly _onDidChangeConnections: () => void, private readonly _onDidReportConnectProgress: (progress: IWSLConnectProgress) => void, private readonly _getRemoteAgentHostCommand: () => string | undefined, - private readonly _createTransportDisposable: (connectionId: string, distro: string, handle: WSLAgentHostConnectionHandle) => IDisposable, + private readonly _createTransportDisposable: (distro: string, handle: WSLAgentHostConnectionHandle) => IDisposable, private readonly _logService: ILogService, ) { super(); @@ -168,34 +170,34 @@ class WSLConnectionFactory extends Disposable implements IRemoteAgentHostConnect stageConfiguration(config: IWSLAgentHostConfig): IRemoteAgentHostEntry { const entry = this._createEntry(config.distro, config.name); - this._stagedConfigurations.set(getEntryAddress(entry), { config, isInitialConnection: true }); + this._stagedConfigurations.set(getEntryAddress(entry), config); this._storeEntry(entry); return entry; } stageEntry(distro: string, name: string, userInitiated = true): IRemoteAgentHostEntry { const entry = this._createEntry(distro, name); - this._stagedConfigurations.set(getEntryAddress(entry), { - config: { distro, name, remoteAgentHostCommand: this._getRemoteAgentHostCommand(), userInitiated }, - isInitialConnection: false, - }); + this._stagedConfigurations.set(getEntryAddress(entry), { distro, name, remoteAgentHostCommand: this._getRemoteAgentHostCommand(), userInitiated }); this._storeEntry(entry); return entry; } + clearStagedConfiguration(address: string): void { + this._stagedConfigurations.delete(address); + } + async createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { if (entry.connection.type !== RemoteAgentHostEntryType.WSL) { throw new Error(`WSL factory cannot create a ${entry.connection.type} connection.`); } const address = getEntryAddress(entry); - let stagedConnection = this._stagedConfigurations.get(address); - this._stagedConfigurations.delete(address); - let config = stagedConnection?.config ?? { + let config = this._stagedConfigurations.get(address) ?? { distro: entry.connection.distro, name: entry.name, remoteAgentHostCommand: this._getRemoteAgentHostCommand(), }; + this._stagedConfigurations.delete(address); let userInitiated = config.userInitiated ?? options.userInitiated; if (!userInitiated) { try { @@ -206,23 +208,19 @@ class WSLConnectionFactory extends Disposable implements IRemoteAgentHostConnect throw err; } this._stagedConfigurations.delete(address); - stagedConnection = userStagedConnection; - config = stagedConnection.config; + config = userStagedConnection; userInitiated = config.userInitiated ?? options.userInitiated; } // A user action may have arrived while the background precondition ran. const userStagedConnection = this._stagedConfigurations.get(address); if (userStagedConnection) { this._stagedConfigurations.delete(address); - stagedConnection = userStagedConnection; - config = stagedConnection.config; + config = userStagedConnection; userInitiated = config.userInitiated ?? options.userInitiated; } } - const result = stagedConnection?.isInitialConnection - ? await this._mainService.connect({ ...config, userInitiated }) - : await this._mainService.reconnect(config.distro, config.name, config.remoteAgentHostCommand, userInitiated); + const result = await this._mainService.connect({ ...config, userInitiated }, WSLConnectionMode.Replace); this._logService.trace(`[WSLRemoteAgentHost] WSL relay established, connectionId=${result.connectionId}`); return this._setupConnection(result, config.remoteAgentHostCommand); } @@ -267,21 +265,15 @@ class WSLConnectionFactory extends Disposable implements IRemoteAgentHostConnect } } - private _setupConnection(result: IWSLConnectResult, remoteAgentHostCommand: string | undefined): IRemoteAgentHostCreatedConnection { + private async _setupConnection(result: IWSLConnectResult, remoteAgentHostCommand: string | undefined): Promise { const existing = this._connections.get(result.connectionId); if (existing) { - if (this._remoteAgentHostService.getConnection(result.address)) { - this._logService.trace(`[WSLRemoteAgentHost] Returning existing connection handle for ${result.address}, connectionId=${result.connectionId}`); - return this._createConnection(result, remoteAgentHostCommand, existing); - } - this._logService.info(`[WSLRemoteAgentHost] Replacing stale connection handle for ${result.address}, connectionId=${result.connectionId}`); - this._connections.delete(result.connectionId); - existing.fireClose(); - existing.dispose(); - this._onDidChangeConnections(); + await this._mainService.disconnect(result.distro); + throw new Error(`WSL shared process reused relay ${result.connectionId} while creating a fresh protocol client for ${result.address}.`); } const handle = new WSLAgentHostConnectionHandle( + result.connectionId, result.distro, result.address, result.name, @@ -308,15 +300,21 @@ class WSLConnectionFactory extends Disposable implements IRemoteAgentHostConnect message: localize('wslProgressHandshake', "Establishing connection to {0}...", result.name), }); const completionObserver = this._observeSuccessfulConnection(result); - const transportDisposable = this._createTransportDisposable(result.connectionId, result.distro, handle); + const transportDisposable = this._createTransportDisposable(result.distro, handle); try { return { - connection: this._relayClientFactory.createClient(this._mainService, result.connectionId, result.address, result, remoteAgentHostCommand), + connection: this._relayClientFactory.createClient(this._mainService, result.connectionId, result.address, result, remoteAgentHostCommand, connectionId => { + if (this._connections.get(handle.connectionId) === handle) { + this._connections.delete(handle.connectionId); + handle.replaceConnectionId(connectionId); + this._connections.set(handle.connectionId, handle); + this._onDidChangeConnections(); + } + }), transportDisposable: toDisposable(() => { completionObserver.dispose(); transportDisposable.dispose(); }), - reconnectTransfersTransportOwnership: true, }; } catch (err) { completionObserver.dispose(); @@ -385,7 +383,7 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA () => this._onDidChangeConnections.fire(), progress => this._onDidReportLocalConnectProgress.fire(progress), () => this._getRemoteAgentHostCommand(), - (connectionId, distro, handle) => this._createTransportDisposable(connectionId, distro, handle), + (distro, handle) => this._createTransportDisposable(distro, handle), this._logService, )); this._register(this._remoteAgentHostService.registerConnectionFactory(this._connectionFactory)); @@ -394,19 +392,8 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA this._logService.info(`[WSLRemoteAgentHost] onDidCloseConnection: connectionId=${connectionId}`); const handle = this._connections.get(connectionId); if (handle) { - this._connections.delete(connectionId); - handle.fireClose(); - handle.dispose(); - this._onDidChangeConnections.fire(); - - // Defense-in-depth: also signal the protocol client directly. - // ReconnectingRelayTransport normally observes `onDidRelayClose` - // (fired from the same shared-process code path as this - // event) and calls back into the client. If that IPC - // delivery is missed for any reason, the renderer-side - // client would stay in `Connected` until its liveness - // watchdog fires — which can take hours when the renderer - // is backgrounded and Chromium throttles `setTimeout`. + // Keep the logical handle while its protocol client restores the + // relay. The reestablish callback rekeys it to the new relay id. this._remoteAgentHostService.notifyConnectionClosed(handle.localAddress); } })); @@ -438,8 +425,12 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA const entry = this._connectionFactory.stageConfiguration(this._augmentConfig({ ...config, userInitiated: config.userInitiated ?? true })); const address = getEntryAddress(entry); this._logService.info(`[WSLRemoteAgentHost] Connecting to distro ${config.distro}`); - this._remoteAgentHostService.reconnect(address, true); - await this._remoteAgentHostService.waitForConnection(address); + this._remoteAgentHostService.ensureConnection(address, true); + try { + await this._remoteAgentHostService.waitForConnection(address); + } finally { + this._connectionFactory.clearStagedConfiguration(address); + } return this._getConnectionHandle(address); } @@ -457,7 +448,11 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA const address = getEntryAddress(entry); this._logService.info(`[WSLRemoteAgentHost] Reconnecting to distro ${distro}`); this._remoteAgentHostService.reconnect(address, userInitiated); - await this._remoteAgentHostService.waitForConnection(address); + try { + await this._remoteAgentHostService.waitForConnection(address); + } finally { + this._connectionFactory.clearStagedConfiguration(address); + } return this._getConnectionHandle(address); } @@ -504,8 +499,9 @@ export class WSLRemoteAgentHostService extends Disposable implements IWSLRemoteA * reconciliation), this tears down the renderer-side handle and the * shared-process WSL relay together so neither is leaked. */ - private _createTransportDisposable(connectionId: string, distro: string, handle: WSLAgentHostConnectionHandle): IDisposable { + private _createTransportDisposable(distro: string, handle: WSLAgentHostConnectionHandle): IDisposable { return toDisposable(() => { + const connectionId = handle.connectionId; if (this._connections.get(connectionId) === handle) { this._connections.delete(connectionId); this._onDidChangeConnections.fire(); @@ -540,6 +536,7 @@ class WSLAgentHostConnectionHandle extends Disposable implements IWSLAgentHostCo private _closedByMain = false; constructor( + private _connectionId: string, readonly distro: string, readonly localAddress: string, readonly name: string, @@ -554,6 +551,14 @@ class WSLAgentHostConnectionHandle extends Disposable implements IWSLAgentHostCo })); } + get connectionId(): string { + return this._connectionId; + } + + replaceConnectionId(connectionId: string): void { + this._connectionId = connectionId; + } + /** Called by the service when the main process signals connection closure. */ fireClose(): void { this._closedByMain = true; diff --git a/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts b/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts index 6aaa253288ad3..ed3592d612733 100644 --- a/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts +++ b/src/vs/platform/agentHost/node/sshRemoteAgentHostService.ts @@ -14,6 +14,7 @@ import { Disposable, DisposableMap, toDisposable } from '../../../base/common/li import { raceTimeout } from '../../../base/common/async.js'; import { CancellationError } from '../../../base/common/errors.js'; import { URI } from '../../../base/common/uri.js'; +import { generateUuid } from '../../../base/common/uuid.js'; import { localize } from '../../../nls.js'; import { ILogService } from '../../log/common/log.js'; import { IProductService } from '../../product/common/productService.js'; @@ -21,6 +22,7 @@ import { ITelemetryService, TelemetryConfiguration } from '../../telemetry/commo import { ISSHRemoteAgentHostMainService, SSHAuthMethod, + SSHConnectionMode, computeSSHConnectionKey, type ISSHAgentHostConfig, type ISSHAgentHostConfigSanitized, @@ -791,12 +793,12 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem return this._nativeRequire; } - async connect(config: ISSHAgentHostConfig, replaceRelay?: boolean): Promise { + async connect(config: ISSHAgentHostConfig, mode: SSHConnectionMode = SSHConnectionMode.Reuse): Promise { const connectionKey = computeSSHConnectionKey(config); const existing = this._connections.get(connectionKey); if (existing) { - if (replaceRelay) { + if (mode === SSHConnectionMode.ReplaceRelay) { // Tear down the old relay and create a fresh one, following // the same dispose-and-recreate pattern as TunnelAgentHostMainService. // The SSH client is detached so only the WebSocket relay is closed. @@ -816,7 +818,7 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem // Create fresh relay and connection. If relay creation fails, // clean up the detached SSH client so it doesn't leak. - const connectionId = connectionKey; + const connectionId = generateUuid(); try { let conn: SSHConnection | undefined; // eslint-disable-line prefer-const // Bound the relay creation: a silently dead SSH client @@ -875,21 +877,25 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem } } - return { - connectionId: existing.connectionId, - address: existing.address, - name: existing.name, - connectionToken: existing.connectionToken, - config: existing.config, - sshConfigHost: config.sshConfigHost, - serverType: existing.serverType, - instanceId: existing.instanceId, - primary: true, - lifecycle: existing.lifecycle, - }; + if (mode === SSHConnectionMode.Reuse) { + return { + connectionId: existing.connectionId, + address: existing.address, + name: existing.name, + connectionToken: existing.connectionToken, + config: existing.config, + sshConfigHost: config.sshConfigHost, + serverType: existing.serverType, + instanceId: existing.instanceId, + primary: true, + lifecycle: existing.lifecycle, + }; + } + + existing.dispose(); } - this._logService.info(`${LOG_PREFIX} ${replaceRelay ? 'Reconnecting' : 'Connecting'} to ${connectionKey}`); + this._logService.info(`${LOG_PREFIX} ${mode === SSHConnectionMode.Reuse ? 'Connecting' : 'Replacing connection'} to ${connectionKey}`); const displayHost = config.sshConfigHost ?? `${config.username}@${config.host}`; let sshClient: SSHClient | undefined; @@ -1073,7 +1079,7 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem // 4. Connect to the exact selected/spawned endpoint via WebSocket relay. reportProgress(localize('sshProgressForwarding', "Connecting to remote agent host...")); - const connectionId = connectionKey; + const connectionId = generateUuid(); let conn: SSHConnection | undefined; // eslint-disable-line prefer-const let relay: { send: (data: string) => void; close: () => void }; try { @@ -1176,7 +1182,7 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem } } - async reconnect(sshConfigHost: string, name: string, remoteAgentHostCommand?: string, agentForward?: boolean, userInitiated?: boolean, preferredAgentLocation?: RemoteAgentHostLocationPreference): Promise { + async reconnect(sshConfigHost: string, name: string, mode: SSHConnectionMode = SSHConnectionMode.ReplaceConnection, remoteAgentHostCommand?: string, agentForward?: boolean, userInitiated?: boolean, preferredAgentLocation?: RemoteAgentHostLocationPreference): Promise { this._logService.info(`${LOG_PREFIX} Reconnecting via SSH config host: ${sshConfigHost} (userInitiated=${userInitiated ?? true})`); const resolved = await this.resolveSSHConfig(sshConfigHost); @@ -1203,7 +1209,7 @@ export class SSHRemoteAgentHostMainService extends Disposable implements ISSHRem agentForward: agentForward && resolved.forwardAgent ? true : undefined, userInitiated, preferredAgentLocation, - }, /* replaceRelay */ true); + }, mode); } async listSSHConfigHosts(): Promise { diff --git a/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts b/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts index 1745e65b921a4..b28ad6ff4b2d0 100644 --- a/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts +++ b/src/vs/platform/agentHost/node/wslRemoteAgentHostService.ts @@ -23,6 +23,7 @@ import { type IWSLConnectProgress, type IWSLConnectResult, type IWSLDistro, + WSLConnectionMode, } from '../common/wslRemoteAgentHost.js'; import { resolveRemotePlatform } from './sshRemoteAgentHostHelpers.js'; import { @@ -173,30 +174,31 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem } } - connect(config: IWSLAgentHostConfig): Promise { + connect(config: IWSLAgentHostConfig, mode: WSLConnectionMode = WSLConnectionMode.Reuse): Promise { const distro = validateDistroName(config.distro); - // Idempotent: a second `connect` for an already-live distro returns - // the existing connection so the renderer-side `_setupConnection` - // reuses its handle (it dedupes by `connectionId`). Picking - // "WSL..." → same distro should be a no-op, not an error. const existingId = this._distroToConnectionId.get(distro); if (existingId) { const existing = this._connections.get(existingId); if (existing) { - return Promise.resolve({ - connectionId: existing.connectionId, - address: existing.address, - distro: existing.distro, - name: existing.name, - connectionToken: existing.connectionToken, - }); + if (mode === WSLConnectionMode.Reuse) { + return Promise.resolve({ + connectionId: existing.connectionId, + address: existing.address, + distro: existing.distro, + name: existing.name, + connectionToken: existing.connectionToken, + }); + } + this._closeConnection(existingId); } } const existingPendingConnect = this._pendingConnects.get(distro); if (existingPendingConnect) { - return existingPendingConnect; + return mode === WSLConnectionMode.Reuse + ? existingPendingConnect + : existingPendingConnect.then(() => this.connect(config, WSLConnectionMode.Replace)); } // Reserve synchronously, before _connectUnguarded reaches its first @@ -460,14 +462,7 @@ export class WSLRemoteAgentHostMainService extends Disposable implements IWSLRem } async reconnect(distro: string, name: string, remoteAgentHostCommand?: string, userInitiated?: boolean): Promise { - const existingId = this._distroToConnectionId.get(distro); - if (existingId) { - this._closeConnection(existingId); - } - // A pending connection is already a fresh bootstrap. Joining it avoids - // starting a competing downloader; callers that reconnect after it - // fails receive that failure and a subsequent reconnect starts anew. - return this.connect({ distro, name, remoteAgentHostCommand, userInitiated }); + return this.connect({ distro, name, remoteAgentHostCommand, userInitiated }, WSLConnectionMode.Replace); } async relaySend(connectionId: string, message: string): Promise { diff --git a/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts index 29be6db51e0e2..00a6928b6000c 100644 --- a/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts @@ -106,7 +106,7 @@ class TestConnectionFactory extends Disposable implements IRemoteAgentHostConnec readonly entries: IObservable; private readonly _entries = observableValue(this, []); - private readonly _createdConnections = new Map(); + private readonly _createdConnections = new Map[]>(); private readonly _failures = new Map(); private readonly _onDidCreateConnection = this._register(new Emitter()); readonly onDidCreateConnection = this._onDidCreateConnection.event; @@ -117,18 +117,29 @@ class TestConnectionFactory extends Disposable implements IRemoteAgentHostConnec this.entries = this._entries; } - stage(entry: IRemoteAgentHostEntry, connection: MockProtocolClient, transportDisposable?: IDisposable, reconnectTransfersTransportOwnership = false): void { + stage(entry: IRemoteAgentHostEntry, connection: MockProtocolClient, transportDisposable?: IDisposable): void { const address = getEntryAddress(entry); const createdConnections = this._createdConnections.get(address) ?? []; - createdConnections.push({ + createdConnections.push(Promise.resolve({ connection: connection as unknown as IRemoteAgentHostProtocolClient, transportDisposable, - reconnectTransfersTransportOwnership, - }); + })); + this._createdConnections.set(address, createdConnections); + this._entries.set([...this._entries.get(), entry], undefined); + } + + stagePending(entry: IRemoteAgentHostEntry, connection: Promise): void { + const address = getEntryAddress(entry); + const createdConnections = this._createdConnections.get(address) ?? []; + createdConnections.push(connection); this._createdConnections.set(address, createdConnections); this._entries.set([...this._entries.get(), entry], undefined); } + clearEntries(): void { + this._entries.set([], undefined); + } + /** Stages a factory-level rejection, as a failed precondition check would produce. */ stageFailure(entry: IRemoteAgentHostEntry, error: Error): void { const address = getEntryAddress(entry); @@ -753,11 +764,11 @@ suite('RemoteAgentHostService', () => { } } - async function reconnectStagedConnection(factory: TestConnectionFactory, entry: IRemoteAgentHostEntry, client: MockProtocolClient, transportDisposable?: IDisposable, reconnectTransfersTransportOwnership = false): Promise { + async function reconnectStagedConnection(factory: TestConnectionFactory, entry: IRemoteAgentHostEntry, client: MockProtocolClient, transportDisposable?: IDisposable): Promise { // Capture the target before staging: `reconnect` dials asynchronously and // may already have created the connection by the time we start waiting. const expectedConnectionCount = factory.createdConnectionCount + 1; - factory.stage(entry, client, transportDisposable, reconnectTransfersTransportOwnership); + factory.stage(entry, client, transportDisposable); service.reconnect(getEntryAddress(entry)); const wait = service.waitForConnection(getEntryAddress(entry)); await waitForFactoryConnection(factory, expectedConnectionCount); @@ -787,11 +798,12 @@ suite('RemoteAgentHostService', () => { automaticCreates: 1, }); + const queuedUserClient = new MockProtocolClient('cloud:reconnect-budget'); + factory.stage(entry, queuedUserClient); service.reconnect(address, true); - // The user request joins the in-flight dial rather than starting a - // second one, but still restores the budget so a later failure is - // retried instead of being reported as exhausted. + // The explicit request is serialized behind the in-flight dial, so + // there is still only one factory call until that dial settles. assert.deepStrictEqual({ automaticAttempts: internals._reconnectAttempts.get(address), pendingReconnectCreates: factory.createdConnectionCount, @@ -803,6 +815,8 @@ suite('RemoteAgentHostService', () => { const automaticWait = service.waitForConnection(address); await waitForFactoryConnection(factory, 1); automaticClient.connectDeferred.complete(); + await waitForFactoryConnection(factory, 2); + queuedUserClient.connectDeferred.complete(); await automaticWait; const automaticDelays: number[] = []; @@ -830,11 +844,81 @@ suite('RemoteAgentHostService', () => { assert.strictEqual(internals._reconnectAttempts.get(address), undefined); const userWait = service.waitForConnection(address); - await waitForFactoryConnection(factory, 2); + await waitForFactoryConnection(factory, 3); userClient.connectDeferred.complete(); await userWait; }); + test('ensure retains a connecting or connected protocol client', async () => { + const factory = createFactory(); + const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:ensure-retains'); + const address = getEntryAddress(entry); + const client = new MockProtocolClient(address); + factory.stage(entry, client); + service.ensureConnection(address); + await waitForFactoryConnection(factory, 1); + + service.ensureConnection(address); + assert.strictEqual(factory.createdConnectionCount, 1); + + client.connectDeferred.complete(); + await service.waitForConnection(address); + service.ensureConnection(address); + + assert.strictEqual(factory.createdConnectionCount, 1); + }); + + test('factory failure preserves a queued reconnect marker through failure notification', async () => { + const factory = createFactory(); + const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:queued-after-factory-failure'); + const address = getEntryAddress(entry); + const firstFactoryResult = new DeferredPromise(); + factory.stagePending(entry, firstFactoryResult.p); + service.reconnect(address, false); + await waitForFactoryConnection(factory, 1); + const waiting = service.waitForConnection(address); + + const replacement = new MockProtocolClient(address); + factory.stage(entry, replacement); + service.reconnect(address, true); + const listener = disposables.add(service.onDidChangeConnections(() => { + if (service.connections.find(connection => connection.address === address)?.status.kind === 'disconnected') { + service.ensureConnection(address); + } + })); + + firstFactoryResult.error(new Error('factory failed')); + await waitForFactoryConnection(factory, 2); + replacement.connectDeferred.complete(); + const connected = await waiting; + listener.dispose(); + + assert.deepStrictEqual({ + createdConnectionCount: factory.createdConnectionCount, + connectedAddress: connected.address, + }, { + createdConnectionCount: 2, + connectedAddress: address, + }); + }); + + test('waitForConnection rejects when an in-flight factory result is discarded', async () => { + const factory = createFactory(); + const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:discarded'); + const address = getEntryAddress(entry); + const factoryResult = new DeferredPromise(); + const client = new MockProtocolClient(address); + factory.stagePending(entry, factoryResult.p); + service.reconnect(address); + await waitForFactoryConnection(factory, 1); + const waiting = service.waitForConnection(address); + + factory.clearEntries(); + factoryResult.complete({ connection: client as unknown as IRemoteAgentHostProtocolClient }); + + await assert.rejects(waiting, /discarded because it is no longer active/); + }); + test('surfaces a protocol reconnect backoff deadline', async () => { const factory = createFactory(); const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:backoff'); @@ -892,7 +976,7 @@ suite('RemoteAgentHostService', () => { }); }); - test('falls back to a fresh dial when a retained entry has no client', async () => { + test('ensure starts a fresh dial when a retained entry has no client', async () => { const factory = createFactory(RemoteAgentHostEntryType.WSL); const entry: IRemoteAgentHostEntry = { name: 'Ubuntu', @@ -906,7 +990,7 @@ suite('RemoteAgentHostService', () => { const client = new MockProtocolClient(address); factory.stage(entry, client); - service.reconnectNow(address); + service.ensureConnection(address); await waitForFactoryConnection(factory, 1); client.connectDeferred.complete(); await waitForConnected(); @@ -926,6 +1010,10 @@ suite('RemoteAgentHostService', () => { client.connectDeferred.error(new InitialAuthenticationError(new Error('Unsupported protocol version'))); await changed; await assert.rejects(() => wait, /Initial authentication failed/); + const replacement = new MockProtocolClient('cloud:incompatible'); + factory.stage(entry, replacement); + service.ensureConnection(getEntryAddress(entry)); + await assert.rejects(() => service.waitForConnection(getEntryAddress(entry)), /Initial authentication failed/); const upgradeResult = await service.triggerServerUpgrade('cloud:incompatible', '_vscodeUpgrade'); @@ -940,6 +1028,13 @@ suite('RemoteAgentHostService', () => { upgradeCalls: ['_vscodeUpgrade'], upgradeResult: { ok: true, upgradeStarted: true }, }); + assert.strictEqual(factory.createdConnectionCount, 1, 'ensure must retain the incompatible client and transport'); + + service.reconnect(getEntryAddress(entry)); + await waitForFactoryConnection(factory, 2); + replacement.connectDeferred.complete(); + await service.waitForConnection(getEntryAddress(entry)); + assert.strictEqual(factory.createdConnectionCount, 2, 'explicit reconnect must replace the incompatible client'); }); test('retains a client-less disconnected entry when the factory rejects before a client exists', async () => { @@ -1016,7 +1111,12 @@ suite('RemoteAgentHostService', () => { connection: { type: RemoteAgentHostEntryType.WSL, address: 'wsl:Ubuntu', distro: 'Ubuntu' }, }; const client = new MockProtocolClient('wsl:Ubuntu'); - await reconnectStagedConnection(factory, entry, client); + factory.stage(entry, client); + service.ensureConnection('wsl:Ubuntu'); + await waitForFactoryConnection(factory, 1); + client.connectDeferred.complete(); + await service.waitForConnection('wsl:Ubuntu'); + await new Promise(resolve => setTimeout(resolve, 0)); const changed = Event.toPromise(service.onDidChangeConnections); client.fireClose(AgentHostTransportFailureReason.HostNotRunning); @@ -1025,7 +1125,13 @@ suite('RemoteAgentHostService', () => { const afterClose = service.connections.find(connection => connection.address === 'wsl:Ubuntu')?.status; // A host that comes back must stop claiming it is not running. - await reconnectStagedConnection(factory, entry, new MockProtocolClient('wsl:Ubuntu')); + const replacement = new MockProtocolClient('wsl:Ubuntu'); + const expectedConnectionCount = factory.createdConnectionCount + 1; + factory.stage(entry, replacement); + service.ensureConnection('wsl:Ubuntu'); + await waitForFactoryConnection(factory, expectedConnectionCount); + replacement.connectDeferred.complete(); + await service.waitForConnection('wsl:Ubuntu'); assert.deepStrictEqual({ afterClose, @@ -1048,16 +1154,16 @@ suite('RemoteAgentHostService', () => { assert.strictEqual(service.getConnection('cloud:remove'), undefined); }); - test('does not dispose a previous transport when a replacement takes ownership', async () => { + test('disposes the previous transport before creating a replacement', async () => { const factory = createFactory(); const entry = cloudSandboxEntry('Cloud Sandbox', 'cloud:replacement'); const t1 = makeTransportDisposable(); - await reconnectStagedConnection(factory, entry, new MockProtocolClient('cloud:replacement'), t1.disposable, true); + await reconnectStagedConnection(factory, entry, new MockProtocolClient('cloud:replacement'), t1.disposable); const t2 = makeTransportDisposable(); - await reconnectStagedConnection(factory, entry, new MockProtocolClient('cloud:replacement'), t2.disposable, true); + await reconnectStagedConnection(factory, entry, new MockProtocolClient('cloud:replacement'), t2.disposable); - assert.strictEqual(t1.disposed(), false, 'previous transport disposable is not run on replacement'); + assert.strictEqual(t1.disposed(), true, 'previous transport disposable runs before replacement'); assert.strictEqual(t2.disposed(), false, 'new transport disposable is still alive'); await service.removeRemoteAgentHost('cloud:replacement'); diff --git a/src/vs/platform/agentHost/test/electron-browser/sshRemoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/electron-browser/sshRemoteAgentHostService.test.ts index 5653791e2d7ca..490dd2577f054 100644 --- a/src/vs/platform/agentHost/test/electron-browser/sshRemoteAgentHostService.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/sshRemoteAgentHostService.test.ts @@ -30,6 +30,8 @@ import { SSHHostKeyTrustService } from '../../browser/sshHostKeyTrustService.js' import { InMemoryStorageService, IStorageService } from '../../../storage/common/storage.js'; import { SSHAuthMethod, + SSHConnectionMode, + computeSSHConnectionKey, type ISSHAgentHostConfig, type ISSHConnectResult, type ISSHEndpointCandidate, @@ -151,17 +153,19 @@ class MockSSHMainService { readonly disconnectCalls: string[] = []; readonly connectCalls: ISSHAgentHostConfig[] = []; - readonly reconnectCalls: Array<{ sshConfigHost: string; name: string; remoteAgentHostCommand?: string; agentForward?: boolean; userInitiated?: boolean; preferredAgentLocation?: RemoteAgentHostLocationPreference }> = []; + readonly connectModes: SSHConnectionMode[] = []; + readonly reconnectCalls: Array<{ sshConfigHost: string; name: string; mode: SSHConnectionMode; remoteAgentHostCommand?: string; agentForward?: boolean; userInitiated?: boolean; preferredAgentLocation?: RemoteAgentHostLocationPreference }> = []; private _nextConnectionId = 1; connectResult: Partial | undefined; - async connect(config: ISSHAgentHostConfig): Promise { + async connect(config: ISSHAgentHostConfig, mode: SSHConnectionMode): Promise { this.connectCalls.push(config); + this.connectModes.push(mode); const connectionId = this.connectResult?.connectionId ?? `conn-${this._nextConnectionId++}`; return { connectionId, - address: this.connectResult?.address ?? `ssh:${config.host}`, + address: this.connectResult?.address ?? computeSSHConnectionKey(config), name: config.name, connectionToken: 'test-token', config: { host: config.host, username: config.username, authMethod: config.authMethod, name: config.name, sshConfigHost: config.sshConfigHost }, @@ -173,8 +177,8 @@ class MockSSHMainService { }; } - async reconnect(sshConfigHost: string, name: string, remoteAgentHostCommand?: string, agentForward?: boolean, userInitiated?: boolean, preferredAgentLocation?: RemoteAgentHostLocationPreference): Promise { - this.reconnectCalls.push({ sshConfigHost, name, remoteAgentHostCommand, agentForward, userInitiated, preferredAgentLocation }); + async reconnect(sshConfigHost: string, name: string, mode: SSHConnectionMode, remoteAgentHostCommand?: string, agentForward?: boolean, userInitiated?: boolean, preferredAgentLocation?: RemoteAgentHostLocationPreference): Promise { + this.reconnectCalls.push({ sshConfigHost, name, mode, remoteAgentHostCommand, agentForward, userInitiated, preferredAgentLocation }); return { connectionId: this.connectResult?.connectionId ?? `conn-${this._nextConnectionId++}`, address: this.connectResult?.address ?? `ssh:${sshConfigHost}`, @@ -241,13 +245,7 @@ function asChannel(target: object): IChannel { /** Drives registered connection factories like RemoteAgentHostService. */ class MockRemoteAgentHostService extends Disposable { readonly added: Array<{ address: string; status?: RemoteAgentHostConnectionStatus; transport?: IDisposable }> = []; - private readonly _entries = new Map void }; status: RemoteAgentHostConnectionStatus }>(); - // Holds transport disposables from prior service-owned connections that - // were replaced by a later connection for the same address. - // Production deliberately does NOT run them at replacement time (doing - // so would call _mainService.disconnect on the brand-new tunnel and - // kill it). They are released when the service itself is disposed. - private readonly _abandonedTransports: IDisposable[] = []; + private readonly _entries = new Map void }; status: RemoteAgentHostConnectionStatus; error?: Error }>(); private readonly _onDidChangeConnections = this._register(new Emitter()); readonly onDidChangeConnections = this._onDidChangeConnections.event; private readonly _connectionWaits = new Map>(); @@ -272,11 +270,23 @@ class MockRemoteAgentHostService extends Disposable { void this._createAndConnect(address, userInitiated); } + ensureConnection(address: string, userInitiated = true): void { + const entry = this._entries.get(address); + if (entry && !RemoteAgentHostConnectionStatus.isDisconnected(entry.status)) { + return; + } + void this._createAndConnect(address, userInitiated); + } + async waitForConnection(address: string): Promise { const existing = this.connections.find(connection => connection.address === address && RemoteAgentHostConnectionStatus.isConnected(connection.status)); if (existing) { return existing; } + const error = this._entries.get(address)?.error; + if (error) { + throw error; + } let wait = this._connectionWaits.get(address); if (!wait) { wait = new DeferredPromise(); @@ -296,9 +306,7 @@ class MockRemoteAgentHostService extends Disposable { const previous = this._entries.get(address); if (previous) { previous.client.dispose?.(); - if (previous.transport) { - this._abandonedTransports.push(previous.transport); - } + previous.transport?.dispose(); this._entries.delete(address); } @@ -306,7 +314,11 @@ class MockRemoteAgentHostService extends Disposable { const created = await factory.createConnection(entry, { userInitiated }); const added = { address, status: RemoteAgentHostConnectionStatus.connecting, transport: created.transportDisposable }; this.added.push(added); - const managed = { client: created.connection as { dispose?: () => void }, transport: created.transportDisposable, status: RemoteAgentHostConnectionStatus.connecting }; + const managed: { client: { dispose?: () => void }; transport?: IDisposable; status: RemoteAgentHostConnectionStatus; error?: Error } = { + client: created.connection as { dispose?: () => void }, + transport: created.transportDisposable, + status: RemoteAgentHostConnectionStatus.connecting, + }; this._entries.set(address, managed); this._onDidChangeConnections.fire(); try { @@ -318,6 +330,7 @@ class MockRemoteAgentHostService extends Disposable { const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]); if (incompatible) { managed.status = incompatible; + managed.error = err instanceof Error ? err : new Error(String(err)); added.status = managed.status; } else { this._entries.delete(address); @@ -383,11 +396,6 @@ class MockRemoteAgentHostService extends Disposable { wait.error(new Error('Mock remote agent host service disposed.')); } this._connectionWaits.clear(); - // Release abandoned transports from prior registrations as well. - for (const t of this._abandonedTransports) { - t.dispose(); - } - this._abandonedTransports.length = 0; super.dispose(); } } @@ -480,6 +488,7 @@ suite('SSHRemoteAgentHostService (renderer)', () => { let configurationService: TestConfigurationService; let notificationService: CapturingNotificationService; let createdClients: MockProtocolClient[]; + let reestablishRelays: Array<() => Promise<{ connectionId: string }>>; let waitForClient: (index: number) => Promise; let service: SSHRemoteAgentHostService; let instantiationService: TestInstantiationService; @@ -493,6 +502,7 @@ suite('SSHRemoteAgentHostService (renderer)', () => { disposables.add({ dispose: () => mainService.dispose() }); remoteAgentHostService = disposables.add(new MockRemoteAgentHostService()); createdClients = []; + reestablishRelays = []; const sharedProcessService: Partial = { getChannel: () => asChannel(mainService), @@ -528,7 +538,8 @@ suite('SSHRemoteAgentHostService (renderer)', () => { }; instantiationService.stub(ISSHRelayClientFactory, { - createClient: (_mainService: ISSHRemoteAgentHostMainService, _connectionId: string, _address: string) => { + createClient: (_mainService: ISSHRemoteAgentHostMainService, _connectionId: string, _address: string, reestablish: () => Promise<{ connectionId: string }>) => { + reestablishRelays.push(reestablish); const c = new MockProtocolClient(); disposables.add(c); const index = createdClients.length; @@ -572,6 +583,7 @@ suite('SSHRemoteAgentHostService (renderer)', () => { stored: readSSHRemoteAgentHostEntries(sshStorageService), connectionCount: service.connections.length, handleAddress: handle.localAddress, + connectModes: mainService.connectModes, }, { managedConnection: 1, address: 'ssh:remote.example', @@ -595,6 +607,75 @@ suite('SSHRemoteAgentHostService (renderer)', () => { }], connectionCount: 1, handleAddress: 'ssh:remote.example', + connectModes: [SSHConnectionMode.ReplaceConnection], + }); + }); + + test('repeated alias connect reuses the live protocol client and relay', async () => { + const firstConnect = service.connect(sampleConfig); + await awaitClientThenResolve(0); + const first = await firstConnect; + + const second = await service.connect(sampleConfig); + + assert.deepStrictEqual({ + sameHandle: first === second, + protocolClients: createdClients.length, + mainConnects: mainService.connectCalls.length, + }, { + sameHandle: true, + protocolClients: 1, + mainConnects: 1, + }); + }); + + test('repeated credential-based connect reuses the live protocol client and relay', async () => { + const config: ISSHAgentHostConfig = { + ...sampleConfig, + sshConfigHost: undefined, + authMethod: SSHAuthMethod.Password, + password: 'secret', + }; + const firstConnect = service.connect(config); + await awaitClientThenResolve(0); + const first = await firstConnect; + + const second = await service.connect(config); + + assert.deepStrictEqual({ + sameHandle: first === second, + address: second.localAddress, + protocolClients: createdClients.length, + mainConnects: mainService.connectCalls.length, + }, { + sameHandle: true, + address: 'user@remote.example:22', + protocolClients: 1, + mainConnects: 1, + }); + }); + + test('soft relay replacement rekeys the renderer handle to the fresh relay id', async () => { + mainService.connectResult = { connectionId: 'conn-1' }; + const connect = service.connect(sampleConfig); + await awaitClientThenResolve(0); + await connect; + + (mainService as unknown as { _onDidCloseConnection: Emitter })._onDidCloseConnection.fire('conn-1'); + mainService.connectResult = { connectionId: 'conn-2' }; + const reestablished = await reestablishRelays[0](); + remoteAgentHostService.removeEntry('ssh:remote.example'); + + assert.deepStrictEqual({ + reestablished, + reconnectModes: mainService.reconnectCalls.map(call => call.mode), + connections: service.connections.length, + disconnectCalls: mainService.disconnectCalls, + }, { + reestablished: { connectionId: 'conn-2' }, + reconnectModes: [SSHConnectionMode.ReplaceRelay], + connections: 0, + disconnectCalls: ['conn-2'], }); }); @@ -734,11 +815,14 @@ suite('SSHRemoteAgentHostService (renderer)', () => { )); await assert.rejects(connectPromise, /Unsupported protocol version/); + await assert.rejects(service.connect(sampleConfig), /Unsupported protocol version/); assert.deepStrictEqual({ added: remoteAgentHostService.added.map(({ address, status }) => ({ address, status })), connections: service.connections.map(connection => connection.localAddress), disconnectCalls: mainService.disconnectCalls, + protocolClients: createdClients.length, + mainConnects: mainService.connectCalls.length, }, { added: [{ address: 'ssh:remote.example', @@ -746,6 +830,8 @@ suite('SSHRemoteAgentHostService (renderer)', () => { }], connections: ['ssh:remote.example'], disconnectCalls: [], + protocolClients: 1, + mainConnects: 1, }); }); @@ -780,9 +866,8 @@ suite('SSHRemoteAgentHostService (renderer)', () => { assert.deepStrictEqual({ clientCount: createdClients.length, added: remoteAgentHostService.added.map(({ address, status }) => ({ address, statusKind: status?.kind })), - // The replaceRelay path keeps the SSH tunnel alive — we must not - // have asked the main service to disconnect it. disconnectCalls: mainService.disconnectCalls, + reconnectModes: mainService.reconnectCalls.map(call => call.mode), // Exactly one renderer-side handle for the address. connections: service.connections.map(connection => connection.localAddress), }, { @@ -791,7 +876,8 @@ suite('SSHRemoteAgentHostService (renderer)', () => { { address: 'ssh:remote.example', statusKind: 'incompatible' }, { address: 'ssh:remote.example', statusKind: 'connected' }, ], - disconnectCalls: [], + disconnectCalls: ['conn-stable'], + reconnectModes: [SSHConnectionMode.ReplaceConnection], connections: ['ssh:remote.example'], }); }); @@ -844,7 +930,7 @@ suite('SSHRemoteAgentHostService (renderer)', () => { assert.strictEqual(remoteAgentHostService.added.length, 2, 'each connect produces a fresh service-owned connection'); }); - test('main-process onDidCloseConnection cleans up renderer handle without double-disconnecting', async () => { + test('main-process onDidCloseConnection retains the logical handle for soft reconnect', async () => { const connectPromise = service.connect(sampleConfig); await awaitClientThenResolve(0); await connectPromise; @@ -856,12 +942,10 @@ suite('SSHRemoteAgentHostService (renderer)', () => { // emitter that the renderer subscribed to. (mainService as unknown as { _onDidCloseConnection: Emitter })._onDidCloseConnection.fire('conn-1'); - assert.strictEqual(service.connections.length, 0, 'handle dropped on main close'); - // Removing the (already-gone) entry shouldn't trigger another disconnect call. + assert.strictEqual(service.connections.length, 1, 'handle retained while the protocol client reconnects'); remoteAgentHostService.removeEntry('ssh:remote.example'); - // One disconnect from the transport disposable is fine; we just want to make - // sure we're not at risk of issuing a second one against a stale id. - assert.ok(mainService.disconnectCalls.length <= 1, 'no duplicate disconnect against a stale connectionId'); + assert.deepStrictEqual(mainService.disconnectCalls, ['conn-1']); + assert.strictEqual(service.connections.length, 0); }); // --- SSH failover notification: editor-owned → standalone on an unattended reconnect --- @@ -890,10 +974,10 @@ suite('SSHRemoteAgentHostService (renderer)', () => { await c1; assert.deepStrictEqual(notificationService.infoMessages, [], 'no notification on initial connect'); - // The SSH tunnel drops and the renderer-side handle is cleaned up. - // This disconnect cleanup must NOT erase the last-known server type. + // The SSH tunnel drops, but the logical handle stays available for the + // protocol client's soft-reconnect attempt. fireMainProcessClose('conn-1'); - assert.strictEqual(service.connections.length, 0); + assert.strictEqual(service.connections.length, 1); // A silent/background reconnect (userInitiated: false) lands on a // standalone endpoint instead of the editor-owned one. @@ -991,19 +1075,16 @@ suite('SSHRemoteAgentHostService (renderer)', () => { assert.deepStrictEqual(notificationService.infoMessages, []); }); - test('a duplicate setup reconnects through the service without notifying', async () => { + test('a duplicate setup reuses the service-owned protocol client without notifying', async () => { mainService.connectResult = { connectionId: 'conn-1', serverType: 'editor' }; const c1 = service.connect(sampleConfig); await awaitClientThenResolve(0); await c1; - // Reconnecting through the shared service replaces its protocol client - // and performs a new handshake without treating it as a failover. - const c2 = service.connect(sampleConfig); - await awaitClientThenResolve(1); - await c2; + const c2 = await service.connect(sampleConfig); - assert.strictEqual(createdClients.length, 2, 'the service owns a new protocol client for the reconnect'); + assert.strictEqual(createdClients.length, 1, 'the service retains the initialized protocol client'); + assert.strictEqual(c2.localAddress, 'ssh:remote.example'); assert.deepStrictEqual(notificationService.infoMessages, []); }); }); diff --git a/src/vs/platform/agentHost/test/electron-browser/wslRemoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/electron-browser/wslRemoteAgentHostService.test.ts index 9c9fcbf67f051..b093d2b9531e1 100644 --- a/src/vs/platform/agentHost/test/electron-browser/wslRemoteAgentHostService.test.ts +++ b/src/vs/platform/agentHost/test/electron-browser/wslRemoteAgentHostService.test.ts @@ -22,12 +22,17 @@ class MockWSLMainService { } class MockRemoteAgentHostService { + readonly ensureCalls: Array<{ readonly address: string; readonly userInitiated: boolean }> = []; readonly reconnectCalls: Array<{ readonly address: string; readonly userInitiated: boolean }> = []; registerConnectionFactory(_factory: IRemoteAgentHostConnectionFactory) { return toDisposable(() => undefined); } + ensureConnection(address: string, userInitiated = true): void { + this.ensureCalls.push({ address, userInitiated }); + } + reconnect(address: string, userInitiated = true): void { this.reconnectCalls.push({ address, userInitiated }); } @@ -83,6 +88,18 @@ suite('WSLRemoteAgentHostService (renderer)', () => { teardown(() => disposables.clear()); ensureNoDisposablesAreLeakedInTestSuite(); + test('connect ensures a connection without forcing replacement', async () => { + await assert.rejects(() => service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }), /not established/); + + assert.deepStrictEqual({ + ensureCalls: remoteAgentHostService.ensureCalls, + reconnectCalls: remoteAgentHostService.reconnectCalls, + }, { + ensureCalls: [{ address: 'wsl:Ubuntu', userInitiated: true }], + reconnectCalls: [], + }); + }); + test('forwards whether reconnect was user-initiated', async () => { await assert.rejects(() => service.reconnect('Ubuntu', 'Ubuntu'), /not established/); await assert.rejects(() => service.reconnect('Ubuntu', 'Ubuntu', false), /not established/); diff --git a/src/vs/platform/agentHost/test/node/sshRemoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/node/sshRemoteAgentHostService.test.ts index 9143912f2b1c9..6b96925b9856e 100644 --- a/src/vs/platform/agentHost/test/node/sshRemoteAgentHostService.test.ts +++ b/src/vs/platform/agentHost/test/node/sshRemoteAgentHostService.test.ts @@ -15,7 +15,7 @@ import { IProductService } from '../../../product/common/productService.js'; import { TelemetryConfiguration } from '../../../telemetry/common/telemetry.js'; import { NullTelemetryService } from '../../../telemetry/common/telemetryUtils.js'; import { AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, type AgentHostEndpointAddress, type IAgentHostEndpointMetadata } from '../../common/agentHostEndpointRegistry.js'; -import { SSHAuthMethod, type ISSHAgentHostConfig, type ISSHConnectProgress, type ISSHEndpointSelection, type ISSHEndpointSelectionRequest, type ISSHKeyboardInteractivePrompt, type ISSHKeyboardInteractiveRequest } from '../../common/sshRemoteAgentHost.js'; +import { SSHAuthMethod, SSHConnectionMode, type ISSHAgentHostConfig, type ISSHConnectProgress, type ISSHEndpointSelection, type ISSHEndpointSelectionRequest, type ISSHKeyboardInteractivePrompt, type ISSHKeyboardInteractiveRequest } from '../../common/sshRemoteAgentHost.js'; import { SSHRemoteAgentHostMainService, makeAuthHandler, type SSHAuthAttempt } from '../../node/sshRemoteAgentHostService.js'; import type { AnyAuthMethod, AuthenticationType, ConnectConfig } from 'ssh2'; @@ -481,7 +481,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { const config = makeConfig({ sshConfigHost: 'myalias' }); const result1 = await service.connect(config); - assert.strictEqual(result1.connectionId, 'ssh:myalias'); + assert.strictEqual(result1.address, 'ssh:myalias'); assert.strictEqual(result1.sshConfigHost, 'myalias'); assert.strictEqual(result1.lifecycle, 'external'); assert.strictEqual(service.startCalled, 0); @@ -505,13 +505,36 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { // Reconnect — creates fresh relay on existing SSH tunnel; does not // rerun endpoint discovery/selection (see connect()'s replaceRelay path). - const result2 = await service.reconnect('myalias', 'test-agent'); - assert.strictEqual(result2.connectionId, result1.connectionId); + const result2 = await service.reconnect('myalias', 'test-agent', SSHConnectionMode.ReplaceRelay); + assert.notStrictEqual(result2.connectionId, result1.connectionId); assert.strictEqual(result2.connectionToken, result1.connectionToken); assert.strictEqual(result2.lifecycle, result1.lifecycle); assert.strictEqual(service.relayCalled, 2); // fresh relay }); + test('replace connection reruns SSH bootstrap and endpoint selection', async () => { + service.execResponses = [ + ...discoveryResponses([makeEndpoint({ type: 'standalone', pid: 1234, instanceId: 'inst-1' })]), + ...discoveryResponses([makeEndpoint({ type: 'standalone', pid: 5678, instanceId: 'inst-2' })]), + ]; + + const config = makeConfig({ sshConfigHost: 'myalias' }); + const first = await service.connect(config); + const replacement = await service.connect(config, SSHConnectionMode.ReplaceConnection); + + assert.deepStrictEqual({ + connectionIdsDiffer: first.connectionId !== replacement.connectionId, + sshClients: service.mockClients.length, + relays: service.relayCalled, + selectedInstance: replacement.instanceId, + }, { + connectionIdsDiffer: true, + sshClients: 2, + relays: 2, + selectedInstance: 'inst-2', + }); + }); + test('reconnect does not fire onDidRelayClose for superseded relay', async () => { service.execResponses = discoveryResponses([makeEndpoint({ type: 'standalone', pid: 1234, instanceId: 'inst-1' })]); @@ -522,7 +545,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { disposables.add(service.onDidRelayClose(id => closeEvents.push(id))); // Reconnect replaces the relay — old relay close should be suppressed - await service.reconnect('myalias', 'test-agent'); + await service.reconnect('myalias', 'test-agent', SSHConnectionMode.ReplaceRelay); // Simulate the old relay's close event firing asynchronously service.simulateOldRelayClose(); @@ -543,7 +566,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { // simulating a WebSocket that fires 'close' synchronously on ws.close(). service.makePreviousRelaySyncClose(); - await service.reconnect('myalias', 'test-agent'); + await service.reconnect('myalias', 'test-agent', SSHConnectionMode.ReplaceRelay); assert.deepStrictEqual(closeEvents, []); }); @@ -551,7 +574,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { service.execResponses = discoveryResponses([makeEndpoint({ type: 'standalone', pid: 1234, instanceId: 'inst-1' })]); const result = await service.connect(makeConfig({ sshConfigHost: 'myhost' })); - assert.strictEqual(result.connectionId, 'ssh:myhost'); + assert.strictEqual(result.address, 'ssh:myhost'); assert.strictEqual(result.sshConfigHost, 'myhost'); }); @@ -563,7 +586,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { const result = await service.connect(makeConfig({ remoteAgentHostCommand: '/custom/agent --port 0', })); - assert.strictEqual(result.connectionId, 'testuser@10.0.0.1:22'); + assert.strictEqual(result.address, 'testuser@10.0.0.1:22'); assert.strictEqual(result.serverType, undefined); assert.strictEqual(result.instanceId, 'override'); assert.strictEqual(result.lifecycle, 'managed'); @@ -774,7 +797,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { const events: ISSHEndpointSelectionRequest[] = []; disposables.add(service.onDidRequestEndpointSelection(r => events.push(r))); - const result = await service.reconnect('myhost', 'test-host', undefined, undefined, /* userInitiated */ false); + const result = await service.reconnect('myhost', 'test-host', SSHConnectionMode.ReplaceConnection, undefined, undefined, /* userInitiated */ false); assert.deepStrictEqual(events, [], 'cold-start silent reconnect() must never fire an endpoint-selection request'); assert.strictEqual(result.serverType, 'standalone'); @@ -792,7 +815,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { const result = await service.withEndpointSelectionResponse( { kind: 'candidate', type: 'editor', pid: 300, instanceId: 'editor-1' }, - () => service.reconnect('myhost', 'test-host', undefined, undefined, /* userInitiated */ true), + () => service.reconnect('myhost', 'test-host', SSHConnectionMode.ReplaceConnection, undefined, undefined, /* userInitiated */ true), ); assert.ok(seenCandidates, 'user-initiated reconnect() must still show the picker when an editor entry exists'); @@ -921,7 +944,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { // userInitiated: true would normally still prompt when an editor is // live (see the contrasting test above) — a stored preference must // pre-empt that entirely. - const result = await service.reconnect('myhost', 'test-host', undefined, undefined, /* userInitiated */ true, /* preferredAgentLocation */ 'editor'); + const result = await service.reconnect('myhost', 'test-host', SSHConnectionMode.ReplaceConnection, undefined, undefined, /* userInitiated */ true, /* preferredAgentLocation */ 'editor'); assert.deepStrictEqual(events, []); assert.strictEqual(result.serverType, 'editor'); @@ -1010,7 +1033,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { remoteAgentHostCommand: '/agent', })); assert.strictEqual(service.startCalled, 1); - assert.strictEqual(result2.connectionId, result.connectionId); + assert.notStrictEqual(result2.connectionId, result.connectionId); }); test('fires onDidChangeConnections on connect and disconnect', async () => { @@ -1165,7 +1188,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { const r2 = await service.reconnect('myhost', 'test-host'); // Should have created a fresh SSH client (not reused the old one) assert.strictEqual(service.mockClients.length, 2); - assert.strictEqual(r2.connectionId, r1.connectionId); + assert.notStrictEqual(r2.connectionId, r1.connectionId); }); // --- Progress events --- @@ -1463,7 +1486,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { port: 2222, remoteAgentHostCommand: '/agent', })); - assert.strictEqual(result.connectionId, 'testuser@192.168.1.1:2222'); + assert.strictEqual(result.address, 'testuser@192.168.1.1:2222'); }); test('defaults to port 22 in connection key', async () => { @@ -1471,7 +1494,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { host: '192.168.1.1', remoteAgentHostCommand: '/agent', })); - assert.strictEqual(result.connectionId, 'testuser@192.168.1.1:22'); + assert.strictEqual(result.address, 'testuser@192.168.1.1:22'); }); // --- Reconnect preserves connection token from initial connect --- @@ -1481,10 +1504,10 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { const original = await service.connect(makeConfig({ sshConfigHost: 'myhost' })); - const reconnected = await service.reconnect('myhost', 'new-name'); + const reconnected = await service.reconnect('myhost', 'new-name', SSHConnectionMode.ReplaceRelay); assert.strictEqual(reconnected.connectionToken, original.connectionToken); assert.strictEqual(reconnected.address, original.address); - assert.strictEqual(reconnected.connectionId, original.connectionId); + assert.notStrictEqual(reconnected.connectionId, original.connectionId); }); // --- Relay messages from superseded relay are still routed (not gated) --- @@ -1498,7 +1521,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { disposables.add(service.onDidRelayMessage(msg => messages.push(msg))); // Reconnect replaces the relay - await service.reconnect('myhost', 'test-host'); + const reconnected = await service.reconnect('myhost', 'test-host', SSHConnectionMode.ReplaceRelay); // Simulate a message arriving from the OLD relay (index 0) service.simulateRelayMessage('stale-message', 0); @@ -1508,7 +1531,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { // Both messages arrive — message suppression is deliberately NOT done assert.deepStrictEqual(messages, [ { connectionId: result.connectionId, data: 'stale-message' }, - { connectionId: result.connectionId, data: 'fresh-message' }, + { connectionId: reconnected.connectionId, data: 'fresh-message' }, ]); }); @@ -1517,7 +1540,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { test('reconnect cleans up SSH client when relay recreation fails', async () => { service.execResponses = discoveryResponses([makeEndpoint({ type: 'standalone', pid: 1234, instanceId: 'inst-1' })]); - await service.connect(makeConfig({ sshConfigHost: 'myhost' })); + const original = await service.connect(makeConfig({ sshConfigHost: 'myhost' })); const originalClient = service.mockClients[0]; assert.strictEqual(originalClient.ended, false); @@ -1533,14 +1556,15 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { disposables.add(service.onDidCloseConnection(id => closeEvents.push(id))); await assert.rejects( - () => service.reconnect('myhost', 'test-host'), + () => service.reconnect('myhost', 'test-host', SSHConnectionMode.ReplaceRelay), /relay failed/, ); // SSH client should have been cleaned up despite the failure assert.strictEqual(originalClient.ended, true); // Close event should have fired to notify the renderer - assert.deepStrictEqual(closeEvents, ['ssh:myhost']); + assert.strictEqual(closeEvents.length, 1); + assert.notStrictEqual(closeEvents[0], original.connectionId); }); test('reconnect rejects with timeout when relay creation hangs (silently dead SSH client)', async () => { @@ -1552,7 +1576,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { // reload, since the shared-process state survives. service.execResponses = discoveryResponses([makeEndpoint({ type: 'standalone', pid: 1234, instanceId: 'inst-1' })]); - await service.connect(makeConfig({ sshConfigHost: 'myhost' })); + const original = await service.connect(makeConfig({ sshConfigHost: 'myhost' })); const originalClient = service.mockClients[0]; assert.strictEqual(originalClient.ended, false); @@ -1565,7 +1589,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { disposables.add(service.onDidCloseConnection(id => closeEvents.push(id))); await assert.rejects( - () => service.reconnect('myhost', 'test-host'), + () => service.reconnect('myhost', 'test-host', SSHConnectionMode.ReplaceRelay), /timed out|timeout/i, 'reconnect should reject (with a timeout error) instead of hanging when relay creation never settles' ); @@ -1576,7 +1600,8 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { assert.strictEqual(originalClient.ended, true, 'dead SSH client should be ended'); // Close event should have fired so the renderer's contribution sees // the reconnect attempt resolved (even as a failure) and can retry. - assert.deepStrictEqual(closeEvents, ['ssh:myhost']); + assert.strictEqual(closeEvents.length, 1); + assert.notStrictEqual(closeEvents[0], original.connectionId); }); // --- Reconnect cleans up old SSH client listeners --- @@ -1594,7 +1619,7 @@ suite('SSHRemoteAgentHostMainService - connect flow', () => { assert.ok(errorListenersBefore > 0, 'should have error listeners after connect'); // Reconnect replaces the SSHConnection — old listeners should be removed - await service.reconnect('myhost', 'test-host'); + await service.reconnect('myhost', 'test-host', SSHConnectionMode.ReplaceRelay); // Listener count should not grow — old ones removed, new ones added assert.strictEqual(client.closeListenerCount, closeListenersBefore); diff --git a/src/vs/platform/agentHost/test/node/wslRemoteAgentHostService.test.ts b/src/vs/platform/agentHost/test/node/wslRemoteAgentHostService.test.ts index e0bea7779f913..559a8334e77c6 100644 --- a/src/vs/platform/agentHost/test/node/wslRemoteAgentHostService.test.ts +++ b/src/vs/platform/agentHost/test/node/wslRemoteAgentHostService.test.ts @@ -12,7 +12,7 @@ import { runWithFakedTimers } from '../../../../base/test/common/timeTravelSched import { NullLogService } from '../../../log/common/log.js'; import type { IProductService } from '../../../product/common/productService.js'; import { NullTelemetryService } from '../../../telemetry/common/telemetryUtils.js'; -import type { IWSLConnectProgress, IWSLConnectResult } from '../../common/wslRemoteAgentHost.js'; +import { WSLConnectionMode, type IWSLConnectProgress, type IWSLConnectResult } from '../../common/wslRemoteAgentHost.js'; import { WSLRemoteAgentHostMainService } from '../../node/wslRemoteAgentHostService.js'; import type WebSocket from 'ws'; @@ -129,6 +129,60 @@ suite('WSL Remote Agent Host Service', () => { ); }); + test('replace creates a fresh relay for a live distro', async () => { + const service = disposables.add(createService()); + const first = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }); + service.resolvePlatform(); + await Promise.resolve(); + service.children[0].emitStdout('ws://127.0.0.1:3000?tkn=first\n'); + const firstResult = await first; + + const replacement = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }, WSLConnectionMode.Replace); + while (service.children.length < 2) { + await Promise.resolve(); + } + service.children[1].emitStdout('ws://127.0.0.1:3001?tkn=second\n'); + const replacementResult = await replacement; + + assert.deepStrictEqual({ + childCount: service.children.length, + firstChildKillCalls: service.children[0].killCalls, + connectionIdsDiffer: firstResult.connectionId !== replacementResult.connectionId, + replacementToken: replacementResult.connectionToken, + }, { + childCount: 2, + firstChildKillCalls: 1, + connectionIdsDiffer: true, + replacementToken: 'second', + }); + }); + + test('replace waits for an in-flight bootstrap and then creates a fresh relay', async () => { + const service = disposables.add(createService()); + const first = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }); + const replacement = service.connect({ distro: 'Ubuntu', name: 'Ubuntu' }, WSLConnectionMode.Replace); + + service.resolvePlatform(); + await Promise.resolve(); + service.children[0].emitStdout('ws://127.0.0.1:3000?tkn=first\n'); + const firstResult = await first; + while (service.children.length < 2) { + await Promise.resolve(); + } + service.children[1].emitStdout('ws://127.0.0.1:3001?tkn=second\n'); + const replacementResult = await replacement; + + assert.deepStrictEqual({ + childCount: service.children.length, + connectionIdsDiffer: firstResult.connectionId !== replacementResult.connectionId, + replacementToken: replacementResult.connectionToken, + }, { + childCount: 2, + connectionIdsDiffer: true, + replacementToken: 'second', + }); + }); + test('accepts initial bootstrap output after the output-idle budget', async () => { return runWithFakedTimers({ useFakeTimers: true, maxTaskCount: 10_000 }, async () => { const service = disposables.add(createService()); 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 453ccc1b46115..ff225d2fd0c70 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 @@ -50,8 +50,14 @@ The remote Agent Host service owns protocol connection construction, handshake c `RemoteAgentHostContribution` owns the workbench integration for a live connection: remote filesystem browsing, agent and model discovery, terminals, authentication, and connection-scoped listener disposal. -Transport-specific callers own discovery, on-demand staging, credentials, and connection leases. They stage -their context by address, request an explicit reconnect, and wait for the service to report the connection. +Transport-specific callers own discovery, on-demand staging, credentials, and connection leases. Ordinary +open and ensure operations retain the protocol client for any live transport, including one that is still +connecting, reconnecting, or retained after an incompatible handshake. Explicit recovery and preferred-location +changes replace the connection, serializing behind an in-flight attempt when necessary. + +Every connection factory invocation constructs a new protocol client over a fresh, uninitialized transport. +If a lower process retains a relay for the same address, the factory must explicitly replace it before returning. +A retained relay may be reused only by the protocol client that already owns its initialized lifecycle. The provider exposes connection state through `IAgentHostSessionsProvider` and delegates protocol operations to the live connection. Disconnecting clears live state without manufacturing successful operation results. diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts index e9266c4eff42b..7b9cba13813b8 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/browserTunnelAgentHostService.ts @@ -102,6 +102,11 @@ class BrowserTunnelConnectionFactory extends Disposable implements IRemoteAgentH } } + clearStagedConnection(address: string): void { + this._stagedUserInitiated.delete(address); + this._stagedAuthProviders.delete(address); + } + createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); @@ -293,14 +298,30 @@ export class BrowserTunnelAgentHostService extends Disposable implements ITunnel } async connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + await this._connect(tunnel, authProvider, options, false); + } + + async reconnect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + await this._connect(tunnel, authProvider, options, true); + } + + private async _connect(tunnel: ITunnelInfo, authProvider: 'github' | 'microsoft' | undefined, options: { readonly userInitiated?: boolean } | undefined, reconnect: boolean): Promise { if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { throw new Error('Remote agent host connections are not enabled.'); } const entry = this._connectionFactory.stageTunnel(tunnel, authProvider, options?.userInitiated ?? true); const address = getEntryAddress(entry); - this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); - await this._remoteAgentHostService.waitForConnection(address); + if (reconnect) { + this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); + } else { + this._remoteAgentHostService.ensureConnection(address, options?.userInitiated ?? true); + } + try { + await this._remoteAgentHostService.waitForConnection(address); + } finally { + this._connectionFactory.clearStagedConnection(address); + } } private async _createConnection(entry: IRemoteAgentHostEntry, authProvider: 'github' | 'microsoft' | undefined, options: IRemoteAgentHostConnectOptions): Promise { diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/entryDrivenProviderContribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/entryDrivenProviderContribution.ts index 5c32a10c622f2..e2bb02e1d9ca0 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/entryDrivenProviderContribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/entryDrivenProviderContribution.ts @@ -17,6 +17,7 @@ import { watchForIncompatibleNotifications } from './remoteHostOptions.js'; /** Options supplied by a remote-host kind when creating its sessions provider. */ export interface IEntryDrivenProviderOptions { readonly connectOnDemand?: () => Promise; + readonly reconnectOnDemand?: () => Promise; readonly disconnectOnDemand?: () => Promise; readonly onDidReportConnectProgress?: Event; readonly autoConnect?: IAgentHostAutoConnect; @@ -98,6 +99,7 @@ export abstract class EntryDrivenProviderContribution extends Disposable { address, name, connectOnDemand: options.connectOnDemand, + reconnectOnDemand: options.reconnectOnDemand, disconnectOnDemand: options.disconnectOnDemand, onDidReportConnectProgress: options.onDidReportConnectProgress, autoConnect: options.autoConnect, diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts index 43e432e21f9fa..397cab08ab7ec 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts @@ -70,6 +70,8 @@ export interface IRemoteAgentHostSessionsProviderConfig { readonly preferenceKey?: string; /** Optional hook to establish a connection on demand (e.g. tunnel relay). */ readonly connectOnDemand?: () => Promise; + /** Optional hook to explicitly replace the current connection on demand. */ + readonly reconnectOnDemand?: () => Promise; /** Optional hook to tear down the active connection on demand (e.g. tunnel relay). */ readonly disconnectOnDemand?: () => Promise; /** Optional progress messages during on-demand connect. */ @@ -181,6 +183,7 @@ export class RemoteAgentHostSessionsProvider extends BaseAgentHostSessionsProvid private readonly _onDidChangeResourceLabelHomes = Event.any(this._onDidChangeSessionsImmediately, this._onDidChangeDraftSessions.event); private readonly _connectionAuthority: string; private readonly _connectOnDemand: (() => Promise) | undefined; + private readonly _reconnectOnDemand: (() => Promise) | undefined; private readonly _disconnectOnDemand: (() => Promise) | undefined; private readonly _sessionSchemeAlias: IAgentHostSessionSchemeAlias | undefined; private readonly _omitHostFromWorkspaceLabel: boolean; @@ -226,6 +229,7 @@ export class RemoteAgentHostSessionsProvider extends BaseAgentHostSessionsProvid this._connectionAuthority = agentHostAuthority(config.address); this._connectOnDemand = config.connectOnDemand; + this._reconnectOnDemand = config.reconnectOnDemand; this._disconnectOnDemand = config.disconnectOnDemand; this._sessionSchemeAlias = config.sessionSchemeAlias; this._omitHostFromWorkspaceLabel = config.omitHostFromWorkspaceLabel === true; @@ -481,10 +485,14 @@ export class RemoteAgentHostSessionsProvider extends BaseAgentHostSessionsProvid /** * Establish (or re-establish) the connection for this host on demand. - * Tunnel-backed providers use their relay hook; other providers fall - * back to the generic remote agent host reconnect path. + * Transport-backed providers use their explicit reconnect hook; other + * on-demand providers fall back to their ensure hook. */ async connect(): Promise { + if (this._reconnectOnDemand) { + await this._reconnectOnDemand(); + return; + } if (this._connectOnDemand) { await this._connectOnDemand(); return; diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.ts index d8e7432dee824..e1149d95aaf16 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/sshAgentHost.contribution.ts @@ -8,7 +8,7 @@ import { isCancellationError } from '../../../../../base/common/errors.js'; import { StopWatch } from '../../../../../base/common/stopwatch.js'; import { type IRemoteAgentHostEntry, IRemoteAgentHostService, type IRemoteAgentHostSSHConnection, getEntryAddress, getEntryTypeConfig, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId } from '../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { computeReconnectDelay } from '../../../../../platform/agentHost/common/reconnectPolicy.js'; -import { computeSSHConnectionKey, isSSHHostKeyDeniedError, ISSHRemoteAgentHostService, SSHAuthMethod } from '../../../../../platform/agentHost/common/sshRemoteAgentHost.js'; +import { computeSSHConnectionKey, isSSHHostKeyDeniedError, ISSHRemoteAgentHostService } from '../../../../../platform/agentHost/common/sshRemoteAgentHost.js'; import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js'; import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; import { ILogService } from '../../../../../platform/log/common/log.js'; @@ -98,7 +98,8 @@ export class SSHAgentHostContribution extends ManagedReconnectAgentHostContribut const connection = entry.connection; const address = getEntryAddress(entry); return { - connectOnDemand: () => this._connectSSHOnDemand(connection, entry.name, address), + connectOnDemand: () => this._ensureSSHOnDemand(address), + reconnectOnDemand: () => this._connectSSHOnDemand(connection, entry.name, address), disconnectOnDemand: () => this._disconnectSSHOnDemand(connection), preferenceKey: computeSSHConnectionKey({ sshConfigHost: connection.sshConfigHost, @@ -109,21 +110,20 @@ export class SSHAgentHostContribution extends ManagedReconnectAgentHostContribut }; } + private async _ensureSSHOnDemand(address: string): Promise { + this._remoteAgentHostService.ensureConnection(address, true); + await this._remoteAgentHostService.waitForConnection(address); + } + private async _connectSSHOnDemand(connection: IRemoteAgentHostSSHConnection, name: string, address: string): Promise { const sshConfigHost = connection.sshConfigHost; if (!sshConfigHost) { const stopwatch = StopWatch.create(false); try { - await this._sshService.connect({ - host: connection.hostName, - port: connection.port, - username: connection.user ?? connection.hostName, - authMethod: SSHAuthMethod.Agent, - name, - userInitiated: true, - }); + this._remoteAgentHostService.reconnect(address, true); + await this._remoteAgentHostService.waitForConnection(address); logSSHConnectAttempt(this._telemetryService, { - operation: 'connect', + operation: 'reconnect', userInitiated: true, attempt: 1, durationMs: stopwatch.elapsed(), @@ -132,7 +132,7 @@ export class SSHAgentHostContribution extends ManagedReconnectAgentHostContribut }); } catch (err) { logSSHConnectAttempt(this._telemetryService, { - operation: 'connect', + operation: 'reconnect', userInitiated: true, attempt: 1, durationMs: stopwatch.elapsed(), diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts index 01c7a10e0287e..3c75f3a0a93a8 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/tunnelAgentHost.contribution.ts @@ -220,7 +220,8 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc RemoteAgentHostSessionsProvider, { address, name, - connectOnDemand: () => this._connectTunnel(address, { userInitiated: true }), + connectOnDemand: () => this._connectTunnel(address, { userInitiated: true, reconnect: false }), + reconnectOnDemand: () => this._connectTunnel(address, { userInitiated: true, reconnect: true }), disconnectOnDemand: () => this._disconnectTunnel(address), }, ); @@ -284,11 +285,25 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc * Establish a relay connection to a cached tunnel. Called on demand * when the user invokes the browse action on an online-but-not-connected tunnel. */ - private _connectTunnel(address: string, options: { readonly userInitiated: boolean }): Promise { + private _connectTunnel(address: string, options: { readonly userInitiated: boolean; readonly reconnect: boolean }): Promise { const existing = this._pendingConnects.get(address); if (existing) { - return existing; + if (!options.reconnect) { + return existing; + } + const queued = existing.then( + () => this._runTunnelConnection(address, options), + () => this._runTunnelConnection(address, options), + ); + this._trackPendingConnection(address, queued); + return queued; } + const pending = this._runTunnelConnection(address, options); + this._trackPendingConnection(address, pending); + return pending; + } + + private async _runTunnelConnection(address: string, options: { readonly userInitiated: boolean; readonly reconnect: boolean }): Promise { this._diagnosticsService.recordHostAction(address, 'connect', options.userInitiated); const tunnelId = address.slice(TUNNEL_ADDRESS_PREFIX.length); @@ -297,49 +312,61 @@ export class TunnelAgentHostContribution extends Disposable implements IWorkbenc } const cached = this._tunnelService.getCachedTunnels().find(t => t.tunnelId === tunnelId); const attemptStart = Date.now(); - const promise = (async () => { - let handle: { close(): void } | undefined; - const timer = options.userInitiated && cached ? setTimeout(() => { - handle = this._notificationService.notify({ - severity: Severity.Info, - message: nls.localize('tunnelConnecting', "Connecting to tunnel '{0}'...", cached.name), - progress: { infinite: true }, - }); - }, 1000) : undefined; - - try { - if (!cached || this._isHostedTunnel(cached)) { - return; - } - const tunnelInfo: ITunnelInfo = { - tunnelId: cached.tunnelId, - clusterId: cached.clusterId, - name: cached.name, - tags: [], - // Legacy cache fallback, not a real capability claim. - protocolVersion: cached.protocolVersion ?? TUNNEL_MIN_PROTOCOL_VERSION, - hostConnectionCount: 0, - }; + let handle: { close(): void } | undefined; + const timer = options.userInitiated && cached ? setTimeout(() => { + handle = this._notificationService.notify({ + severity: Severity.Info, + message: nls.localize('tunnelConnecting', "Connecting to tunnel '{0}'...", cached.name), + progress: { infinite: true }, + }); + }, 1000) : undefined; + + try { + if (!cached || this._isHostedTunnel(cached)) { + return; + } + const tunnelInfo: ITunnelInfo = { + tunnelId: cached.tunnelId, + clusterId: cached.clusterId, + name: cached.name, + tags: [], + // Legacy cache fallback, not a real capability claim. + protocolVersion: cached.protocolVersion ?? TUNNEL_MIN_PROTOCOL_VERSION, + hostConnectionCount: 0, + }; + if (options.reconnect) { + await this._tunnelService.reconnect(tunnelInfo, cached.authProvider, { userInitiated: options.userInitiated }); + } else { await this._tunnelService.connect(tunnelInfo, cached.authProvider, { userInitiated: options.userInitiated }); - logTunnelConnectAttempt(this._telemetryService, { isReconnect: false, attempt: 1, durationMs: Date.now() - attemptStart, success: true }); - logTunnelConnectResolved(this._telemetryService, { isReconnect: false, totalAttempts: 1, totalDurationMs: Date.now() - attemptStart, success: true }); - } catch (err) { - this._logService.warn(`[TunnelAgentHost] Connect to ${cached?.name ?? address} failed:`, err); - logTunnelConnectAttempt(this._telemetryService, { isReconnect: false, attempt: 1, durationMs: Date.now() - attemptStart, success: false, errorCategory: 'other' }); - logTunnelConnectResolved(this._telemetryService, { isReconnect: false, totalAttempts: 1, totalDurationMs: Date.now() - attemptStart, success: false }); - throw err; - } finally { - if (timer !== undefined) { - clearTimeout(timer); - } - handle?.close(); - this._pendingConnects.delete(address); - this._updateConnectionStatuses(); } - })(); + logTunnelConnectAttempt(this._telemetryService, { isReconnect: options.reconnect, attempt: 1, durationMs: Date.now() - attemptStart, success: true }); + logTunnelConnectResolved(this._telemetryService, { isReconnect: options.reconnect, totalAttempts: 1, totalDurationMs: Date.now() - attemptStart, success: true }); + } catch (err) { + this._logService.warn(`[TunnelAgentHost] Connect to ${cached?.name ?? address} failed:`, err); + logTunnelConnectAttempt(this._telemetryService, { isReconnect: options.reconnect, attempt: 1, durationMs: Date.now() - attemptStart, success: false, errorCategory: 'other' }); + logTunnelConnectResolved(this._telemetryService, { isReconnect: options.reconnect, totalAttempts: 1, totalDurationMs: Date.now() - attemptStart, success: false }); + throw err; + } finally { + if (timer !== undefined) { + clearTimeout(timer); + } + handle?.close(); + } + } - this._pendingConnects.set(address, promise); - return promise; + private _trackPendingConnection(address: string, pending: Promise): void { + this._pendingConnects.set(address, pending); + void pending.then( + () => this._completePendingConnection(address, pending), + () => this._completePendingConnection(address, pending), + ); + } + + private _completePendingConnection(address: string, pending: Promise): void { + if (this._pendingConnects.get(address) === pending) { + this._pendingConnects.delete(address); + this._updateConnectionStatuses(); + } } /** diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts index 3880b9a809a1b..e561caf96f88b 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.contribution.ts @@ -78,6 +78,10 @@ class BrowserTunnelAgentHostServiceSelector extends Disposable implements ITunne return this._delegate.connect(tunnel, authProvider, options); } + reconnect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + return this._delegate.reconnect(tunnel, authProvider, options); + } + get canDeleteTunnels(): boolean { return this._delegate.canDeleteTunnels; } diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts index 26f89e59303e7..a80eb6e4950e2 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/webTunnelAgentHostService.ts @@ -89,6 +89,11 @@ class WebTunnelConnectionFactory extends Disposable implements IRemoteAgentHostC } } + clearStagedConnection(address: string): void { + this._stagedUserInitiated.delete(address); + this._stagedAuthProviders.delete(address); + } + createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); @@ -231,14 +236,30 @@ export class WebTunnelAgentHostService extends Disposable implements ITunnelAgen // Connection (via embedder) async connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + await this._connect(tunnel, authProvider, options, false); + } + + async reconnect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + await this._connect(tunnel, authProvider, options, true); + } + + private async _connect(tunnel: ITunnelInfo, authProvider: 'github' | 'microsoft' | undefined, options: { readonly userInitiated?: boolean } | undefined, reconnect: boolean): Promise { if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { throw new Error('Remote agent host connections are not enabled.'); } const entry = this._connectionFactory.stageTunnel(tunnel, authProvider, options?.userInitiated ?? true); const address = getEntryAddress(entry); - this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); - await this._remoteAgentHostService.waitForConnection(address); + if (reconnect) { + this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); + } else { + this._remoteAgentHostService.ensureConnection(address, options?.userInitiated ?? true); + } + try { + await this._remoteAgentHostService.waitForConnection(address); + } finally { + this._connectionFactory.clearStagedConnection(address); + } } private async _createConnection(entry: IRemoteAgentHostEntry, _options: IRemoteAgentHostConnectOptions): Promise { diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts index c6df198dbef18..f0310964732cf 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/browser/wslAgentHost.contribution.ts @@ -125,13 +125,19 @@ export class WSLAgentHostContribution extends ManagedReconnectAgentHostContribut } const { distro, address } = entry.connection; return { - connectOnDemand: () => this._connectWSLOnDemand(distro, entry.name, address), + connectOnDemand: () => this._ensureWSLOnDemand(address), + reconnectOnDemand: () => this._connectWSLOnDemand(distro, entry.name, address), disconnectOnDemand: () => this._disconnectWSLOnDemand(distro, address), onDidReportConnectProgress: this._wslService.onDidReportConnectProgress, autoConnect: this._autoConnect, }; } + private async _ensureWSLOnDemand(address: string): Promise { + this._remoteAgentHostService.ensureConnection(address, true); + await this._remoteAgentHostService.waitForConnection(address); + } + private async _connectWSLOnDemand(distro: string, name: string, address: string): Promise { while (true) { const inFlight = this._pendingReconnects.get(distro); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts index 189d803a15896..a0be75e874ea1 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/electron-browser/tunnelAgentHostServiceImpl.ts @@ -119,6 +119,11 @@ class TunnelConnectionFactory extends Disposable implements IRemoteAgentHostConn } } + clearStagedConnection(address: string): void { + this._stagedUserInitiated.delete(address); + this._stagedAuthProviders.delete(address); + } + createConnection(entry: IRemoteAgentHostEntry, options: IRemoteAgentHostConnectOptions): Promise { if (entry.connection.type !== RemoteAgentHostEntryType.Tunnel) { throw new Error(`Tunnel factory cannot create a ${entry.connection.type} connection.`); @@ -224,14 +229,30 @@ export class TunnelAgentHostService extends Disposable implements ITunnelAgentHo } async connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + await this._connect(tunnel, authProvider, options, false); + } + + async reconnect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + await this._connect(tunnel, authProvider, options, true); + } + + private async _connect(tunnel: ITunnelInfo, authProvider: 'github' | 'microsoft' | undefined, options: { readonly userInitiated?: boolean } | undefined, reconnect: boolean): Promise { if (!this._configurationService.getValue(RemoteAgentHostsEnabledSettingId)) { throw new Error('Remote agent host connections are not enabled.'); } const entry = this._connectionFactory.stageTunnel(tunnel, authProvider, options?.userInitiated ?? true); const address = getEntryAddress(entry); - this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); - await this._remoteAgentHostService.waitForConnection(address); + if (reconnect) { + this._remoteAgentHostService.reconnect(address, options?.userInitiated ?? true); + } else { + this._remoteAgentHostService.ensureConnection(address, options?.userInitiated ?? true); + } + try { + await this._remoteAgentHostService.waitForConnection(address); + } finally { + this._connectionFactory.clearStagedConnection(address); + } } private async _createConnection(entry: IRemoteAgentHostEntry, authProvider: 'github' | 'microsoft' | undefined, options: IRemoteAgentHostConnectOptions): Promise { diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/browserTunnelAgentHostService.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/browserTunnelAgentHostService.test.ts index fbf0f03a47799..fdbbbf84279d3 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/browserTunnelAgentHostService.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/browserTunnelAgentHostService.test.ts @@ -10,7 +10,7 @@ import { type ITunnelApplicationConfig } from '../../../../../../base/common/pro import { mock } from '../../../../../../base/test/common/mock.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; import { IRemoteAgentHostLocationPreferenceService } from '../../../../../../platform/agentHost/common/remoteAgentHostLocationPreference.js'; -import { IRemoteAgentHostConnectionFactory, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IRemoteAgentHostConnectionFactory, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { type ITunnelConnectResult, type ITunnelGatewaySelection, type ITunnelGatewaySelectionSession, type ITunnelInfo } from '../../../../../../platform/agentHost/common/tunnelAgentHost.js'; import { resolveGatewaySelection, type IGatewaySelectionRequest } from '../../../../../../platform/agentHost/common/tunnelGatewaySelection.js'; import type { ITunnelDuplexStream } from '../../../../../../platform/agentHost/common/tunnelMessageSocket.js'; @@ -105,18 +105,40 @@ class FakeConnector implements ITunnelAgentHostConnector { } } -function createRemoteAgentHostService(): IRemoteAgentHostService { - return new class extends mock() { - override registerConnectionFactory(_factory: IRemoteAgentHostConnectionFactory) { - return { dispose() { } }; - } - }(); +class TestRemoteAgentHostService extends mock() { + readonly ensureCalls: Array<{ address: string; userInitiated: boolean }> = []; + readonly reconnectCalls: Array<{ address: string; userInitiated: boolean }> = []; + + override registerConnectionFactory(_factory: IRemoteAgentHostConnectionFactory) { + return { dispose() { } }; + } + + override ensureConnection(address: string, userInitiated = true): void { + this.ensureCalls.push({ address, userInitiated }); + } + + override reconnect(address: string, userInitiated = true): void { + this.reconnectCalls.push({ address, userInitiated }); + } + + override async waitForConnection(address: string) { + return { + address, + name: address, + status: RemoteAgentHostConnectionStatus.connected, + }; + } +} + +function createRemoteAgentHostService(): TestRemoteAgentHostService { + return new TestRemoteAgentHostService(); } function createBrowserTunnelService( store: Pick, sessions: readonly AuthenticationSession[], listTunnels: () => Promise, + remoteAgentHostService: IRemoteAgentHostService = createRemoteAgentHostService(), ): BrowserTunnelAgentHostService { class FakeManagementClient implements IDevTunnelsWebManagementClient { constructor(_userAgent: string, _apiVersion: object, _userTokenCallback: () => Promise) { @@ -155,7 +177,7 @@ function createBrowserTunnelService( const configurationService = new TestConfigurationService({ [RemoteAgentHostsEnabledSettingId]: true }); return store.add(new BrowserTunnelAgentHostService( - createRemoteAgentHostService(), + remoteAgentHostService, new NullLogService(), store.add(new TestInstantiationService()), configurationService, @@ -207,6 +229,30 @@ suite('BrowserTunnelAgentHostService', () => { await assert.rejects(service.listTunnels(), /enumeration failed/); }); + test('browser tunnel connect ensures while reconnect explicitly replaces and refreshes metadata', async () => { + const remoteAgentHostService = createRemoteAgentHostService(); + const service = createBrowserTunnelService(store, [], async () => [], remoteAgentHostService); + + await service.connect(tunnel, 'github', { userInitiated: true }); + await service.reconnect({ ...tunnel, clusterId: 'updated-cluster', name: 'Updated tunnel' }, 'microsoft', { userInitiated: true }); + + assert.deepStrictEqual({ + ensureCalls: remoteAgentHostService.ensureCalls, + reconnectCalls: remoteAgentHostService.reconnectCalls, + cached: service.getCachedTunnels(), + }, { + ensureCalls: [{ address: 'tunnel:tunnel-id', userInitiated: true }], + reconnectCalls: [{ address: 'tunnel:tunnel-id', userInitiated: true }], + cached: [{ + tunnelId: 'tunnel-id', + clusterId: 'updated-cluster', + name: 'Updated tunnel', + protocolVersion: 6, + authProvider: 'microsoft', + }], + }); + }); + test('rejects embedder tunnel discovery failures', async () => { const discoveryProvider = new class extends mock() { override async listTunnels(): Promise { @@ -229,6 +275,40 @@ suite('BrowserTunnelAgentHostService', () => { await assert.rejects(service.listTunnels(), /authentication failed/); }); + test('web tunnel connect ensures while reconnect explicitly replaces and refreshes metadata', async () => { + const remoteAgentHostService = createRemoteAgentHostService(); + const service = store.add(new WebTunnelAgentHostService( + remoteAgentHostService, + new class extends mock() { + override readonly options = { tunnelDiscoveryProvider: new class extends mock() { }() }; + }(), + new NullLogService(), + store.add(new TestInstantiationService()), + new TestConfigurationService({ [RemoteAgentHostsEnabledSettingId]: true }), + new class extends mock() { }(), + store.add(new InMemoryStorageService()), + )); + + await service.connect(tunnel, 'github', { userInitiated: true }); + await service.reconnect({ ...tunnel, clusterId: 'updated-cluster', name: 'Updated tunnel' }, 'microsoft', { userInitiated: true }); + + assert.deepStrictEqual({ + ensureCalls: remoteAgentHostService.ensureCalls, + reconnectCalls: remoteAgentHostService.reconnectCalls, + cached: service.getCachedTunnels(), + }, { + ensureCalls: [{ address: 'tunnel:tunnel-id', userInitiated: true }], + reconnectCalls: [{ address: 'tunnel:tunnel-id', userInitiated: true }], + cached: [{ + tunnelId: 'tunnel-id', + clusterId: 'updated-cluster', + name: 'Updated tunnel', + protocolVersion: 6, + authProvider: 'microsoft', + }], + }); + }); + test('completes the version-six gateway selection returned by the browser picker', async () => { const connector = new FakeConnector({ selectionId: 'selection-id', diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts index 169fb1cea3b67..7e1fb51f4d092 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts @@ -242,7 +242,7 @@ function createSession(id: string, opts?: { provider?: string; summary?: string; }; } -function createProvider(disposables: DisposableStore, connection: MockAgentConnection, overrides?: { address?: string; preferenceKey?: string; connectionName?: string | undefined; sendRequest?: (resource: URI, message: string, options?: IChatSendRequestOptions) => Promise; openSession?: boolean; storageService?: IStorageService; localAgentHostService?: IAgentHostService; noConnection?: boolean; connectOnDemand?: () => Promise; isWebPlatform?: boolean; workspaceTrusted?: boolean; omitHostFromWorkspaceLabel?: boolean; workspaceTypeIcon?: ThemeIcon; sessionSchemeAlias?: IAgentHostSessionSchemeAlias; defaultChangesetKind?: IRemoteAgentHostSessionsProviderConfig['defaultChangesetKind']; sessionResolutionPolicies?: Array<{ authority: string; policy: IAgentHostSessionResolutionPolicy }>; devContainerWorktreeScope?: string; readOnlyWhenDisconnected?: boolean; ctor?: typeof RemoteAgentHostSessionsProvider; labelService?: ILabelService; defaultDirectory?: string }): RemoteAgentHostSessionsProvider { +function createProvider(disposables: DisposableStore, connection: MockAgentConnection, overrides?: { address?: string; preferenceKey?: string; connectionName?: string | undefined; sendRequest?: (resource: URI, message: string, options?: IChatSendRequestOptions) => Promise; openSession?: boolean; storageService?: IStorageService; localAgentHostService?: IAgentHostService; noConnection?: boolean; connectOnDemand?: () => Promise; reconnectOnDemand?: () => Promise; isWebPlatform?: boolean; workspaceTrusted?: boolean; omitHostFromWorkspaceLabel?: boolean; workspaceTypeIcon?: ThemeIcon; sessionSchemeAlias?: IAgentHostSessionSchemeAlias; defaultChangesetKind?: IRemoteAgentHostSessionsProviderConfig['defaultChangesetKind']; sessionResolutionPolicies?: Array<{ authority: string; policy: IAgentHostSessionResolutionPolicy }>; devContainerWorktreeScope?: string; readOnlyWhenDisconnected?: boolean; ctor?: typeof RemoteAgentHostSessionsProvider; labelService?: ILabelService; defaultDirectory?: string }): RemoteAgentHostSessionsProvider { const instantiationService = disposables.add(new TestInstantiationService()); instantiationService.stub(IFileDialogService, {}); @@ -303,6 +303,7 @@ function createProvider(disposables: DisposableStore, connection: MockAgentConne preferenceKey: overrides?.preferenceKey, name: overrides !== undefined && Object.prototype.hasOwnProperty.call(overrides, 'connectionName') ? overrides.connectionName ?? '' : 'Test Host', connectOnDemand: overrides?.connectOnDemand, + reconnectOnDemand: overrides?.reconnectOnDemand, omitHostFromWorkspaceLabel: overrides?.omitHostFromWorkspaceLabel, workspaceTypeIcon: overrides?.workspaceTypeIcon, sessionSchemeAlias: overrides?.sessionSchemeAlias, @@ -574,6 +575,20 @@ suite('RemoteAgentHostSessionsProvider', () => { }); }); + test('uses the explicit reconnect hook instead of the ordinary ensure hook', async () => { + let ensureCalls = 0; + let reconnectCalls = 0; + const provider = createProvider(disposables, connection, { + noConnection: true, + connectOnDemand: async () => { ensureCalls++; }, + reconnectOnDemand: async () => { reconnectCalls++; }, + }); + + await provider.connect(); + + assert.deepStrictEqual({ ensureCalls, reconnectCalls }, { ensureCalls: 0, reconnectCalls: 1 }); + }); + test('remoteLocationPreferenceKey defaults to the live address when no stable preference key is given (e.g. tunnels/WSL)', () => { const provider = createProvider(disposables, connection, { address: 'tunnel:abc123' }); assert.strictEqual(provider.remoteLocationPreferenceKey, 'tunnel:abc123'); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts index 9e4e9fe759877..cd03d8a574889 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/browser/tunnelAgentHost.contribution.test.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { DeferredPromise } from '../../../../../../base/common/async.js'; import { Emitter, Event } from '../../../../../../base/common/event.js'; import { Disposable, IDisposable, toDisposable } from '../../../../../../base/common/lifecycle.js'; import { observableValue } from '../../../../../../base/common/observable.js'; @@ -86,6 +87,8 @@ class StubTunnelService extends Disposable implements ITunnelAgentHostService { /** Records every `connect()` call for assertions on the `userInitiated` threading. */ readonly connectCalls: Array<{ tunnel: ITunnelInfo; authProvider: string | undefined; options: { readonly userInitiated?: boolean } | undefined }> = []; + readonly reconnectCalls: Array<{ tunnel: ITunnelInfo; authProvider: string | undefined; options: { readonly userInitiated?: boolean } | undefined }> = []; + connectBarrier: DeferredPromise | undefined; readonly disconnectCalls: string[] = []; setCached(tunnels: ICachedTunnel[]): void { @@ -134,6 +137,11 @@ class StubTunnelService extends Disposable implements ITunnelAgentHostService { async connect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { this.connectCalls.push({ tunnel, authProvider, options }); + await this.connectBarrier?.p; + } + + async reconnect(tunnel: ITunnelInfo, authProvider?: 'github' | 'microsoft', options?: { readonly userInitiated?: boolean }): Promise { + this.reconnectCalls.push({ tunnel, authProvider, options }); } async disconnect(address: string): Promise { this.disconnectCalls.push(address); } @@ -329,7 +337,7 @@ suite('TunnelAgentHostContribution', () => { assert.deepStrictEqual(providersService.getProviders(), []); }); - test('on-demand connect threads userInitiated to tunnelService.connect', async () => { + test('explicit reconnect is serialized behind an on-demand tunnel connect', async () => { const tunnelService = store.add(new StubTunnelService()); const remoteService = store.add(new StubRemoteAgentHostService()); const providersService = store.add(new StubSessionsProvidersService()); @@ -357,18 +365,28 @@ suite('TunnelAgentHostContribution', () => { // Access the private on-demand orchestration method via a typed seam. const testable = contribution as unknown as { - _connectTunnel(address: string, options: { readonly userInitiated: boolean }): Promise; + _connectTunnel(address: string, options: { readonly userInitiated: boolean; readonly reconnect: boolean }): Promise; }; tunnelService.dismissTunnel(tunnelId); - await testable._connectTunnel(address, { userInitiated: true }); + tunnelService.connectBarrier = new DeferredPromise(); + const connect = testable._connectTunnel(address, { userInitiated: true, reconnect: false }); + while (tunnelService.connectCalls.length === 0) { + await Promise.resolve(); + } + const reconnect = testable._connectTunnel(address, { userInitiated: true, reconnect: true }); + assert.strictEqual(tunnelService.reconnectCalls.length, 0); + tunnelService.connectBarrier.complete(); + await Promise.all([connect, reconnect]); assert.deepStrictEqual({ dismissed: tunnelService.isTunnelDismissed(tunnelId), connectCalls: tunnelService.connectCalls.map(call => call.options?.userInitiated), + reconnectCalls: tunnelService.reconnectCalls.map(call => call.options?.userInitiated), providers: providersService.getProviders().map(provider => provider.id), }, { dismissed: false, connectCalls: [true], + reconnectCalls: [true], providers: [`agenthost-${address}`], }); }); diff --git a/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts b/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts index dc648fd9f187b..415afcce11334 100644 --- a/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts +++ b/src/vs/sessions/contrib/providers/remoteAgentHost/test/electron-browser/tunnelAgentHostServiceImpl.test.ts @@ -8,7 +8,7 @@ import { Event } from '../../../../../../base/common/event.js'; import type { IChannel } from '../../../../../../base/parts/ipc/common/ipc.js'; import { mock } from '../../../../../../base/test/common/mock.js'; import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js'; -import { IRemoteAgentHostConnectionFactory, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; +import { IRemoteAgentHostConnectionFactory, IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId } from '../../../../../../platform/agentHost/common/remoteAgentHostService.js'; import { IRemoteAgentHostLocationPreferenceService } from '../../../../../../platform/agentHost/common/remoteAgentHostLocationPreference.js'; import { ITunnelGatewayInventory } from '../../../../../../platform/agentHost/common/tunnelAgentHost.js'; import { IConfigurationService } from '../../../../../../platform/configuration/common/configuration.js'; @@ -81,6 +81,73 @@ suite('TunnelAgentHostService discovery', () => { await assert.rejects(service.listTunnels({ silent: true }), /No authentication is available to enumerate tunnels/); }); + + test('desktop tunnel connect ensures while reconnect explicitly replaces and refreshes metadata', async () => { + const remoteAgentHostService = new class extends mock() { + readonly ensureCalls: Array<{ address: string; userInitiated: boolean }> = []; + readonly reconnectCalls: Array<{ address: string; userInitiated: boolean }> = []; + + override registerConnectionFactory(_factory: IRemoteAgentHostConnectionFactory) { + return { dispose() { } }; + } + + override ensureConnection(address: string, userInitiated = true): void { + this.ensureCalls.push({ address, userInitiated }); + } + + override reconnect(address: string, userInitiated = true): void { + this.reconnectCalls.push({ address, userInitiated }); + } + + override async waitForConnection(address: string) { + return { address, name: address, status: RemoteAgentHostConnectionStatus.connected }; + } + }(); + const instantiationService = store.add(new TestInstantiationService()); + instantiationService.stub(ISharedProcessService, new class extends mock() { + override getChannel(): IChannel { + return new class extends mock() { }(); + } + }()); + instantiationService.stub(IRemoteAgentHostService, remoteAgentHostService); + instantiationService.stub(ILogService, new NullLogService()); + instantiationService.stub(IConfigurationService, new TestConfigurationService({ [RemoteAgentHostsEnabledSettingId]: true })); + instantiationService.stub(IAuthenticationService, new class extends mock() { }()); + instantiationService.stub(IProductService, TestProductService); + instantiationService.stub(IStorageService, store.add(new InMemoryStorageService())); + instantiationService.stub(IEnvironmentService, new class extends mock() { }()); + instantiationService.stub(IRemoteAgentHostLocationPreferenceService, new class extends mock() { }()); + instantiationService.stub(IDialogService, new class extends mock() { }()); + instantiationService.stub(INotificationService, new class extends mock() { }()); + const service = store.add(instantiationService.createInstance(TunnelAgentHostService)); + const tunnel = { + tunnelId: 'tunnel-id', + clusterId: 'cluster-id', + name: 'Remote tunnel', + tags: ['vscode-server-launcher', 'protocolv6'], + protocolVersion: 6, + hostConnectionCount: 1, + }; + + await service.connect(tunnel, 'github', { userInitiated: true }); + await service.reconnect({ ...tunnel, clusterId: 'updated-cluster', name: 'Updated tunnel' }, 'microsoft', { userInitiated: true }); + + assert.deepStrictEqual({ + ensureCalls: remoteAgentHostService.ensureCalls, + reconnectCalls: remoteAgentHostService.reconnectCalls, + cached: service.getCachedTunnels(), + }, { + ensureCalls: [{ address: 'tunnel:tunnel-id', userInitiated: true }], + reconnectCalls: [{ address: 'tunnel:tunnel-id', userInitiated: true }], + cached: [{ + tunnelId: 'tunnel-id', + clusterId: 'updated-cluster', + name: 'Updated tunnel', + protocolVersion: 6, + authProvider: 'microsoft', + }], + }); + }); }); suite('tunnelAgentHostServiceImpl - gateway selection', () => {