Support session configuration controls on remote AHP hosts

src/vs/sessions/contrib/providers/agentHost/test/browser/agentHostModePicker.fixture.ts

@@ -12,7 +12,7 @@ import { mock } from '../../../../../../base/test/common/mock.js';
 import { ActionWidgetService, IActionWidgetService } from '../../../../../../platform/actionWidget/browser/actionWidget.js';
 import { IAgentHostEnablementService } from '../../../../../../platform/agentHost/common/agentHostEnablementService.js';
 import { getAgentHostCopilotSandboxSettingId, IAgentConnection, IAgentHostNetworkDiagnosticsInfo, IAgentHostService } from '../../../../../../platform/agentHost/common/agentService.js';
-import { IAgentHostConnectionsService } from '../../../../../../platform/agentHost/common/agentHostConnectionsService.js';
+import { AMBIENT_AGENT_HOST_AUTHORITY, IAgentHostConnectionsService } from '../../../../../../platform/agentHost/common/agentHostConnectionsService.js';
 import { getAgentHostOperatingSystem } from '../../../../../../platform/agentHost/common/agentHostOperatingSystem.js';
 import { IAgentSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js';
 import { ComponentToState, StateComponents } from '../../../../../../platform/agentHost/common/state/sessionState.js';
@@ -25,6 +25,7 @@ import { IContextViewService } from '../../../../../../platform/contextview/brow
 import { ContextViewService } from '../../../../../../platform/contextview/browser/contextViewService.js';
 import { ILayoutService } from '../../../../../../platform/layout/browser/layoutService.js';
 import { AgentHostChatInputPicker } from '../../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostChatInputPicker.js';
+import { toAgentHostBackendSessionUri } from '../../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionUri.js';
 import { IAgentHostSessionWorkingDirectoryResolver } from '../../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionWorkingDirectoryResolver.js';
 import { IAgentHostNewSessionFolderService } from '../../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostNewSessionFolderService.js';
 import { IAgentHostUntitledProvisionalSessionService } from '../../../../../../workbench/contrib/chat/browser/agentSessions/agentHost/agentHostUntitledProvisionalSessionService.js';
@@ -158,13 +159,15 @@ async function render(context: ComponentFixtureContext, mode: string, permission
 	if (editor) {
 		const state = new class extends mock<SessionState>() {
 			override readonly config = config;
+			override readonly provider = 'copilotcli';
 		}();
 		const changed = disposableStore.add(new Emitter<SessionState>());
 		const subscriptions: { [K in StateComponents]?: IAgentSubscription<ComponentToState[K]> } = {
 			[StateComponents.Session]: { value: state, verifiedValue: state, onDidChange: changed.event, onWillApplyAction: Event.None, onDidApplyAction: Event.None },
 		};
 		const hostService = new class extends mock<IAgentHostService>() {
 			override readonly onAgentHostStart = Event.None;
+			override readonly onDidNotification = Event.None;
 			override getNetworkDiagnosticsInfo(): Promise<IAgentHostNetworkDiagnosticsInfo> {
 				return connection.getNetworkDiagnosticsInfo();
 			}
@@ -183,6 +186,14 @@ async function render(context: ComponentFixtureContext, mode: string, permission
 			}
 		}();
 		instantiationService.set(IAgentHostService, hostService);
+		instantiationService.stub(IAgentHostConnectionsService, {
+			ambientConnection: hostService,
+			onDidChangeSessionResolution: Event.None,
+			resolveSessionResource: sessionResource => {
+				const backendSession = toAgentHostBackendSessionUri(sessionResource);
+				return backendSession ? { connection: hostService, backendSession, connectionAuthority: AMBIENT_AGENT_HOST_AUTHORITY } : undefined;
+			},
+		});
 		instantiationService.stub(IAgentHostSessionWorkingDirectoryResolver, { resolve: () => undefined });
 		instantiationService.stub(IAgentHostNewSessionFolderService, { getFolder: () => undefined, getDefaultFolder: () => undefined });
 		instantiationService.stub(IAgentHostUntitledProvisionalSessionService, { onDidChange: Event.None, get: () => undefined, getResolvedConfig: () => config, refreshResolvedConfig: async () => { } });

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostGenericConfigChips.ts

@@ -6,21 +6,23 @@
 import * as dom from '../../../../../../base/browser/dom.js';
 import { CancellationTokenSource } from '../../../../../../base/common/cancellation.js';
 import { Disposable, DisposableMap, IDisposable, MutableDisposable } from '../../../../../../base/common/lifecycle.js';
+import { isEqual } from '../../../../../../base/common/resources.js';
 import { URI } from '../../../../../../base/common/uri.js';
 import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js';
-import { IAgentHostService } from '../../../../../../platform/agentHost/common/agentService.js';
+import { IAgentHostConnectionsService, IAgentHostSessionResolution } from '../../../../../../platform/agentHost/common/agentHostConnectionsService.js';
 import type { ResolveSessionConfigResult, SessionConfigPropertySchema } from '../../../../../../platform/agentHost/common/state/protocol/commands.js';
 import type { SessionState } from '../../../../../../platform/agentHost/common/state/protocol/state.js';
 import { StateComponents } from '../../../../../../platform/agentHost/common/state/sessionState.js';
 import { type IAgentSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js';
 import { isUntitledChatSession } from '../../../common/model/chatUri.js';
 import type { IChatWidget } from '../../chat.js';
-import { AgentHostChatInputPicker, isClaimedByDedicatedPicker } from './agentHostChatInputPicker.js';
+import { AgentHostChatInputPicker, isGenericConfigPickerProperty } from './agentHostChatInputPicker.js';
 import { IAgentHostSessionWorkingDirectoryResolver } from './agentHostSessionWorkingDirectoryResolver.js';
 import { IAgentHostNewSessionFolderService } from './agentHostNewSessionFolderService.js';
 import { IAgentHostUntitledProvisionalSessionService } from './agentHostUntitledProvisionalSessionService.js';
 import { IWorkspaceContextService } from '../../../../../../platform/workspace/common/workspace.js';
-import { toAgentHostBackendSessionUri } from './agentHostSessionUri.js';
+import { resolveAgentHostChatSession, toAgentHostBackendSessionUri } from './agentHostSessionUri.js';
+import { retrySessionConfigSubscriptionOnCreation } from './agentHostSessionConfigSubscription.js';
 
 /**
  * Direct-render chip lane for agent-host session-config properties that are
@@ -39,14 +41,10 @@ export class AgentHostGenericConfigChips extends Disposable {
 	private readonly _chips = this._register(new DisposableMap<string>());
 	private readonly _chipElements = new Map<string, HTMLElement>();
 
-	/**
-	 * Subscription to the active session's backend state. Maintained for the
-	 * lifetime of any one (sessionResource, backendSession) pair; replaced
-	 * via {@link _reattach} when the active session changes.
-	 */
-	private readonly _subRef = this._register(new MutableDisposable<IDisposable & {
+	/** Subscription to the active session, replaced when its resource or owning connection changes. */
+	private readonly _subRef = this._register(new MutableDisposable<IDisposable & IAgentHostSessionResolution & {
 		readonly sub: IAgentSubscription<SessionState>;
-		readonly backendSession: URI;
+		readonly sessionResource: URI;
 	}>());
 
 	private _initialResolved: { readonly sessionResource: URI; readonly result: ResolveSessionConfigResult } | undefined;
@@ -55,14 +53,15 @@ export class AgentHostGenericConfigChips extends Disposable {
 	constructor(
 		private readonly _widget: IChatWidget,
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
-		@IAgentHostService private readonly _agentHostService: IAgentHostService,
+		@IAgentHostConnectionsService private readonly _connectionsService: IAgentHostConnectionsService,
 		@IAgentHostUntitledProvisionalSessionService private readonly _provisional: IAgentHostUntitledProvisionalSessionService,
 		@IAgentHostSessionWorkingDirectoryResolver private readonly _workingDirectoryResolver: IAgentHostSessionWorkingDirectoryResolver,
 		@IWorkspaceContextService private readonly _workspaceContextService: IWorkspaceContextService,
 		@IAgentHostNewSessionFolderService private readonly _newSessionFolderService: IAgentHostNewSessionFolderService,
 	) {
 		super();
 		this._register(this._widget.onDidChangeViewModel(() => this._reattach()));
+		this._register(this._connectionsService.onDidChangeSessionResolution(() => this._reattach()));
 		this._register(this._provisional.onDidChange((sessionResource: URI) => {
 			const current = this._widget.viewModel?.sessionResource;
 			if (current && current.toString() === sessionResource.toString()) {
@@ -84,36 +83,43 @@ export class AgentHostGenericConfigChips extends Disposable {
 	private _reattach(): void {
 		const sessionResource = this._widget.viewModel?.sessionResource;
 		const provisionalBackend = sessionResource ? this._provisional.get(sessionResource) : undefined;
-		const backendSession = provisionalBackend
-			?? (sessionResource ? toAgentHostBackendSessionUri(sessionResource) : undefined);
+		const resolution = sessionResource ? resolveAgentHostChatSession(sessionResource, provisionalBackend, this._connectionsService) : undefined;
 
-		if (!sessionResource || !backendSession) {
+		if (!sessionResource || !resolution) {
 			this._subRef.clear();
 			this._initialResolved = undefined;
 			this._cancelInitialResolve();
 			this._sync();
 			return;
 		}
 
-		if (isUntitledChatSession(sessionResource) && !provisionalBackend) {
+		const localBackend = toAgentHostBackendSessionUri(sessionResource);
+		if (localBackend && isUntitledChatSession(sessionResource) && !provisionalBackend) {
 			this._subRef.clear();
 			if (!this._initialResolved || this._initialResolved.sessionResource.toString() !== sessionResource.toString()) {
 				this._initialResolved = undefined;
-				void this._refreshInitialResolved(sessionResource, backendSession);
+				void this._refreshInitialResolved(sessionResource, localBackend);
 			}
 			this._sync();
 			return;
 		}
 
 		this._initialResolved = undefined;
 		this._cancelInitialResolve();
-		const ref = this._agentHostService.getSubscription(StateComponents.Session, backendSession, 'AgentHostGenericConfigChips');
+		const current = this._subRef.value;
+		if (current && !(current.sub.value instanceof Error) && isEqual(current.sessionResource, sessionResource) && current.connection === resolution.connection && isEqual(current.backendSession, resolution.backendSession)) {
+			this._sync();
+			return;
+		}
+		const ref = resolution.connection.getSubscription(StateComponents.Session, resolution.backendSession, 'AgentHostGenericConfigChips');
 		const sub = ref.object;
 		const listener = sub.onDidChange(() => this._sync());
+		const creationListener = retrySessionConfigSubscriptionOnCreation(resolution.connection, resolution.backendSession, sub, () => this._reattach());
 		this._subRef.value = {
+			...resolution,
 			sub,
-			backendSession,
-			dispose: () => { listener.dispose(); ref.dispose(); },
+			sessionResource,
+			dispose: () => { creationListener.dispose(); listener.dispose(); ref.dispose(); },
 		};
 		this._sync();
 	}
@@ -128,7 +134,7 @@ export class AgentHostGenericConfigChips extends Disposable {
 		const cts = new CancellationTokenSource();
 		this._initialResolveCts.value = cts;
 		try {
-			const result = await this._agentHostService.resolveSessionConfig({
+			const result = await this._connectionsService.ambientConnection.resolveSessionConfig({
 				provider: backendSession.scheme,
 				workingDirectory: this._readWorkingDirectory(),
 			});
@@ -176,10 +182,12 @@ export class AgentHostGenericConfigChips extends Disposable {
 			return;
 		}
 		const entries = this._readSchemaProperties();
+		const sessionResource = this._widget.viewModel?.sessionResource;
+		const isStartedSession = !!sessionResource && !(isUntitledChatSession(sessionResource) && toAgentHostBackendSessionUri(sessionResource));
 		const desired = new Set<string>();
 		if (entries) {
 			for (const [property, schema] of entries) {
-				if (isClaimedByDedicatedPicker(property, schema)) {
+				if (!isGenericConfigPickerProperty(property, schema, isStartedSession)) {
 					continue;
 				}
 				desired.add(property);
@@ -206,15 +214,15 @@ export class AgentHostGenericConfigChips extends Disposable {
 			// chips' container — required so the secondary-toolbar styling
 			// in `chat.css` (height, padding, chevron) applies here too.
 			const slot = dom.append(this._container, dom.$('.agent-host-generic-chip-slot.chat-input-picker-item'));
-			chip.render(slot);
-			this._chipElements.set(property, slot);
 			this._chips.set(property, {
 				dispose: () => {
 					chip.dispose();
 					slot.remove();
 					this._chipElements.delete(property);
 				},
 			});
+			this._chipElements.set(property, slot);
+			chip.render(slot);
 		}
 	}
 }

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionConfigSubscription.ts

@@ -0,0 +1,40 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { DisposableStore, IDisposable } from '../../../../../../base/common/lifecycle.js';
+import { isEqual } from '../../../../../../base/common/resources.js';
+import { URI } from '../../../../../../base/common/uri.js';
+import { IAgentConnection } from '../../../../../../platform/agentHost/common/agentService.js';
+import { IAgentSubscription } from '../../../../../../platform/agentHost/common/state/agentSubscription.js';
+import { SessionState } from '../../../../../../platform/agentHost/common/state/protocol/state.js';
+import { NotificationType } from '../../../../../../platform/agentHost/common/state/sessionActions.js';
+
+/** Retries an early subscription failure once the host announces that the session exists. */
+export function retrySessionConfigSubscriptionOnCreation(connection: IAgentConnection, session: URI, subscription: IAgentSubscription<SessionState>, retry: () => void): IDisposable {
+	const store = new DisposableStore();
+	let creationAnnounced = false;
+	const retryIfFailed = () => {
+		if (creationAnnounced && subscription.value instanceof Error) {
+			creationAnnounced = false;
+			// Let every consumer observe the error before retrying disposes the shared subscription.
+			queueMicrotask(() => {
+				if (!store.isDisposed) {
+					retry();
+				}
+			});
+		}
+	};
+	store.add(connection.onDidNotification(notification => {
+		if (notification.type === NotificationType.SessionAdded && isEqual(URI.parse(notification.summary.resource), session)) {
+			creationAnnounced = subscription.value === undefined || subscription.value instanceof Error;
+			retryIfFailed();
+		}
+	}));
+	store.add(subscription.onDidChange(() => creationAnnounced = false));
+	if (subscription.onDidError) {
+		store.add(subscription.onDidError(retryIfFailed));
+	}
+	return store;
+}

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionUri.ts

@@ -4,6 +4,15 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { URI } from '../../../../../../base/common/uri.js';
+import { AMBIENT_AGENT_HOST_AUTHORITY, IAgentHostConnectionsService, IAgentHostSessionResolution } from '../../../../../../platform/agentHost/common/agentHostConnectionsService.js';
+
+/** Preserves local provisional resources while resolving all other sessions through their owning host. */
+export function resolveAgentHostChatSession(sessionResource: URI, provisionalBackend: URI | undefined, connectionsService: IAgentHostConnectionsService): IAgentHostSessionResolution | undefined {
+	const resolution = connectionsService.resolveSessionResource(sessionResource);
+	return resolution && provisionalBackend && resolution.connectionAuthority === AMBIENT_AGENT_HOST_AUTHORITY
+		? { ...resolution, backendSession: provisionalBackend }
+		: resolution;
+}
 
 export function toAgentHostBackendSessionUri(sessionResource: URI): URI | undefined {
 	const scheme = sessionResource.scheme;