diff --git a/src/vs/platform/agentHost/node/agentService.ts b/src/vs/platform/agentHost/node/agentService.ts index f32c48995c322..89416efd670f4 100644 --- a/src/vs/platform/agentHost/node/agentService.ts +++ b/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; } diff --git a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts index e8546189e922b..a17d394ef628b 100644 --- a/src/vs/platform/agentHost/node/copilot/copilotAgent.ts +++ b/src/vs/platform/agentHost/node/copilot/copilotAgent.ts @@ -313,7 +313,7 @@ export class CopilotAgent extends Disposable implements IAgent { } async getSessionCustomizations(session: URI): Promise { - return this._plugins.getSessionCustomizations(await this._getSessionCustomizationDirectory(session)); + return this._plugins.getSessionCustomizationsSettled(await this._getSessionCustomizationDirectory(session)); } private async _getSessionCustomizationDirectory(session: URI): Promise { @@ -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 { + 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. */ diff --git a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts b/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts index a75464c3d7dfe..60341acaa7161 100644 --- a/src/vs/sessions/contrib/providers/agentHost/browser/baseAgentHostSessionsProvider.ts +++ b/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 | undefined; + private _subscription: IReference> | 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(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) { diff --git a/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts b/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts index 1dfc9ac7636c8..6ea4a91f3f66b 100644 --- a/src/vs/sessions/contrib/providers/agentHost/test/browser/localAgentHostSessionsProvider.test.ts +++ b/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', () => { diff --git a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts b/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts index 0204dd4740bb7..baded1399deb5 100644 --- a/src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostSessionHandler.ts +++ b/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. diff --git a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts index 706862fcc56f0..ab5ba9436077a 100644 --- a/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostChatContribution.test.ts @@ -4426,7 +4426,7 @@ suite('AgentHostChatContribution', () => { assert.strictEqual(ac.activeClient.customizations?.[0].uri, 'file:///plugin-b'); }); - test('does not dispatch activeClientChanged when an existing session is restored', async () => { + test('does not dispatch activeClientChanged when an existing session is restored and this client is already active', async () => { const { instantiationService, agentHostService } = createTestServices(disposables); const sessionResource = AgentSession.uri('copilot', 'existing-session'); const summary: SessionSummary = { @@ -4441,8 +4441,9 @@ suite('AgentHostChatContribution', () => { ...createSessionState(summary), lifecycle: SessionLifecycle.Ready, activeClient: { - clientId: 'other-client', + clientId: agentHostService.clientId, tools: [], + customizations: [], }, }); @@ -4465,8 +4466,8 @@ suite('AgentHostChatContribution', () => { ); }); - test('dispatches activeClientChanged before sending a message when another client is active', async () => { - const { instantiationService, agentHostService, chatAgentService } = createTestServices(disposables); + test('dispatches activeClientChanged when restoring a session where another client is active', async () => { + const { instantiationService, agentHostService } = createTestServices(disposables); const sessionResource = AgentSession.uri('copilot', 'existing-session'); const summary: SessionSummary = { resource: sessionResource.toString(), @@ -4495,17 +4496,16 @@ suite('AgentHostChatContribution', () => { connectionAuthority: 'local', })); - const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, chatAgentService, disposables, { sessionResource }); - fire({ type: 'session/turnComplete', session, turnId } as SessionAction); - await turnPromise; + const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => chatSession.dispose())); const activeClientActions = agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientChanged'); assert.strictEqual(activeClientActions.length, 1); assert.strictEqual(activeClientActions[0].channel, sessionResource.toString()); }); - test('dispatches activeClientChanged before sending a message when current client customizations are stale', async () => { - const { instantiationService, agentHostService, chatAgentService } = createTestServices(disposables); + test('dispatches activeClientChanged when restoring a session where current client customizations are stale', async () => { + const { instantiationService, agentHostService } = createTestServices(disposables); const customizations = observableValue('customizations', [ { uri: 'file:///plugin-new', displayName: 'Plugin New' }, ]); @@ -4539,9 +4539,8 @@ suite('AgentHostChatContribution', () => { customizations, })); - const { turnPromise, session, turnId, fire } = await startTurn(sessionHandler, agentHostService, chatAgentService, disposables, { sessionResource }); - fire({ type: 'session/turnComplete', session, turnId } as SessionAction); - await turnPromise; + const chatSession = await sessionHandler.provideChatSessionContent(sessionResource, CancellationToken.None); + disposables.add(toDisposable(() => chatSession.dispose())); const activeClientActions = agentHostService.dispatchedActions.filter(d => d.action.type === 'session/activeClientChanged'); assert.strictEqual(activeClientActions.length, 1);