agentHost: Restore selected model across sessions

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

@@ -44,6 +44,7 @@ export interface IAgentSessionMetadata {
 	readonly project?: IAgentSessionProjectInfo;
 	readonly summary?: string;
 	readonly status?: SessionStatus;
+	readonly model?: string;
 	readonly workingDirectory?: URI;
 	readonly isRead?: boolean;
 	readonly isDone?: boolean;

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

@@ -180,7 +180,7 @@ export class AgentService extends Disposable implements IAgentService {
 		const withStatus = result.map(s => {
 			const liveState = this._stateManager.getSessionState(s.session.toString());
 			if (liveState) {
-				return { ...s, status: liveState.summary.status };
+				return { ...s, status: liveState.summary.status, model: liveState.summary.model ?? s.model };
 			}
 			return s;
 		});
@@ -232,6 +232,7 @@ export class AgentService extends Disposable implements IAgentService {
 				createdAt: Date.now(),
 				modifiedAt: Date.now(),
 				...(created.project ? { project: { uri: created.project.uri.toString(), displayName: created.project.displayName } } : {}),
+				model: config?.model,
 				workingDirectory: config.workingDirectory?.toString(),
 			};
 			const state = this._stateManager.createSession(summary);
@@ -247,6 +248,7 @@ export class AgentService extends Disposable implements IAgentService {
 				createdAt: Date.now(),
 				modifiedAt: Date.now(),
 				...(created.project ? { project: { uri: created.project.uri.toString(), displayName: created.project.displayName } } : {}),
+				model: config?.model,
 				workingDirectory: config?.workingDirectory?.toString(),
 			};
 			const state = this._stateManager.createSession(summary);
@@ -462,6 +464,7 @@ export class AgentService extends Disposable implements IAgentService {
 			createdAt: meta.startTime,
 			modifiedAt: meta.modifiedTime,
 			...(meta.project ? { project: { uri: meta.project.uri.toString(), displayName: meta.project.displayName } } : {}),
+			model: meta.model,
 			workingDirectory: meta.workingDirectory?.toString(),
 			isRead,
 			isDone,

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

@@ -191,7 +191,8 @@ export class CopilotAgent extends Disposable implements IAgent {
 		const projectByContext = new Map<string, Promise<IAgentSessionProjectInfo | undefined>>();
 		const result: IAgentSessionMetadata[] = await Promise.all(sessions.map(async s => {
 			const session = AgentSession.uri(this.id, s.sessionId);
-			let { project, resolved } = await this._readSessionProject(session);
+			const metadata = await this._readStoredSessionMetadata(session);
+			let { project, resolved } = metadata;
 			if (!resolved) {
 				project = await this._resolveSessionProject(s.context, projectLimiter, projectByContext);
 				this._storeSessionProjectResolution(session, project);
@@ -202,7 +203,8 @@ export class CopilotAgent extends Disposable implements IAgent {
 				modifiedTime: s.modifiedTime.getTime(),
 				...(project ? { project } : {}),
 				summary: s.summary,
-				workingDirectory: typeof s.context?.cwd === 'string' ? URI.file(s.context.cwd) : undefined,
+				model: metadata.model,
+				workingDirectory: typeof s.context?.cwd === 'string' ? URI.file(s.context.cwd) : metadata.workingDirectory,
 			};
 		}));
 		this._logService.info(`[Copilot] Found ${result.length} sessions`);
@@ -258,7 +260,7 @@ export class CopilotAgent extends Disposable implements IAgent {
 				const session = agentSession.sessionUri;
 				this._logService.info(`[Copilot] Forked session created: ${session.toString()}`);
 				const project = await projectFromCopilotContext({ cwd: config.workingDirectory?.fsPath }, this._gitService);
-				this._storeSessionMetadata(session, undefined, config.workingDirectory, project, true);
+				this._storeSessionMetadata(session, config.model, config.workingDirectory, project, true);
 				return { session, ...(project ? { project } : {}) };
 			});
 		}
@@ -573,11 +575,12 @@ export class CopilotAgent extends Disposable implements IAgent {
 		const parsedPlugins = await this._plugins.getAppliedPlugins();
 
 		const sessionUri = AgentSession.uri(this.id, sessionId);
+		const storedMetadata = await this._readSessionMetadata(sessionUri);
 		const sessionMetadata = await client.getSessionMetadata(sessionId).catch(err => {
 			this._logService.warn(`[Copilot:${sessionId}] getSessionMetadata failed`, err);
 			return undefined;
 		});
-		const workingDirectory = typeof sessionMetadata?.context?.cwd === 'string' ? URI.file(sessionMetadata.context.cwd) : undefined;
+		const workingDirectory = typeof sessionMetadata?.context?.cwd === 'string' ? URI.file(sessionMetadata.context.cwd) : storedMetadata.workingDirectory;
 		const shellManager = this._instantiationService.createInstance(ShellManager, sessionUri);
 		const sessionConfig = this._buildSessionConfig(parsedPlugins, shellManager);
 
@@ -598,13 +601,12 @@ export class CopilotAgent extends Disposable implements IAgent {
 				}
 
 				this._logService.warn(`[Copilot:${sessionId}] Resume failed (session not found in SDK), recreating`);
-				const metadata = await this._readSessionMetadata(sessionUri);
 				const raw = await client.createSession({
 					...config,
 					sessionId,
 					streaming: true,
-					model: metadata.model,
-					workingDirectory: metadata.workingDirectory?.fsPath,
+					model: storedMetadata.model,
+					workingDirectory: workingDirectory?.fsPath,
 				});
 
 				return new CopilotSessionWrapper(raw);
@@ -717,19 +719,21 @@ export class CopilotAgent extends Disposable implements IAgent {
 		}
 	}
 
-	private async _readSessionProject(session: URI): Promise<{ project?: IAgentSessionProjectInfo; resolved: boolean }> {
+	private async _readStoredSessionMetadata(session: URI): Promise<{ model?: string; workingDirectory?: URI; project?: IAgentSessionProjectInfo; resolved: boolean }> {
 		const ref = await this._sessionDataService.tryOpenDatabase(session);
 		if (!ref) {
 			return { resolved: false };
 		}
 		try {
-			const [resolved, uri, displayName] = await Promise.all([
+			const [model, cwd, resolved, uri, displayName] = await Promise.all([
+				ref.object.getMetadata(CopilotAgent._META_MODEL),
+				ref.object.getMetadata(CopilotAgent._META_CWD),
 				ref.object.getMetadata(CopilotAgent._META_PROJECT_RESOLVED),
 				ref.object.getMetadata(CopilotAgent._META_PROJECT_URI),
 				ref.object.getMetadata(CopilotAgent._META_PROJECT_DISPLAY_NAME),
 			]);
 			const project = uri && displayName ? { uri: URI.parse(uri), displayName } : undefined;
-			return { project, resolved: resolved === 'true' || project !== undefined };
+			return { model, workingDirectory: cwd ? URI.parse(cwd) : undefined, project, resolved: resolved === 'true' || project !== undefined };
 		} finally {
 			ref.dispose();
 		}

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

@@ -387,6 +387,7 @@ export class ProtocolServerHandler extends Disposable {
 				createdAt: s.startTime,
 				modifiedAt: s.modifiedTime,
 				...(s.project ? { project: { uri: s.project.uri.toString(), displayName: s.project.displayName } } : {}),
+				model: s.model,
 				workingDirectory: s.workingDirectory?.toString(),
 				isRead: s.isRead,
 				isDone: s.isDone,

src/vs/platform/agentHost/test/node/protocol/sessionFeatures.integrationTest.ts

@@ -6,7 +6,7 @@
 import assert from 'assert';
 import { timeout } from '../../../../../base/common/async.js';
 import { ISubscribeResult } from '../../../common/state/protocol/commands.js';
-import type { IResponsePartAction, ISessionAddedNotification, ITitleChangedAction } from '../../../common/state/sessionActions.js';
+import type { IModelChangedAction, IResponsePartAction, ISessionAddedNotification, ITitleChangedAction } from '../../../common/state/sessionActions.js';
 import { PROTOCOL_VERSION } from '../../../common/state/sessionCapabilities.js';
 import type { IListSessionsResult, INotificationBroadcastParams } from '../../../common/state/sessionProtocol.js';
 import { PendingMessageKind, ResponsePartKind, type ISessionState } from '../../../common/state/sessionState.js';
@@ -118,6 +118,51 @@ suite('Protocol WebSocket — Session Features', function () {
 		assert.strictEqual(session.title, 'Persisted Title');
 	});
 
+	// ---- Session model --------------------------------------------------------
+
+	test('session model flows through create, subscribe, listSessions, and modelChanged', async function () {
+		this.timeout(10_000);
+
+		await client.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-model-summary' });
+
+		const sessionUri = nextSessionUri();
+		await client.call('createSession', { session: sessionUri, provider: 'mock', model: 'mock-model' });
+
+		const addedNotif = await client.waitForNotification(n =>
+			n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionAdded'
+		);
+		const addedSession = (addedNotif.params as INotificationBroadcastParams).notification as ISessionAddedNotification;
+		assert.strictEqual(addedSession.summary.model, 'mock-model');
+		const createdSessionUri = addedSession.summary.resource;
+
+		const initialSnapshot = await client.call<ISubscribeResult>('subscribe', { resource: createdSessionUri });
+		const initialState = initialSnapshot.snapshot.state as ISessionState;
+		assert.strictEqual(initialState.summary.model, 'mock-model');
+
+		const initialList = await client.call<IListSessionsResult>('listSessions');
+		assert.strictEqual(initialList.items.find(s => s.resource === createdSessionUri)?.model, 'mock-model');
+
+		client.notify('dispatchAction', {
+			clientSeq: 1,
+			action: {
+				type: 'session/modelChanged',
+				session: createdSessionUri,
+				model: 'mock-model-2',
+			},
+		});
+
+		const modelNotif = await client.waitForNotification(n => isActionNotification(n, 'session/modelChanged'));
+		const modelAction = getActionEnvelope(modelNotif).action as IModelChangedAction;
+		assert.strictEqual(modelAction.model, 'mock-model-2');
+
+		const updatedSnapshot = await client.call<ISubscribeResult>('subscribe', { resource: createdSessionUri });
+		const updatedState = updatedSnapshot.snapshot.state as ISessionState;
+		assert.strictEqual(updatedState.summary.model, 'mock-model-2');
+
+		const updatedList = await client.call<IListSessionsResult>('listSessions');
+		assert.strictEqual(updatedList.items.find(s => s.resource === createdSessionUri)?.model, 'mock-model-2');
+	});
+
 	// ---- Reasoning events ------------------------------------------------------
 
 	test('reasoning events produce reasoning response parts and append actions', async function () {

src/vs/sessions/contrib/copilotChatSessions/browser/copilotChatSessionsActions.ts

@@ -20,7 +20,6 @@ import { IChatInputPickerOptions } from '../../../../workbench/contrib/chat/brow
 import { EnhancedModelPickerActionItem } from '../../../../workbench/contrib/chat/browser/widget/input/modelPickerActionItem2.js';
 import { HoverPosition } from '../../../../base/browser/ui/hover/hoverWidget.js';
 import { IContextKeyService, ContextKeyExpr } from '../../../../platform/contextkey/common/contextkey.js';
-import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
 import { Menus } from '../../../browser/menus.js';
 import { ISessionsManagementService } from '../../../services/sessions/common/sessionsManagement.js';
 import { ISessionsProvidersService } from '../../../services/sessions/browser/sessionsProvidersService.js';
@@ -187,7 +186,6 @@ class CopilotPickerActionViewItemContribution extends Disposable implements IWor
 		@ILanguageModelsService languageModelsService: ILanguageModelsService,
 		@ISessionsManagementService sessionsManagementService: ISessionsManagementService,
 		@ISessionsProvidersService sessionsProvidersService: ISessionsProvidersService,
-		@IStorageService storageService: IStorageService,
 	) {
 		super();
 
@@ -219,15 +217,13 @@ class CopilotPickerActionViewItemContribution extends Disposable implements IWor
 				const delegate: IModelPickerDelegate = {
 					currentModel,
 					setModel: (model: ILanguageModelChatMetadataAndIdentifier) => {
-						currentModel.set(model, undefined);
-						storageService.store('sessions.localModelPicker.selectedModelId', model.identifier, StorageScope.PROFILE, StorageTarget.MACHINE);
 						const session = sessionsManagementService.activeSession.get();
 						if (session) {
 							const provider = sessionsProvidersService.getProviders().find(p => p.id === session.providerId);
 							provider?.setModel(session.sessionId, model.identifier);
 						}
 					},
-					getModels: () => getAvailableModels(languageModelsService, sessionsManagementService),
+					getModels: () => getAvailableModels(languageModelsService, sessionsManagementService.activeSession.get()),
 					useGroupedModelPicker: () => true,
 					showManageModelsAction: () => false,
 					showUnavailableFeatured: () => false,
@@ -240,29 +236,24 @@ class CopilotPickerActionViewItemContribution extends Disposable implements IWor
 				const action = { id: 'sessions.modelPicker', label: '', enabled: true, class: undefined, tooltip: '', run: () => { } };
 				const modelPicker = instantiationService.createInstance(EnhancedModelPickerActionItem, action, delegate, pickerOptions);
 
-				// Initialize with remembered model or first available model
-				const rememberedModelId = storageService.get('sessions.localModelPicker.selectedModelId', StorageScope.PROFILE);
-				const initModel = () => {
-					const models = getAvailableModels(languageModelsService, sessionsManagementService);
+				const updatePickerModel = (session: ISession | undefined, sessionModelId: string | undefined) => {
+					const models = getAvailableModels(languageModelsService, session);
 					modelPicker.setEnabled(models.length > 0);
-					if (!currentModel.get() && models.length > 0) {
-						const remembered = rememberedModelId ? models.find(m => m.identifier === rememberedModelId) : undefined;
-						delegate.setModel(remembered ?? models[0]);
-					}
+					currentModel.set(sessionModelId ? models.find(model => model.identifier === sessionModelId) : undefined, undefined);
 				};
-				initModel();
+				const updatePickerModelFromActiveSession = () => {
+					const session = sessionsManagementService.activeSession.get();
+					updatePickerModel(session, session?.modelId.get());
+				};
+				updatePickerModelFromActiveSession();
 
 				const disposableStore = new DisposableStore();
-				disposableStore.add(languageModelsService.onDidChangeLanguageModels(() => initModel()));
+				disposableStore.add(languageModelsService.onDidChangeLanguageModels(() => updatePickerModelFromActiveSession()));
 
-				// 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 sessionModelId = session?.modelId.read(reader);
+					updatePickerModel(session, sessionModelId);
 				}));
 
 				return new PickerActionViewItem(modelPicker, disposableStore);
@@ -287,9 +278,8 @@ class CopilotPickerActionViewItemContribution extends Disposable implements IWor
 
 function getAvailableModels(
 	languageModelsService: ILanguageModelsService,
-	sessionsManagementService: ISessionsManagementService,
+	session: ISession | undefined,
 ): ILanguageModelChatMetadataAndIdentifier[] {
-	const session = sessionsManagementService.activeSession.get();
 	if (!session) {
 		return [];
 	}