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', () => {