Clone static plot for editor

Author: timtmok

src/vs/workbench/contrib/positronPlots/browser/positronPlots.tsx

@@ -13,14 +13,13 @@ import React, { PropsWithChildren, useCallback, useEffect, useState } from 'reac
 import { IWorkbenchLayoutService } from '../../../services/layout/browser/layoutService.js';
 import { PositronPlotsServices } from './positronPlotsState.js';
 import { PositronPlotsContextProvider } from './positronPlotsContext.js';
-import { HistoryPolicy, IPositronPlotsService, ZoomLevel } from '../../../services/positronPlots/common/positronPlots.js';
+import { HistoryPolicy, IPositronPlotsService, isZoomablePlotClient, ZoomLevel } from '../../../services/positronPlots/common/positronPlots.js';
 import { DisposableStore } from '../../../../base/common/lifecycle.js';
 import { PlotsContainer } from './components/plotsContainer.js';
 import { ActionBars } from './components/actionBars.js';
 import { INotificationService } from '../../../../platform/notification/common/notification.js';
 import { PositronPlotsViewPane } from './positronPlotsView.js';
 import { IPreferencesService } from '../../../services/preferences/common/preferences.js';
-import { PlotClientInstance } from '../../../services/languageRuntime/common/languageRuntimePlotClient.js';
 
 /**
  * PositronPlotsProps interface.
@@ -66,18 +65,16 @@ export const PositronPlots = (props: PropsWithChildren<PositronPlotsProps>) => {
 	}, [props.positronPlotsService.positronPlotInstances.length, props.reactComponentContainer.height, props.reactComponentContainer.width]);
 
 	const zoomHandler = (zoom: number) => {
-		setZoom(zoom);
 		const currentPlotId = props.positronPlotsService.selectedPlotId;
 		if (!currentPlotId) {
 			return;
 		}
 
 		const plot = props.positronPlotsService.positronPlotInstances.find(plot => plot.id === currentPlotId);
-		if (!plot || !(plot instanceof PlotClientInstance)) {
-			return;
+		if (isZoomablePlotClient(plot)) {
+			// Update the zoom level in the plot metadata.
+			plot.zoomLevel = zoom;
 		}
-		// Update the zoom level in the plot metadata.
-		plot.zoomLevel = zoom;
 	};
 
 	// Hooks.
@@ -142,30 +139,29 @@ export const PositronPlots = (props: PropsWithChildren<PositronPlotsProps>) => {
 
 	useEffect(() => {
 		// Set the initial zoom level for the current plot.
-		const currentPlot = props.positronPlotsService.selectedPlotId;
-
-		if (currentPlot) {
-			const plot = props.positronPlotsService.positronPlotInstances.find(plot => plot.id === currentPlot);
-			if (plot) {
-				setZoom(plot.metadata.zoom_level || ZoomLevel.Fit);
-			}
-		}
-	}, [props.positronPlotsService.positronPlotInstances, props.positronPlotsService.selectedPlotId]);
+		const disposableStore = new DisposableStore();
 
-	useEffect(() => {
-		const disposable = props.positronPlotsService.onDidSelectPlot((selectedPlotId) => {
-			const plot = props.positronPlotsService.positronPlotInstances.find(plot => plot.id === selectedPlotId);
-			if (plot) {
-				setZoom(plot.metadata.zoom_level || ZoomLevel.Fit);
-			} else {
-				setZoom(ZoomLevel.Fit);
+		disposableStore.add(props.positronPlotsService.onDidSelectPlot(plotId => {
+			const currentPlot = props.positronPlotsService.selectedPlotId;
+
+			if (currentPlot) {
+				const plot = props.positronPlotsService.positronPlotInstances.find(plot => plot.id === currentPlot);
+				if (isZoomablePlotClient(plot)) {
+					disposableStore.add(plot.onDidChangeZoomLevel((zoomLevel) => {
+						setZoom(zoomLevel);
+					}));
+					setZoom(plot.zoomLevel);
+				} else {
+					setZoom(ZoomLevel.Fit);
+				}
 			}
-		});
+		}));
 
 		return () => {
-			disposable.dispose();
-		};
-	}, [props.positronPlotsService])
+			// Dispose of the disposable store to clean up event handlers.
+			disposableStore.dispose();
+		}
+	}, [props.positronPlotsService]);
 
 	// Render.
 	return (

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

@@ -16,7 +16,7 @@ import { IQuickInputService, IQuickPick, IQuickPickItem } from '../../../../plat
 import { PLOT_IS_ACTIVE_EDITOR } from '../../positronPlotsEditor/browser/positronPlotsEditor.contribution.js';
 import { PositronPlotsEditorInput } from '../../positronPlotsEditor/browser/positronPlotsEditorInput.js';
 import { IEditorService } from '../../../services/editor/common/editorService.js';
-import { IPositronPlotClient, IPositronPlotsService, ZoomLevel } from '../../../services/positronPlots/common/positronPlots.js';
+import { IPositronPlotClient, IPositronPlotsService, isZoomablePlotClient, ZoomLevel } from '../../../services/positronPlots/common/positronPlots.js';
 import { PlotClientInstance } from '../../../services/languageRuntime/common/languageRuntimePlotClient.js';
 import { ThemeIcon } from '../../../../base/common/themables.js';
 import { Uri } from 'vscode';
@@ -664,7 +664,10 @@ abstract class PlotsEditorZoomAction extends Action2 {
 	 */
 	async run(accessor: ServicesAccessor, plotId: Uri): Promise<void> {
 		const plotsService = accessor.get(IPositronPlotsService);
-		plotsService.setEditorPlotZoom(plotId.path, this.zoomLevel);
+		const plotInstance = plotsService.getEditorInstance(plotId.path);
+		if (isZoomablePlotClient(plotInstance)) {
+			plotInstance.zoomLevel = this.zoomLevel;
+		}
 	}
 }
 

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

@@ -890,7 +890,7 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 		const code = this._recentExecutions.get(message.parent_id) ?? '';
 
 		if (message.kind === RuntimeOutputKind.StaticImage) {
-			return new StaticPlotClient(this._storageService, session.sessionId, message, code);
+			return StaticPlotClient.fromMessage(this._storageService, session.sessionId, message, code);
 		} else if (message.kind === RuntimeOutputKind.PlotWidget) {
 			return new NotebookOutputPlotClient(this._notebookOutputWebviewService, session, message, code);
 		}
@@ -1423,7 +1423,10 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 		}
 
 		if (plotClient instanceof StaticPlotClient) {
-			this._editorPlots.set(plotClient.id, plotClient);
+			// Create a copy of the StaticPlotClient for the editor
+			const plotCopy = StaticPlotClient.fromMetadata(this._storageService, plotClient.metadata, plotClient.mimeType, plotClient.data);
+			this._editorPlots.set(plotClient.id, plotCopy);
+			this._register(plotCopy);
 		}
 
 		// Create a new plot client instance for the editor

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

@@ -17,18 +17,21 @@ export class StaticPlotClient extends Disposable implements IPositronPlotClient,
 	public readonly metadata: IPositronPlotMetadata;
 	public readonly mimeType;
 	public readonly data;
+	public readonly code?: string;
 
 	// Zoom level emitter
 	public onDidChangeZoomLevel: Event<ZoomLevel>;
 
 	private readonly _zoomLevelEventEmitter = new Emitter<ZoomLevel>();
 
-	constructor(storageService: IStorageService, sessionId: string, message: ILanguageRuntimeMessageOutput,
-		public readonly code?: string) {
-		super();
+	static fromMetadata(storageService: IStorageService, metadata: IPositronPlotMetadata, mimeType: string, data: string): StaticPlotClient {
+		// Create a new StaticPlotClient instance from the provided metadata, MIME type, and data.
+		return new StaticPlotClient(storageService, metadata.session_id, metadata, mimeType, data);
+	}
 
+	static fromMessage(storageService: IStorageService, sessionId: string, message: ILanguageRuntimeMessageOutput, code?: string): StaticPlotClient {
 		// Create the metadata for the plot.
-		this.metadata = {
+		const metadata = {
 			id: message.id,
 			parent_id: message.parent_id,
 			created: Date.parse(message.when),
@@ -52,8 +55,18 @@ export class StaticPlotClient extends Disposable implements IPositronPlotClient,
 		}
 
 		// Save the MIME type and data for the image.
-		this.mimeType = imageKey;
+		const mimeType = imageKey;
+
+		return new StaticPlotClient(storageService, sessionId, metadata, mimeType, data);
+	}
+
+	private constructor(storageService: IStorageService, sessionId: string, metadata: IPositronPlotMetadata, mimeType: string, data: string) {
+		super();
+
+		this.metadata = metadata;
+		this.mimeType = mimeType;
 		this.data = data;
+		this.code = metadata.code;
 
 		// Set up the zoom level event emitter.
 		this.onDidChangeZoomLevel = this._zoomLevelEventEmitter.event;