Skip to content

Commit 9cd2d09

Browse files
roblourensroblourens
committed
Start supporting custom chat modes
1 parent c1ed063 commit 9cd2d09

File tree

12 files changed

+270
-38
lines changed

12 files changed

+270
-38
lines changed

‎src/vs/platform/actions/common/actions.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,7 @@ export class MenuId {
237237
static readonly ChatInput = new MenuId('ChatInput');
238238
static readonly ChatInputSide = new MenuId('ChatInputSide');
239239
static readonly ChatModelPicker = new MenuId('ChatModelPicker');
240+
static readonly ChatModePicker = new MenuId('ChatModePicker');
240241
static readonly ChatEditingWidgetToolbar = new MenuId('ChatEditingWidgetToolbar');
241242
static readonly ChatEditingEditorContent = new MenuId('ChatEditingEditorContent');
242243
static readonly ChatEditingEditorHunk = new MenuId('ChatEditingEditorHunk');

‎src/vs/workbench/contrib/chat/browser/actions/chatExecuteActions.ts

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,10 @@ import { IDialogService } from '../../../../../platform/dialogs/common/dialogs.j
1717
import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
1818
import { KeybindingWeight } from '../../../../../platform/keybinding/common/keybindingsRegistry.js';
1919
import { ChatContextKeys } from '../../common/chatContextKeys.js';
20+
import { ChatMode2, IChatMode, validateChatMode2 } from '../../common/chatModes.js';
2021
import { chatVariableLeader } from '../../common/chatParserTypes.js';
2122
import { IChatService } from '../../common/chatService.js';
22-
import { ChatAgentLocation, ChatConfiguration, ChatMode, validateChatMode } from '../../common/constants.js';
23+
import { ChatAgentLocation, ChatConfiguration, ChatMode, } from '../../common/constants.js';
2324
import { ILanguageModelChatMetadata } from '../../common/languageModels.js';
2425
import { ILanguageModelToolsService } from '../../common/languageModelToolsService.js';
2526
import { IChatWidget, IChatWidgetService } from '../chat.js';
@@ -90,7 +91,7 @@ export class ChatSubmitAction extends SubmitAction {
9091
export const ToggleAgentModeActionId = 'workbench.action.chat.toggleAgentMode';
9192

9293
export interface IToggleChatModeArgs {
93-
mode: ChatMode;
94+
mode: IChatMode | ChatMode;
9495
}
9596

9697
class ToggleChatModeAction extends Action2 {
@@ -142,32 +143,32 @@ class ToggleChatModeAction extends Action2 {
142143
const arg = args.at(0) as IToggleChatModeArgs | undefined;
143144
const chatSession = context.chatWidget.viewModel?.model;
144145
const requestCount = chatSession?.getRequests().length ?? 0;
145-
const switchToMode = validateChatMode(arg?.mode) ?? this.getNextMode(context.chatWidget, requestCount, configurationService);
146+
const switchToMode = validateChatMode2(arg?.mode) ?? this.getNextMode(context.chatWidget, requestCount, configurationService);
146147

147-
if (switchToMode === context.chatWidget.input.currentMode) {
148+
if (switchToMode.id === context.chatWidget.input.currentMode2.id) {
148149
return;
149150
}
150151

151-
const chatModeCheck = await instaService.invokeFunction(handleModeSwitch, context.chatWidget.input.currentMode, switchToMode, requestCount, context.editingSession);
152+
const chatModeCheck = await instaService.invokeFunction(handleModeSwitch, context.chatWidget.input.currentMode, switchToMode.kind, requestCount, context.editingSession);
152153
if (!chatModeCheck) {
153154
return;
154155
}
155156

156-
context.chatWidget.input.setChatMode(switchToMode);
157+
context.chatWidget.input.setChatMode2(switchToMode);
157158

158159
if (chatModeCheck.needToClearSession) {
159160
await commandService.executeCommand(ACTION_ID_NEW_CHAT);
160161
}
161162
}
162163

163-
private getNextMode(chatWidget: IChatWidget, requestCount: number, configurationService: IConfigurationService): ChatMode {
164-
const modes = [ChatMode.Ask];
164+
private getNextMode(chatWidget: IChatWidget, requestCount: number, configurationService: IConfigurationService): IChatMode {
165+
const modes = [ChatMode2.Ask];
165166
if (configurationService.getValue(ChatConfiguration.Edits2Enabled) || requestCount === 0) {
166-
modes.push(ChatMode.Edit);
167+
modes.push(ChatMode2.Edit);
167168
}
168-
modes.push(ChatMode.Agent);
169+
modes.push(ChatMode2.Agent);
169170

170-
const modeIndex = modes.indexOf(chatWidget.input.currentMode);
171+
const modeIndex = modes.findIndex(mode => mode.id === chatWidget.input.currentMode2.id);
171172
const newMode = modes[(modeIndex + 1) % modes.length];
172173
return newMode;
173174
}

‎src/vs/workbench/contrib/chat/browser/actions/promptActions/chatModeActions.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ import { PromptsConfig } from '../../../../../../platform/prompts/common/config.
1212
import { PromptFilePickers } from './dialogs/askToSelectPrompt/promptFilePickers.js';
1313
import { ServicesAccessor } from '../../../../../../editor/browser/editorExtensions.js';
1414
import { ContextKeyExpr } from '../../../../../../platform/contextkey/common/contextkey.js';
15-
import { Action2, registerAction2 } from '../../../../../../platform/actions/common/actions.js';
15+
import { Action2, MenuId, registerAction2 } from '../../../../../../platform/actions/common/actions.js';
1616
import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js';
1717
import { PromptsType } from '../../../../../../platform/prompts/common/prompts.js';
1818
import { IOpenerService } from '../../../../../../platform/opener/common/opener.js';
@@ -27,10 +27,16 @@ class ManageModeAction extends Action2 {
2727
super({
2828
id: MANAGE_CUSTOM_MODE_ACTION_ID,
2929
title: localize2('manage-mode.capitalized', "Manage Custom Chat Modes..."),
30+
shortTitle: localize('manage-mode', "Manage Modes..."),
3031
icon: Codicon.bookmark,
3132
f1: true,
3233
precondition: ContextKeyExpr.and(PromptsConfig.enabledCtx, ChatContextKeys.enabled),
3334
category: CHAT_CATEGORY,
35+
menu: [
36+
{
37+
id: MenuId.ChatModePicker,
38+
}
39+
]
3440
});
3541
}
3642

‎src/vs/workbench/contrib/chat/browser/chat.contribution.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@ import './promptSyntax/contributions/createPromptCommand/createPromptCommand.js'
108108
import { ChatViewsWelcomeHandler } from './viewsWelcome/chatViewsWelcomeHandler.js';
109109
import { registerAction2 } from '../../../../platform/actions/common/actions.js';
110110
import product from '../../../../platform/product/common/product.js';
111+
import { ChatModeService, IChatModeService } from '../common/chatModes.js';
111112

112113
// Register configuration
113114
const configurationRegistry = Registry.as<IConfigurationRegistry>(ConfigurationExtensions.Configuration);
@@ -721,6 +722,7 @@ registerSingleton(ILanguageModelIgnoredFilesService, LanguageModelIgnoredFilesSe
721722
registerSingleton(IChatEntitlementService, ChatEntitlementService, InstantiationType.Delayed);
722723
registerSingleton(IPromptsService, PromptsService, InstantiationType.Delayed);
723724
registerSingleton(IChatContextPickService, ChatContextPickService, InstantiationType.Delayed);
725+
registerSingleton(IChatModeService, ChatModeService, InstantiationType.Delayed);
724726

725727
registerWorkbenchContribution2(ChatEditingNotebookFileSystemProviderContrib.ID, ChatEditingNotebookFileSystemProviderContrib, WorkbenchPhase.BlockStartup);
726728

‎src/vs/workbench/contrib/chat/browser/chatInputPart.ts

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ import { ChatContextKeys } from '../common/chatContextKeys.js';
7474
import { IChatEditingSession } from '../common/chatEditingService.js';
7575
import { ChatEntitlement, IChatEntitlementService } from '../common/chatEntitlementService.js';
7676
import { IChatRequestVariableEntry, isElementVariableEntry, isImageVariableEntry, isNotebookOutputVariableEntry, isPasteVariableEntry, isSCMHistoryItemVariableEntry } from '../common/chatModel.js';
77+
import { IChatMode, ChatMode2, BuiltinChatMode } from '../common/chatModes.js';
7778
import { IChatFollowup } from '../common/chatService.js';
7879
import { IChatVariablesService } from '../common/chatVariables.js';
7980
import { IChatResponseViewModel } from '../common/chatViewModel.js';
@@ -315,11 +316,15 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
315316
private _onDidChangeCurrentChatMode: Emitter<void>;
316317
readonly onDidChangeCurrentChatMode: Event<void>;
317318

318-
private _currentMode: ChatMode;
319+
private _currentMode: IChatMode;
319320
public get currentMode(): ChatMode {
320-
return this._currentMode === ChatMode.Agent && !this.agentService.hasToolsAgent ?
321+
return this._currentMode.kind === ChatMode.Agent && !this.agentService.hasToolsAgent ?
321322
ChatMode.Edit :
322-
this._currentMode;
323+
this._currentMode.kind;
324+
}
325+
326+
public get currentMode2(): IChatMode {
327+
return this._currentMode;
323328
}
324329

325330
private cachedDimensions: dom.Dimension | undefined;
@@ -414,7 +419,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
414419
this._onDidChangeCurrentLanguageModel = this._register(new Emitter<ILanguageModelChatMetadataAndIdentifier>());
415420
this._onDidChangeCurrentChatMode = this._register(new Emitter<void>());
416421
this.onDidChangeCurrentChatMode = this._onDidChangeCurrentChatMode.event;
417-
this._currentMode = ChatMode.Ask;
422+
this._currentMode = ChatMode2.Ask;
418423
this.inputUri = URI.parse(`${ChatInputPart.INPUT_SCHEME}:input-${ChatInputPart._counter++}`);
419424
this._chatEditsActionsDisposables = this._register(new DisposableStore());
420425
this._chatEditsDisposables = this._register(new DisposableStore());
@@ -468,7 +473,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
468473

469474
this.initSelectedModel();
470475

471-
this._register(this.onDidChangeCurrentChatMode(() => this.accessibilityService.alert(this._currentMode)));
476+
this._register(this.onDidChangeCurrentChatMode(() => this.accessibilityService.alert(this._currentMode.kind)));
472477
this._register(this._onDidChangeCurrentLanguageModel.event(() => {
473478
if (this._currentLanguageModel?.metadata.name) {
474479
this.accessibilityService.alert(this._currentLanguageModel.metadata.name);
@@ -557,7 +562,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
557562
}
558563

559564
mode = validateChatMode(mode) ?? ChatMode.Ask;
560-
this._currentMode = mode;
565+
this._currentMode = new BuiltinChatMode(mode);
561566
this.chatMode.set(mode);
562567
this._onDidChangeCurrentChatMode.fire();
563568

@@ -566,6 +571,20 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
566571
}
567572
}
568573

574+
setChatMode2(mode: IChatMode, storeSelection = true): void {
575+
if (!this.options.supportsChangingModes) {
576+
return;
577+
}
578+
579+
this._currentMode = mode;
580+
this.chatMode.set(mode.kind);
581+
this._onDidChangeCurrentChatMode.fire();
582+
583+
if (storeSelection) {
584+
this.storageService.store(GlobalLastChatModeKey, mode, StorageScope.APPLICATION, StorageTarget.USER);
585+
}
586+
}
587+
569588
private modelSupportedForDefaultAgent(model: ILanguageModelChatMetadataAndIdentifier): boolean {
570589
// Probably this logic could live in configuration on the agent, or somewhere else, if it gets more complex
571590
if (this.currentMode === ChatMode.Agent || (this.currentMode === ChatMode.Edit && this.configurationService.getValue(ChatConfiguration.Edits2Enabled))) {
@@ -678,7 +697,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
678697
}
679698
}
680699

681-
if (typeof defaultLanguageModelTreatment === 'string' && this._currentMode === ChatMode.Agent) {
700+
if (typeof defaultLanguageModelTreatment === 'string' && this._currentMode.kind === ChatMode.Agent) {
682701
this.storageService.store(storageKey, true, StorageScope.WORKSPACE, StorageTarget.MACHINE);
683702
this.logService.trace(`Applying default language model from experiment: ${defaultLanguageModelTreatment}`);
684703
this.setExpModelOrWait(defaultLanguageModelTreatment);
@@ -861,7 +880,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
861880
}
862881

863882
validateCurrentMode(): void {
864-
if (!this.agentService.hasToolsAgent && this._currentMode === ChatMode.Agent) {
883+
if (!this.agentService.hasToolsAgent && this._currentMode.kind === ChatMode.Agent) {
865884
this.setChatMode(ChatMode.Edit);
866885
}
867886
}
@@ -1079,7 +1098,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
10791098
}
10801099
} else if (action.id === ToggleAgentModeActionId && action instanceof MenuItemAction) {
10811100
const delegate: IModePickerDelegate = {
1082-
getMode: () => this.currentMode,
1101+
getMode: () => this.currentMode2,
10831102
onDidChangeMode: this._onDidChangeCurrentChatMode.event
10841103
};
10851104
return this.instantiationService.createInstance(ModePickerActionItem, action, delegate);

‎src/vs/workbench/contrib/chat/browser/chatWidget.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1311,7 +1311,9 @@ export class ChatWidget extends Disposable implements IChatWidget {
13111311
this.input.validateCurrentMode();
13121312

13131313
let userSelectedTools: Record<string, boolean> | undefined;
1314-
if (this.input.currentMode === ChatMode.Agent) {
1314+
if (this.input.currentMode2.customTools) {
1315+
userSelectedTools = this.toolsService.toEnablementMap(this.input.currentMode2.customTools);
1316+
} else if (this.input.currentMode === ChatMode.Agent) {
13151317
userSelectedTools = {};
13161318
for (const [tool, enablement] of this.inputPart.selectedToolsModel.asEnablementMap()) {
13171319
userSelectedTools[tool.id] = enablement;

‎src/vs/workbench/contrib/chat/browser/languageModelToolsService.ts

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -455,6 +455,29 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
455455
}
456456
}
457457

458+
toEnablementMap(toolOrToolsetNames: Iterable<string>): Record<string, boolean> {
459+
const toolOrToolset = new Set<string>(toolOrToolsetNames);
460+
const result: Record<string, boolean> = {};
461+
for (const tool of this._tools.values()) {
462+
if (tool.data.toolReferenceName && toolOrToolset.has(tool.data.toolReferenceName) || toolOrToolset.has(tool.data.id)) {
463+
result[tool.data.id] = true;
464+
}
465+
}
466+
467+
for (const toolSet of this._toolSets) {
468+
if (toolOrToolset.has(toolSet.toolReferenceName)) {
469+
result[toolSet.toolReferenceName] = true;
470+
}
471+
for (const tool of toolSet.getTools()) {
472+
if (toolOrToolset.has(tool.id)) {
473+
result[tool.id] = true;
474+
}
475+
}
476+
}
477+
478+
return result;
479+
}
480+
458481
private readonly _toolSets = new ObservableSet<ToolSet>();
459482

460483
readonly toolSets: IObservable<Iterable<ToolSet>> = this._toolSets.observable;

‎src/vs/workbench/contrib/chat/browser/modelPicker/modePickerActionItem.ts

Lines changed: 49 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,22 @@ import { IAction } from '../../../../../base/common/actions.js';
99
import { Event } from '../../../../../base/common/event.js';
1010
import { IDisposable } from '../../../../../base/common/lifecycle.js';
1111
import { ActionWidgetDropdownActionViewItem } from '../../../../../platform/actions/browser/actionWidgetDropdownActionViewItem.js';
12-
import { MenuItemAction } from '../../../../../platform/actions/common/actions.js';
12+
import { getFlatActionBarActions } from '../../../../../platform/actions/browser/menuEntryActionViewItem.js';
13+
import { IMenuService, MenuId, MenuItemAction } from '../../../../../platform/actions/common/actions.js';
1314
import { IActionWidgetService } from '../../../../../platform/actionWidget/browser/actionWidget.js';
14-
import { IActionWidgetDropdownActionProvider, IActionWidgetDropdownOptions } from '../../../../../platform/actionWidget/browser/actionWidgetDropdown.js';
15+
import { IActionWidgetDropdownAction, IActionWidgetDropdownActionProvider, IActionWidgetDropdownOptions } from '../../../../../platform/actionWidget/browser/actionWidgetDropdown.js';
1516
import { IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js';
1617
import { IKeybindingService } from '../../../../../platform/keybinding/common/keybinding.js';
1718
import { IChatAgentService } from '../../common/chatAgents.js';
19+
import { IChatMode, IChatModeService } from '../../common/chatModes.js';
1820
import { ChatAgentLocation, ChatMode, modeToString } from '../../common/constants.js';
21+
import { IPromptsService } from '../../common/promptSyntax/service/types.js';
1922
import { getOpenChatActionIdForMode } from '../actions/chatActions.js';
2023
import { IToggleChatModeArgs } from '../actions/chatExecuteActions.js';
2124

2225
export interface IModePickerDelegate {
2326
onDidChangeMode: Event<void>;
24-
getMode(): ChatMode;
27+
getMode(): IChatMode;
2528
}
2629

2730
export class ModePickerActionItem extends ActionWidgetDropdownActionViewItem {
@@ -31,53 +34,83 @@ export class ModePickerActionItem extends ActionWidgetDropdownActionViewItem {
3134
@IActionWidgetService actionWidgetService: IActionWidgetService,
3235
@IChatAgentService chatAgentService: IChatAgentService,
3336
@IKeybindingService keybindingService: IKeybindingService,
34-
@IContextKeyService contextKeyService: IContextKeyService,
37+
@IContextKeyService private readonly contextKeyService: IContextKeyService,
38+
@IPromptsService promptsService: IPromptsService,
39+
@IChatModeService chatModeService: IChatModeService,
40+
@IMenuService private readonly menuService: IMenuService
3541
) {
36-
const makeAction = (mode: ChatMode): IAction => ({
42+
const makeAction = (mode: ChatMode, includeCategory: boolean): IActionWidgetDropdownAction => ({
3743
...action,
3844
id: getOpenChatActionIdForMode(mode),
3945
label: modeToString(mode),
4046
class: undefined,
4147
enabled: true,
42-
checked: delegate.getMode() === mode,
48+
checked: delegate.getMode().id === mode,
4349
tooltip: chatAgentService.getDefaultAgent(ChatAgentLocation.Panel, mode)?.description ?? action.tooltip,
4450
run: async () => {
4551
const result = await action.run({ mode } satisfies IToggleChatModeArgs);
4652
this.renderLabel(this.element!);
4753
return result;
48-
}
54+
},
55+
category: includeCategory ? { label: 'Standard', order: 0 } : undefined
56+
});
57+
58+
const makeActionFromCustomMode = (mode: IChatMode): IActionWidgetDropdownAction => ({
59+
...action,
60+
id: getOpenChatActionIdForMode(mode.name as ChatMode),
61+
label: mode.name,
62+
class: undefined,
63+
enabled: true,
64+
checked: delegate.getMode().id === mode.id,
65+
tooltip: mode.description ?? chatAgentService.getDefaultAgent(ChatAgentLocation.Panel, mode.kind)?.description ?? action.tooltip,
66+
run: async () => {
67+
const result = await action.run({ mode } satisfies IToggleChatModeArgs);
68+
this.renderLabel(this.element!);
69+
return result;
70+
},
71+
category: { label: 'Custom', order: 1 }
4972
});
5073

5174
const actionProvider: IActionWidgetDropdownActionProvider = {
5275
getActions: () => {
53-
const agentStateActions = [
54-
makeAction(ChatMode.Edit),
55-
];
56-
if (chatAgentService.hasToolsAgent) {
57-
agentStateActions.push(makeAction(ChatMode.Agent));
76+
const modes = chatModeService.getModes();
77+
const hasCustomModes = modes.custom && modes.custom.length > 0;
78+
const agentStateActions: IActionWidgetDropdownAction[] = modes.builtin.map(mode => makeAction(mode.kind, !!hasCustomModes));
79+
if (modes.custom) {
80+
agentStateActions.push(...modes.custom.map(mode => makeActionFromCustomMode(mode)));
5881
}
5982

60-
agentStateActions.unshift(makeAction(ChatMode.Ask));
6183
return agentStateActions;
6284
}
6385
};
6486

65-
const modelPickerActionWidgetOptions: Omit<IActionWidgetDropdownOptions, 'label' | 'labelRenderer'> = {
87+
const modePickerActionWidgetOptions: Omit<IActionWidgetDropdownOptions, 'label' | 'labelRenderer'> = {
6688
actionProvider,
89+
actionBarActionProvider: {
90+
getActions: () => this.getModePickerActionBarActions()
91+
},
6792
showItemKeybindings: true
6893
};
6994

70-
super(action, modelPickerActionWidgetOptions, actionWidgetService, keybindingService, contextKeyService);
95+
super(action, modePickerActionWidgetOptions, actionWidgetService, keybindingService, contextKeyService);
7196

7297
this._register(delegate.onDidChangeMode(() => this.renderLabel(this.element!)));
7398
}
7499

100+
private getModePickerActionBarActions(): IAction[] {
101+
const menuActions = this.menuService.createMenu(MenuId.ChatModePicker, this.contextKeyService);
102+
const menuContributions = getFlatActionBarActions(menuActions.getActions({ renderShortTitle: true }));
103+
menuActions.dispose();
104+
105+
return menuContributions;
106+
}
107+
75108
protected override renderLabel(element: HTMLElement): IDisposable | null {
76109
if (!this.element) {
77110
return null;
78111
}
79112
this.setAriaLabelAttributes(element);
80-
const state = modeToString(this.delegate.getMode());
113+
const state = this.delegate.getMode().name;
81114
dom.reset(element, dom.$('span.chat-model-label', undefined, state), ...renderLabelWithIcons(`$(chevron-down)`));
82115
return null;
83116
}

‎src/vs/workbench/contrib/chat/common/chatContextKeys.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,10 @@ export namespace ChatContextKeys {
7777
export const Tools = {
7878
toolsCount: new RawContextKey<number>('toolsCount', 0, { type: 'number', description: localize('toolsCount', "The count of tools available in the chat.") })
7979
};
80+
81+
export const Modes = {
82+
hasCustomChatModes: new RawContextKey<boolean>('chatHasCustomChatModes', false, { type: 'boolean', description: localize('chatHasCustomChatModes', "True when the chat has custom chat modes available.") }),
83+
};
8084
}
8185

8286
export namespace ChatContextKeyExprs {

0 commit comments

Comments
 (0)