Skip to content

Commit 59dfdc4

Browse files
roblourensroblourens
committed
agentHost: Restore selected model across sessions (Written by Copilot)
1 parent 78e3397 commit 59dfdc4

13 files changed

Lines changed: 308 additions & 41 deletions

File tree

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ export interface IAgentSessionMetadata {
4444
readonly project?: IAgentSessionProjectInfo;
4545
readonly summary?: string;
4646
readonly status?: SessionStatus;
47+
readonly model?: string;
4748
readonly workingDirectory?: URI;
4849
readonly isRead?: boolean;
4950
readonly isDone?: boolean;

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ export class AgentService extends Disposable implements IAgentService {
180180
const withStatus = result.map(s => {
181181
const liveState = this._stateManager.getSessionState(s.session.toString());
182182
if (liveState) {
183-
return { ...s, status: liveState.summary.status };
183+
return { ...s, status: liveState.summary.status, model: liveState.summary.model ?? s.model };
184184
}
185185
return s;
186186
});
@@ -230,6 +230,7 @@ export class AgentService extends Disposable implements IAgentService {
230230
createdAt: Date.now(),
231231
modifiedAt: Date.now(),
232232
...(created.project ? { project: { uri: created.project.uri.toString(), displayName: created.project.displayName } } : {}),
233+
model: config?.model,
233234
workingDirectory: config.workingDirectory?.toString(),
234235
};
235236
const state = this._stateManager.createSession(summary);
@@ -244,6 +245,7 @@ export class AgentService extends Disposable implements IAgentService {
244245
createdAt: Date.now(),
245246
modifiedAt: Date.now(),
246247
...(created.project ? { project: { uri: created.project.uri.toString(), displayName: created.project.displayName } } : {}),
248+
model: config?.model,
247249
workingDirectory: config?.workingDirectory?.toString(),
248250
};
249251
this._stateManager.createSession(summary);
@@ -423,6 +425,7 @@ export class AgentService extends Disposable implements IAgentService {
423425
createdAt: meta.startTime,
424426
modifiedAt: meta.modifiedTime,
425427
...(meta.project ? { project: { uri: meta.project.uri.toString(), displayName: meta.project.displayName } } : {}),
428+
model: meta.model,
426429
workingDirectory: meta.workingDirectory?.toString(),
427430
isRead,
428431
isDone,

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,7 @@ export class CopilotAgent extends Disposable implements IAgent {
166166
const projectByContext = new Map<string, Promise<IAgentSessionProjectInfo | undefined>>();
167167
const result: IAgentSessionMetadata[] = await Promise.all(sessions.map(async s => {
168168
const session = AgentSession.uri(this.id, s.sessionId);
169+
const metadata = await this._readSessionMetadata(session);
169170
let { project, resolved } = await this._readSessionProject(session);
170171
if (!resolved) {
171172
project = await this._resolveSessionProject(s.context, projectLimiter, projectByContext);
@@ -177,6 +178,7 @@ export class CopilotAgent extends Disposable implements IAgent {
177178
modifiedTime: s.modifiedTime.getTime(),
178179
...(project ? { project } : {}),
179180
summary: s.summary,
181+
model: metadata.model,
180182
workingDirectory: typeof s.context?.cwd === 'string' ? URI.file(s.context.cwd) : undefined,
181183
};
182184
}));
@@ -233,7 +235,7 @@ export class CopilotAgent extends Disposable implements IAgent {
233235
const session = agentSession.sessionUri;
234236
this._logService.info(`[Copilot] Forked session created: ${session.toString()}`);
235237
const project = await projectFromCopilotContext({ cwd: config.workingDirectory?.fsPath });
236-
this._storeSessionMetadata(session, undefined, config.workingDirectory, project, true);
238+
this._storeSessionMetadata(session, config.model, config.workingDirectory, project, true);
237239
return { session, ...(project ? { project } : {}) };
238240
});
239241
}
@@ -467,11 +469,12 @@ export class CopilotAgent extends Disposable implements IAgent {
467469
const parsedPlugins = await this._plugins.getAppliedPlugins();
468470

469471
const sessionUri = AgentSession.uri(this.id, sessionId);
472+
const storedMetadata = await this._readSessionMetadata(sessionUri);
470473
const sessionMetadata = await client.getSessionMetadata(sessionId).catch(err => {
471474
this._logService.warn(`[Copilot:${sessionId}] getSessionMetadata failed`, err);
472475
return undefined;
473476
});
474-
const workingDirectory = typeof sessionMetadata?.context?.cwd === 'string' ? URI.file(sessionMetadata.context.cwd) : undefined;
477+
const workingDirectory = typeof sessionMetadata?.context?.cwd === 'string' ? URI.file(sessionMetadata.context.cwd) : storedMetadata.workingDirectory;
475478
const shellManager = this._instantiationService.createInstance(ShellManager, sessionUri);
476479
const sessionConfig = this._buildSessionConfig(parsedPlugins, shellManager);
477480

@@ -492,13 +495,12 @@ export class CopilotAgent extends Disposable implements IAgent {
492495
}
493496

494497
this._logService.warn(`[Copilot:${sessionId}] Resume failed (session not found in SDK), recreating`);
495-
const metadata = await this._readSessionMetadata(sessionUri);
496498
const raw = await client.createSession({
497499
...config,
498500
sessionId,
499501
streaming: true,
500-
model: metadata.model,
501-
workingDirectory: metadata.workingDirectory?.fsPath,
502+
model: storedMetadata.model,
503+
workingDirectory: workingDirectory?.fsPath,
502504
});
503505

504506
return new CopilotSessionWrapper(raw);

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -386,6 +386,7 @@ export class ProtocolServerHandler extends Disposable {
386386
createdAt: s.startTime,
387387
modifiedAt: s.modifiedTime,
388388
...(s.project ? { project: { uri: s.project.uri.toString(), displayName: s.project.displayName } } : {}),
389+
model: s.model,
389390
workingDirectory: s.workingDirectory?.toString(),
390391
isRead: s.isRead,
391392
isDone: s.isDone,

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

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import assert from 'assert';
77
import { timeout } from '../../../../../base/common/async.js';
88
import { ISubscribeResult } from '../../../common/state/protocol/commands.js';
9-
import type { IResponsePartAction, ISessionAddedNotification, ITitleChangedAction } from '../../../common/state/sessionActions.js';
9+
import type { IModelChangedAction, IResponsePartAction, ISessionAddedNotification, ITitleChangedAction } from '../../../common/state/sessionActions.js';
1010
import { PROTOCOL_VERSION } from '../../../common/state/sessionCapabilities.js';
1111
import type { IListSessionsResult, INotificationBroadcastParams } from '../../../common/state/sessionProtocol.js';
1212
import { PendingMessageKind, ResponsePartKind, type ISessionState } from '../../../common/state/sessionState.js';
@@ -118,6 +118,51 @@ suite('Protocol WebSocket — Session Features', function () {
118118
assert.strictEqual(session.title, 'Persisted Title');
119119
});
120120

121+
// ---- Session model --------------------------------------------------------
122+
123+
test('session model flows through create, subscribe, listSessions, and modelChanged', async function () {
124+
this.timeout(10_000);
125+
126+
await client.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-model-summary' });
127+
128+
const sessionUri = nextSessionUri();
129+
await client.call('createSession', { session: sessionUri, provider: 'mock', model: 'mock-model' });
130+
131+
const addedNotif = await client.waitForNotification(n =>
132+
n.method === 'notification' && (n.params as INotificationBroadcastParams).notification.type === 'notify/sessionAdded'
133+
);
134+
const addedSession = (addedNotif.params as INotificationBroadcastParams).notification as ISessionAddedNotification;
135+
assert.strictEqual(addedSession.summary.model, 'mock-model');
136+
const createdSessionUri = addedSession.summary.resource;
137+
138+
const initialSnapshot = await client.call<ISubscribeResult>('subscribe', { resource: createdSessionUri });
139+
const initialState = initialSnapshot.snapshot.state as ISessionState;
140+
assert.strictEqual(initialState.summary.model, 'mock-model');
141+
142+
const initialList = await client.call<IListSessionsResult>('listSessions');
143+
assert.strictEqual(initialList.items.find(s => s.resource === createdSessionUri)?.model, 'mock-model');
144+
145+
client.notify('dispatchAction', {
146+
clientSeq: 1,
147+
action: {
148+
type: 'session/modelChanged',
149+
session: createdSessionUri,
150+
model: 'mock-model-2',
151+
},
152+
});
153+
154+
const modelNotif = await client.waitForNotification(n => isActionNotification(n, 'session/modelChanged'));
155+
const modelAction = getActionEnvelope(modelNotif).action as IModelChangedAction;
156+
assert.strictEqual(modelAction.model, 'mock-model-2');
157+
158+
const updatedSnapshot = await client.call<ISubscribeResult>('subscribe', { resource: createdSessionUri });
159+
const updatedState = updatedSnapshot.snapshot.state as ISessionState;
160+
assert.strictEqual(updatedState.summary.model, 'mock-model-2');
161+
162+
const updatedList = await client.call<IListSessionsResult>('listSessions');
163+
assert.strictEqual(updatedList.items.find(s => s.resource === createdSessionUri)?.model, 'mock-model-2');
164+
});
165+
121166
// ---- Reasoning events ------------------------------------------------------
122167

123168
test('reasoning events produce reasoning response parts and append actions', async function () {

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

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ import { IChatInputPickerOptions } from '../../../../workbench/contrib/chat/brow
2020
import { EnhancedModelPickerActionItem } from '../../../../workbench/contrib/chat/browser/widget/input/modelPickerActionItem2.js';
2121
import { HoverPosition } from '../../../../base/browser/ui/hover/hoverWidget.js';
2222
import { IContextKeyService, ContextKeyExpr } from '../../../../platform/contextkey/common/contextkey.js';
23-
import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
2423
import { Menus } from '../../../browser/menus.js';
2524
import { ISessionsManagementService } from '../../../services/sessions/common/sessionsManagement.js';
2625
import { ISessionsProvidersService } from '../../../services/sessions/browser/sessionsProvidersService.js';
@@ -187,7 +186,6 @@ class CopilotPickerActionViewItemContribution extends Disposable implements IWor
187186
@ILanguageModelsService languageModelsService: ILanguageModelsService,
188187
@ISessionsManagementService sessionsManagementService: ISessionsManagementService,
189188
@ISessionsProvidersService sessionsProvidersService: ISessionsProvidersService,
190-
@IStorageService storageService: IStorageService,
191189
) {
192190
super();
193191

@@ -219,8 +217,6 @@ class CopilotPickerActionViewItemContribution extends Disposable implements IWor
219217
const delegate: IModelPickerDelegate = {
220218
currentModel,
221219
setModel: (model: ILanguageModelChatMetadataAndIdentifier) => {
222-
currentModel.set(model, undefined);
223-
storageService.store('sessions.localModelPicker.selectedModelId', model.identifier, StorageScope.PROFILE, StorageTarget.MACHINE);
224220
const session = sessionsManagementService.activeSession.get();
225221
if (session) {
226222
const provider = sessionsProvidersService.getProviders().find(p => p.id === session.providerId);
@@ -240,29 +236,24 @@ class CopilotPickerActionViewItemContribution extends Disposable implements IWor
240236
const action = { id: 'sessions.modelPicker', label: '', enabled: true, class: undefined, tooltip: '', run: () => { } };
241237
const modelPicker = instantiationService.createInstance(EnhancedModelPickerActionItem, action, delegate, pickerOptions);
242238

243-
// Initialize with remembered model or first available model
244-
const rememberedModelId = storageService.get('sessions.localModelPicker.selectedModelId', StorageScope.PROFILE);
245-
const initModel = () => {
239+
const updatePickerModel = (session: ISession | undefined, sessionModelId: string | undefined) => {
246240
const models = getAvailableModels(languageModelsService, sessionsManagementService);
247241
modelPicker.setEnabled(models.length > 0);
248-
if (!currentModel.get() && models.length > 0) {
249-
const remembered = rememberedModelId ? models.find(m => m.identifier === rememberedModelId) : undefined;
250-
delegate.setModel(remembered ?? models[0]);
251-
}
242+
currentModel.set(sessionModelId ? models.find(model => model.identifier === sessionModelId) : undefined, undefined);
252243
};
253-
initModel();
244+
const updatePickerModelFromActiveSession = () => {
245+
const session = sessionsManagementService.activeSession.get();
246+
updatePickerModel(session, session?.modelId.get());
247+
};
248+
updatePickerModelFromActiveSession();
254249

255250
const disposableStore = new DisposableStore();
256-
disposableStore.add(languageModelsService.onDidChangeLanguageModels(() => initModel()));
251+
disposableStore.add(languageModelsService.onDidChangeLanguageModels(() => updatePickerModelFromActiveSession()));
257252

258-
// When the active session changes, push the selected model to the new session
259253
disposableStore.add(autorun(reader => {
260254
const session = sessionsManagementService.activeSession.read(reader);
261-
const model = currentModel.read(reader);
262-
if (session && model) {
263-
const provider = sessionsProvidersService.getProviders().find(p => p.id === session.providerId);
264-
provider?.setModel(session.sessionId, model.identifier);
265-
}
255+
const sessionModelId = session?.modelId.read(reader);
256+
updatePickerModel(session, sessionModelId);
266257
}));
267258

268259
return new PickerActionViewItem(modelPicker, disposableStore);

‎src/vs/sessions/contrib/localAgentHost/browser/localAgentHostSessionsProvider.ts‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ class LocalSessionAdapter implements ISession {
6060
readonly updatedAt: ISettableObservable<Date>;
6161
readonly status = observableValue<SessionStatus>('status', SessionStatus.Completed);
6262
readonly changes = observableValue<readonly IChatSessionFileChange[]>('changes', []);
63-
readonly modelId = observableValue<string | undefined>('modelId', undefined);
63+
readonly modelId: ISettableObservable<string | undefined>;
6464
readonly mode = observableValue<{ readonly id: string; readonly kind: string } | undefined>('mode', undefined);
6565
readonly loading = observableValue('loading', false);
6666
readonly isArchived = observableValue('isArchived', false);
@@ -89,6 +89,7 @@ class LocalSessionAdapter implements ISession {
8989
this.createdAt = new Date(metadata.startTime);
9090
this.title = observableValue('title', metadata.summary ?? `Session ${rawId.substring(0, 8)}`);
9191
this.updatedAt = observableValue('updatedAt', new Date(metadata.modifiedTime));
92+
this.modelId = observableValue<string | undefined>('modelId', metadata.model ? `${logicalSessionType}:${metadata.model}` : undefined);
9293
this.lastTurnEnd = observableValue('lastTurnEnd', metadata.modifiedTime ? new Date(metadata.modifiedTime) : undefined);
9394
this.description = observableValue('description', new MarkdownString().appendText(localize('localAgentHostDescription', "Local")));
9495
this.workspace = observableValue('workspace', LocalAgentHostSessionsProvider.buildWorkspace(metadata.project, metadata.workingDirectory));
@@ -155,6 +156,12 @@ class LocalSessionAdapter implements ISession {
155156
didChange = true;
156157
}
157158

159+
const modelId = metadata.model ? `${this.sessionType}:${metadata.model}` : undefined;
160+
if (modelId !== this.modelId.get()) {
161+
this.modelId.set(modelId, undefined);
162+
didChange = true;
163+
}
164+
158165
return didChange;
159166
}
160167
}
@@ -209,6 +216,7 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
209216
private _selectedModelId: string | undefined;
210217
private _currentNewSession: ISession | undefined;
211218
private _currentNewSessionStatus: ISettableObservable<SessionStatus> | undefined;
219+
private _currentNewSessionModelId: ISettableObservable<string | undefined> | undefined;
212220

213221
private _cacheInitialized = false;
214222

@@ -245,6 +253,8 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
245253
this._refreshSessions();
246254
} else if (e.action.type === ActionType.SessionTitleChanged && isSessionAction(e.action)) {
247255
this._handleTitleChanged(e.action.session, e.action.title);
256+
} else if (e.action.type === ActionType.SessionModelChanged && isSessionAction(e.action)) {
257+
this._handleModelChanged(e.action.session, e.action.model);
248258
} else if (e.action.type === ActionType.SessionIsReadChanged && isSessionAction(e.action)) {
249259
this._handleIsReadChanged(e.action.session, e.action.isRead);
250260
} else if (e.action.type === ActionType.SessionIsDoneChanged && isSessionAction(e.action)) {
@@ -376,6 +386,7 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
376386

377387
this._currentNewSession = undefined;
378388
this._selectedModelId = undefined;
389+
this._currentNewSessionModelId = undefined;
379390

380391
const defaultType = this.sessionTypes[0];
381392
if (!defaultType) {
@@ -429,6 +440,7 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
429440
};
430441
this._currentNewSession = session;
431442
this._currentNewSessionStatus = status;
443+
this._currentNewSessionModelId = modelId;
432444
return session;
433445
}
434446

@@ -458,6 +470,18 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
458470
setModel(sessionId: string, modelId: string): void {
459471
if (this._currentNewSession?.sessionId === sessionId) {
460472
this._selectedModelId = modelId;
473+
this._currentNewSessionModelId?.set(modelId, undefined);
474+
return;
475+
}
476+
477+
const rawId = this._rawIdFromChatId(sessionId);
478+
const cached = rawId ? this._sessionCache.get(rawId) : undefined;
479+
if (cached && rawId) {
480+
cached.modelId.set(modelId, undefined);
481+
this._onDidChangeSessions.fire({ added: [], removed: [], changed: [cached] });
482+
const rawModelId = modelId.startsWith(`${cached.sessionType}:`) ? modelId.substring(cached.sessionType.length + 1) : modelId;
483+
const action = { type: ActionType.SessionModelChanged as const, session: AgentSession.uri(cached.agentProvider, rawId).toString(), model: rawModelId };
484+
this._agentHostService.dispatch(action);
461485
}
462486
}
463487

@@ -581,11 +605,13 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
581605

582606
this._selectedModelId = undefined;
583607
this._currentNewSessionStatus = undefined;
608+
this._currentNewSessionModelId = undefined;
584609

585610
try {
586611
const committedSession = await this._waitForNewSession(existingKeys);
587612
if (committedSession) {
588613
this._currentNewSession = undefined;
614+
this._currentNewSessionModelId = undefined;
589615
this._onDidReplaceSession.fire({ from: newSession, to: committedSession });
590616
return committedSession;
591617
}
@@ -596,6 +622,7 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
596622
}
597623

598624
this._currentNewSession = undefined;
625+
this._currentNewSessionModelId = undefined;
599626
return newSession;
600627
}
601628

@@ -676,7 +703,7 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
676703
}
677704
}
678705

679-
private _handleSessionAdded(summary: { resource: string; provider: string; title: string; createdAt: number; modifiedAt: number; project?: { uri: string; displayName: string }; workingDirectory?: string; isRead?: boolean; isDone?: boolean }): void {
706+
private _handleSessionAdded(summary: { resource: string; provider: string; title: string; createdAt: number; modifiedAt: number; project?: { uri: string; displayName: string }; model?: string; workingDirectory?: string; isRead?: boolean; isDone?: boolean }): void {
680707
const sessionUri = URI.parse(summary.resource);
681708
const rawId = AgentSession.id(sessionUri);
682709
if (this._sessionCache.has(rawId)) {
@@ -692,6 +719,7 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
692719
modifiedTime: summary.modifiedAt,
693720
summary: summary.title,
694721
...(summary.project ? { project: { uri: URI.parse(summary.project.uri), displayName: summary.project.displayName } } : {}),
722+
model: summary.model,
695723
workingDirectory: workingDir,
696724
isRead: summary.isRead,
697725
isDone: summary.isDone,
@@ -725,6 +753,16 @@ export class LocalAgentHostSessionsProvider extends Disposable implements ISessi
725753
}
726754
}
727755

756+
private _handleModelChanged(session: string, model: string): void {
757+
const rawId = AgentSession.id(session);
758+
const cached = this._sessionCache.get(rawId);
759+
const modelId = cached ? `${cached.sessionType}:${model}` : undefined;
760+
if (cached && cached.modelId.get() !== modelId) {
761+
cached.modelId.set(modelId, undefined);
762+
this._onDidChangeSessions.fire({ added: [], removed: [], changed: [cached] });
763+
}
764+
}
765+
728766
private _handleIsReadChanged(session: string, isRead: boolean): void {
729767
const rawId = AgentSession.id(session);
730768
const cached = this._sessionCache.get(rawId);

0 commit comments

Comments
 (0)