From 2f1297db89a664b0675bbf19e4c5bb7a335001a1 Mon Sep 17 00:00:00 2001 From: Don Jayamanne Date: Sun, 15 Feb 2026 18:21:41 +1100 Subject: [PATCH] Add model vendor filtering for chat session model picker (#295124) * Add model vendor filtering for chat session model picker * Udpates * Refactor chat session model handling to use targetChatSessionType for model filtering * Updates * Updates --- .../api/common/extHostLanguageModels.ts | 1 + .../browser/actions/chatExecuteActions.ts | 4 +- .../chatSessions/chatSessions.contribution.ts | 10 ++ .../browser/widget/input/chatInputPart.ts | 112 ++++++++++++++++-- .../chat/common/actions/chatContextKeys.ts | 6 + .../chat/common/chatSessionsService.ts | 6 + .../contrib/chat/common/languageModels.ts | 6 + .../test/common/mockChatSessionsService.ts | 4 + .../vscode.proposed.chatProvider.d.ts | 9 ++ 9 files changed, 150 insertions(+), 8 deletions(-) diff --git a/src/vs/workbench/api/common/extHostLanguageModels.ts b/src/vs/workbench/api/common/extHostLanguageModels.ts index f4383ce38d2..18c2045825c 100644 --- a/src/vs/workbench/api/common/extHostLanguageModels.ts +++ b/src/vs/workbench/api/common/extHostLanguageModels.ts @@ -228,6 +228,7 @@ export class ExtHostLanguageModels implements ExtHostLanguageModelsShape { isDefaultForLocation, isUserSelectable: m.isUserSelectable, statusIcon: m.statusIcon, + targetChatSessionType: m.targetChatSessionType, modelPickerCategory: m.category ?? DEFAULT_MODEL_PICKER_CATEGORY, capabilities: m.capabilities ? { vision: m.capabilities.imageInput, diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts index 2192b382727..bccd2696f44 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts @@ -417,7 +417,9 @@ export class OpenModelPickerAction extends Action2 { group: 'navigation', when: ContextKeyExpr.and( - ChatContextKeys.lockedToCodingAgent.negate(), + ContextKeyExpr.or( + ChatContextKeys.lockedToCodingAgent.negate(), + ChatContextKeys.chatSessionHasTargetedModels), ContextKeyExpr.or( ContextKeyExpr.equals(ChatContextKeys.location.key, ChatAgentLocation.Chat), ContextKeyExpr.equals(ChatContextKeys.location.key, ChatAgentLocation.EditorInline), diff --git a/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts b/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts index 1bb18e32332..cb8b452684a 100644 --- a/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chatSessions/chatSessions.contribution.ts @@ -207,6 +207,11 @@ const extensionPoint = ExtensionsRegistry.registerExtensionPoint; private agentSessionTypeKey: IContextKey; private chatSessionHasCustomAgentTarget: IContextKey; + private chatSessionHasTargetedModels: IContextKey; private modelWidget: ModelPickerActionItem | undefined; private modeWidget: ModePickerActionItem | undefined; private sessionTargetWidget: SessionTypePickerActionItem | undefined; @@ -473,6 +474,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge private _emptyInputState: ObservableMemento; private _chatSessionIsEmpty = false; private _pendingDelegationTarget: AgentSessionProviders | undefined = undefined; + private _currentSessionType: string | undefined = undefined; constructor( // private readonly editorOptions: ChatEditorOptions, // TODO this should be used @@ -576,6 +578,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge } } this.chatSessionHasCustomAgentTarget = ChatContextKeys.chatSessionHasCustomAgentTarget.bindTo(contextKeyService); + this.chatSessionHasTargetedModels = ChatContextKeys.chatSessionHasTargetedModels.bindTo(contextKeyService); this.history = this._register(this.instantiationService.createInstance(ChatHistoryNavigator, this.location)); @@ -618,8 +621,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge // We've changed models and the current one is no longer available. Select a new one const selectedModel = this._currentLanguageModel ? this.getModels().find(m => m.identifier === this._currentLanguageModel.get()?.identifier) : undefined; - const selectedModelNotAvailable = this._currentLanguageModel && (!selectedModel?.metadata.isUserSelectable); - if (!this.currentLanguageModel || selectedModelNotAvailable) { + if (!this.currentLanguageModel || !selectedModel) { this.setCurrentLanguageModelToDefault(); } })); @@ -680,10 +682,18 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge } private getSelectedModelStorageKey(): string { + const sessionType = this._currentSessionType; + if (sessionType && this.hasModelsTargetingSessionType()) { + return `chat.currentLanguageModel.${this.location}.${sessionType}`; + } return `chat.currentLanguageModel.${this.location}`; } private getSelectedModelIsDefaultStorageKey(): string { + const sessionType = this._currentSessionType; + if (sessionType && this.hasModelsTargetingSessionType()) { + return `chat.currentLanguageModel.${this.location}.${sessionType}.isDefault`; + } return `chat.currentLanguageModel.${this.location}.isDefault`; } @@ -922,11 +932,16 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge } } - // Sync selected model + // Sync selected model - validate it belongs to the current session's model pool if (state?.selectedModel) { const lm = this._currentLanguageModel.get(); if (!lm || lm.identifier !== state.selectedModel.identifier) { - this.setCurrentLanguageModel(state.selectedModel); + if (this.isModelValidForCurrentSession(state.selectedModel)) { + this.setCurrentLanguageModel(state.selectedModel); + } else { + // Model from state doesn't belong to this session's pool - use default + this.setCurrentLanguageModelToDefault(); + } } } @@ -992,7 +1007,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge private checkModelSupported(): void { const lm = this._currentLanguageModel.get(); - if (lm && (!this.modelSupportedForDefaultAgent(lm) || !this.modelSupportedForInlineChat(lm))) { + if (lm && (!this.modelSupportedForDefaultAgent(lm) || !this.modelSupportedForInlineChat(lm) || !this.isModelValidForCurrentSession(lm))) { this.setCurrentLanguageModelToDefault(); } } @@ -1049,12 +1064,79 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge this.storageService.store(CachedLanguageModelsKey, models, StorageScope.APPLICATION, StorageTarget.MACHINE); } models.sort((a, b) => a.metadata.name.localeCompare(b.metadata.name)); - return models.filter(entry => entry.metadata?.isUserSelectable && this.modelSupportedForDefaultAgent(entry) && this.modelSupportedForInlineChat(entry)); + + const sessionType = this.getCurrentSessionType(); + if (sessionType) { + // Session has a specific chat session type - show only models that target + // this session type, if any such models exist. + const targeted = models.filter(entry => entry.metadata?.targetChatSessionType === sessionType); + if (targeted.length > 0) { + return targeted; + } + } + + // No session type or no targeted models - show general models (those without + // a targetChatSessionType) filtered by the standard criteria. + return models.filter(entry => !entry.metadata?.targetChatSessionType && entry.metadata?.isUserSelectable && this.modelSupportedForDefaultAgent(entry) && this.modelSupportedForInlineChat(entry)); + } + + /** + * Get the chat session type for the current session, if any. + * Uses the delegate or `getChatSessionFromInternalUri` to determine the session type. + */ + private getCurrentSessionType(): string | undefined { + const delegateSessionType = this.options.sessionTypePickerDelegate?.getActiveSessionProvider?.(); + if (delegateSessionType) { + return delegateSessionType; + } + const sessionResource = this._widget?.viewModel?.model.sessionResource; + const ctx = sessionResource ? this.chatService.getChatSessionFromInternalUri(sessionResource) : undefined; + return ctx?.chatSessionType; + } + + /** + * Check if any registered models target the current session type. + * This is used to set the context key that controls model picker visibility. + */ + private hasModelsTargetingSessionType(): boolean { + const sessionType = this.getCurrentSessionType(); + if (!sessionType) { + return false; + } + return this.languageModelsService.getLanguageModelIds().some(modelId => { + const metadata = this.languageModelsService.lookupLanguageModel(modelId); + return metadata?.targetChatSessionType === sessionType; + }); + } + + /** + * Check if a model is valid for the current session's model pool. + * If the session has targeted models, the model must target this session type. + * If no models target this session, the model must not have a targetChatSessionType. + */ + private isModelValidForCurrentSession(model: ILanguageModelChatMetadataAndIdentifier): boolean { + if (this.hasModelsTargetingSessionType()) { + // Session has targeted models - model must match + return model.metadata.targetChatSessionType === this.getCurrentSessionType(); + } + // No targeted models - model must not be session-specific + return !model.metadata.targetChatSessionType; + } + + /** + * Validate that the current model belongs to the current session's pool. + * Called when switching sessions to prevent cross-contamination. + */ + private checkModelInSessionPool(): void { + const lm = this._currentLanguageModel.get(); + if (lm && !this.isModelValidForCurrentSession(lm)) { + this.setCurrentLanguageModelToDefault(); + } } private setCurrentLanguageModelToDefault() { const allModels = this.getModels(); - const defaultModel = allModels.find(m => m.metadata.isDefaultForLocation[this.location]) || allModels.find(m => m.metadata.isUserSelectable); + const defaultModel = allModels.find(m => m.metadata.isDefaultForLocation[this.location]) || allModels[0]; if (defaultModel) { this.setCurrentLanguageModel(defaultModel); } @@ -1439,6 +1521,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge const customAgentTarget = ctx && this.chatSessionsService.getCustomAgentTargetForSessionType(ctx.chatSessionType); this.chatSessionHasCustomAgentTarget.set(customAgentTarget !== Target.Undefined); + // Check if this session type requires custom models + const requiresCustomModels = ctx && this.chatSessionsService.requiresCustomModelsForSessionType(ctx.chatSessionType); + this.chatSessionHasTargetedModels.set(!!requiresCustomModels); + // Handle agent option from session - set initial mode if (customAgentTarget) { const agentOption = this.chatSessionsService.getSessionOption(ctx.chatSessionResource, agentOptionId); @@ -1761,6 +1847,18 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge this.tryUpdateWidgetController(); this.updateContextUsageWidget(); this.clearQuestionCarousel(); + + // Track the current session type and re-initialize model selection + // when the session type changes (different session types may have + // different model pools via targetChatSessionType). + const newSessionType = this.getCurrentSessionType(); + if (newSessionType !== this._currentSessionType) { + this._currentSessionType = newSessionType; + this.initSelectedModel(); + } + + // Validate that the current model belongs to the new session's pool + this.checkModelInSessionPool(); })); let elements; diff --git a/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts b/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts index c8f138be639..5e45b28b9de 100644 --- a/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts +++ b/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts @@ -58,6 +58,12 @@ export namespace ChatContextKeys { * which means the mode picker should be shown with filtered custom agents. */ export const chatSessionHasCustomAgentTarget = new RawContextKey('chatSessionHasCustomAgentTarget', false, { type: 'boolean', description: localize('chatSessionHasCustomAgentTarget', "True when the chat session has a customAgentTarget defined to filter modes.") }); + /** + * True when the current chat session has models that specifically target it + * via `targetChatSessionType`, which means the model picker should be shown + * even when the widget is locked to a coding agent. + */ + export const chatSessionHasTargetedModels = new RawContextKey('chatSessionHasTargetedModels', false, { type: 'boolean', description: localize('chatSessionHasTargetedModels', "True when the chat session has language models that target it via targetChatSessionType.") }); export const agentSupportsAttachments = new RawContextKey('agentSupportsAttachments', false, { type: 'boolean', description: localize('agentSupportsAttachments', "True when the chat agent supports attachments.") }); export const withinEditSessionDiff = new RawContextKey('withinEditSessionDiff', false, { type: 'boolean', description: localize('withinEditSessionDiff', "True when the chat widget dispatches to the edit session chat.") }); export const filePartOfEditSession = new RawContextKey('filePartOfEditSession', false, { type: 'boolean', description: localize('filePartOfEditSession', "True when the chat widget is within a file with an edit session.") }); diff --git a/src/vs/workbench/contrib/chat/common/chatSessionsService.ts b/src/vs/workbench/contrib/chat/common/chatSessionsService.ts index 6afc7b19a9d..989e005c0b5 100644 --- a/src/vs/workbench/contrib/chat/common/chatSessionsService.ts +++ b/src/vs/workbench/contrib/chat/common/chatSessionsService.ts @@ -85,6 +85,7 @@ export interface IChatSessionsExtensionPoint { readonly capabilities?: IChatAgentAttachmentCapabilities; readonly commands?: IChatSessionCommandContribution[]; readonly canDelegate?: boolean; + readonly isReadOnly?: boolean; /** * When set, the chat session will show a filtered mode picker with custom agents * that have a matching `target` property. This enables contributed chat sessions @@ -92,6 +93,7 @@ export interface IChatSessionsExtensionPoint { * Custom agents without a `target` property are also shown in all filtered lists */ readonly customAgentTarget?: Target; + readonly requiresCustomModels?: boolean; } export interface IChatSessionItem { @@ -273,6 +275,10 @@ export interface IChatSessionsService { */ getCustomAgentTargetForSessionType(chatSessionType: string): Target; + /** + * Returns whether the session type requires custom models. When true, the model picker should show filtered custom models. + */ + requiresCustomModelsForSessionType(chatSessionType: string): boolean; onDidChangeOptionGroups: Event; getOptionGroupsForSessionType(chatSessionType: string): IChatSessionProviderOptionGroup[] | undefined; diff --git a/src/vs/workbench/contrib/chat/common/languageModels.ts b/src/vs/workbench/contrib/chat/common/languageModels.ts index f01c341dd9a..f77edfbc6b4 100644 --- a/src/vs/workbench/contrib/chat/common/languageModels.ts +++ b/src/vs/workbench/contrib/chat/common/languageModels.ts @@ -195,6 +195,12 @@ export interface ILanguageModelChatMetadata { readonly agentMode?: boolean; readonly editTools?: ReadonlyArray; }; + /** + * When set, this model is only shown in the model picker for the specified chat session type. + * Models with this property are excluded from the general model picker and only appear + * when the user is in a session matching this type. + */ + readonly targetChatSessionType?: string; } export namespace ILanguageModelChatMetadata { diff --git a/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts b/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts index 097866e2e4d..932cbb036b0 100644 --- a/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts +++ b/src/vs/workbench/contrib/chat/test/common/mockChatSessionsService.ts @@ -205,6 +205,10 @@ export class MockChatSessionsService implements IChatSessionsService { return this.contributions.find(c => c.type === chatSessionType)?.customAgentTarget ?? Target.Undefined; } + requiresCustomModelsForSessionType(chatSessionType: string): boolean { + return this.contributions.find(c => c.type === chatSessionType)?.requiresCustomModels ?? false; + } + getContentProviderSchemes(): string[] { return Array.from(this.contentProviders.keys()); } diff --git a/src/vscode-dts/vscode.proposed.chatProvider.d.ts b/src/vscode-dts/vscode.proposed.chatProvider.d.ts index 48a3f1048fe..1ded6ac9ba7 100644 --- a/src/vscode-dts/vscode.proposed.chatProvider.d.ts +++ b/src/vscode-dts/vscode.proposed.chatProvider.d.ts @@ -65,6 +65,15 @@ declare module 'vscode' { readonly category?: { label: string; order: number }; readonly statusIcon?: ThemeIcon; + + /** + * When set, this model is only shown in the model picker for the specified chat session type. + * Models with this property are excluded from the general model picker and only appear + * when the user is in a session matching this type. + * + * The value must match a `type` declared in a `chatSessions` extension contribution. + */ + readonly targetChatSessionType?: string; } export interface LanguageModelChatCapabilities {