chat: refactor input state to be central (#278683)

* chat: add additional timing information on the model

- Add a public created `timestamp` on the IChatModel
- Add a started `timestamp` and optional `completedAt` timestamp on
  the IChatResponseModel
- Make `isPendingConfirmation<{startedWaitingAt: number}> | undefined`
  to encode the time when the response started waiting for confirmation.
- Add a `confirmationAdjustedTimestamp` that can be used to reflect the
  duration a chat response was waiting for user input vs working.

* update snapshots

* wip

* fix mode picker changing wrong input when focused

* wip

* make dynamic variables happy

* tidy

Author: connor4312

src/vs/workbench/api/browser/mainThreadChatAgents2.ts

@@ -147,16 +147,14 @@ export class MainThreadChatAgents2 extends Disposable implements MainThreadChatA
 
 	$transferActiveChatSession(toWorkspace: UriComponents): void {
 		const widget = this._chatWidgetService.lastFocusedWidget;
-		const sessionId = widget?.viewModel?.model.sessionId;
-		if (!sessionId) {
+		const model = widget?.viewModel?.model;
+		if (!model) {
 			this._logService.error(`MainThreadChat#$transferActiveChatSession: No active chat session found`);
 			return;
 		}
 
-		const inputValue = widget?.inputEditor.getValue() ?? '';
 		const location = widget.location;
-		const mode = widget.input.currentModeKind;
-		this._chatService.transferChatSession({ sessionId, inputValue, location, mode }, URI.revive(toWorkspace));
+		this._chatService.transferChatSession({ sessionId: model.sessionId, inputState: model.inputModel.state.get(), location }, URI.revive(toWorkspace));
 	}
 
 	async $registerAgent(handle: number, extension: ExtensionIdentifier, id: string, metadata: IExtensionChatAgentMetadata, dynamicProps: IDynamicChatAgentProps | undefined): Promise<void> {

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

@@ -8,6 +8,7 @@ import { KeyCode, KeyMod } from '../../../../../base/common/keyCodes.js';
 import { basename } from '../../../../../base/common/resources.js';
 import { ThemeIcon } from '../../../../../base/common/themables.js';
 import { assertType } from '../../../../../base/common/types.js';
+import { URI } from '../../../../../base/common/uri.js';
 import { ServicesAccessor } from '../../../../../editor/browser/editorExtensions.js';
 import { EditorContextKeys } from '../../../../../editor/common/editorContextKeys.js';
 import { localize, localize2 } from '../../../../../nls.js';
@@ -273,6 +274,7 @@ export const ToggleAgentModeActionId = 'workbench.action.chat.toggleAgentMode';
 
 export interface IToggleChatModeArgs {
 	modeId: ChatModeKind | string;
+	sessionResource: URI | undefined;
 }
 
 type ChatModeChangeClassification = {
@@ -319,23 +321,30 @@ class ToggleChatModeAction extends Action2 {
 		const instaService = accessor.get(IInstantiationService);
 		const modeService = accessor.get(IChatModeService);
 		const telemetryService = accessor.get(ITelemetryService);
+		const chatWidgetService = accessor.get(IChatWidgetService);
 
-		const context = getEditingSessionContext(accessor, args);
-		if (!context?.chatWidget) {
+		const arg = args.at(0) as IToggleChatModeArgs | undefined;
+		let widget: IChatWidget | undefined;
+		if (arg?.sessionResource) {
+			widget = chatWidgetService.getWidgetBySessionResource(arg.sessionResource);
+		} else {
+			widget = getEditingSessionContext(accessor, args)?.chatWidget;
+		}
+
+		if (!widget) {
 			return;
 		}
 
-		const arg = args.at(0) as IToggleChatModeArgs | undefined;
-		const chatSession = context.chatWidget.viewModel?.model;
+		const chatSession = widget.viewModel?.model;
 		const requestCount = chatSession?.getRequests().length ?? 0;
-		const switchToMode = (arg && modeService.findModeById(arg.modeId)) ?? this.getNextMode(context.chatWidget, requestCount, configurationService, modeService);
+		const switchToMode = (arg && modeService.findModeById(arg.modeId)) ?? this.getNextMode(widget, requestCount, configurationService, modeService);
 
-		const currentMode = context.chatWidget.input.currentModeObs.get();
+		const currentMode = widget.input.currentModeObs.get();
 		if (switchToMode.id === currentMode.id) {
 			return;
 		}
 
-		const chatModeCheck = await instaService.invokeFunction(handleModeSwitch, context.chatWidget.input.currentModeKind, switchToMode.kind, requestCount, context.editingSession);
+		const chatModeCheck = await instaService.invokeFunction(handleModeSwitch, widget.input.currentModeKind, switchToMode.kind, requestCount, widget.viewModel?.model.editingSession);
 		if (!chatModeCheck) {
 			return;
 		}
@@ -356,7 +365,7 @@ class ToggleChatModeAction extends Action2 {
 			handoffsCount
 		});
 
-		context.chatWidget.input.setChatMode(switchToMode.id);
+		widget.input.setChatMode(switchToMode.id);
 
 		if (chatModeCheck.needToClearSession) {
 			await commandService.executeCommand(ACTION_ID_NEW_CHAT);

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

@@ -131,13 +131,15 @@ async function executeMoveToAction(accessor: ServicesAccessor, moveTo: MoveToNew
 		return;
 	}
 
-	// Save off the state before clearing
-	const viewState = widget.getViewState();
+	// Save off the session resource before clearing
 	const resourceToOpen = widget.viewModel.sessionResource;
 
+	// Todo: can possibly go away with https://github.com/microsoft/vscode/pull/278476
+	const modelInputState = existingWidget.getViewState();
+
 	await widget.clear();
 
-	const options: IChatEditorOptions = { pinned: true, viewState, auxiliary };
+	const options: IChatEditorOptions = { pinned: true, modelInputState, auxiliary };
 	await editorService.openEditor({ resource: resourceToOpen, options }, moveTo === MoveToNewLocation.Window ? AUX_WINDOW_GROUP : ACTIVE_GROUP);
 }
 
@@ -150,9 +152,15 @@ async function moveToSidebar(accessor: ServicesAccessor): Promise<void> {
 	const chatEditorInput = chatEditor?.input;
 	let view: ChatViewPane;
 	if (chatEditor instanceof ChatEditor && chatEditorInput instanceof ChatEditorInput && chatEditorInput.sessionResource) {
+		const previousViewState = chatEditor.widget.getViewState();
 		await editorService.closeEditor({ editor: chatEditor.input, groupId: editorGroupService.activeGroup.id });
 		view = await viewsService.openView(ChatViewId) as ChatViewPane;
-		await view.loadSession(chatEditorInput.sessionResource, chatEditor.getViewState());
+
+		// Todo: can possibly go away with https://github.com/microsoft/vscode/pull/278476
+		const newModel = await view.loadSession(chatEditorInput.sessionResource);
+		if (previousViewState && newModel && !newModel.inputModel.state.get()) {
+			newModel.inputModel.setState(previousViewState);
+		}
 	} else {
 		view = await viewsService.openView(ChatViewId) as ChatViewPane;
 	}

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

@@ -15,7 +15,7 @@ import { IContextKeyService } from '../../../../platform/contextkey/common/conte
 import { createDecorator } from '../../../../platform/instantiation/common/instantiation.js';
 import { PreferredGroup } from '../../../services/editor/common/editorService.js';
 import { IChatAgentAttachmentCapabilities, IChatAgentCommand, IChatAgentData } from '../common/chatAgents.js';
-import { IChatResponseModel } from '../common/chatModel.js';
+import { IChatResponseModel, IChatModelInputState } from '../common/chatModel.js';
 import { IChatMode } from '../common/chatModes.js';
 import { IParsedChatRequest } from '../common/chatParserTypes.js';
 import { CHAT_PROVIDER_ID } from '../common/chatParticipantContribTypes.js';
@@ -25,7 +25,7 @@ import { ChatAgentLocation, ChatModeKind } from '../common/constants.js';
 import { ChatAttachmentModel } from './chatAttachmentModel.js';
 import { IChatEditorOptions } from './chatEditor.js';
 import { ChatInputPart } from './chatInputPart.js';
-import { ChatWidget, IChatViewState, IChatWidgetContrib } from './chatWidget.js';
+import { ChatWidget, IChatWidgetContrib } from './chatWidget.js';
 import { ICodeBlockActionContext } from './codeBlockPart.js';
 
 export const IChatWidgetService = createDecorator<IChatWidgetService>('chatWidgetService');
@@ -223,6 +223,7 @@ export interface IChatWidget {
 	readonly input: ChatInputPart;
 	readonly attachmentModel: ChatAttachmentModel;
 	readonly locationData?: IChatLocationData;
+	readonly contribs: readonly IChatWidgetContrib[];
 
 	readonly supportsChangingModes: boolean;
 
@@ -254,7 +255,7 @@ export interface IChatWidget {
 	getFileTreeInfosForResponse(response: IChatResponseViewModel): IChatFileTreeInfo[];
 	getLastFocusedFileTreeForResponse(response: IChatResponseViewModel): IChatFileTreeInfo | undefined;
 	clear(): Promise<void>;
-	getViewState(): IChatViewState;
+	getViewState(): IChatModelInputState | undefined;
 	lockToCodingAgent(name: string, displayName: string, agentId?: string): void;
 
 	delegateScrollFromMouseWheelEvent(event: IMouseWheelEvent): void;

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

@@ -14,26 +14,30 @@ import { IContextKeyService, IScopedContextKeyService } from '../../../../platfo
 import { IEditorOptions } from '../../../../platform/editor/common/editor.js';
 import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
 import { ServiceCollection } from '../../../../platform/instantiation/common/serviceCollection.js';
-import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
+import { IStorageService } from '../../../../platform/storage/common/storage.js';
 import { ITelemetryService } from '../../../../platform/telemetry/common/telemetry.js';
 import { editorBackground, editorForeground, inputBackground } from '../../../../platform/theme/common/colorRegistry.js';
 import { IThemeService } from '../../../../platform/theme/common/themeService.js';
 import { EditorPane } from '../../../browser/parts/editor/editorPane.js';
 import { IEditorOpenContext } from '../../../common/editor.js';
-import { Memento } from '../../../common/memento.js';
 import { EDITOR_DRAG_AND_DROP_BACKGROUND } from '../../../common/theme.js';
 import { IEditorGroup } from '../../../services/editor/common/editorGroupsService.js';
 import { ChatContextKeys } from '../common/chatContextKeys.js';
-import { IChatModel, IExportableChatData, ISerializableChatData } from '../common/chatModel.js';
-import { CHAT_PROVIDER_ID } from '../common/chatParticipantContribTypes.js';
+import { IChatModel, IChatModelInputState, IExportableChatData, ISerializableChatData } from '../common/chatModel.js';
 import { IChatService } from '../common/chatService.js';
 import { IChatSessionsService, localChatSessionType } from '../common/chatSessionsService.js';
 import { ChatAgentLocation, ChatModeKind } from '../common/constants.js';
 import { clearChatEditor } from './actions/chatClear.js';
 import { ChatEditorInput } from './chatEditorInput.js';
-import { ChatWidget, IChatViewState } from './chatWidget.js';
+import { ChatWidget } from './chatWidget.js';
 
 export interface IChatEditorOptions extends IEditorOptions {
+	/**
+	 * Input state of the model when the editor is opened. Currently needed since
+	 * new sessions are not persisted but may go away with
+	 * https://github.com/microsoft/vscode/pull/278476 as input state is stored on the model.
+	 */
+	modelInputState?: IChatModelInputState;
 	target?: { data: IExportableChatData | ISerializableChatData };
 	title?: {
 		preferred?: string;
@@ -52,8 +56,6 @@ export class ChatEditor extends EditorPane {
 		return this._scopedContextKeyService;
 	}
 
-	private _memento: Memento<IChatViewState> | undefined;
-	private _viewState: IChatViewState | undefined;
 	private dimension = new dom.Dimension(0, 0);
 	private _loadingContainer: HTMLElement | undefined;
 	private _editorContainer: HTMLElement | undefined;
@@ -63,7 +65,7 @@ export class ChatEditor extends EditorPane {
 		@ITelemetryService telemetryService: ITelemetryService,
 		@IThemeService themeService: IThemeService,
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
-		@IStorageService private readonly storageService: IStorageService,
+		@IStorageService storageService: IStorageService,
 		@IChatSessionsService private readonly chatSessionsService: IChatSessionsService,
 		@IContextKeyService private readonly contextKeyService: IContextKeyService,
 		@IChatService private readonly chatService: IChatService,
@@ -231,8 +233,11 @@ export class ChatEditor extends EditorPane {
 				this.hideLoadingInChatWidget();
 			}
 
-			const viewState = options?.viewState ?? input.options.viewState;
-			this.updateModel(editorModel.model, viewState);
+			if (options?.modelInputState) {
+				editorModel.model.inputModel.setState(options.modelInputState);
+			}
+
+			this.updateModel(editorModel.model);
 
 			if (isContributedChatSession && options?.title?.preferred && input.sessionResource) {
 				this.chatService.setChatSessionTitle(input.sessionResource, options.title.preferred);
@@ -243,27 +248,8 @@ export class ChatEditor extends EditorPane {
 		}
 	}
 
-	private updateModel(model: IChatModel, viewState?: IChatViewState): void {
-		this._memento = new Memento('interactive-session-editor-' + CHAT_PROVIDER_ID, this.storageService);
-		this._viewState = viewState ?? this._memento.getMemento(StorageScope.WORKSPACE, StorageTarget.MACHINE);
-		this.widget.setModel(model, { ...this._viewState });
-	}
-
-	protected override saveState(): void {
-		this.widget?.saveState();
-
-		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.inputState = widgetViewState.inputState;
-			this._memento.saveMemento();
-		}
-	}
-
-	override getViewState(): object | undefined {
-		return { ...this._viewState };
+	private updateModel(model: IChatModel): void {
+		this.widget.setModel(model);
 	}
 
 	override layout(dimension: dom.Dimension, position?: dom.IDomPosition | undefined): void {