agentHost: Register sessions provider independently of having a connection (#305915)

* agentHost: Register sessions provider independently of having a connection

* fix

src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHost.contribution.ts

@@ -9,11 +9,12 @@ import * as nls from '../../../../nls.js';
 import { AgentHostFileSystemProvider } from '../../../../platform/agentHost/common/agentHostFileSystemProvider.js';
 import { AGENT_HOST_LABEL_FORMATTER, AGENT_HOST_SCHEME, agentHostAuthority, fromAgentHostUri } from '../../../../platform/agentHost/common/agentHostUri.js';
 import { type AgentProvider, type IAgentConnection } from '../../../../platform/agentHost/common/agentService.js';
-import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId } from '../../../../platform/agentHost/common/remoteAgentHostService.js';
+import { IRemoteAgentHostConnectionInfo, IRemoteAgentHostEntry, IRemoteAgentHostService, RemoteAgentHostsEnabledSettingId, RemoteAgentHostsSettingId } from '../../../../platform/agentHost/common/remoteAgentHostService.js';
 import { isSessionAction } from '../../../../platform/agentHost/common/state/sessionActions.js';
 import { SessionClientState } from '../../../../platform/agentHost/common/state/sessionClientState.js';
 import { ROOT_STATE_URI, type IAgentInfo, type IRootState } from '../../../../platform/agentHost/common/state/sessionState.js';
 import { Extensions as ConfigurationExtensions, IConfigurationRegistry } from '../../../../platform/configuration/common/configurationRegistry.js';
+import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
 import { IDefaultAccountService } from '../../../../platform/defaultAccount/common/defaultAccount.js';
 import { IFileService } from '../../../../platform/files/common/files.js';
 import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
@@ -68,6 +69,10 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
 	/** Per-connection state: client state + per-agent registrations. */
 	private readonly _connections = this._register(new DisposableMap<string, ConnectionState>());
 
+	/** Per-address sessions providers, registered for all configured entries. */
+	private readonly _providerStores = this._register(new DisposableMap<string, DisposableStore>());
+	private readonly _providerInstances = new Map<string, RemoteAgentHostSessionsProvider>();
+
 	/** Maps sanitized authority strings back to original addresses. */
 	private readonly _fsProvider: AgentHostFileSystemProvider;
 
@@ -83,6 +88,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
 		@IFileService private readonly _fileService: IFileService,
 		@ISessionsProvidersService private readonly _sessionsProvidersService: ISessionsProvidersService,
 		@ILabelService private readonly _labelService: ILabelService,
+		@IConfigurationService private readonly _configurationService: IConfigurationService,
 	) {
 		super();
 
@@ -94,17 +100,76 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
 		// Display agent-host URIs with the original file path
 		this._register(this._labelService.registerFormatter(AGENT_HOST_LABEL_FORMATTER));
 
+		// Reconcile providers when configured entries change
+		this._register(this._configurationService.onDidChangeConfiguration(e => {
+			if (e.affectsConfiguration(RemoteAgentHostsSettingId) || e.affectsConfiguration(RemoteAgentHostsEnabledSettingId)) {
+				this._reconcile();
+			}
+		}));
+
 		// Reconcile when connections change (added/removed/reconnected)
 		this._register(this._remoteAgentHostService.onDidChangeConnections(() => {
-			this._reconcileConnections();
+			this._reconcile();
 		}));
 
 		// Push auth token whenever the default account or sessions change
 		this._register(this._defaultAccountService.onDidChangeDefaultAccount(() => this._authenticateAllConnections()));
 		this._register(this._authenticationService.onDidChangeSessions(() => this._authenticateAllConnections()));
 
-		// Initial setup for already-connected remotes
+		// Initial setup for configured entries and connected remotes
+		this._reconcile();
+	}
+
+	private _reconcile(): void {
+		this._reconcileProviders();
 		this._reconcileConnections();
+
+		// Ensure every live connection is wired to its provider.
+		// This covers the case where a provider was recreated (e.g. name
+		// change) while a connection for that address already existed.
+		for (const [address, connState] of this._connections) {
+			const provider = this._providerInstances.get(address);
+			if (provider) {
+				const connectionInfo = this._remoteAgentHostService.connections.find(c => c.address === address);
+				provider.setConnection(connState.loggedConnection, connectionInfo?.defaultDirectory);
+			}
+		}
+	}
+
+	private _reconcileProviders(): void {
+		const enabled = this._configurationService.getValue<boolean>(RemoteAgentHostsEnabledSettingId);
+		const entries = enabled ? this._remoteAgentHostService.configuredEntries : [];
+		const desiredAddresses = new Set(entries.map(e => e.address));
+
+		// Remove providers no longer configured
+		for (const [address] of this._providerStores) {
+			if (!desiredAddresses.has(address)) {
+				this._providerStores.deleteAndDispose(address);
+			}
+		}
+
+		// Add or recreate providers for configured entries
+		for (const entry of entries) {
+			const existing = this._providerInstances.get(entry.address);
+			if (existing && existing.label !== (entry.name || entry.address)) {
+				// Name changed — recreate since ISessionsProvider.label is readonly
+				this._providerStores.deleteAndDispose(entry.address);
+			}
+			if (!this._providerStores.has(entry.address)) {
+				this._createProvider(entry);
+			}
+		}
+	}
+
+	private _createProvider(entry: IRemoteAgentHostEntry): void {
+		const store = new DisposableStore();
+		const provider = this._instantiationService.createInstance(
+			RemoteAgentHostSessionsProvider, { address: entry.address, name: entry.name });
+		store.add(provider);
+		store.add(this._sessionsProvidersService.registerProvider(provider));
+		this._providerInstances.set(entry.address, provider);
+		store.add(toDisposable(() => this._providerInstances.delete(entry.address)));
+		this._providerStores.set(entry.address, store);
 	}
 
 	private _reconcileConnections(): void {
@@ -114,6 +179,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
 		for (const [address] of this._connections) {
 			if (!currentAddresses.has(address)) {
 				this._logService.info(`[RemoteAgentHost] Removing contribution for ${address}`);
+				this._providerInstances.get(address)?.clearConnection();
 				this._connections.deleteAndDispose(address);
 			}
 		}
@@ -125,6 +191,7 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
 				// If the name changed, tear down and re-register with new name
 				if (existing.name !== connectionInfo.name) {
 					this._logService.info(`[RemoteAgentHost] Name changed for ${connectionInfo.address}: ${existing.name} -> ${connectionInfo.name}`);
+					this._providerInstances.get(connectionInfo.address)?.clearConnection();
 					this._connections.deleteAndDispose(connectionInfo.address);
 					this._setupConnection(connectionInfo);
 				}
@@ -184,12 +251,8 @@ export class RemoteAgentHostContribution extends Disposable implements IWorkbenc
 		// Authenticate with this new connection
 		this._authenticateWithConnection(loggedConnection);
 
-		// Register a single sessions provider for the entire connection.
-		// It handles all agents discovered on this connection.
-		const sessionsProvider = this._instantiationService.createInstance(
-			RemoteAgentHostSessionsProvider, { connectionInfo, connection: loggedConnection });
-		store.add(sessionsProvider);
-		store.add(this._sessionsProvidersService.registerProvider(sessionsProvider));
+		// Wire connection to existing sessions provider
+		this._providerInstances.get(address)?.setConnection(loggedConnection, connectionInfo.defaultDirectory);
 	}
 
 	private _handleRootStateChange(address: string, loggedConnection: LoggingAgentConnection, rootState: IRootState): void {

src/vs/sessions/contrib/remoteAgentHost/browser/remoteAgentHostSessionsProvider.ts

@@ -6,7 +6,7 @@
 import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js';
 import { Codicon } from '../../../../base/common/codicons.js';
 import { Emitter, Event } from '../../../../base/common/event.js';
-import { Disposable } from '../../../../base/common/lifecycle.js';
+import { Disposable, DisposableStore, toDisposable } from '../../../../base/common/lifecycle.js';
 import { basename } from '../../../../base/common/resources.js';
 import { ISettableObservable, observableValue } from '../../../../base/common/observable.js';
 import { ThemeIcon } from '../../../../base/common/themables.js';
@@ -18,6 +18,7 @@ import { AGENT_HOST_SCHEME, agentHostAuthority, toAgentHostUri } from '../../../
 import { AgentSession, type IAgentConnection, type IAgentSessionMetadata } from '../../../../platform/agentHost/common/agentService.js';
 import { isSessionAction } from '../../../../platform/agentHost/common/state/sessionActions.js';
 import { IFileDialogService } from '../../../../platform/dialogs/common/dialogs.js';
+import { INotificationService } from '../../../../platform/notification/common/notification.js';
 import { ChatViewPaneTarget, IChatWidgetService } from '../../../../workbench/contrib/chat/browser/chat.js';
 import { IChatSendRequestOptions, IChatService } from '../../../../workbench/contrib/chat/common/chatService/chatService.js';
 import { IChatSessionFileChange, IChatSessionsService } from '../../../../workbench/contrib/chat/common/chatSessionsService.js';
@@ -26,11 +27,10 @@ import { ILanguageModelsService } from '../../../../workbench/contrib/chat/commo
 import { ISessionChangeEvent, ISendRequestOptions, ISessionsBrowseAction, ISessionsProvider, ISessionType } from '../../sessions/browser/sessionsProvider.js';
 import { CopilotCLISessionType } from '../../sessions/browser/sessionTypes.js';
 import { ISessionData, ISessionPullRequest, ISessionWorkspace, SessionStatus } from '../../sessions/common/sessionData.js';
-import { IRemoteAgentHostConnectionInfo } from '../../../../platform/agentHost/common/remoteAgentHostService.js';
 
 export interface IRemoteAgentHostSessionsProviderConfig {
-	readonly connectionInfo: IRemoteAgentHostConnectionInfo;
-	readonly connection: IAgentConnection;
+	readonly address: string;
+	readonly name: string;
 }
 
 /**
@@ -128,8 +128,10 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 	/** Selected model for the current new session. */
 	private _selectedModelId: string | undefined;
 
-	private readonly _connectionInfo: IRemoteAgentHostConnectionInfo;
-	private readonly _connection: IAgentConnection;
+	private _connection: IAgentConnection | undefined;
+	private _defaultDirectory: string | undefined;
+	private readonly _connectionListeners = this._register(new DisposableStore());
+	private readonly _disconnectListeners = this._register(new DisposableStore());
 	private readonly _connectionAuthority: string;
 
 	constructor(
@@ -139,13 +141,12 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 		@IChatService private readonly _chatService: IChatService,
 		@IChatWidgetService private readonly _chatWidgetService: IChatWidgetService,
 		@ILanguageModelsService private readonly _languageModelsService: ILanguageModelsService,
+		@INotificationService private readonly _notificationService: INotificationService,
 	) {
 		super();
 
-		this._connectionInfo = config.connectionInfo;
-		this._connection = config.connection;
-		this._connectionAuthority = agentHostAuthority(config.connectionInfo.address);
-		const displayName = config.connectionInfo.name || config.connectionInfo.address;
+		this._connectionAuthority = agentHostAuthority(config.address);
+		const displayName = config.name || config.address;
 
 		this.id = `agenthost-${this._connectionAuthority}`;
 		this.label = displayName;
@@ -159,9 +160,23 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 			providerId: this.id,
 			execute: () => this._browseForFolder(),
 		}];
+	}
 
-		// Listen for session notifications from the connection
-		this._register(this._connection.onDidNotification(n => {
+	// -- Connection Management --
+
+	/**
+	 * Wire a live connection to this provider, enabling session operations and folder browsing.
+	 */
+	setConnection(connection: IAgentConnection, defaultDirectory?: string): void {
+		if (this._connection === connection && this._defaultDirectory === defaultDirectory) {
+			return;
+		}
+
+		this._connectionListeners.clear();
+		this._connection = connection;
+		this._defaultDirectory = defaultDirectory;
+
+		this._connectionListeners.add(connection.onDidNotification(n => {
 			if (n.type === 'notify/sessionAdded') {
 				this._handleSessionAdded(n.summary);
 			} else if (n.type === 'notify/sessionRemoved') {
@@ -169,13 +184,35 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 			}
 		}));
 
-		// Refresh on turnComplete actions for metadata updates (title, timing)
-		this._register(this._connection.onDidAction(e => {
+		this._connectionListeners.add(connection.onDidAction(e => {
 			if (e.action.type === 'session/turnComplete' && isSessionAction(e.action)) {
 				const cts = new CancellationTokenSource();
 				this._refreshSessions(cts.token).finally(() => cts.dispose());
 			}
 		}));
+
+		// Always refresh sessions when a connection is (re)established
+		const cts = new CancellationTokenSource();
+		this._cacheInitialized = true;
+		this._refreshSessions(cts.token).finally(() => cts.dispose());
+	}
+
+	/**
+	 * Clear the connection, e.g. when the remote host disconnects.
+	 * Retains the provider registration so it remains visible in the UI.
+	 */
+	clearConnection(): void {
+		this._connectionListeners.clear();
+		this._disconnectListeners.clear();
+		this._connection = undefined;
+		this._defaultDirectory = undefined;
+
+		const removed = Array.from(this._sessionCache.values());
+		this._sessionCache.clear();
+		this._cacheInitialized = false;
+		if (removed.length > 0) {
+			this._onDidChangeSessions.fire({ added: [], removed, changed: [] });
+		}
 	}
 
 	// -- Workspaces --
@@ -229,6 +266,10 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 	}
 
 	createNewSession(workspace: ISessionWorkspace): ISessionData {
+		if (!this._connection) {
+			throw new Error(localize('notConnectedSession', "Cannot create session: not connected to remote agent host '{0}'.", this.label));
+		}
+
 		const workspaceUri = workspace.repositories[0]?.uri;
 		if (!workspaceUri) {
 			throw new Error('Workspace has no repository URI');
@@ -291,7 +332,7 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 	async deleteSession(chatId: string): Promise<void> {
 		const rawId = this._rawIdFromChatId(chatId);
 		const cached = rawId ? this._sessionCache.get(rawId) : undefined;
-		if (cached && rawId) {
+		if (cached && rawId && this._connection) {
 			await this._connection.disposeSession(AgentSession.uri(cached.agentProvider, rawId));
 			this._sessionCache.delete(rawId);
 			this._onDidChangeSessions.fire({ added: [], removed: [cached], changed: [] });
@@ -311,6 +352,10 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 	}
 
 	async sendRequest(chatId: string, options: ISendRequestOptions): Promise<ISessionData> {
+		if (!this._connection) {
+			throw new Error(localize('notConnectedSend', "Cannot send request: not connected to remote agent host '{0}'.", this.label));
+		}
+
 		const session = this._currentNewSession;
 		if (!session || session.id !== chatId) {
 			throw new Error(`Session '${chatId}' not found or not a new session`);
@@ -387,6 +432,9 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 	}
 
 	private async _refreshSessions(_token: unknown): Promise<void> {
+		if (!this._connection) {
+			return;
+		}
 		try {
 			const sessions = await this._connection.listSessions();
 			const currentKeys = new Set<string>();
@@ -428,6 +476,7 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 	/**
 	 * Wait for a new session to appear in the cache that wasn't present before.
 	 * Tries an immediate refresh, then listens for the session-added notification.
+	 * Rejects if the connection is cleared before a session appears.
 	 */
 	private async _waitForNewSession(existingKeys: Set<string>): Promise<ISessionData | undefined> {
 		// First, try an immediate refresh
@@ -439,17 +488,28 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 		}
 
 		// If not found yet, wait for the next onDidChangeSessions event
-		return new Promise<ISessionData | undefined>(resolve => {
+		// or reject if the connection is cleared.
+		return new Promise<ISessionData | undefined>((resolve, reject) => {
+			let settled = false;
 			const listener = this._onDidChangeSessions.event(e => {
 				const newSession = e.added.find(s => {
 					const rawId = s.resource.path.substring(1);
 					return !existingKeys.has(rawId);
 				});
 				if (newSession) {
+					settled = true;
 					listener.dispose();
+					disconnectListener.dispose();
 					resolve(newSession);
 				}
 			});
+			const disconnectListener = toDisposable(() => {
+				if (!settled) {
+					listener.dispose();
+					reject(new Error('Connection lost while waiting for session'));
+				}
+			});
+			this._disconnectListeners.add(disconnectListener);
 		});
 	}
 
@@ -500,8 +560,12 @@ export class RemoteAgentHostSessionsProvider extends Disposable implements ISess
 	// -- Private: Browse --
 
 	private async _browseForFolder(): Promise<ISessionWorkspace | undefined> {
-		const authority = agentHostAuthority(this._connectionInfo.address);
-		const defaultUri = agentHostUri(authority, this._connectionInfo.defaultDirectory ?? '/');
+		if (!this._connection) {
+			this._notificationService.error(localize('notConnected', "Unable to connect to remote agent host '{0}'.", this.label));
+			return undefined;
+		}
+
+		const defaultUri = agentHostUri(this._connectionAuthority, this._defaultDirectory ?? '/');
 
 		try {
 			const selected = await this._fileDialogService.showOpenDialog({

src/vs/sessions/contrib/remoteAgentHost/test/browser/remoteAgentHostSessionsProvider.test.ts

@@ -10,6 +10,7 @@ import { URI } from '../../../../../base/common/uri.js';
 import { mock } from '../../../../../base/test/common/mock.js';
 import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js';
 import { IFileDialogService } from '../../../../../platform/dialogs/common/dialogs.js';
+import { INotificationService } from '../../../../../platform/notification/common/notification.js';
 import { TestInstantiationService } from '../../../../../platform/instantiation/test/common/instantiationServiceMock.js';
 import { AgentSession, type IAgentConnection, type IAgentSessionMetadata } from '../../../../../platform/agentHost/common/agentService.js';
 import type { IActionEnvelope, INotification } from '../../../../../platform/agentHost/common/state/sessionActions.js';
@@ -83,6 +84,7 @@ function createProvider(disposables: DisposableStore, connection: MockAgentConne
 	const instantiationService = disposables.add(new TestInstantiationService());
 
 	instantiationService.stub(IFileDialogService, {});
+	instantiationService.stub(INotificationService, { error: () => { } });
 	instantiationService.stub(IChatSessionsService, {
 		getChatSessionContribution: () => ({ type: 'remote-test-copilot', name: 'test', displayName: 'Test', description: 'test', icon: undefined }),
 		getOrCreateChatSession: async () => ({ onWillDispose: () => ({ dispose() { } }), sessionResource: URI.from({ scheme: 'test' }), history: [], dispose() { } }),
@@ -99,15 +101,13 @@ function createProvider(disposables: DisposableStore, connection: MockAgentConne
 	});
 
 	const config: IRemoteAgentHostSessionsProviderConfig = {
-		connectionInfo: {
-			address: overrides?.address ?? 'localhost:4321',
-			name: overrides !== undefined && Object.prototype.hasOwnProperty.call(overrides, 'connectionName') ? overrides.connectionName ?? '' : 'Test Host',
-			clientId: 'test-client',
-		},
-		connection,
+		address: overrides?.address ?? 'localhost:4321',
+		name: overrides !== undefined && Object.prototype.hasOwnProperty.call(overrides, 'connectionName') ? overrides.connectionName ?? '' : 'Test Host',
 	};
 
-	return disposables.add(instantiationService.createInstance(RemoteAgentHostSessionsProvider, config));
+	const provider = disposables.add(instantiationService.createInstance(RemoteAgentHostSessionsProvider, config));
+	provider.setConnection(connection);
+	return provider;
 }
 
 function fireSessionAdded(connection: MockAgentConnection, rawId: string, opts?: { provider?: string; title?: string; workingDirectory?: string }): void {