De-contribify the chat attachments (#230451)

* De-contribify the chat attachments
It was done this way originally to match "dynamic variables" and fit the flow of how extra metadata is attached to chat history. But it didn't really make sense for attachments to be considered a contribution given how they're rendered and used now. And I might change this for dynamic variables too anyway. So now we can add some input state that comes directly from the input part and not a contrib.
Fix microsoft/vscode-copilot#8678

* Simplify

src/vs/workbench/contrib/chat/browser/actions/chatClearActions.ts

@@ -20,7 +20,6 @@ import { ChatViewPane } from '../chatViewPane.js';
 import { CONTEXT_IN_CHAT_SESSION, CONTEXT_CHAT_ENABLED, CONTEXT_CHAT_EDITING_PARTICIPANT_REGISTERED } from '../../common/chatContextKeys.js';
 import { IViewsService } from '../../../../services/views/common/viewsService.js';
 import { ChatAgentLocation } from '../../common/chatAgents.js';
-import { ChatContextAttachments } from '../chatWidget.js';
 import { isChatViewTitleActionContext } from '../../common/chatActions.js';
 
 export const ACTION_ID_NEW_CHAT = `workbench.action.chat.newChat`;
@@ -142,7 +141,7 @@ export function registerNewChatActions() {
 				const widget = widgetService.getWidgetBySessionId(context.sessionId);
 				if (widget) {
 					widget.clear();
-					widget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.setContext(true, ...[]);
+					widget.attachmentModel.clear();
 					widget.focusInput();
 				}
 			} else {
@@ -157,7 +156,7 @@ export function registerNewChatActions() {
 
 				announceChatCleared(accessibilitySignalService);
 				widget.clear();
-				widget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.setContext(true, ...[]);
+				widget.attachmentModel.clear();
 				widget.focusInput();
 			}
 		}

src/vs/workbench/contrib/chat/browser/actions/chatContextActions.ts

@@ -26,7 +26,7 @@ import { AnythingQuickAccessProviderRunOptions } from '../../../../../platform/q
 import { IQuickInputService, IQuickPickItem, IQuickPickItemWithResource, IQuickPickSeparator, QuickPickItem } from '../../../../../platform/quickinput/common/quickInput.js';
 import { CHAT_CATEGORY } from './chatActions.js';
 import { IChatWidget, IChatWidgetService, IQuickChatService, showChatView } from '../chat.js';
-import { ChatContextAttachments, isQuickChat } from '../chatWidget.js';
+import { isQuickChat } from '../chatWidget.js';
 import { ChatAgentLocation, IChatAgentService } from '../../common/chatAgents.js';
 import { CONTEXT_CHAT_LOCATION, CONTEXT_IN_CHAT_INPUT } from '../../common/chatContextKeys.js';
 import { IChatEditingService } from '../../common/chatEditingService.js';
@@ -324,7 +324,7 @@ export class AttachContextAction extends Action2 {
 			}
 		}
 
-		widget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.setContext(false, ...toAttach);
+		widget.attachmentModel.addContext(...toAttach);
 	}
 
 	override async run(accessor: ServicesAccessor, ...args: any[]): Promise<void> {
@@ -477,7 +477,7 @@ export class AttachContextAction extends Action2 {
 			additionPicks: quickPickItems,
 			filter: (item: IChatContextQuickPickItem | IQuickPickSeparator) => {
 				// Avoid attaching the same context twice
-				const attachedContext = widget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.getContext() ?? new Set();
+				const attachedContext = widget.attachmentModel.getAttachmentIDs();
 
 				if ('kind' in item && item.kind === 'image') {
 					return !attachedContext.has(item.id);

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

@@ -65,6 +65,7 @@ import { ChatResponseAccessibleView } from './chatResponseAccessibleView.js';
 import { ChatVariablesService } from './chatVariables.js';
 import { ChatWidgetService } from './chatWidget.js';
 import { ChatCodeBlockContextProviderService } from './codeBlockContextProviderService.js';
+import './chatAttachmentModel.js';
 import './contrib/chatInputCompletions.js';
 import './contrib/chatInputEditorContrib.js';
 import './contrib/chatInputEditorHover.js';

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

@@ -22,6 +22,7 @@ import { CHAT_PROVIDER_ID } from '../common/chatParticipantContribTypes.js';
 import { IChatRequestViewModel, IChatResponseViewModel, IChatViewModel } from '../common/chatViewModel.js';
 import { IViewsService } from '../../../services/views/common/viewsService.js';
 import { ChatInputPart } from './chatInputPart.js';
+import { ChatAttachmentModel } from './chatAttachmentModel.js';
 
 export const IChatWidgetService = createDecorator<IChatWidgetService>('chatWidgetService');
 
@@ -162,6 +163,7 @@ export interface IChatWidget {
 	lastSelectedAgent: IChatAgentData | undefined;
 	readonly scopedContextKeyService: IContextKeyService;
 	readonly input: ChatInputPart;
+	readonly attachmentModel: ChatAttachmentModel;
 
 	getContrib<T extends IChatWidgetContrib>(id: string): T | undefined;
 	reveal(item: ChatTreeItem): void;
@@ -182,7 +184,6 @@ export interface IChatWidget {
 	getCodeBlockInfosForResponse(response: IChatResponseViewModel): IChatCodeBlockInfo[];
 	getFileTreeInfosForResponse(response: IChatResponseViewModel): IChatFileTreeInfo[];
 	getLastFocusedFileTreeForResponse(response: IChatResponseViewModel): IChatFileTreeInfo | undefined;
-	setContext(overwrite: boolean, ...context: IChatRequestVariableEntry[]): void;
 	clear(): void;
 	getViewState(): IChatViewState;
 }

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

@@ -0,0 +1,51 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Emitter } from '../../../../base/common/event.js';
+import { Disposable } from '../../../../base/common/lifecycle.js';
+import { IChatRequestVariableEntry } from '../common/chatModel.js';
+
+export class ChatAttachmentModel extends Disposable {
+	private _attachments = new Map<string, IChatRequestVariableEntry>();
+	get attachments(): ReadonlyArray<IChatRequestVariableEntry> {
+		return Array.from(this._attachments.values());
+	}
+
+	private _onDidChangeContext = this._register(new Emitter<void>());
+	readonly onDidChangeContext = this._onDidChangeContext.event;
+
+	get size(): number {
+		return this._attachments.size;
+	}
+
+	getAttachmentIDs() {
+		return new Set(this._attachments.keys());
+	}
+
+	clear(): void {
+		this._attachments.clear();
+		this._onDidChangeContext.fire();
+	}
+
+	delete(variableEntryId: string) {
+		this._attachments.delete(variableEntryId);
+		this._onDidChangeContext.fire();
+	}
+
+	addContext(...attachments: IChatRequestVariableEntry[]) {
+		for (const attachment of attachments) {
+			if (!this._attachments.has(attachment.id)) {
+				this._attachments.set(attachment.id, attachment);
+			}
+		}
+
+		this._onDidChangeContext.fire();
+	}
+
+	clearAndSetContext(...attachments: IChatRequestVariableEntry[]) {
+		this.clear();
+		this.addContext(...attachments);
+	}
+}

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

@@ -95,7 +95,7 @@ export class ChatDragAndDrop extends Themable {
 		e.preventDefault();
 
 		// Make sure to attach only new contexts
-		const currentContextIds = new Set(Array.from(this.inputPart.attachedContext).map(context => context.id));
+		const currentContextIds = this.inputPart.attachmentModel.getAttachmentIDs();
 		const filteredContext = [];
 		for (const context of contexts) {
 			if (!currentContextIds.has(context.id)) {
@@ -104,7 +104,7 @@ export class ChatDragAndDrop extends Themable {
 			}
 		}
 
-		this.inputPart.attachContext(false, ...filteredContext);
+		this.inputPart.attachmentModel.addContext(...filteredContext);
 	}
 
 	private updateDropFeedback(e: DragEvent, dropType: ChatDragAndDropType | undefined): void {

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

@@ -14,7 +14,6 @@ import { EditorActivation } from '../../../../platform/editor/common/editor.js';
 import { IEditorService } from '../../../services/editor/common/editorService.js';
 import { CHAT_EDITING_MULTI_DIFF_SOURCE_RESOLVER_SCHEME, chatEditingResourceContextKey, chatEditingWidgetFileStateContextKey, decidedChatEditingResourceContextKey, IChatEditingService, IChatEditingSession, WorkingSetEntryState } from '../common/chatEditingService.js';
 import { IChatWidget, IChatWidgetService } from './chat.js';
-import { ChatContextAttachments } from './chatWidget.js';
 
 abstract class WorkingSetAction extends Action2 {
 	run(accessor: ServicesAccessor, ...args: any[]) {
@@ -66,13 +65,13 @@ registerAction2(class RemoveFileFromWorkingSet extends WorkingSetAction {
 		const resourceSet = new ResourceSet(uris);
 		const newContext = [];
 
-		for (const context of chatWidget.input.attachedContext) {
+		for (const context of chatWidget.input.attachmentModel.attachments) {
 			if (!URI.isUri(context.value) || !context.isFile || !resourceSet.has(context.value)) {
 				newContext.push(context);
 			}
 		}
 
-		chatWidget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.setContext(true, ...newContext);
+		chatWidget.attachmentModel.clearAndSetContext(...newContext);
 	}
 });
 

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

@@ -37,7 +37,7 @@ export class ChatImageDropAndPaste extends Disposable {
 			return;
 		}
 
-		const currentContextIds = new Set(Array.from(this.inputPart.attachedContext).map(context => context.id));
+		const currentContextIds = this.inputPart.attachmentModel.getAttachmentIDs();
 		const filteredContext = [];
 
 		if (!currentContextIds.has(context.id)) {
@@ -45,7 +45,7 @@ export class ChatImageDropAndPaste extends Disposable {
 			filteredContext.push(context);
 		}
 
-		this.inputPart.attachContext(false, ...filteredContext);
+		this.inputPart.attachmentModel.addContext(...filteredContext);
 	}
 }
 
@@ -84,4 +84,3 @@ export function isImage(array: Uint8Array): boolean {
 		signature.every((byte, index) => array[index] === byte)
 	);
 }
-

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

@@ -68,10 +68,11 @@ import { ChatEditingSessionState, IChatEditingSession, WorkingSetEntryState } fr
 import { IChatRequestVariableEntry } from '../common/chatModel.js';
 import { IChatFollowup } from '../common/chatService.js';
 import { IChatResponseViewModel } from '../common/chatViewModel.js';
-import { IChatHistoryEntry, IChatWidgetHistoryService } from '../common/chatWidgetHistoryService.js';
+import { IChatHistoryEntry, IChatInputState, IChatWidgetHistoryService } from '../common/chatWidgetHistoryService.js';
 import { ILanguageModelChatMetadata, ILanguageModelsService } from '../common/languageModels.js';
 import { CancelAction, ChatModelPickerActionId, ChatSubmitSecondaryAgentAction, IChatExecuteActionContext, SubmitAction } from './actions/chatExecuteActions.js';
 import { IChatWidget } from './chat.js';
+import { ChatAttachmentModel } from './chatAttachmentModel.js';
 import { IDisposableReference } from './chatContentParts/chatCollections.js';
 import { CollapsibleListPool, IChatCollapsibleListItem } from './chatContentParts/chatReferencesContentPart.js';
 import { ChatEditingAcceptAllAction, ChatEditingDiscardAllAction, ChatEditingShowChangesAction } from './chatEditingActions.js';
@@ -93,26 +94,6 @@ interface IChatInputPartOptions {
 	editorOverflowWidgetsDomNode?: HTMLElement;
 }
 
-// TODO: @justschen
-// class ChatInputContentProvider extends Disposable implements ITextModelContentProvider {
-// 	constructor(
-// 		textModelService: ITextModelService,
-// 		private readonly modelService: IModelService,
-// 		private readonly languageService: ILanguageService,
-// 	) {
-// 		super();
-// 		this._register(textModelService.registerTextModelContentProvider(ChatInputPart.INPUT_SCHEME, this));
-// 	}
-
-// 	async provideTextContent(resource: URI): Promise<ITextModel | null> {
-// 		const existing = this.modelService.getModel(resource);
-// 		if (existing) {
-// 			return existing;
-// 		}
-// 		return this.modelService.createModel('', this.languageService.createById('chatSessionInput'), resource);
-// 	}
-// }
-
 export class ChatInputPart extends Disposable implements IHistoryNavigationWidget {
 	static readonly INPUT_SCHEME = 'chatSessionInput';
 	private static _counter = 0;
@@ -135,12 +116,12 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	private _onDidAcceptFollowup = this._register(new Emitter<{ followup: IChatFollowup; response: IChatResponseViewModel | undefined }>());
 	readonly onDidAcceptFollowup = this._onDidAcceptFollowup.event;
 
-	public get attachedContext(): ReadonlySet<IChatRequestVariableEntry> {
-		return this._attachedContext;
+	private readonly _attachmentModel: ChatAttachmentModel;
+	public get attachmentModel(): ChatAttachmentModel {
+		return this._attachmentModel;
 	}
 
 	private _indexOfLastAttachedContextDeletedWithKeyboard: number = -1;
-	private readonly _attachedContext = new Set<IChatRequestVariableEntry>();
 
 	private readonly _onDidChangeVisibility = this._register(new Emitter<boolean>());
 	private readonly _contextResourceLabels = this.instantiationService.createInstance(ResourceLabels, { onDidChangeVisibility: this._onDidChangeVisibility.event });
@@ -182,7 +163,6 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	private history: HistoryNavigator2<IChatHistoryEntry>;
 	private historyNavigationBackwardsEnablement!: IContextKey<boolean>;
 	private historyNavigationForewardsEnablement!: IContextKey<boolean>;
-	private inHistoryNavigation = false;
 	private inputModel: ITextModel | undefined;
 	private inputEditorHasText: IContextKey<boolean>;
 	private chatCursorAtTop: IContextKey<boolean>;
@@ -218,11 +198,13 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		return edits;
 	}
 
+	private readonly getInputState: () => IChatInputState;
+
 	constructor(
 		// private readonly editorOptions: ChatEditorOptions, // TODO this should be used
 		private readonly location: ChatAgentLocation,
 		private readonly options: IChatInputPartOptions,
-		private readonly getInputState: () => any,
+		getContribsInputState: () => any,
 		@IChatWidgetHistoryService private readonly historyService: IChatWidgetHistoryService,
 		@IModelService private readonly modelService: IModelService,
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
@@ -239,7 +221,14 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 	) {
 		super();
 
-
+		this._attachmentModel = this._register(new ChatAttachmentModel());
+		this.getInputState = (): IChatInputState => {
+			// Get input state from widget contribs, merge with attachments
+			return {
+				...getContribsInputState(),
+				chatContextAttachments: this._attachmentModel.attachments,
+			};
+		};
 		this.inputEditorMaxHeight = this.options.renderStyle === 'compact' ? INPUT_EDITOR_MAX_HEIGHT / 3 : INPUT_EDITOR_MAX_HEIGHT;
 
 		this.inputEditorHasText = CONTEXT_CHAT_INPUT_HAS_TEXT.bindTo(contextKeyService);
@@ -295,29 +284,14 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		return localize('chatInput', "Chat Input");
 	}
 
-	updateState(inputState: Object): void {
-		if (this.inHistoryNavigation) {
-			return;
-		}
-
-		const newEntry = { text: this._inputEditor.getValue(), state: inputState };
-
-		if (this.history.isAtEnd()) {
-			// The last history entry should always be the current input value
-			this.history.replaceLast(newEntry);
-		} else {
-			// Added a reference while in the middle of history navigation, it's a new entry
-			this.history.replaceLast(newEntry);
-			this.history.resetCursor();
-		}
-	}
-
 	initForNewChatModel(state: IChatViewState): void {
 		this.history = this.loadHistory();
 		this.history.add({
 			text: state.inputValue ?? this.history.current().text,
 			state: state.inputState ?? this.getInputState()
 		});
+		const attachments = state.inputState?.chatContextAttachments ?? [];
+		this._attachmentModel.clearAndSetContext(...attachments);
 
 		if (state.inputValue) {
 			this.setValue(state.inputValue, false);
@@ -342,7 +316,6 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 				});
 			}
 		}
-
 	}
 
 	logInputHistory(): void {
@@ -390,11 +363,11 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		const historyEntry = previous ?
 			this.history.previous() : this.history.next();
 
-		aria.status(historyEntry.text);
+		const historyAttachments = historyEntry.state?.chatContextAttachments ?? [];
+		this._attachmentModel.clearAndSetContext(...historyAttachments);
 
-		this.inHistoryNavigation = true;
+		aria.status(historyEntry.text);
 		this.setValue(historyEntry.text, true);
-		this.inHistoryNavigation = false;
 
 		this._onDidLoadInputState.fire(historyEntry.state);
 
@@ -455,7 +428,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		}
 
 		// Clear attached context, fire event to clear input state, and clear the input editor
-		this._attachedContext.clear();
+		this.attachmentModel.clear();
 		this._onDidLoadInputState.fire({});
 		if (this.accessibilityService.isScreenReaderOptimized() && isMacintosh) {
 			this._acceptInputForVoiceover();
@@ -478,28 +451,6 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		this._inputEditor.focus();
 	}
 
-	attachContext(overwrite: boolean, ...contentReferences: IChatRequestVariableEntry[]): void {
-		const removed = [];
-		if (overwrite) {
-			removed.push(...Array.from(this._attachedContext));
-			this._attachedContext.clear();
-		}
-
-		if (contentReferences.length > 0) {
-			for (const reference of contentReferences) {
-				this._attachedContext.add(reference);
-			}
-		}
-
-		if (removed.length > 0 || contentReferences.length > 0) {
-			this.initAttachedContext(this.attachedContextContainer);
-
-			if (!overwrite) {
-				this._onDidChangeContext.fire({ removed, added: contentReferences });
-			}
-		}
-	}
-
 	render(container: HTMLElement, initialValue: string, widget: IChatWidget) {
 		let elements;
 		if (this.options.renderStyle === 'compact') {
@@ -540,6 +491,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		const toolbarsContainer = elements.inputToolbars;
 		this.chatEditingSessionWidgetContainer = elements.chatEditingSessionWidgetContainer;
 		this.initAttachedContext(this.attachedContextContainer);
+		this._register(this._attachmentModel.onDidChangeContext(() => this.initAttachedContext(this.attachedContextContainer)));
 		this.renderChatEditingSessionState(null, undefined, widget);
 
 		const inputScopedContextKeyService = this._register(this.contextKeyService.createScoped(inputContainer));
@@ -710,11 +662,11 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		dom.clearNode(container);
 		this.attachedContextDisposables.clear();
 		const hoverDelegate = this.attachedContextDisposables.add(createInstantHoverDelegate());
-		dom.setVisibility(Boolean(this.attachedContext.size), this.attachedContextContainer);
-		if (!this.attachedContext.size) {
+		dom.setVisibility(Boolean(this.attachmentModel.size), this.attachedContextContainer);
+		if (!this.attachmentModel.size) {
 			this._indexOfLastAttachedContextDeletedWithKeyboard = -1;
 		}
-		[...this.attachedContext.values()].forEach(async (attachment, index) => {
+		this.attachmentModel.attachments.forEach(async (attachment, index) => {
 			if (attachment.isFile && this.location === ChatAgentLocation.EditingSession) {
 				return;
 			}
@@ -802,14 +754,14 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		});
 
 		// If this item is rendering in place of the last attached context item, focus the clear button so the user can continue deleting attached context items with the keyboard
-		if (index === Math.min(this._indexOfLastAttachedContextDeletedWithKeyboard, this.attachedContext.size - 1)) {
+		if (index === Math.min(this._indexOfLastAttachedContextDeletedWithKeyboard, this.attachmentModel.size - 1)) {
 			clearButton.focus();
 		}
 
 		this.attachedContextDisposables.add(clearButton);
 		clearButton.icon = Codicon.close;
 		const disp = clearButton.onDidClick((e) => {
-			this._attachedContext.delete(attachment);
+			this._attachmentModel.delete(attachment.id);
 			disp.dispose();
 
 			// Set focus to the next attached context item if deletion was triggered by a keystroke (vs a mouse click)
@@ -820,7 +772,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 				}
 			}
 
-			if (this._attachedContext.size === 0) {
+			if (this._attachmentModel.size === 0) {
 				this.focus();
 			}
 
@@ -876,7 +828,7 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 				kind: 'reference',
 			};
 		}) ?? [];
-		for (const attachment of this._attachedContext) {
+		for (const attachment of this.attachmentModel.attachments) {
 			if (attachment.isFile && URI.isUri(attachment.value) && !modifiedFiles.has(attachment.value)) {
 				entries.unshift({
 					reference: attachment.value,
@@ -1101,6 +1053,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		};
 	}
 
+	getViewState(): IChatInputState {
+		return this.getInputState();
+	}
+
 	saveState(): void {
 		this.saveCurrentValue(this.getInputState());
 		const inputHistory = [...this.history];

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

@@ -19,7 +19,6 @@ import { ChatRequestDynamicVariablePart, ChatRequestToolPart, ChatRequestVariabl
 import { IChatContentReference } from '../common/chatService.js';
 import { IChatRequestVariableValue, IChatVariableData, IChatVariableResolver, IChatVariableResolverProgress, IChatVariablesService, IDynamicVariable } from '../common/chatVariables.js';
 import { IChatWidgetService, showChatView } from './chat.js';
-import { ChatContextAttachments } from './chatWidget.js';
 import { ChatDynamicVariableModel } from './contrib/chatDynamicVariables.js';
 
 interface IChatData {
@@ -176,7 +175,7 @@ export class ChatVariablesService implements IChatVariablesService {
 		if (key === 'file' && typeof value !== 'string') {
 			const uri = URI.isUri(value) ? value : value.uri;
 			const range = 'range' in value ? value.range : undefined;
-			widget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.setContext(false, { value, id: uri.toString() + (range?.toString() ?? ''), name: basename(uri.path), isFile: true, isDynamic: true });
+			widget.attachmentModel.addContext({ value, id: uri.toString() + (range?.toString() ?? ''), name: basename(uri.path), isFile: true, isDynamic: true });
 			return;
 		}
 
@@ -185,6 +184,6 @@ export class ChatVariablesService implements IChatVariablesService {
 			return;
 		}
 
-		widget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.setContext(false, { ...resolved.data, value });
+		widget.attachmentModel.addContext({ ...resolved.data, value });
 	}
 }

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

@@ -245,8 +245,9 @@ export class ChatViewPane extends ViewPane {
 
 	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;
+		for (const [key, value] of Object.entries(newViewState)) {
+			// Assign all props to the memento so they get saved
+			(this.viewState as any)[key] = value;
+		}
 	}
 }

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

@@ -33,15 +33,17 @@ import { TerminalChatController } from '../../terminal/terminalContribChatExport
 import { ChatAgentLocation, IChatAgentCommand, IChatAgentData, IChatAgentService, IChatWelcomeMessageContent, isChatWelcomeMessageContent } from '../common/chatAgents.js';
 import { CONTEXT_CHAT_INPUT_HAS_AGENT, CONTEXT_CHAT_LOCATION, CONTEXT_CHAT_REQUEST_IN_PROGRESS, CONTEXT_IN_CHAT_SESSION, CONTEXT_IN_QUICK_CHAT, CONTEXT_LAST_ITEM_ID, CONTEXT_PARTICIPANT_SUPPORTS_MODEL_PICKER, CONTEXT_RESPONSE_FILTERED } from '../common/chatContextKeys.js';
 import { ChatEditingSessionState, IChatEditingService, IChatEditingSession } from '../common/chatEditingService.js';
-import { IChatModel, IChatRequestVariableEntry, IChatResponseModel, isChatRequestVariableEntry } from '../common/chatModel.js';
+import { IChatModel, IChatRequestVariableEntry, IChatResponseModel } from '../common/chatModel.js';
 import { ChatRequestAgentPart, IParsedChatRequest, chatAgentLeader, chatSubcommandLeader, formatChatQuestion } from '../common/chatParserTypes.js';
 import { ChatRequestParser } from '../common/chatRequestParser.js';
 import { IChatFollowup, IChatLocationData, IChatService } from '../common/chatService.js';
 import { IChatSlashCommandService } from '../common/chatSlashCommands.js';
 import { ChatViewModel, IChatResponseViewModel, isRequestVM, isResponseVM } from '../common/chatViewModel.js';
+import { IChatInputState } from '../common/chatWidgetHistoryService.js';
 import { CodeBlockModelCollection } from '../common/codeBlockModelCollection.js';
 import { ChatTreeItem, IChatAccessibilityService, IChatCodeBlockInfo, IChatFileTreeInfo, IChatListItemRendererOptions, IChatWidget, IChatWidgetService, IChatWidgetViewContext, IChatWidgetViewOptions } from './chat.js';
 import { ChatAccessibilityProvider } from './chatAccessibilityProvider.js';
+import { ChatAttachmentModel } from './chatAttachmentModel.js';
 import { ChatDragAndDrop } from './chatDragAndDrop.js';
 import { ChatImageDropAndPaste } from './chatImagePaste.js';
 import { ChatInputPart } from './chatInputPart.js';
@@ -56,7 +58,6 @@ function revealLastElement(list: WorkbenchObjectTree<any>) {
 	list.scrollTop = list.scrollHeight - list.renderHeight;
 }
 
-export type IChatInputState = Record<string, any>;
 export interface IChatViewState {
 	inputValue?: string;
 	inputState?: IChatInputState;
@@ -286,7 +287,6 @@ export class ChatWidget extends Disposable implements IChatWidget {
 						currentEditSession = undefined;
 					}
 				}
-
 			}));
 		}
 
@@ -385,6 +385,10 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		return this.inputPart.contentHeight + this.tree.contentHeight;
 	}
 
+	get attachmentModel(): ChatAttachmentModel {
+		return this.inputPart.attachmentModel;
+	}
+
 	render(parent: HTMLElement): void {
 		const viewId = 'viewId' in this.viewContext ? this.viewContext.viewId : undefined;
 		this.editorOptions = this._register(this.instantiationService.createInstance(ChatEditorOptions, viewId, this.styles.listForeground, this.styles.inputEditorBackground, this.styles.resultEditorBackground));
@@ -768,6 +772,12 @@ export class ChatWidget extends Disposable implements IChatWidget {
 			}
 			this._onDidChangeContentHeight.fire();
 		}));
+		this._register(this.inputPart.attachmentModel.onDidChangeContext(() => {
+			if (this.chatEditingService.currentEditingSession && this.chatEditingService.currentEditingSession?.chatSessionId === this.viewModel?.sessionId) {
+				// TODO still needed? Do this inside input part and fire onDidChangeHeight?
+				this.renderChatEditingSessionState(this.chatEditingService.currentEditingSession);
+			}
+		}));
 		this._register(this.inputEditor.onDidChangeModelContent(() => {
 			this.parsedChatRequest = undefined;
 			this.updateChatInputContext();
@@ -910,7 +920,7 @@ export class ChatWidget extends Disposable implements IChatWidget {
 				location: this.location,
 				locationData: this._location.resolveData?.(),
 				parserContext: { selectedAgent: this._lastSelectedAgent },
-				attachedContext: [...this.inputPart.attachedContext.values()]
+				attachedContext: [...this.attachmentModel.attachments]
 			});
 
 			if (result) {
@@ -934,14 +944,6 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		return undefined;
 	}
 
-	setContext(overwrite: boolean, ...contentReferences: IChatRequestVariableEntry[]) {
-		this.inputPart.attachContext(overwrite, ...contentReferences);
-
-		if (this.chatEditingService.currentEditingSession && this.chatEditingService.currentEditingSession?.chatSessionId === this.viewModel?.sessionId) {
-			this.renderChatEditingSessionState(this.chatEditingService.currentEditingSession);
-		}
-	}
-
 	getCodeBlockInfosForResponse(response: IChatResponseViewModel): IChatCodeBlockInfo[] {
 		return this.renderer.getCodeBlockInfosForResponse(response);
 	}
@@ -1112,8 +1114,8 @@ export class ChatWidget extends Disposable implements IChatWidget {
 	getViewState(): IChatViewState {
 		return {
 			inputValue: this.getInput(),
-			inputState: this.collectInputState(),
-			selectedLanguageModelId: this.inputPart.currentLanguageModel
+			inputState: this.inputPart.getViewState(),
+			selectedLanguageModelId: this.inputPart.currentLanguageModel,
 		};
 	}
 
@@ -1170,66 +1172,3 @@ export class ChatWidgetService implements IChatWidgetService {
 		);
 	}
 }
-
-export class ChatContextAttachments extends Disposable implements IChatWidgetContrib {
-
-	private _attachedContext = new Map<string, IChatRequestVariableEntry>();
-
-	public static readonly ID = 'chatContextAttachments';
-
-	get id() {
-		return ChatContextAttachments.ID;
-	}
-
-	constructor(readonly widget: IChatWidget) {
-		super();
-
-		this._register(this.widget.onDidChangeContext(({ removed, added }) => {
-			removed?.forEach(attachment => this._attachedContext.delete(attachment.id));
-			added?.forEach(attachment => {
-				if (!this._attachedContext.has(attachment.id)) {
-					this._attachedContext.set(attachment.id, attachment);
-				}
-			});
-		}));
-
-		this._register(this.widget.onDidSubmitAgent(() => {
-			this._clearAttachedContext();
-		}));
-	}
-
-	getInputState(): IChatRequestVariableEntry[] {
-		return [...this._attachedContext.values()];
-	}
-
-	setInputState(s: unknown): void {
-		const attachments = Array.isArray(s) ? s.filter(isChatRequestVariableEntry) : [];
-		this.setContext(true, ...attachments);
-	}
-
-	getContext() {
-		return new Set(this._attachedContext.keys());
-	}
-
-	setContext(overwrite: boolean, ...attachments: IChatRequestVariableEntry[]) {
-		if (overwrite) {
-			this._attachedContext.clear();
-		}
-
-		const newAttachments = [];
-		for (const attachment of attachments) {
-			if (!this._attachedContext.has(attachment.id)) {
-				this._attachedContext.set(attachment.id, attachment);
-				newAttachments.push(attachment);
-			}
-		}
-
-		this.widget.setContext(overwrite, ...newAttachments);
-	}
-
-	private _clearAttachedContext() {
-		this._attachedContext.clear();
-	}
-}
-
-ChatWidget.CONTRIBS.push(ChatContextAttachments);

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

@@ -8,11 +8,18 @@ import { createDecorator } from '../../../../platform/instantiation/common/insta
 import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
 import { Memento } from '../../../common/memento.js';
 import { ChatAgentLocation } from './chatAgents.js';
+import { IChatRequestVariableEntry } from './chatModel.js';
 import { CHAT_PROVIDER_ID } from './chatParticipantContribTypes.js';
 
 export interface IChatHistoryEntry {
 	text: string;
-	state?: any;
+	state?: IChatInputState;
+}
+
+/** The collected input state of ChatWidget contribs + attachments */
+export interface IChatInputState {
+	[key: string]: any;
+	chatContextAttachments?: ReadonlyArray<IChatRequestVariableEntry>;
 }
 
 export const IChatWidgetHistoryService = createDecorator<IChatWidgetHistoryService>('IChatWidgetHistoryService');

src/vs/workbench/contrib/inlineChat/browser/inlineChatWidget.ts

@@ -487,7 +487,7 @@ export class InlineChatWidget {
 	}
 
 	reset() {
-		this._chatWidget.setContext(true);
+		this._chatWidget.attachmentModel.clear();
 		this._chatWidget.saveState();
 
 		reset(this._elements.statusLabel);

src/vs/workbench/contrib/notebook/browser/controller/chat/notebook.chat.contribution.ts

@@ -21,7 +21,6 @@ import { IWorkbenchContribution, registerWorkbenchContribution2, WorkbenchPhase
 import { IEditorService } from '../../../../../services/editor/common/editorService.js';
 import { IChatWidget, IChatWidgetService } from '../../../../chat/browser/chat.js';
 import { ChatInputPart } from '../../../../chat/browser/chatInputPart.js';
-import { ChatContextAttachments } from '../../../../chat/browser/chatWidget.js';
 import { ChatDynamicVariableModel } from '../../../../chat/browser/contrib/chatDynamicVariables.js';
 import { computeCompletionRanges } from '../../../../chat/browser/contrib/chatInputCompletions.js';
 import { ChatAgentLocation, IChatAgentService } from '../../../../chat/common/chatAgents.js';
@@ -219,13 +218,13 @@ export class SelectAndInsertKernelVariableAction extends Action2 {
 				icon: codiconsLibrary.variable,
 			});
 		} else {
-			widget.getContrib<ChatContextAttachments>(ChatContextAttachments.ID)?.setContext(false, ...[{
+			widget.attachmentModel.addContext({
 				id: 'vscode.notebook.variable',
 				name: variableName,
 				value: variableName,
 				icon: codiconsLibrary.variable,
 				isDynamic: true
-			}]);
+			});
 		}
 	}
 }