Enhance session customization handling and simplify sync logic

src/vs/platform/agentHost/node/agentService.ts

@@ -482,6 +482,33 @@ export class AgentService extends Disposable implements IAgentService {
 			this._attachGitState(session, created.workingDirectory ?? config?.workingDirectory);
 		}
 
+		// Publish the current customization set for the new session now that
+		// the state manager knows about it. Without this, host/global
+		// customizations — and the custom agents they contribute — are
+		// invisible to subscribers (e.g. the agent picker) until something
+		// else triggers a republish (`RootConfigChanged`, plugin reload, or
+		// the first message's `setClientCustomizations`). The provider's
+		// `createSession` cannot do this itself because it runs *before*
+		// `_stateManager.createSession` above, so any action it dispatches
+		// is dropped as "Action for unknown session".
+		if (provider.getSessionCustomizations) {
+			void provider.getSessionCustomizations(session).then(customizations => {
+				if (customizations.length === 0) {
+					return;
+				}
+				const sessionKey = session.toString();
+				if (!this._stateManager.getSessionState(sessionKey)) {
+					return;
+				}
+				this._stateManager.dispatchServerAction(sessionKey, {
+					type: ActionType.SessionCustomizationsChanged,
+					customizations: [...customizations],
+				});
+			}).catch(err => {
+				this._logService.error('[AgentService] createSession: failed to publish initial customizations', err);
+			});
+		}
+
 		return session;
 	}
 

src/vs/platform/agentHost/node/copilot/copilotAgent.ts

@@ -313,7 +313,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 	}
 
 	async getSessionCustomizations(session: URI): Promise<readonly SessionCustomization[]> {
-		return this._plugins.getSessionCustomizations(await this._getSessionCustomizationDirectory(session));
+		return this._plugins.getSessionCustomizationsSettled(await this._getSessionCustomizationDirectory(session));
 	}
 
 	private async _getSessionCustomizationDirectory(session: URI): Promise<URI | undefined> {
@@ -1949,6 +1949,32 @@ class PluginController extends Disposable {
 		return result;
 	}
 
+	/**
+	 * Settled variant of {@link getSessionCustomizations}: awaits the
+	 * in-flight host sync, the in-flight client sync, and (when a directory
+	 * is supplied) the session-discovered entry's initial scan + parse
+	 * before snapshotting the customization list.
+	 *
+	 * Callers that publish customizations into session state at session
+	 * creation time MUST use this — the synchronous variant can return an
+	 * empty list for a brand-new working directory because
+	 * {@link SessionDiscoveredEntry} kicks off its `_refresh()` in its
+	 * constructor without anyone awaiting it.
+	 */
+	public async getSessionCustomizationsSettled(directory: URI | undefined): Promise<readonly SessionCustomization[]> {
+		const entry = directory ? this._getOrCreateSessionEntry(directory) : undefined;
+		await Promise.all([
+			this._hostSync.catch(err => {
+				this._logService.warn('[Copilot:PluginController] Host customization update failed', err);
+			}),
+			this._clientSync.catch(err => {
+				this._logService.warn('[Copilot:PluginController] Customization sync failed', err);
+			}),
+			entry?.whenSettled(),
+		]);
+		return this.getSessionCustomizations(directory);
+	}
+
 	/**
 	 * Returns the current parsed plugins, awaiting any pending sync.
 	 */

src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts

@@ -23,6 +23,7 @@ import { ResolveSessionConfigResult } from '../../../../../platform/agentHost/co
 import { AgentSelection, CustomizationAgentRef, ModelSelection, SessionStatus as ProtocolSessionStatus, RootConfigState, RootState, SessionState, SessionSummary, type ChangesetSummary } from '../../../../../platform/agentHost/common/state/protocol/state.js';
 import { ActionType, isSessionAction, NotificationType } from '../../../../../platform/agentHost/common/state/sessionActions.js';
 import { readSessionGitState, ROOT_STATE_URI, SessionMeta, StateComponents, type ISessionGitState } from '../../../../../platform/agentHost/common/state/sessionState.js';
+import type { IAgentSubscription } from '../../../../../platform/agentHost/common/state/agentSubscription.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';
@@ -469,6 +470,16 @@ interface INewSessionConstructionContext {
 	 * list as the committed session that replaces it.
 	 */
 	readonly instantiationService: IInstantiationService;
+	/**
+	 * Forwards `SessionState` snapshots from the eagerly-held wire
+	 * subscription back to the provider. `state === undefined` is a
+	 * cleanup sentinel emitted by {@link NewSession.dispose} on the
+	 * close-without-graduation path so the provider can drop any cached
+	 * entry it accumulated for this session. The graduation path skips
+	 * this sentinel because the running-session subscription pipeline
+	 * takes over ownership of the same `sessionId` key.
+	 */
+	readonly onSessionState?: (sessionId: string, state: SessionState | undefined) => void;
 }
 
 /**
@@ -526,7 +537,16 @@ class NewSession extends Disposable {
 	/** Connection used to create the backend session, captured for `disposeSession` on tear-down. */
 	private _connection: IAgentConnection | undefined;
 	/** Held state subscription. Set after the wire `createSession` resolves. */
-	private _subscription: IReference<unknown> | undefined;
+	private _subscription: IReference<IAgentSubscription<SessionState>> | undefined;
+	/**
+	 * `onDidChange` listener for {@link _subscription}. Forwards every
+	 * `SessionState` snapshot to the provider via {@link _onSessionState}
+	 * so the new session's customizations (and any other state) reach
+	 * `_lastSessionStates` while the session is still Untitled. Detached
+	 * in {@link graduate} (handoff) and {@link dispose} (close-without-send).
+	 */
+	private readonly _stateListener = this._register(new MutableDisposable());
+	private readonly _onSessionState: ((sessionId: string, state: SessionState | undefined) => void) | undefined;
 
 	private readonly _logService: ILogService;
 	private readonly _providerId: string;
@@ -541,6 +561,7 @@ class NewSession extends Disposable {
 		this.agentProvider = ctx.sessionType.id;
 		this._providerId = ctx.providerId;
 		this._logService = ctx.logService;
+		this._onSessionState = ctx.onSessionState;
 
 		const resource = URI.from({ scheme: ctx.resourceScheme, path: `/${generateUuid()}` });
 		this._status = observableValue<SessionStatus>(this, SessionStatus.Untitled);
@@ -742,7 +763,23 @@ class NewSession extends Disposable {
 			// handler refcounts the same subscription via `getSubscription`
 			// when chat content opens, so when we release this ref on
 			// graduation the wire-level refcount stays positive.
-			this._subscription = connection.getSubscription(StateComponents.Session, backendUri);
+			const ref = connection.getSubscription(StateComponents.Session, backendUri);
+			this._subscription = ref;
+
+			// Forward `SessionState` updates back to the provider so
+			// `_lastSessionStates` (and therefore `getCustomAgents`) becomes
+			// populated for this still-Untitled session. Seed once from the
+			// cached value, then attach a listener for subsequent deltas.
+			const onSessionState = this._onSessionState;
+			if (onSessionState) {
+				const initial = ref.object.value;
+				if (initial && !(initial instanceof Error)) {
+					onSessionState(this.sessionId, initial);
+				}
+				this._stateListener.value = ref.object.onDidChange(state => {
+					onSessionState(this.sessionId, state);
+				});
+			}
 		})();
 	}
 
@@ -752,6 +789,12 @@ class NewSession extends Disposable {
 	 * graduated into a real running session.
 	 */
 	graduate(): void {
+		// Detach the new-session listener BEFORE releasing the subscription.
+		// Both code paths (this one and the running-session pipeline) write
+		// `_lastSessionStates` under the same `sessionId` key, so detaching
+		// here hands ownership cleanly to `_ensureSessionStateSubscription`
+		// without a transient empty-read window or a duplicate writer.
+		this._stateListener.clear();
 		this._subscription?.dispose();
 		this._subscription = undefined;
 		this._backendUri = undefined;
@@ -763,6 +806,18 @@ class NewSession extends Disposable {
 		// Bump the seq so any in-flight resolveConfig discards itself.
 		this._configRequestSeq++;
 
+		// Detach the state listener BEFORE firing the cleanup sentinel so
+		// a racing `onDidChange` cannot re-populate `_lastSessionStates`
+		// after we have asked the provider to delete the entry. Then fire
+		// the sentinel so the provider drops the cached snapshot. Only
+		// fires when a listener was actually wired (i.e. `eagerCreate`
+		// reached the post-`createSession` branch).
+		const hadListener = !!this._stateListener.value;
+		this._stateListener.clear();
+		if (hadListener) {
+			this._onSessionState?.(this.sessionId, undefined);
+		}
+
 		this._subscription?.dispose();
 		this._subscription = undefined;
 
@@ -1120,6 +1175,9 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement
 			logService: this._logService,
 			initialConfigValues: this._initialNewSessionConfig(),
 			instantiationService: this._instantiationService,
+			onSessionState: (id, state) => state === undefined
+				? this._handleNewSessionStateGone(id)
+				: this._handleNewSessionStateUpdate(id, state),
 		});
 		this._newSession = newSession;
 		this._onDidChangeSessionConfig.fire(newSession.sessionId);
@@ -1730,6 +1788,34 @@ export abstract class BaseAgentHostSessionsProvider extends Disposable implement
 		this._applySessionMetaFromState(sessionId, state);
 	}
 
+	/**
+	 * NewSession variant of {@link _applySessionStateUpdate}: writes the
+	 * customizations subset (the only one the agent picker reads) and
+	 * fires `_onDidChangeCustomAgents` when it changes. Skips
+	 * {@link _seedRunningConfigFromState} (NewSession owns its own config
+	 * via `NewSession._config`) and {@link _applySessionMetaFromState}
+	 * (which only applies to cached running sessions).
+	 */
+	private _handleNewSessionStateUpdate(sessionId: string, state: SessionState): void {
+		const previous = this._lastSessionStates.get(sessionId);
+		this._lastSessionStates.set(sessionId, state);
+		if (previous?.customizations !== state.customizations || previous?.activeClient?.customizations !== state.activeClient?.customizations) {
+			this._onDidChangeCustomAgents.fire();
+		}
+	}
+
+	/**
+	 * Cleanup sentinel from {@link NewSession.dispose}: drops the cached
+	 * `_lastSessionStates` entry the new session contributed. Fires
+	 * `_onDidChangeCustomAgents` so any open picker re-reads and falls
+	 * back to the empty list rather than rendering stale agents.
+	 */
+	private _handleNewSessionStateGone(sessionId: string): void {
+		if (this._lastSessionStates.delete(sessionId)) {
+			this._onDidChangeCustomAgents.fire();
+		}
+	}
+
 	private _applySessionMetaFromState(sessionId: string, state: SessionState): void {
 		const rawId = this._rawIdFromChatId(sessionId);
 		if (!rawId) {

src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts

@@ -849,6 +849,104 @@ suite('LocalAgentHostSessionsProvider', () => {
 		assert.strictEqual(fired, afterFirstCustomization, 'expected event NOT to fire when customizations are unchanged');
 	});
 
+	test('NewSession forwards SessionState into _lastSessionStates so the picker sees customizations before first message', async () => {
+		const provider = createProvider(disposables, agentHost);
+		const sessionTypeId = provider.sessionTypes[0].id;
+		const session = provider.createNewSession(URI.parse('file:///home/user/proj'), sessionTypeId);
+		await timeout(0); // let eagerCreate complete and the subscription seed
+
+		const rawId = session.resource.path.substring(1);
+
+		let fired = 0;
+		disposables.add(provider.onDidChangeCustomAgents(() => { fired++; }));
+
+		// Push a SessionState carrying customizations as if the host had
+		// resolved them and dispatched a SessionCustomizationsChanged.
+		const customizations = [{
+			customization: { uri: 'plugin://new-session', displayName: 'p' },
+			enabled: true,
+			status: CustomizationStatus.Loaded,
+			agents: [
+				{ uri: 'agent://reviewer', name: 'reviewer' },
+				{ uri: 'agent://triage', name: 'triage' },
+			],
+		}];
+		const state: SessionState = {
+			summary: {
+				resource: AgentSession.uri(sessionTypeId, rawId).toString(),
+				provider: sessionTypeId,
+				title: '',
+				status: ProtocolSessionStatus.Idle,
+				createdAt: 0,
+				modifiedAt: 0,
+			},
+			lifecycle: SessionLifecycle.Ready,
+			turns: [],
+			customizations,
+		};
+		agentHost.setSessionState(rawId, sessionTypeId, state);
+
+		assert.deepStrictEqual(provider.getCustomAgents(session.sessionId), [
+			{ uri: 'agent://reviewer', name: 'reviewer' },
+			{ uri: 'agent://triage', name: 'triage' },
+		]);
+		assert.ok(fired > 0, 'expected onDidChangeCustomAgents to fire when SessionState arrives');
+
+		// A second update with a different customizations identity should
+		// re-fire and update the picker.
+		const after = fired;
+		agentHost.setSessionState(rawId, sessionTypeId, {
+			...state,
+			customizations: [{
+				...customizations[0],
+				agents: [{ uri: 'agent://only', name: 'only' }],
+			}],
+		});
+		assert.deepStrictEqual(provider.getCustomAgents(session.sessionId), [
+			{ uri: 'agent://only', name: 'only' },
+		]);
+		assert.ok(fired > after, 'expected onDidChangeCustomAgents to fire again on a second update');
+	});
+
+	test('NewSession dispose clears _lastSessionStates entry and fires onDidChangeCustomAgents', async () => {
+		const provider = createProvider(disposables, agentHost);
+		const sessionTypeId = provider.sessionTypes[0].id;
+		const first = provider.createNewSession(URI.parse('file:///home/user/a'), sessionTypeId);
+		await timeout(0);
+
+		const rawId = first.resource.path.substring(1);
+		agentHost.setSessionState(rawId, sessionTypeId, {
+			summary: {
+				resource: AgentSession.uri(sessionTypeId, rawId).toString(),
+				provider: sessionTypeId,
+				title: '',
+				status: ProtocolSessionStatus.Idle,
+				createdAt: 0,
+				modifiedAt: 0,
+			},
+			lifecycle: SessionLifecycle.Ready,
+			turns: [],
+			customizations: [{
+				customization: { uri: 'plugin://x', displayName: 'p' },
+				enabled: true,
+				status: CustomizationStatus.Loaded,
+				agents: [{ uri: 'agent://x', name: 'x' }],
+			}],
+		});
+		assert.strictEqual(provider.getCustomAgents(first.sessionId).length, 1);
+
+		let fired = 0;
+		disposables.add(provider.onDidChangeCustomAgents(() => { fired++; }));
+
+		// Trigger disposal of the first NewSession by creating a second one
+		// (the MutableDisposable holding `_newSession` disposes the previous).
+		provider.createNewSession(URI.parse('file:///home/user/b'), sessionTypeId);
+		await timeout(0);
+
+		assert.deepStrictEqual(provider.getCustomAgents(first.sessionId), []);
+		assert.ok(fired > 0, 'expected onDidChangeCustomAgents to fire on NewSession dispose');
+	});
+
 	// ---- Session lifecycle -------
 
 	test('createNewSession returns session with correct fields', () => {

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

@@ -693,6 +693,13 @@ export class AgentHostSessionHandler extends Disposable implements IChatSessionC
 		if (!isNewSession) {
 			this._ensurePendingMessageSubscription(sessionResource, resolvedSession);
 
+			// Claim active client up-front so the host syncs this client's
+			// customizations into session state immediately. Without this,
+			// session-scoped customizations (and the agents derived from
+			// them) wouldn't land until the user sent their first message,
+			// leaving the agent picker empty on session open.
+			this._ensureActiveClientForMessage(resolvedSession);
+
 			// If there are historical turns with file edits, eagerly create
 			// the editing session once the ChatModel is available so that
 			// edit pills render with diff info on session restore.