Fix plot client disposed too early

Use editor input dispose to determine the editor is finished with the plot client

src/vs/workbench/contrib/positronPlots/browser/positronPlotsService.ts

@@ -47,7 +47,6 @@ import { WebviewPlotClient } from 'vs/workbench/contrib/positronPlots/browser/we
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { URI } from 'vs/base/common/uri';
 import { PositronPlotCommProxy } from 'vs/workbench/services/languageRuntime/common/positronPlotCommProxy';
-import { EditorCloseContext } from 'vs/workbench/common/editor';
 
 /** The maximum number of recent executions to store. */
 const MaxRecentExecutions = 10;
@@ -1051,20 +1050,6 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 			}),
 		});
 
-		this._editorService.onDidCloseEditor(editor => {
-			// Dispose the plot client for the editor when the editor matches the plot ID
-			// and the editor was closed, not moved or replaced.
-			if (editorPane?.getId() === editor.editor.editorId
-				&& editor.editor.resource?.path === plotId
-				&& editor.context === EditorCloseContext.UNKNOWN) {
-				const editorPlot = this._editorPlots.get(plotId);
-				if (editorPlot instanceof PlotClientInstance) {
-					this._editorPlots.delete(plotId);
-					this.unregisterPlotClient(editorPlot);
-				}
-			}
-		});
-
 		if (!editorPane) {
 			throw new Error('Failed to open editor');
 		}
@@ -1074,7 +1059,7 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 		return this._editorPlots.get(id);
 	}
 
-	private unregisterPlotClient(plotClient: IPositronPlotClient) {
+	public unregisterPlotClient(plotClient: IPositronPlotClient) {
 		if (plotClient instanceof PlotClientInstance) {
 			const plotId = plotClient.id;
 			const plotClients = this._plotClientsByComm.get(plotId);

src/vs/workbench/contrib/positronPlotsEditor/browser/positronPlotsEditorInput.ts

@@ -6,6 +6,7 @@
 import { URI } from 'vs/base/common/uri';
 import { IUntypedEditorInput } from 'vs/workbench/common/editor';
 import { EditorInput } from 'vs/workbench/common/editor/editorInput';
+import { IPositronPlotsService } from 'vs/workbench/services/positronPlots/common/positronPlots';
 
 export class PositronPlotsEditorInput extends EditorInput {
 	static readonly TypeID: string = 'workbench.input.positronPlots';
@@ -16,9 +17,14 @@ export class PositronPlotsEditorInput extends EditorInput {
 
 	constructor(
 		readonly resource: URI,
+		@IPositronPlotsService private readonly _positronPlotsService: IPositronPlotsService
 	) { super(); }
 
 	override dispose(): void {
+		const plotClient = this._positronPlotsService.getEditorInstance(this.resource.path);
+		if (plotClient) {
+			this._positronPlotsService.unregisterPlotClient(plotClient);
+		}
 		super.dispose();
 	}
 

src/vs/workbench/services/positronPlots/common/positronPlots.ts

@@ -172,6 +172,8 @@ export interface IPositronPlotsService {
 
 	getEditorInstance(id: string): IPositronPlotClient | undefined;
 
+	unregisterPlotClient(plotClient: IPositronPlotClient): void;
+
 	/**
 	 * Placeholder for service initialization.
 	 */