Restoring chat model if opening changes

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

@@ -8,7 +8,7 @@ import { localize, localize2 } from '../../../../../nls.js';
 import { getAgentChangesSummary, IAgentSession } from './agentSessionsModel.js';
 import { Action, IAction } from '../../../../../base/common/actions.js';
 import { ActionViewItem, IActionViewItemOptions } from '../../../../../base/browser/ui/actionbar/actionViewItems.js';
-import { CommandsRegistry, ICommandService } from '../../../../../platform/commands/common/commands.js';
+import { ICommandService } from '../../../../../platform/commands/common/commands.js';
 import { EventHelper, h, hide, show } from '../../../../../base/browser/dom.js';
 import { assertReturnsDefined } from '../../../../../base/common/types.js';
 import { Action2, MenuId } from '../../../../../platform/actions/common/actions.js';
@@ -17,7 +17,6 @@ import { ServicesAccessor } from '../../../../../editor/browser/editorExtensions
 import { ViewAction } from '../../../../browser/parts/views/viewPane.js';
 import { AGENT_SESSIONS_VIEW_ID, AgentSessionProviders, IAgentSessionsControl } from './agentSessions.js';
 import { AgentSessionsView } from './agentSessionsView.js';
-import { URI } from '../../../../../base/common/uri.js';
 import { IChatService } from '../../common/chatService.js';
 import { ChatContextKeys } from '../../common/chatContextKeys.js';
 import { IMarshalledChatSessionContext } from '../actions/chatSessionActions.js';
@@ -27,6 +26,9 @@ import { ACTIVE_GROUP, AUX_WINDOW_GROUP, PreferredGroup, SIDE_GROUP } from '../.
 import { IViewDescriptorService, ViewContainerLocation } from '../../../../common/views.js';
 import { getPartByLocation } from '../../../../services/views/browser/viewsService.js';
 import { IWorkbenchLayoutService } from '../../../../services/layout/browser/layoutService.js';
+import { getChatSessionType } from '../../common/chatUri.js';
+import { ChatAgentLocation } from '../../common/constants.js';
+import { CancellationToken } from '../../../../../base/common/cancellation.js';
 
 //#region Session Title Actions
 
@@ -100,7 +102,8 @@ export class AgentSessionDiffActionViewItem extends ActionViewItem {
 	constructor(
 		action: IAction,
 		options: IActionViewItemOptions,
-		@ICommandService private readonly commandService: ICommandService
+		@ICommandService private readonly commandService: ICommandService,
+		@IChatService private readonly chatService: IChatService
 	) {
 		super(null, action, options);
 	}
@@ -152,20 +155,34 @@ export class AgentSessionDiffActionViewItem extends ActionViewItem {
 
 	override onClick(event: MouseEvent): void {
 		EventHelper.stop(event, true);
-
-		const session = this.action.getSession();
-
-		this.commandService.executeCommand(`agentSession.${session.providerType}.openChanges`, this.action.getSession().resource);
+		this.openChanges();
+
+	}
+	async openChanges(): Promise<void> {
+		const resource = this.action.getSession().resource;
+		const type = getChatSessionType(resource);
+		let session = this.chatService.getSession(resource);
+		if (!session) {
+			const chatModelRef = type === AgentSessionProviders.Local ? await this.chatService.getOrRestoreSession(resource) : await this.chatService.loadSessionForResource(resource, ChatAgentLocation.Chat, CancellationToken.None);
+			session = chatModelRef?.object;
+			const disposable = session?.editingSession?.onDidEditorPaneClose(() => {
+				chatModelRef?.dispose();
+				disposable?.dispose();
+			});
+			if (disposable) {
+				this._register(disposable);
+			}
+		}
+		if (session) {
+			if (type === AgentSessionProviders.Local) {
+				await session?.editingSession?.show();
+			} else {
+				this.commandService.executeCommand(`agentSession.${type}.openChanges`, resource);
+			}
+		}
 	}
 }
 
-CommandsRegistry.registerCommand(`agentSession.${AgentSessionProviders.Local}.openChanges`, async (accessor: ServicesAccessor, resource: URI) => {
-	const chatService = accessor.get(IChatService);
-
-	const session = chatService.getSession(resource);
-	session?.editingSession?.show();
-});
-
 //#endregion
 
 //#region Session Context Actions

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

@@ -175,6 +175,12 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
 		return this._onDidDispose.event;
 	}
 
+	private readonly _onDidEditorPaneClose = new Emitter<void>();
+	get onDidEditorPaneClose() {
+		this._assertNotDisposed();
+		return this._onDidEditorPaneClose.event;
+	}
+
 	constructor(
 		readonly chatSessionResource: URI,
 		readonly isGlobalEditingSession: boolean,
@@ -435,6 +441,14 @@ export class ChatEditingSession extends Disposable implements IChatEditingSessio
 		}, this._instantiationService);
 
 		this._editorPane = await this._editorGroupsService.activeGroup.openEditor(input, { pinned: true, activation: EditorActivation.ACTIVATE }) as MultiDiffEditor | undefined;
+		const disposable = this._editorGroupsService.activeGroup.onDidCloseEditor(e => {
+			if (e.editor === input) {
+				this._editorPane = undefined;
+				this._onDidEditorPaneClose.fire();
+				disposable.dispose();
+			}
+		});
+		this._register(disposable);
 	}
 
 	private _stopPromise: Promise<void> | undefined;

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

@@ -116,6 +116,7 @@ export interface IChatEditingSession extends IDisposable {
 	readonly isGlobalEditingSession: boolean;
 	readonly chatSessionResource: URI;
 	readonly onDidDispose: Event<void>;
+	readonly onDidEditorPaneClose: Event<void>;
 	readonly state: IObservable<ChatEditingSessionState>;
 	readonly entries: IObservable<readonly IModifiedFileEntry[]>;
 	/** Requests disabled by undo/redo in the session */