Make selected chat model sticky (#229094)

microsoft · Sep 19, 2024 · 1bb1b0e · 1bb1b0e
1 parent 17752c8
commit 1bb1b0e
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 30 deletions.
diff --git a/src/vs/workbench/api/common/extHostLanguageModels.ts b/src/vs/workbench/api/common/extHostLanguageModels.ts
@@ -165,7 +165,7 @@ export class ExtHostLanguageModels implements ExtHostLanguageModelsShape {
 				accountLabel: typeof metadata.auth === 'object' ? metadata.auth.label : undefined
 			};
 		}
-		this._proxy.$registerLanguageModelProvider(handle, `${ExtensionIdentifier.toKey(extension.identifier)}/${handle}/${identifier}`, {
+		this._proxy.$registerLanguageModelProvider(handle, `${ExtensionIdentifier.toKey(extension.identifier)}/${identifier}`, {
 			extension: extension.identifier,
 			id: identifier,
 			vendor: metadata.vendor ?? ExtensionIdentifier.toKey(extension.identifier),

diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatMoveActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatMoveActions.ts
@@ -10,7 +10,7 @@ import { ServicesAccessor } from '../../../../../platform/instantiation/common/i
 import { ActiveEditorContext } from '../../../../common/contextkeys.js';
 import { CHAT_CATEGORY, isChatViewTitleActionContext } from './chatActions.js';
 import { CHAT_VIEW_ID, IChatWidgetService } from '../chat.js';
-import { IChatEditorOptions } from '../chatEditor.js';
+import { ChatEditor, IChatEditorOptions } from '../chatEditor.js';
 import { ChatEditorInput } from '../chatEditorInput.js';
 import { ChatViewPane } from '../chatViewPane.js';
 import { CONTEXT_CHAT_ENABLED } from '../../common/chatContextKeys.js';
@@ -118,12 +118,13 @@ async function moveToSidebar(accessor: ServicesAccessor): Promise<void> {
 	const editorService = accessor.get(IEditorService);
 	const editorGroupService = accessor.get(IEditorGroupsService);
 
-	const chatEditorInput = editorService.activeEditor;
+	const chatEditor = editorService.activeEditorPane;
+	const chatEditorInput = chatEditor?.input;
 	let view: ChatViewPane;
-	if (chatEditorInput instanceof ChatEditorInput && chatEditorInput.sessionId) {
-		await editorService.closeEditor({ editor: chatEditorInput, groupId: editorGroupService.activeGroup.id });
+	if (chatEditor instanceof ChatEditor && chatEditorInput instanceof ChatEditorInput && chatEditorInput.sessionId) {
+		await editorService.closeEditor({ editor: chatEditor.input, groupId: editorGroupService.activeGroup.id });
 		view = await viewsService.openView(CHAT_VIEW_ID) as ChatViewPane;
-		view.loadSession(chatEditorInput.sessionId);
+		view.loadSession(chatEditorInput.sessionId, chatEditor.getViewState());
 	} else {
 		view = await viewsService.openView(CHAT_VIEW_ID) as ChatViewPane;
 	}

diff --git a/src/vs/workbench/contrib/chat/browser/chatEditor.ts b/src/vs/workbench/contrib/chat/browser/chatEditor.ts
@@ -122,11 +122,18 @@ export class ChatEditor extends EditorPane {
 
 		if (this._memento && this._viewState) {
 			const widgetViewState = this.widget.getViewState();
+
+			// Need to set props individually on the memento
 			this._viewState.inputValue = widgetViewState.inputValue;
+			this._viewState.selectedLanguageModelId = widgetViewState.selectedLanguageModelId;
 			this._memento.saveMemento();
 		}
 	}
 
+	override getViewState(): object | undefined {
+		return { ...this._viewState };
+	}
+
 	override layout(dimension: dom.Dimension, position?: dom.IDomPosition | undefined): void {
 		if (this.widget) {
 			this.widget.layout(dimension.height, dimension.width);

diff --git a/src/vs/workbench/contrib/chat/browser/chatInputPart.ts b/src/vs/workbench/contrib/chat/browser/chatInputPart.ts
@@ -11,10 +11,10 @@ import * as aria from '../../../../base/browser/ui/aria/aria.js';
 import { Button } from '../../../../base/browser/ui/button/button.js';
 import { IAction } from '../../../../base/common/actions.js';
 import { Codicon } from '../../../../base/common/codicons.js';
-import { Emitter } from '../../../../base/common/event.js';
+import { Emitter, Event } from '../../../../base/common/event.js';
 import { HistoryNavigator2 } from '../../../../base/common/history.js';
 import { KeyCode } from '../../../../base/common/keyCodes.js';
-import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js';
+import { Disposable, DisposableStore, IDisposable, MutableDisposable } from '../../../../base/common/lifecycle.js';
 import { basename, dirname } from '../../../../base/common/path.js';
 import { isMacintosh } from '../../../../base/common/platform.js';
 import { URI } from '../../../../base/common/uri.js';
@@ -59,6 +59,7 @@ import { ILanguageModelChatMetadata, ILanguageModelsService } from '../common/la
 import { CancelAction, ChatModelPickerActionId, ChatSubmitSecondaryAgentAction, IChatExecuteActionContext, SubmitAction } from './actions/chatExecuteActions.js';
 import { IChatWidget } from './chat.js';
 import { ChatFollowups } from './chatFollowups.js';
+import { IChatViewState } from './chatWidget.js';
 
 const $ = dom.$;
 
@@ -143,6 +144,8 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	private chatCursorAtTop: IContextKey<boolean>;
 	private inputEditorHasFocus: IContextKey<boolean>;
 
+	private readonly _waitForPersistedLanguageModel = this._register(new MutableDisposable<IDisposable>());
+	private _onDidChangeCurrentLanguageModel = new Emitter<string>();
 	private _currentLanguageModel: string | undefined;
 	get currentLanguageModel() {
 		return this._currentLanguageModel;
@@ -187,7 +190,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		}));
 	}
 
-	private resetCurrentLanguageModel() {
+	private setCurrentLanguageModelToDefault() {
 		const defaultLanguageModel = this.languageModelsService.getLanguageModelIds().find(id => this.languageModelsService.lookupLanguageModel(id)?.isDefault);
 		const hasUserSelectableLanguageModels = this.languageModelsService.getLanguageModelIds().find(id => {
 			const model = this.languageModelsService.lookupLanguageModel(id);
@@ -196,6 +199,13 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		this._currentLanguageModel = hasUserSelectableLanguageModels ? defaultLanguageModel : undefined;
 	}
 
+	private setCurrentLanguageModelByUser(modelId: string) {
+		this._currentLanguageModel = modelId;
+
+		// The user changed the language model, so we don't wait for the persisted option to be registered
+		this._waitForPersistedLanguageModel.clear();
+	}
+
 	private loadHistory(): HistoryNavigator2<IChatHistoryEntry> {
 		const history = this.historyService.getHistory(this.location);
 		if (history.length === 0) {
@@ -231,13 +241,37 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		}
 	}
 
-	initForNewChatModel(inputValue: string | undefined, inputState: Object): void {
+	initForNewChatModel(state: IChatViewState): void {
 		this.history = this.loadHistory();
-		this.history.add({ text: inputValue ?? this.history.current().text, state: inputState });
+		this.history.add({
+			text: state.inputValue ?? this.history.current().text,
+			state: state.inputState ?? this.getInputState()
+		});
 
-		if (inputValue) {
-			this.setValue(inputValue, false);
+		if (state.inputValue) {
+			this.setValue(state.inputValue, false);
 		}
+
+		if (state.selectedLanguageModelId) {
+			const model = this.languageModelsService.lookupLanguageModel(state.selectedLanguageModelId);
+			if (model) {
+				this._currentLanguageModel = state.selectedLanguageModelId;
+				this._onDidChangeCurrentLanguageModel.fire(this._currentLanguageModel);
+			} else {
+				this._waitForPersistedLanguageModel.value = this.languageModelsService.onDidChangeLanguageModels(e => {
+					const persistedModel = e.added?.find(m => m.identifier === state.selectedLanguageModelId);
+					if (persistedModel) {
+						this._waitForPersistedLanguageModel.clear();
+
+						if (persistedModel.metadata.isUserSelectable) {
+							this._currentLanguageModel = state.selectedLanguageModelId;
+							this._onDidChangeCurrentLanguageModel.fire(this._currentLanguageModel!);
+						}
+					}
+				});
+			}
+		}
+
 	}
 
 	logInputHistory(): void {
@@ -504,12 +538,20 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 					}
 				}
 
-				if (!this._currentLanguageModel) {
-					this.resetCurrentLanguageModel();
-				}
+				if (action.id === ChatModelPickerActionId && action instanceof MenuItemAction) {
+					if (!this._currentLanguageModel) {
+						this.setCurrentLanguageModelToDefault();
+					}
 
-				if (action.id === ChatModelPickerActionId && action instanceof MenuItemAction && this._currentLanguageModel) {
-					return this.instantiationService.createInstance(ModelPickerActionViewItem, action, this._currentLanguageModel, modelId => this._currentLanguageModel = modelId);
+					if (this._currentLanguageModel) {
+						const itemDelegate: ModelPickerDelegate = {
+							onDidChangeModel: this._onDidChangeCurrentLanguageModel.event,
+							setModel: (modelId: string) => {
+								this.setCurrentLanguageModelByUser(modelId);
+							}
+						};
+						return this.instantiationService.createInstance(ModelPickerActionViewItem, action, this._currentLanguageModel, itemDelegate);
+					}
 				}
 
 				return undefined;
@@ -772,11 +814,16 @@ class ChatSubmitDropdownActionItem extends DropdownWithPrimaryActionViewItem {
 	}
 }
 
+interface ModelPickerDelegate {
+	onDidChangeModel: Event<string>;
+	setModel(selectedModelId: string): void;
+}
+
 class ModelPickerActionViewItem extends MenuEntryActionViewItem {
 	constructor(
 		action: MenuItemAction,
 		private currentLanguageModel: string,
-		private readonly setModelDelegate: (selectedModelId: string) => void,
+		private readonly delegate: ModelPickerDelegate,
 		@IKeybindingService keybindingService: IKeybindingService,
 		@INotificationService notificationService: INotificationService,
 		@IContextKeyService contextKeyService: IContextKeyService,
@@ -786,6 +833,11 @@ class ModelPickerActionViewItem extends MenuEntryActionViewItem {
 		@IAccessibilityService _accessibilityService: IAccessibilityService
 	) {
 		super(action, undefined, keybindingService, notificationService, contextKeyService, themeService, contextMenuService, _accessibilityService);
+
+		this._register(delegate.onDidChangeModel(modelId => {
+			this.currentLanguageModel = modelId;
+			this.updateLabel();
+		}));
 	}
 
 	override async onClick(event: MouseEvent): Promise<void> {
@@ -817,7 +869,7 @@ class ModelPickerActionViewItem extends MenuEntryActionViewItem {
 				checked: id === this.currentLanguageModel,
 				run: () => {
 					this.currentLanguageModel = id;
-					this.setModelDelegate(id);
+					this.delegate.setModel(id);
 					this.updateLabel();
 				}
 			};

diff --git a/src/vs/workbench/contrib/chat/browser/chatViewPane.ts b/src/vs/workbench/contrib/chat/browser/chatViewPane.ts
@@ -101,7 +101,7 @@ export class ChatViewPane extends ViewPane {
 		};
 	}
 
-	private updateModel(model?: IChatModel | undefined): void {
+	private updateModel(model?: IChatModel | undefined, viewState?: IChatViewState): void {
 		this.modelDisposables.clear();
 
 		model = model ?? (this.chatService.transferredSessionData?.sessionId
@@ -111,8 +111,12 @@ export class ChatViewPane extends ViewPane {
 			throw new Error('Could not start chat session');
 		}
 
-		this._widget.setModel(model, { ...this.viewState });
+		if (viewState) {
+			this.updateViewState(viewState);
+		}
+
 		this.viewState.sessionId = model.sessionId;
+		this._widget.setModel(model, { ...this.viewState });
 	}
 
 	override shouldShowWelcome(): boolean {
@@ -193,13 +197,13 @@ export class ChatViewPane extends ViewPane {
 		this.updateModel(undefined);
 	}
 
-	loadSession(sessionId: string): void {
+	loadSession(sessionId: string, viewState?: IChatViewState): void {
 		if (this.widget.viewModel) {
 			this.chatService.clearSession(this.widget.viewModel.sessionId);
 		}
 
 		const newModel = this.chatService.getOrRestoreSession(sessionId);
-		this.updateModel(newModel);
+		this.updateModel(newModel, viewState);
 	}
 
 	focusInput(): void {
@@ -229,9 +233,10 @@ export class ChatViewPane extends ViewPane {
 		super.saveState();
 	}
 
-	private updateViewState(): void {
-		const widgetViewState = this._widget.getViewState();
-		this.viewState.inputValue = widgetViewState.inputValue;
-		this.viewState.inputState = widgetViewState.inputState;
+	private updateViewState(viewState?: IChatViewState): void {
+		const newViewState = viewState ?? this._widget.getViewState();
+		this.viewState.inputValue = newViewState.inputValue;
+		this.viewState.inputState = newViewState.inputState;
+		this.viewState.selectedLanguageModelId = newViewState.selectedLanguageModelId;
 	}
 }
diff --git a/src/vs/workbench/contrib/chat/browser/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/chatWidget.ts
@@ -52,6 +52,7 @@ export type IChatInputState = Record<string, any>;
 export interface IChatViewState {
 	inputValue?: string;
 	inputState?: IChatInputState;
+	selectedLanguageModelId?: string;
 }
 
 export interface IChatWidgetStyles {
@@ -709,7 +710,7 @@ export class ChatWidget extends Disposable implements IChatWidget {
 			this.viewModel = undefined;
 			this.onDidChangeItems();
 		}));
-		this.inputPart.initForNewChatModel(viewState.inputValue, viewState.inputState ?? this.collectInputState());
+		this.inputPart.initForNewChatModel(viewState);
 		this.contribs.forEach(c => {
 			if (c.setInputState && viewState.inputState?.[c.id]) {
 				c.setInputState(viewState.inputState?.[c.id]);
@@ -990,7 +991,11 @@ export class ChatWidget extends Disposable implements IChatWidget {
 	}
 
 	getViewState(): IChatViewState {
-		return { inputValue: this.getInput(), inputState: this.collectInputState() };
+		return {
+			inputValue: this.getInput(),
+			inputState: this.collectInputState(),
+			selectedLanguageModelId: this.inputPart.currentLanguageModel
+		};
 	}
 
 	private updateChatInputContext() {