agentHost: keep SSH connection registered after incompatible handshake

src/vs/platform/agentHost/browser/remoteAgentHostServiceImpl.ts

@@ -284,7 +284,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo
 		return connection;
 	}
 
-	async addManagedConnection(entry: IRemoteAgentHostEntry, connection: IAgentConnection, transportDisposable?: IDisposable): Promise<IRemoteAgentHostConnectionInfo> {
+	async addManagedConnection(entry: IRemoteAgentHostEntry, connection: IAgentConnection, transportDisposable?: IDisposable, status = RemoteAgentHostConnectionStatus.connected): Promise<IRemoteAgentHostConnectionInfo> {
 		if (!this._configurationService.getValue<boolean>(RemoteAgentHostsEnabledSettingId)) {
 			throw new Error('Remote agent host connections are not enabled.');
 		}
@@ -311,7 +311,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo
 		// Create a connection entry wrapping the pre-connected client
 		const protocolClient = connection as RemoteAgentHostProtocolClient;
 		store.add(protocolClient);
-		const connEntry: IConnectionEntry = { store, client: protocolClient, transportDisposable, connected: true, status: RemoteAgentHostConnectionStatus.connected };
+		const connEntry: IConnectionEntry = { store, client: protocolClient, transportDisposable, connected: RemoteAgentHostConnectionStatus.isConnected(status), status };
 		this._entries.set(address, connEntry);
 		this._names.set(address, entry.name);
 		this._registeredEntries.set(address, entry);
@@ -340,7 +340,7 @@ export class RemoteAgentHostService extends Disposable implements IRemoteAgentHo
 			name: entry.name,
 			clientId: protocolClient.clientId,
 			defaultDirectory: protocolClient.defaultDirectory,
-			status: RemoteAgentHostConnectionStatus.connected,
+			status,
 		};
 	}
 

src/vs/platform/agentHost/common/remoteAgentHostService.ts

@@ -252,8 +252,13 @@ export interface IRemoteAgentHostService {
 	 * Callers should put any teardown that needs to happen on entry removal
 	 * (e.g. closing the shared-process tunnel, dropping renderer-side handles)
 	 * into this disposable, so a single removal path tears down the whole stack.
+	 *
+	 * `status` defaults to `connected`. Pass `incompatible` when the managed
+	 * transport is alive but the protocol handshake rejected the client version;
+	 * this keeps recovery actions (such as server upgrade) addressable without
+	 * exposing the connection as ready for session traffic.
 	 */
-	addManagedConnection(entry: IRemoteAgentHostEntry, connection: IAgentConnection, transportDisposable?: IDisposable): Promise<IRemoteAgentHostConnectionInfo>;
+	addManagedConnection(entry: IRemoteAgentHostEntry, connection: IAgentConnection, transportDisposable?: IDisposable, status?: RemoteAgentHostConnectionStatus): Promise<IRemoteAgentHostConnectionInfo>;
 
 	/**
 	 * Force the protocol client at `address` (if any) to treat its

src/vs/platform/agentHost/electron-browser/sshRemoteAgentHostServiceImpl.ts

@@ -13,13 +13,14 @@ import { IConfigurationService } from '../../configuration/common/configuration.
 import { IEnvironmentService } from '../../environment/common/environment.js';
 import { ISharedProcessService } from '../../ipc/electron-browser/services.js';
 import { ProxyChannel } from '../../../base/parts/ipc/common/ipc.js';
-import { IRemoteAgentHostService, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../common/remoteAgentHostService.js';
+import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostEntryType, RemoteAgentHostsEnabledSettingId } from '../common/remoteAgentHostService.js';
 import { createDecorator, IInstantiationService } from '../../instantiation/common/instantiation.js';
 import { IQuickInputService } from '../../quickinput/common/quickInput.js';
 import { AhpJsonlLogger } from '../common/ahpJsonlLogger.js';
 import { AgentHostAhpJsonlLoggingSettingId } from '../common/agentService.js';
 import { SSHRelayTransport } from './sshRelayTransport.js';
 import { RemoteAgentHostProtocolClient } from '../browser/remoteAgentHostProtocolClient.js';
+import { PROTOCOL_VERSION } from '../common/state/protocol/version/registry.js';
 import {
 	ISSHRemoteAgentHostService,
 	SSH_REMOTE_AGENT_HOST_CHANNEL,
@@ -185,25 +186,57 @@ export class SSHRemoteAgentHostService extends Disposable implements ISSHRemoteA
 	private async _setupConnection(result: ISSHConnectResult): Promise<ISSHAgentHostConnection> {
 		const existing = this._connections.get(result.connectionId);
 		if (existing) {
-			this._logService.trace('[SSHRemoteAgentHost] Returning existing connection handle');
-			return existing;
+			// Reuse the existing handle only if the managed entry is still
+			// in a usable state. After a `reconnect` that replaced the
+			// underlying SSH relay (e.g. following a CLI-driven server
+			// upgrade), the previous protocol client is bound to a
+			// torn-down transport and — if its handshake had failed with
+			// `incompatible` — will never re-handshake on its own. Drop
+			// the stale local state and fall through to a fresh
+			// handshake; the subsequent `addManagedConnection` call
+			// disposes the stale protocol client by replacing the entry.
+			if (this._remoteAgentHostService.getConnection(result.address)) {
+				this._logService.trace('[SSHRemoteAgentHost] Returning existing connection handle');
+				return existing;
+			}
+			this._logService.info(`[SSHRemoteAgentHost] Replacing stale connection handle for ${result.address}`);
+			this._connections.delete(result.connectionId);
+			// Mark closed-by-main so disposing the handle does NOT call
+			// disconnect() — the main service kept the SSH client alive
+			// across `replaceRelay`, and we'd kill the brand-new tunnel
+			// otherwise.
+			existing.fireClose();
+			existing.dispose();
+			this._onDidChangeConnections.fire();
 		}
-
-		let protocolClient: RemoteAgentHostProtocolClient | undefined;
-		let handle: SSHAgentHostConnectionHandle | undefined;
 		let registeredHandle = false;
+		const protocolClient = this._createRelayClient(result);
+		let status = RemoteAgentHostConnectionStatus.connected;
+		let connectError: unknown;
 		try {
-			protocolClient = this._createRelayClient(result);
 			await protocolClient.connect();
 			this._logService.trace('[SSHRemoteAgentHost] Protocol handshake completed');
+		} catch (err) {
+			const incompatible = RemoteAgentHostConnectionStatus.fromConnectError(err, [PROTOCOL_VERSION]);
+			if (!RemoteAgentHostConnectionStatus.isIncompatible(incompatible)) {
+				this._logService.error('[SSHRemoteAgentHost] Connection setup failed', err);
+				protocolClient.dispose();
+				this._mainService.disconnect(result.connectionId).catch(() => { /* best effort */ });
+				throw err;
+			}
+			this._logService.warn(`[SSHRemoteAgentHost] Incompatible with ${result.address}: ${incompatible.message}`);
+			status = incompatible;
+			connectError = err;
+		}
 
-			handle = new SSHAgentHostConnectionHandle(
-				result.config,
-				result.address,
-				result.name,
-				() => this._mainService.disconnect(result.connectionId),
-			);
+		const handle = new SSHAgentHostConnectionHandle(
+			result.config,
+			result.address,
+			result.name,
+			() => this._mainService.disconnect(result.connectionId),
+		);
 
+		try {
 			this._connections.set(result.connectionId, handle);
 			registeredHandle = true;
 			this._onDidChangeConnections.fire();
@@ -219,20 +252,24 @@ export class SSHRemoteAgentHostService extends Disposable implements ISSHRemoteA
 					user: result.config.username || undefined,
 					port: result.config.port,
 				},
-			}, protocolClient, this._createTransportDisposable(result.connectionId, handle));
-
-			return handle;
+			}, protocolClient, this._createTransportDisposable(result.connectionId, handle), status);
 		} catch (err) {
 			this._logService.error('[SSHRemoteAgentHost] Connection setup failed', err);
 			if (registeredHandle && this._connections.get(result.connectionId) === handle) {
 				this._connections.delete(result.connectionId);
 				this._onDidChangeConnections.fire();
 			}
-			handle?.dispose();
-			protocolClient?.dispose();
+			handle.dispose();
+			protocolClient.dispose();
 			this._mainService.disconnect(result.connectionId).catch(() => { /* best effort */ });
 			throw err;
 		}
+
+		if (connectError) {
+			throw connectError;
+		}
+
+		return handle;
 	}
 
 	/**

src/vs/platform/agentHost/test/electron-browser/remoteAgentHostService.test.ts

@@ -46,6 +46,7 @@ class MockProtocolClient extends Disposable {
 	readonly connectionState = 'connecting' as const;
 	readonly initializeResult = undefined;
 	readonly telemetryCapabilities = undefined;
+	readonly triggerVscodeUpgradeCalls: string[] = [];
 
 	public connectDeferred = new DeferredPromise<void>();
 
@@ -57,6 +58,11 @@ class MockProtocolClient extends Disposable {
 		return this.connectDeferred.p;
 	}
 
+	async triggerVscodeUpgrade(method: string) {
+		this.triggerVscodeUpgradeCalls.push(method);
+		return { ok: true, upgradeStarted: true };
+	}
+
 	fireClose(): void {
 		this._onDidClose.fire();
 	}
@@ -517,6 +523,38 @@ suite('RemoteAgentHostService', () => {
 			);
 		}
 
+		test('keeps incompatible managed connection addressable for server upgrade', async () => {
+			const mockClient = disposables.add(new MockProtocolClient('ssh:remote.example'));
+			await service.addManagedConnection(
+				{
+					name: 'SSH Host',
+					connection: {
+						type: RemoteAgentHostEntryType.SSH,
+						address: 'ssh:remote.example',
+						sshConfigHost: 'remote',
+						hostName: 'remote.example',
+					},
+				},
+				mockClient as unknown as Parameters<typeof service.addManagedConnection>[1],
+				undefined,
+				RemoteAgentHostConnectionStatus.incompatible('Unsupported protocol version', ['0.3.0'], ['^0.2.0'], '_vscodeUpgrade'),
+			);
+
+			const upgradeResult = await service.triggerServerUpgrade('ssh:remote.example', '_vscodeUpgrade');
+
+			assert.deepStrictEqual({
+				status: service.connections[0].status,
+				connectedConnection: service.getConnection('ssh:remote.example'),
+				upgradeCalls: mockClient.triggerVscodeUpgradeCalls,
+				upgradeResult,
+			}, {
+				status: RemoteAgentHostConnectionStatus.incompatible('Unsupported protocol version', ['0.3.0'], ['^0.2.0'], '_vscodeUpgrade'),
+				connectedConnection: undefined,
+				upgradeCalls: ['_vscodeUpgrade'],
+				upgradeResult: { ok: true, upgradeStarted: true },
+			});
+		});
+
 		test('disposes transportDisposable when entry is removed via removeRemoteAgentHost', async () => {
 			const t = makeTransportDisposable();
 			await addManaged('Managed', 'managed:1234', t.disposable);

src/vs/platform/agentHost/test/electron-browser/sshRemoteAgentHostService.test.ts

@@ -16,8 +16,9 @@ import { IConfigurationService } from '../../../configuration/common/configurati
 
 import { ISharedProcessService } from '../../../ipc/electron-browser/services.js';
 import { IQuickInputService } from '../../../quickinput/common/quickInput.js';
-import { IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId } from '../../common/remoteAgentHostService.js';
+import { IRemoteAgentHostService, RemoteAgentHostConnectionStatus, RemoteAgentHostsEnabledSettingId } from '../../common/remoteAgentHostService.js';
 import type { IAgentConnection } from '../../common/agentService.js';
+import { AHP_UNSUPPORTED_PROTOCOL_VERSION, ProtocolError } from '../../common/state/sessionProtocol.js';
 import type {
 	ISSHAgentHostConfig,
 	ISSHConnectResult,
@@ -26,6 +27,7 @@ import type {
 	ISSHResolvedConfig,
 	ISSHRemoteAgentHostMainService,
 } from '../../common/sshRemoteAgentHost.js';
+import { PROTOCOL_VERSION } from '../../common/state/protocol/version/registry.js';
 import { ISSHRelayClientFactory, SSHRemoteAgentHostService } from '../../electron-browser/sshRemoteAgentHostServiceImpl.js';
 import { RemoteAgentHostProtocolClient } from '../../browser/remoteAgentHostProtocolClient.js';
 
@@ -85,8 +87,8 @@ class MockSSHMainService {
 	async reconnect(sshConfigHost: string, name: string): Promise<ISSHConnectResult> {
 		this.reconnectCalls.push({ sshConfigHost, name });
 		return {
-			connectionId: `conn-${this._nextConnectionId++}`,
-			address: `ssh:${sshConfigHost}`,
+			connectionId: this.connectResult?.connectionId ?? `conn-${this._nextConnectionId++}`,
+			address: this.connectResult?.address ?? `ssh:${sshConfigHost}`,
 			name,
 			connectionToken: 'test-token',
 			config: { host: sshConfigHost, username: 'u', authMethod: 0 as never, name, sshConfigHost },
@@ -140,14 +142,36 @@ function asChannel(target: object): IChannel {
 
 /** Captures addManagedConnection calls so tests can inspect transportDisposable. */
 class MockRemoteAgentHostService extends Disposable {
-	readonly added: Array<{ address: string; transport?: IDisposable }> = [];
-	private readonly _entries = new Map<string, { transport?: IDisposable; client: { dispose?: () => void } }>();
-
-	async addManagedConnection(entry: { name: string; connection: { address?: string; sshConfigHost?: string } }, client: IAgentConnection, transportDisposable?: IDisposable): Promise<unknown> {
+	readonly added: Array<{ address: string; status?: RemoteAgentHostConnectionStatus; transport?: IDisposable }> = [];
+	private readonly _entries = new Map<string, { transport?: IDisposable; client: { dispose?: () => void }; status: RemoteAgentHostConnectionStatus }>();
+	// Holds transport disposables from prior registrations that were
+	// replaced by a later `addManagedConnection` 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[] = [];
+
+	async addManagedConnection(entry: { name: string; connection: { address?: string; sshConfigHost?: string } }, client: IAgentConnection, transportDisposable?: IDisposable, status: RemoteAgentHostConnectionStatus = RemoteAgentHostConnectionStatus.connected): Promise<unknown> {
 		const address = entry.connection.address ?? `ssh:${entry.connection.sshConfigHost}`;
-		this.added.push({ address, transport: transportDisposable });
-		this._entries.set(address, { client: client as { dispose?: () => void }, transport: transportDisposable });
-		return { address, name: entry.name, clientId: 'mock', defaultDirectory: undefined, status: 0 };
+		// Mirror RemoteAgentHostService: re-registering an address replaces
+		// the previous entry and disposes its protocol client (but NOT its
+		// transport disposable — the new entry owns the underlying tunnel).
+		const previous = this._entries.get(address);
+		if (previous) {
+			previous.client.dispose?.();
+			if (previous.transport) {
+				this._abandonedTransports.push(previous.transport);
+			}
+		}
+		this.added.push({ address, status, transport: transportDisposable });
+		this._entries.set(address, { client: client as { dispose?: () => void }, transport: transportDisposable, status });
+		return { address, name: entry.name, clientId: 'mock', defaultDirectory: undefined, status };
+	}
+
+	/** Mirrors IRemoteAgentHostService.getConnection: returns the client only when the entry is connected. */
+	getConnection(address: string): IAgentConnection | undefined {
+		const entry = this._entries.get(address);
+		return entry && RemoteAgentHostConnectionStatus.isConnected(entry.status) ? entry.client as unknown as IAgentConnection : undefined;
 	}
 
 	notifyConnectionClosed(_address: string): void {
@@ -173,6 +197,11 @@ class MockRemoteAgentHostService extends Disposable {
 			e.transport?.dispose();
 		}
 		this._entries.clear();
+		// Release abandoned transports from prior registrations as well.
+		for (const t of this._abandonedTransports) {
+			t.dispose();
+		}
+		this._abandonedTransports.length = 0;
 		super.dispose();
 	}
 }
@@ -268,11 +297,84 @@ suite('SSHRemoteAgentHostService (renderer)', () => {
 
 		assert.strictEqual(remoteAgentHostService.added.length, 1);
 		assert.strictEqual(remoteAgentHostService.added[0].address, 'ssh:remote.example');
+		assert.strictEqual(remoteAgentHostService.added[0].status?.kind, 'connected');
 		assert.ok(remoteAgentHostService.added[0].transport, 'a transport disposable is passed so removal can tear down the SSH tunnel');
 		assert.strictEqual(service.connections.length, 1);
 		assert.strictEqual(handle.localAddress, 'ssh:remote.example');
 	});
 
+	test('incompatible handshake keeps SSH tunnel registered for server upgrade', async () => {
+		const connectPromise = service.connect(sampleConfig);
+		const client = await waitForClient(0);
+		await client.connectDeferred.error(new ProtocolError(
+			AHP_UNSUPPORTED_PROTOCOL_VERSION,
+			'Unsupported protocol version',
+			{ supportedVersions: ['^0.2.0'], _meta: { vscodeUpgradeMethod: '_vscodeUpgrade' } },
+		));
+
+		await assert.rejects(connectPromise, /Unsupported protocol version/);
+
+		assert.deepStrictEqual({
+			added: remoteAgentHostService.added.map(({ address, status }) => ({ address, status })),
+			connections: service.connections.map(connection => connection.localAddress),
+			disconnectCalls: mainService.disconnectCalls,
+		}, {
+			added: [{
+				address: 'ssh:remote.example',
+				status: RemoteAgentHostConnectionStatus.incompatible('Unsupported protocol version', [PROTOCOL_VERSION], ['^0.2.0'], '_vscodeUpgrade'),
+			}],
+			connections: ['ssh:remote.example'],
+			disconnectCalls: [],
+		});
+	});
+
+	test('reconnect after incompatible handshake replaces the stale handle and re-handshakes', async () => {
+		// Pin a stable connectionId so the simulated `replaceRelay` reconnect
+		// returns the same id as the initial connect — that is the real
+		// behavior of SSHRemoteAgentHostMainService.connect(replaceRelay=true).
+		mainService.connectResult = { connectionId: 'conn-stable', address: 'ssh:remote.example' };
+
+		// First connect: handshake rejected as incompatible. Per the existing
+		// fix, this still registers a managed connection in `incompatible`
+		// state so the server-upgrade RPC can reach the host.
+		const firstConnect = service.connect(sampleConfig);
+		const firstClient = await waitForClient(0);
+		await firstClient.connectDeferred.error(new ProtocolError(
+			AHP_UNSUPPORTED_PROTOCOL_VERSION,
+			'Unsupported protocol version',
+			{ supportedVersions: ['^0.2.0'], _meta: { vscodeUpgradeMethod: '_vscodeUpgrade' } },
+		));
+		await assert.rejects(firstConnect, /Unsupported protocol version/);
+
+		// User triggers the server upgrade and then the contribution reconnects.
+		// The reconnect must NOT short-circuit to the stale handle (whose
+		// protocol client is permanently stuck in incompatible state); it must
+		// build a fresh client and complete a fresh handshake against the
+		// upgraded server.
+		const reconnectPromise = service.reconnect('remote.example', 'My Remote');
+		const secondClient = await waitForClient(1);
+		await secondClient.connectDeferred.complete();
+		await reconnectPromise;
+
+		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,
+			// Exactly one renderer-side handle for the address.
+			connections: service.connections.map(connection => connection.localAddress),
+		}, {
+			clientCount: 2,
+			added: [
+				{ address: 'ssh:remote.example', statusKind: 'incompatible' },
+				{ address: 'ssh:remote.example', statusKind: 'connected' },
+			],
+			disconnectCalls: [],
+			connections: ['ssh:remote.example'],
+		});
+	});
+
 	test('disabled setting prevents SSH tunnel connects and reconnects', async () => {
 		configurationService.setRemoteAgentHostsEnabled(false);