agent-host: restore selected model in picker

src/vs/sessions/contrib/chat/browser/agentHost/agentHostModelPicker.ts

@@ -17,7 +17,7 @@ import { type ILanguageModelChatMetadataAndIdentifier, ILanguageModelsService }
 import { type IChatInputPickerOptions } from '../../../../../workbench/contrib/chat/browser/widget/input/chatInputPickerActionItem.js';
 import { ModelPickerActionItem, type IModelPickerDelegate } from '../../../../../workbench/contrib/chat/browser/widget/input/modelPickerActionItem.js';
 import { ActiveSessionProviderIdContext, IsPhoneLayoutContext } from '../../../../common/contextkeys.js';
-import { type ISession } from '../../../../services/sessions/common/session.js';
+import { SessionStatus, type ISession } from '../../../../services/sessions/common/session.js';
 import { ISessionsManagementService } from '../../../../services/sessions/common/sessionsManagement.js';
 import { ISessionsProvidersService } from '../../../../services/sessions/browser/sessionsProvidersService.js';
 import { Menus } from '../../../../browser/menus.js';
@@ -52,7 +52,10 @@ registerAction2(class extends Action2 {
 
 // -- Agent Host Model Picker Contribution --
 
-function getAgentHostModels(
+/**
+ * Gets the language models registered for the active agent-host session resource scheme.
+ */
+export function getAgentHostModels(
 	languageModelsService: ILanguageModelsService,
 	session: ISession | undefined,
 ): ILanguageModelChatMetadataAndIdentifier[] {
@@ -72,7 +75,25 @@ function getAgentHostModels(
 		.filter((m): m is ILanguageModelChatMetadataAndIdentifier => !!m && m.metadata.targetChatSessionType === resourceScheme);
 }
 
-const STORAGE_KEY = 'sessions.agentHostModelPicker.selectedModelId';
+export function agentHostModelPickerStorageKey(resourceScheme: string): string {
+	return `workbench.agentsession.agentHostModelPicker.${resourceScheme}.selectedModelId`;
+}
+
+/**
+ * Resolves the model that should be shown for a session.
+ */
+export function resolveAgentHostModel(
+	models: readonly ILanguageModelChatMetadataAndIdentifier[],
+	sessionModelId: string | undefined,
+	storedModelId: string | undefined,
+): ILanguageModelChatMetadataAndIdentifier | undefined {
+	const sessionModel = sessionModelId ? models.find(model => model.identifier === sessionModelId) : undefined;
+	if (sessionModel) {
+		return sessionModel;
+	}
+
+	return storedModelId ? models.find(model => model.identifier === storedModelId) : undefined;
+}
 
 class AgentHostModelPickerContribution extends Disposable implements IWorkbenchContribution {
 
@@ -96,9 +117,9 @@ class AgentHostModelPickerContribution extends Disposable implements IWorkbenchC
 					currentModel,
 					setModel: (model: ILanguageModelChatMetadataAndIdentifier) => {
 						currentModel.set(model, undefined);
-						storageService.store(STORAGE_KEY, model.identifier, StorageScope.PROFILE, StorageTarget.MACHINE);
 						const session = sessionsManagementService.activeSession.get();
 						if (session) {
+							storageService.store(agentHostModelPickerStorageKey(session.resource.scheme), model.identifier, StorageScope.PROFILE, StorageTarget.MACHINE);
 							const provider = sessionsProvidersService.getProviders().find(p => p.id === session.providerId);
 							provider?.setModel(session.sessionId, model.identifier);
 						}
@@ -115,31 +136,27 @@ class AgentHostModelPickerContribution extends Disposable implements IWorkbenchC
 				const action = { id: 'sessions.agentHost.modelPicker', label: '', enabled: true, class: undefined, tooltip: '', run: () => { } };
 				const modelPicker = instantiationService.createInstance(ModelPickerActionItem, action, delegate, pickerOptions);
 
-				const rememberedModelId = storageService.get(STORAGE_KEY, StorageScope.PROFILE);
-				const initModel = (session: ISession | undefined, sessionModelId: string | undefined) => {
+				const initModel = (session: ISession | undefined, sessionModelId: string | undefined, isUntitled: boolean) => {
 					const models = getAgentHostModels(languageModelsService, session);
 					modelPicker.setEnabled(models.length > 0);
 
-					let resolvedModel = sessionModelId
-						? models.find(model => model.identifier === sessionModelId)
-						: undefined;
-
-					// When no model is explicitly selected, restore the
-					// remembered model or pick the first available one so
-					// the picker shows a real model name instead of the
-					// misleading "Auto" label (the copilot "auto"
-					// pseudo-model is not available in agent host sessions).
-					if (!resolvedModel && models.length > 0) {
-						const remembered = rememberedModelId ? models.find(m => m.identifier === rememberedModelId) : undefined;
-						resolvedModel = remembered ?? models[0];
-						delegate.setModel(resolvedModel);
+					if (!session || models.length === 0) {
+						currentModel.set(undefined, undefined);
+						return;
 					}
 
+					const storedModelId = isUntitled
+						? storageService.get(agentHostModelPickerStorageKey(session.resource.scheme), StorageScope.PROFILE)
+						: undefined;
+					const resolvedModel = resolveAgentHostModel(models, sessionModelId, storedModelId);
 					currentModel.set(resolvedModel, undefined);
+					if (!sessionModelId && isUntitled && resolvedModel) {
+						delegate.setModel(resolvedModel);
+					}
 				};
 				const initModelFromActiveSession = () => {
 					const session = sessionsManagementService.activeSession.get();
-					initModel(session, session?.modelId.get());
+					initModel(session, session?.modelId.get(), session?.status.get() === SessionStatus.Untitled);
 				};
 				initModelFromActiveSession();
 
@@ -149,17 +166,8 @@ class AgentHostModelPickerContribution extends Disposable implements IWorkbenchC
 				disposableStore.add(autorun(reader => {
 					const session = sessionsManagementService.activeSession.read(reader);
 					const sessionModelId = session?.modelId.read(reader);
-					initModel(session, sessionModelId);
-				}));
-
-				// When the active session changes, push the selected model to the new session
-				disposableStore.add(autorun(reader => {
-					const session = sessionsManagementService.activeSession.read(reader);
-					const model = currentModel.read(reader);
-					if (session && model) {
-						const provider = sessionsProvidersService.getProviders().find(p => p.id === session.providerId);
-						provider?.setModel(session.sessionId, model.identifier);
-					}
+					const isUntitled = session?.status.read(reader) === SessionStatus.Untitled;
+					initModel(session, sessionModelId, isUntitled);
 				}));
 
 				return new AgentHostPickerActionViewItem(modelPicker, disposableStore);

src/vs/sessions/contrib/chat/test/browser/agentHostModelPicker.test.ts

@@ -0,0 +1,82 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import assert from 'assert';
+import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../base/test/common/utils.js';
+import { ExtensionIdentifier } from '../../../../../platform/extensions/common/extensions.js';
+import type { ILanguageModelChatMetadataAndIdentifier } from '../../../../../workbench/contrib/chat/common/languageModels.js';
+import { agentHostModelPickerStorageKey, resolveAgentHostModel } from '../../browser/agentHost/agentHostModelPicker.js';
+
+function makeModel(identifier: string): ILanguageModelChatMetadataAndIdentifier {
+	return {
+		identifier,
+		metadata: {
+			extension: new ExtensionIdentifier('test.ext'),
+			id: identifier,
+			name: identifier,
+			vendor: 'copilot',
+			version: '1.0',
+			family: 'copilot',
+			maxInputTokens: 128000,
+			maxOutputTokens: 4096,
+			isDefaultForLocation: {},
+			isUserSelectable: true,
+			modelPickerCategory: undefined,
+			targetChatSessionType: 'agent-host-copilotcli',
+		},
+	};
+}
+
+suite('AgentHostModelPicker', () => {
+	ensureNoDisposablesAreLeakedInTestSuite();
+
+	test('uses resource-scheme-scoped storage keys', () => {
+		assert.strictEqual(
+			agentHostModelPickerStorageKey('agent-host-copilotcli'),
+			'workbench.agentsession.agentHostModelPicker.agent-host-copilotcli.selectedModelId',
+		);
+		assert.strictEqual(
+			agentHostModelPickerStorageKey('remote-localhost__4321-copilotcli'),
+			'workbench.agentsession.agentHostModelPicker.remote-localhost__4321-copilotcli.selectedModelId',
+		);
+	});
+
+	test('uses the current session model from session state', () => {
+		const models = [
+			makeModel('agent-host-copilotcli:other'),
+			makeModel('agent-host-copilotcli:session'),
+		];
+
+		assert.strictEqual(
+			resolveAgentHostModel(models, 'agent-host-copilotcli:session', 'agent-host-copilotcli:other'),
+			models[1],
+		);
+	});
+
+	test('does not synthesize a model for existing sessions without one in state', () => {
+		const models = [
+			makeModel('agent-host-copilotcli:first'),
+		];
+
+		assert.strictEqual(resolveAgentHostModel(models, undefined, undefined), undefined);
+	});
+
+	test('uses the stored model for new untitled sessions with no model yet', () => {
+		const models = [
+			makeModel('agent-host-copilotcli:first'),
+			makeModel('agent-host-copilotcli:stored'),
+		];
+
+		assert.strictEqual(resolveAgentHostModel(models, undefined, 'agent-host-copilotcli:stored'), models[1]);
+	});
+
+	test('does not fall back to the first model for new untitled sessions without stored state', () => {
+		const models = [
+			makeModel('agent-host-copilotcli:first'),
+		];
+
+		assert.strictEqual(resolveAgentHostModel(models, undefined, undefined), undefined);
+	});
+});