Fix using dynamic variables with chat input history navigation (#199133)

* Fix using dynamic variables with chat input history navigation and reloading the window
Instead of just saving string history items, saving and restoring a string input plus

* Don't deepClone variable values that contain URIs

* Remove test code

src/vs/workbench/contrib/chat/browser/chatEditor.ts

@@ -18,7 +18,7 @@ import { EditorPane } from 'vs/workbench/browser/parts/editor/editorPane';
 import { IEditorOpenContext } from 'vs/workbench/common/editor';
 import { Memento } from 'vs/workbench/common/memento';
 import { ChatEditorInput } from 'vs/workbench/contrib/chat/browser/chatEditorInput';
-import { IViewState, ChatWidget } from 'vs/workbench/contrib/chat/browser/chatWidget';
+import { IChatViewState, ChatWidget } from 'vs/workbench/contrib/chat/browser/chatWidget';
 import { IChatModel, ISerializableChatData } from 'vs/workbench/contrib/chat/common/chatModel';
 import { clearChatEditor } from 'vs/workbench/contrib/chat/browser/actions/chatClear';
 
@@ -35,7 +35,7 @@ export class ChatEditor extends EditorPane {
 	}
 
 	private _memento: Memento | undefined;
-	private _viewState: IViewState | undefined;
+	private _viewState: IChatViewState | undefined;
 
 	constructor(
 		@ITelemetryService telemetryService: ITelemetryService,
@@ -99,7 +99,7 @@ export class ChatEditor extends EditorPane {
 
 	private updateModel(model: IChatModel): void {
 		this._memento = new Memento('interactive-session-editor-' + model.sessionId, this.storageService);
-		this._viewState = this._memento.getMemento(StorageScope.WORKSPACE, StorageTarget.MACHINE) as IViewState;
+		this._viewState = this._memento.getMemento(StorageScope.WORKSPACE, StorageTarget.MACHINE) as IChatViewState;
 		this.widget.setModel(model, { ...this._viewState });
 	}
 

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

@@ -51,6 +51,9 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	static readonly INPUT_SCHEME = 'chatSessionInput';
 	private static _counter = 0;
 
+	private _onDidLoadInputState = this._register(new Emitter<any>());
+	readonly onDidLoadInputState = this._onDidLoadInputState.event;
+
 	private _onDidChangeHeight = this._register(new Emitter<void>());
 	readonly onDidChangeHeight = this._onDidChangeHeight.event;
 
@@ -78,7 +81,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		return this._inputEditor;
 	}
 
-	private history: HistoryNavigator<string>;
+	private history: HistoryNavigator<{ text: string; state?: any }>;
 	private historyNavigationBackwardsEnablement!: IContextKey<boolean>;
 	private historyNavigationForewardsEnablement!: IContextKey<boolean>;
 	private inputModel: ITextModel | undefined;
@@ -124,7 +127,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 
 	setState(providerId: string, inputValue: string | undefined): void {
 		this.providerId = providerId;
-		const history = this.historyService.getHistory(providerId).map(h => typeof h === 'object' ? (h as any).text : h);
+		const history = this.historyService.getHistory(providerId);
 		this.history = new HistoryNavigator(history, 50);
 
 		if (typeof inputValue === 'string') {
@@ -147,10 +150,11 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	private navigateHistory(previous: boolean): void {
 		const historyInput = (previous ?
 			(this.history.previous() ?? this.history.first()) : this.history.next())
-			?? '';
+			?? { text: '' };
 
-		aria.status(historyInput);
-		this.setValue(historyInput);
+		aria.status(historyInput.text);
+		this.setValue(historyInput.text);
+		this._onDidLoadInputState.fire(historyInput.state);
 		if (previous) {
 			this._inputEditor.setPosition({ lineNumber: 1, column: 1 });
 		} else {
@@ -181,9 +185,9 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	 * Reset the input and update history.
 	 * @param userQuery If provided, this will be added to the history. Followups and programmatic queries should not be passed.
 	 */
-	async acceptInput(userQuery?: string): Promise<void> {
+	async acceptInput(userQuery?: string, inputState?: any): Promise<void> {
 		if (userQuery) {
-			this.history.add(userQuery);
+			this.history.add({ text: userQuery, state: inputState });
 		}
 
 		if (this.accessibilityService.isScreenReaderOptimized() && isMacintosh) {

src/vs/workbench/contrib/chat/browser/chatViewPane.ts

@@ -22,15 +22,15 @@ import { Memento } from 'vs/workbench/common/memento';
 import { SIDE_BAR_FOREGROUND } from 'vs/workbench/common/theme';
 import { IViewDescriptorService } from 'vs/workbench/common/views';
 import { IChatViewPane } from 'vs/workbench/contrib/chat/browser/chat';
-import { IViewState, ChatWidget } from 'vs/workbench/contrib/chat/browser/chatWidget';
+import { IChatViewState, ChatWidget } from 'vs/workbench/contrib/chat/browser/chatWidget';
 import { IChatModel } from 'vs/workbench/contrib/chat/common/chatModel';
 import { IChatService } from 'vs/workbench/contrib/chat/common/chatService';
 
 export interface IChatViewOptions {
 	readonly providerId: string;
 }
 
-interface IViewPaneState extends IViewState {
+interface IViewPaneState extends IChatViewState {
 	sessionId?: string;
 }
 
@@ -43,7 +43,7 @@ export class ChatViewPane extends ViewPane implements IChatViewPane {
 
 	private modelDisposables = this._register(new DisposableStore());
 	private memento: Memento;
-	private viewState: IViewPaneState;
+	private readonly viewState: IViewPaneState;
 	private didProviderRegistrationFail = false;
 
 	constructor(
@@ -199,6 +199,7 @@ export class ChatViewPane extends ViewPane implements IChatViewPane {
 
 			const widgetViewState = this._widget.getViewState();
 			this.viewState.inputValue = widgetViewState.inputValue;
+			this.viewState.inputState = widgetViewState.inputState;
 			this.memento.saveMemento();
 		}
 

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

@@ -39,9 +39,10 @@ function revealLastElement(list: WorkbenchObjectTree<any>) {
 	list.scrollTop = list.scrollHeight - list.renderHeight;
 }
 
-export interface IViewState {
+export type IChatInputState = Record<string, any>;
+export interface IChatViewState {
 	inputValue?: string;
-	// renderData
+	inputState?: IChatInputState;
 }
 
 export interface IChatWidgetStyles {
@@ -53,6 +54,16 @@ export interface IChatWidgetStyles {
 
 export interface IChatWidgetContrib extends IDisposable {
 	readonly id: string;
+
+	/**
+	 * A piece of state which is related to the input editor of the chat widget
+	 */
+	getInputState?(): any;
+
+	/**
+	 * Called with the result of getInputState when navigating input history.
+	 */
+	setInputState?(s: any): void;
 }
 
 export class ChatWidget extends Disposable implements IChatWidget {
@@ -390,6 +401,13 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		}));
 		this.inputPart.render(container, '', this);
 
+		this._register(this.inputPart.onDidLoadInputState(state => {
+			this.contribs.forEach(c => {
+				if (c.setInputState && state[c.id]) {
+					c.setInputState(state[c.id]);
+				}
+			});
+		}));
 		this._register(this.inputPart.onDidFocus(() => this._onDidFocus.fire()));
 		this._register(this.inputPart.onDidAcceptFollowup(e => {
 			if (!this.viewModel) {
@@ -423,7 +441,7 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		this.container.style.setProperty('--vscode-interactive-session-foreground', this.editorOptions.configuration.foreground?.toString() ?? '');
 	}
 
-	setModel(model: IChatModel, viewState: IViewState): void {
+	setModel(model: IChatModel, viewState: IChatViewState): void {
 		if (!this.container) {
 			throw new Error('Call render() before setModel()');
 		}
@@ -447,6 +465,11 @@ export class ChatWidget extends Disposable implements IChatWidget {
 			this.onDidChangeItems();
 		}));
 		this.inputPart.setState(model.providerId, viewState.inputValue);
+		this.contribs.forEach(c => {
+			if (c.setInputState && viewState.inputState?.[c.id]) {
+				c.setInputState(viewState.inputState?.[c.id]);
+			}
+		});
 
 		if (this.tree) {
 			this.onDidChangeItems();
@@ -497,6 +520,16 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		this._acceptInput({ prefix });
 	}
 
+	private collectInputState(): IChatInputState {
+		const inputState: IChatInputState = {};
+		this.contribs.forEach(c => {
+			if (c.getInputState) {
+				inputState[c.id] = c.getInputState();
+			}
+		});
+		return inputState;
+	}
+
 	private async _acceptInput(opts: { query: string } | { prefix: string } | undefined): Promise<void> {
 		if (this.viewModel) {
 			this._onDidAcceptInput.fire();
@@ -510,7 +543,8 @@ export class ChatWidget extends Disposable implements IChatWidget {
 			const result = await this.chatService.sendRequest(this.viewModel.sessionId, input);
 
 			if (result) {
-				this.inputPart.acceptInput(isUserQuery ? input : undefined);
+				const inputState = this.collectInputState();
+				this.inputPart.acceptInput(isUserQuery ? input : undefined, isUserQuery ? inputState : undefined);
 				result.responseCompletePromise.then(async () => {
 					const responses = this.viewModel?.getItems().filter(isResponseVM);
 					const lastResponse = responses?.[responses.length - 1];
@@ -676,9 +710,9 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		this.inputPart.saveState();
 	}
 
-	getViewState(): IViewState {
+	getViewState(): IChatViewState {
 		this.inputPart.saveState();
-		return { inputValue: this.getInput() };
+		return { inputValue: this.getInput(), inputState: this.collectInputState() };
 	}
 }
 

src/vs/workbench/contrib/chat/browser/contrib/chatDynamicVariables.ts

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { coalesce } from 'vs/base/common/arrays';
 import { IMarkdownString, MarkdownString } from 'vs/base/common/htmlContent';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { basename } from 'vs/base/common/resources';
@@ -24,7 +25,7 @@ export const dynamicVariableDecorationType = 'chat-dynamic-variable';
 export class ChatDynamicVariableModel extends Disposable implements IChatWidgetContrib {
 	public static readonly ID = 'chatDynamicVariableModel';
 
-	private readonly _variables: IDynamicVariable[] = [];
+	private _variables: IDynamicVariable[] = [];
 	get variables(): ReadonlyArray<IDynamicVariable> {
 		return [...this._variables];
 	}
@@ -35,37 +36,59 @@ export class ChatDynamicVariableModel extends Disposable implements IChatWidgetC
 
 	constructor(
 		private readonly widget: IChatWidget,
-		@ILabelService private readonly labelService: ILabelService
+		@ILabelService private readonly labelService: ILabelService,
+		@ILogService private readonly logService: ILogService,
 	) {
 		super();
 		this._register(widget.inputEditor.onDidChangeModelContent(e => {
 			e.changes.forEach(c => {
-				this._variables.forEach((ref, i) => {
+				// Don't mutate entries in _variables, since they will be returned from the getter
+				this._variables = coalesce(this._variables.map(ref => {
 					if (Range.areIntersecting(ref.range, c.range)) {
 						// The reference text was changed, it's broken
-						this._variables.splice(i, 1);
+						return null;
 					} else if (Range.compareRangesUsingStarts(ref.range, c.range) > 0) {
 						const delta = c.text.length - c.rangeLength;
-						ref.range = {
-							startLineNumber: ref.range.startLineNumber,
-							startColumn: ref.range.startColumn + delta,
-							endLineNumber: ref.range.endLineNumber,
-							endColumn: ref.range.endColumn + delta
+						return {
+							...ref,
+							range: {
+								startLineNumber: ref.range.startLineNumber,
+								startColumn: ref.range.startColumn + delta,
+								endLineNumber: ref.range.endLineNumber,
+								endColumn: ref.range.endColumn + delta
+							}
 						};
 					}
-				});
+
+					return ref;
+				}));
 			});
 
-			this.updateReferences();
+			this.updateDecorations();
 		}));
 	}
 
+	getInputState(): any {
+		return this.variables;
+	}
+
+	setInputState(s: any): void {
+		if (!Array.isArray(s)) {
+			// Something went wrong
+			this.logService.warn('ChatDynamicVariableModel.setInputState called with invalid state: ' + JSON.stringify(s));
+			return;
+		}
+
+		this._variables = s;
+		this.updateDecorations();
+	}
+
 	addReference(ref: IDynamicVariable): void {
 		this._variables.push(ref);
-		this.updateReferences();
+		this.updateDecorations();
 	}
 
-	private updateReferences(): void {
+	private updateDecorations(): void {
 		this.widget.inputEditor.setDecorationsByType('chat', dynamicVariableDecorationType, this._variables.map(r => (<IDecorationOptions>{
 			range: r.range,
 			hoverMessage: this.getHoverForReference(r)

src/vs/workbench/contrib/chat/common/chatVariables.ts

@@ -52,6 +52,5 @@ export interface IChatVariableResolveResult {
 
 export interface IDynamicVariable {
 	range: IRange;
-	// data: any; // File details for a file, something else for a different type of thing, is it typed?
 	data: IChatRequestVariableValue[];
 }

src/vs/workbench/contrib/chat/common/chatWidgetHistoryService.ts

@@ -8,19 +8,24 @@ import { createDecorator } from 'vs/platform/instantiation/common/instantiation'
 import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage';
 import { Memento } from 'vs/workbench/common/memento';
 
+export interface IChatHistoryEntry {
+	text: string;
+	state?: any;
+}
+
 export const IChatWidgetHistoryService = createDecorator<IChatWidgetHistoryService>('IChatWidgetHistoryService');
 export interface IChatWidgetHistoryService {
 	_serviceBrand: undefined;
 
 	readonly onDidClearHistory: Event<void>;
 
 	clearHistory(): void;
-	getHistory(providerId: string): string[];
-	saveHistory(providerId: string, history: string[]): void;
+	getHistory(providerId: string): IChatHistoryEntry[];
+	saveHistory(providerId: string, history: IChatHistoryEntry[]): void;
 }
 
 interface IChatHistory {
-	history: { [providerId: string]: string[] };
+	history: { [providerId: string]: IChatHistoryEntry[] };
 }
 
 export class ChatWidgetHistoryService implements IChatWidgetHistoryService {
@@ -36,14 +41,20 @@ export class ChatWidgetHistoryService implements IChatWidgetHistoryService {
 		@IStorageService storageService: IStorageService
 	) {
 		this.memento = new Memento('interactive-session', storageService);
-		this.viewState = this.memento.getMemento(StorageScope.WORKSPACE, StorageTarget.MACHINE) as IChatHistory;
+		const loadedState = this.memento.getMemento(StorageScope.WORKSPACE, StorageTarget.MACHINE) as IChatHistory;
+		for (const provider in loadedState.history) {
+			// Migration from old format
+			loadedState.history[provider] = loadedState.history[provider].map(entry => typeof entry === 'string' ? { text: entry } : entry);
+		}
+
+		this.viewState = loadedState;
 	}
 
-	getHistory(providerId: string): string[] {
+	getHistory(providerId: string): IChatHistoryEntry[] {
 		return this.viewState.history?.[providerId] ?? [];
 	}
 
-	saveHistory(providerId: string, history: string[]): void {
+	saveHistory(providerId: string, history: IChatHistoryEntry[]): void {
 		if (!this.viewState.history) {
 			this.viewState.history = {};
 		}