Add plot thumbnail cache feature so plots pane shows thumbnails after a reload

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

@@ -8,6 +8,7 @@ import React, { useEffect, useState } from 'react';
 
 // Other dependencies.
 import { PlaceholderThumbnail } from './placeholderThumbnail.js';
+import { usePositronPlotsContext } from '../positronPlotsContext.js';
 import { PlotClientInstance } from '../../../../services/languageRuntime/common/languageRuntimePlotClient.js';
 
 /**
@@ -24,21 +25,25 @@ interface DynamicPlotThumbnailProps {
  * @returns The rendered component.
  */
 export const DynamicPlotThumbnail = (props: DynamicPlotThumbnailProps) => {
-
-	const [uri, setUri] = useState('');
-
-	useEffect(() => {
-		// If the plot is already rendered, show the URI; otherwise, wait for
-		// the plot to render.
+	const context = usePositronPlotsContext();
+	const [uri, setUri] = useState(() => {
+		// If the plot is already rendered, set the URI; otherwise, try to use the cached URI until
+		// the plot is rendered.
 		if (props.plotClient.lastRender) {
-			setUri(props.plotClient.lastRender.uri);
+			return props.plotClient.lastRender.uri;
+		} else {
+			return context.positronPlotsService.getCachedPlotThumbnailURI(props.plotClient.id);
 		}
+	});
 
+	useEffect(() => {
 		// When the plot is rendered, update the URI. This can happen multiple times if the plot
 		// is resized.
-		props.plotClient.onDidCompleteRender((result) => {
+		const disposable = props.plotClient.onDidCompleteRender(result => {
 			setUri(result.uri);
 		});
+
+		return () => disposable.dispose();
 	}, [props.plotClient]);
 
 	// If the plot is not yet rendered yet (no URI), show a placeholder;

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

@@ -7,8 +7,9 @@
 import React, { useEffect, useState } from 'react';
 
 // Other dependencies.
-import { PlaceholderThumbnail } from './placeholderThumbnail.js';
 import { WebviewPlotClient } from '../webviewPlotClient.js';
+import { PlaceholderThumbnail } from './placeholderThumbnail.js';
+import { usePositronPlotsContext } from '../positronPlotsContext.js';
 
 /**
  * WebviewPlotThumbnailProps interface.
@@ -25,25 +26,26 @@ interface WebviewPlotThumbnailProps {
  * @returns The rendered component.
  */
 export const WebviewPlotThumbnail = (props: WebviewPlotThumbnailProps) => {
-
-	const [uri, setUri] = useState('');
-
-	useEffect(() => {
-		// If the plot is already rendered, show the URI; otherwise, wait for
-		// the plot to render.
+	const context = usePositronPlotsContext();
+	const [uri, setUri] = useState(() => {
+		// If the plot is already rendered, set the URI; otherwise, try to use the cached URI until
+		// the plot is rendered.
 		if (props.plotClient.thumbnailUri) {
-			setUri(props.plotClient.thumbnailUri);
+			return props.plotClient.thumbnailUri;
+		} else {
+			return context.positronPlotsService.getCachedPlotThumbnailURI(props.plotClient.id);
 		}
+	});
 
+	useEffect(() => {
 		// When the plot is rendered, update the URI. This can happen multiple times if the plot
 		// is resized.
 		const disposable = props.plotClient.onDidRenderThumbnail((result) => {
 			setUri(result);
 		});
-		return () => {
-			disposable.dispose();
-		};
-	}, [props.plotClient]);
+
+		return () => disposable.dispose();
+	}, [context.positronPlotsService, props.plotClient]);
 
 	// If the plot is not yet rendered yet (no URI), show a placeholder;
 	// otherwise, show the rendered thumbnail.

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

@@ -67,6 +67,9 @@ const WebviewPlotInactiveTimeout = 120_000;
 /** Interval in milliseconds at which inactive webview plots are checked. */
 const WebviewPlotInactiveInterval = 1_000;
 
+/** The key used to store the cached plot thumbnail descriptors */
+const CachedPlotThumbnailDescriptorsKey = 'positron.plots.cachedPlotThumbnailDescriptors';
+
 /** The key used to store the preferred history policy */
 const HistoryPolicyStorageKey = 'positron.plots.historyPolicy';
 
@@ -86,12 +89,23 @@ interface DataUri {
 }
 
 /**
- * PositronPlotsService class.
+ * ICachedPlotThumbnailDescriptor interface.
  */
+interface ICachedPlotThumbnailDescriptor {
+	readonly plotClientId: string;
+	readonly thumbnailURI: string;
+}
+
+/**
+* PositronPlotsService class.
+*/
 export class PositronPlotsService extends Disposable implements IPositronPlotsService {
 	/** Needed for service branding in dependency injector. */
 	declare readonly _serviceBrand: undefined;
 
+	/** The map of cached plot thumbnail descriptors. */
+	private readonly _cachedPlotThumbnailDescriptors = new Map<string, ICachedPlotThumbnailDescriptor>();
+
 	/** The list of Positron plots. */
 	private readonly _plots: IPositronPlotClient[] = [];
 
@@ -224,8 +238,7 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 			await this.registerWebviewPlotClient(plotClient);
 		}));
 
-		// When the storage service is about to save state, store the current history policy
-		// and storage policy in the workspace storage.
+		// When the storage service is about to save state, store policies and cached plot thumbnail descriptors.
 		this._register(this._storageService.onWillSaveState(() => {
 			this._storageService.store(
 				HistoryPolicyStorageKey,
@@ -254,6 +267,45 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 					StorageScope.WORKSPACE,
 					StorageTarget.MACHINE);
 			}
+
+			// Enumerate the plot clients and update the cached plot thumbnail descriptors.
+			const keysToDelete: Set<string> = new Set(this._cachedPlotThumbnailDescriptors.keys());
+			this._plots.forEach(plotClient => {
+				keysToDelete.delete(plotClient.id);
+				if (plotClient instanceof PlotClientInstance) {
+					if (plotClient.lastRender?.uri) {
+						this._cachedPlotThumbnailDescriptors.set(plotClient.id, {
+							plotClientId: plotClient.id,
+							thumbnailURI: plotClient.lastRender.uri
+						});
+					}
+				} else if (plotClient instanceof HtmlPlotClient) {
+					if (plotClient.thumbnailUri) {
+						this._cachedPlotThumbnailDescriptors.set(plotClient.id, {
+							plotClientId: plotClient.id,
+							thumbnailURI: plotClient.thumbnailUri
+						});
+					}
+				}
+			});
+
+			// Delete any cached plot thumbnail descriptors that are no longer valid.
+			keysToDelete.forEach(key => this._cachedPlotThumbnailDescriptors.delete(key));
+
+			// Update the cached plot thumbnail descriptors in workspace storage.
+			if (this._cachedPlotThumbnailDescriptors.size) {
+				this._storageService.store(
+					CachedPlotThumbnailDescriptorsKey,
+					JSON.stringify([...this._cachedPlotThumbnailDescriptors.values()]),
+					StorageScope.WORKSPACE,
+					StorageTarget.MACHINE);
+			} else {
+				this._storageService.store(
+					CachedPlotThumbnailDescriptorsKey,
+					undefined,
+					StorageScope.WORKSPACE,
+					StorageTarget.MACHINE);
+			}
 		}));
 
 		// Listen for changes to the dark mode configuration
@@ -331,6 +383,22 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 			this._selectedHistoryPolicy = preferredHistoryPolicy as HistoryPolicy;
 		}
 
+		// Load the cached plot thumbnail descriptors from workspace storage.
+		const cachedPlotThumbnailDescriptorsJSON = this._storageService.get(CachedPlotThumbnailDescriptorsKey, StorageScope.WORKSPACE);
+		if (cachedPlotThumbnailDescriptorsJSON) {
+			try {
+				// Parse the cached plot thumbnail descriptors.
+				const cachedPlotThumbnailDescriptors = JSON.parse(cachedPlotThumbnailDescriptorsJSON) as ICachedPlotThumbnailDescriptor[];
+
+				// Initialize the cached plot thumbnail descriptors.
+				for (const cachedPlotThumbnailDescriptor of cachedPlotThumbnailDescriptors) {
+					this._cachedPlotThumbnailDescriptors.set(cachedPlotThumbnailDescriptor.plotClientId, cachedPlotThumbnailDescriptor);
+				}
+			} catch (error) {
+				this._logService.error(`Error parsing cached plot thumbnail descriptors: ${error}`);
+			}
+		}
+
 		// When a plot is selected, update its last selected time.
 		this._register(this._onDidSelectPlot.event(async (id) => {
 			this._lastSelectedTimeByPlotId.set(id, Date.now());
@@ -427,6 +495,15 @@ export class PositronPlotsService extends Disposable implements IPositronPlotsSe
 		return this._selectedDarkFilterMode;
 	}
 
+	/**
+	 * Gets the cached plot thumbnail URI for a given plot ID.
+	 * @param plotId The plot ID to get the thumbnail URI for.
+	 * @returns The thumbnail URI for the plot, or undefined if not found.
+	 */
+	getCachedPlotThumbnailURI(plotId: string) {
+		return this._cachedPlotThumbnailDescriptors.get(plotId)?.thumbnailURI;
+	}
+
 	/**
 	 * Selects a new sizing policy and fires an event indicating that the policy
 	 * has changed.

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

@@ -164,6 +164,13 @@ export interface IPositronPlotsService {
 	 */
 	readonly onDidChangeSizingPolicy: Event<IPositronPlotSizingPolicy>;
 
+	/**
+	 * Gets the cached plot thumbnail URI for a given plot ID.
+	 * @param plotId The plot ID to get the thumbnail URI for.
+	 * @returns The thumbnail URI for the plot, or undefined if not found.
+	 */
+	getCachedPlotThumbnailURI(plotId: string): string | undefined;
+
 	/**
 	 * Selects the plot with the specified ID.
 	 *

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

@@ -208,6 +208,15 @@ export class TestPositronPlotsService extends Disposable implements IPositronPlo
 	 */
 	readonly onDidChangeSizingPolicy = this._onDidChangeSizingPolicyEmitter.event;
 
+	/**
+	 * Gets the cached plot thumbnail URI for a given plot ID.
+	 * @param plotId The plot ID to get the thumbnail URI for.
+	 * @returns The thumbnail URI for the plot, or undefined if not found.
+	 */
+	getCachedPlotThumbnailURI(plotId: string) {
+		// Noop in test implementation. In a real implementation, this would return the URI of the cached thumbnail.
+		return undefined;
+	}
 
 	/**
 	 * Selects the plot with the specified ID.