Fall back to open Positron notebooks in notebooks.getContext when none is the active editor pane

src/positron-dts/positron.d.ts

@@ -3785,13 +3785,19 @@ declare module 'positron' {
 		/**
 		 * Get context about the active notebook.
 		 *
+		 * When no Positron notebook is the active editor pane (e.g. focus is
+		 * in another editor or view), falls back to the open notebook attached
+		 * to the foreground session, then to the most recently active Positron
+		 * notebook that is still open.
+		 *
 		 * Resolves with `undefined` when no notebook is open. Rejects with an
-		 * actionable error when a notebook is open but in an editor other than
-		 * the Positron Notebook Editor (e.g. the built-in/Jupyter notebook
-		 * editor), since the notebook API only operates on Positron Notebook
-		 * Editor instances; the error message explains how to switch editors.
+		 * actionable error when the active editor holds a notebook in an
+		 * editor other than the Positron Notebook Editor (e.g. the
+		 * built-in/Jupyter notebook editor), since the notebook API only
+		 * operates on Positron Notebook Editor instances; the error message
+		 * explains how to switch editors.
 		 *
-		 * @returns The notebook context, or undefined if no notebook is active
+		 * @returns The notebook context, or undefined if no notebook is open
 		 */
 		export function getContext(): Thenable<NotebookContext | undefined>;
 

src/vs/workbench/api/browser/positron/mainThreadNotebookFeatures.ts

@@ -23,6 +23,8 @@ import { rasterizeSvgToPng } from '../../../contrib/positronNotebook/browser/svg
 import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
 import { POSITRON_NOTEBOOK_ASSISTANT_AUTO_FOLLOW_KEY } from '../../../contrib/positronNotebook/common/positronNotebookConfig.js';
 import { IRuntimeSessionService } from '../../../services/runtimeSession/common/runtimeSessionService.js';
+import { EditorsOrder } from '../../../common/editor.js';
+import { POSITRON_NOTEBOOK_EDITOR_INPUT_ID } from '../../../contrib/positronNotebook/common/positronNotebookCommon.js';
 
 /**
  * Main thread implementation of notebook features for extension host communication.
@@ -90,12 +92,15 @@ export class MainThreadNotebookFeatures implements MainThreadNotebookFeaturesSha
 	}
 
 	/**
-	 * Gets the context information for the currently active notebook.
-	 * @returns The notebook context DTO, or undefined if no notebook is active.
+	 * Gets the context information for the currently active notebook. When no
+	 * Positron notebook is the active editor pane, falls back to the open
+	 * notebook attached to the foreground session, then to the most recently
+	 * active Positron notebook that is still open.
+	 * @returns The notebook context DTO, or undefined if no notebook is open.
 	 */
 	async $getActiveNotebookContext(): Promise<INotebookContextDTO | undefined> {
 		// Use existing helper function instead of service method
-		const instance = getNotebookInstanceFromActiveEditorPane(this._editorService);
+		let instance = getNotebookInstanceFromActiveEditorPane(this._editorService);
 		if (!instance) {
 			// A notebook may be open, but in the built-in editor rather than the
 			// Positron Notebook Editor that this API operates on. Surface an
@@ -106,7 +111,14 @@ export class MainThreadNotebookFeatures implements MainThreadNotebookFeaturesSha
 			if (unsupportedEditorMessage) {
 				throw new Error(unsupportedEditorMessage);
 			}
-			return undefined;
+			// A Positron notebook can be open without being the active editor
+			// pane (a Data Explorer or plot editor took the pane, focus moved
+			// to another tab mid-turn). Fall back to an open notebook so
+			// assistant tools can still find it (#14762).
+			instance = this._getFallbackNotebookInstance();
+			if (!instance) {
+				return undefined;
+			}
 		}
 
 		// Get current state from observables
@@ -703,6 +715,34 @@ export class MainThreadNotebookFeatures implements MainThreadNotebookFeaturesSha
 		return instance.textModel;
 	}
 
+	/**
+	 * Finds an open Positron notebook instance when none is the active editor
+	 * pane. Prefers the notebook attached to the foreground session (what the
+	 * interpreter picker shows, and what assistant notebook mode is keyed on),
+	 * then falls back to the notebook whose editor was most recently active.
+	 * @returns The notebook instance, or undefined if no Positron notebook
+	 * editor is open.
+	 */
+	private _getFallbackNotebookInstance(): IPositronNotebookInstance | undefined {
+		const foregroundNotebookUri = this._runtimeSessionService.foregroundSession?.metadata.notebookUri;
+		if (foregroundNotebookUri) {
+			const instances = this._positronNotebookService.listInstances(foregroundNotebookUri);
+			if (instances.length > 0) {
+				return instances[0];
+			}
+		}
+		for (const { editor } of this._editorService.getEditors(EditorsOrder.MOST_RECENTLY_ACTIVE)) {
+			if (editor.typeId !== POSITRON_NOTEBOOK_EDITOR_INPUT_ID || !editor.resource) {
+				continue;
+			}
+			const instances = this._positronNotebookService.listInstances(editor.resource);
+			if (instances.length > 0) {
+				return instances[0];
+			}
+		}
+		return undefined;
+	}
+
 	/**
 	 * Helper method to get a notebook instance by URI string.
 	 * @param uriString The notebook URI as a string.

src/vs/workbench/api/test/browser/positron/mainThreadNotebookFeatures.vitest.ts

@@ -9,14 +9,22 @@ import { stubInterface } from '../../../../../test/vitest/stubInterface.js';
 import { createTestContainer } from '../../../../../test/vitest/positronTestContainer.js';
 import { encodeBase64, VSBuffer } from '../../../../../base/common/buffer.js';
 import { constObservable } from '../../../../../base/common/observable.js';
+import { URI } from '../../../../../base/common/uri.js';
+import { isEqual } from '../../../../../base/common/resources.js';
 import { IExtHostContext } from '../../../../services/extensions/common/extHostCustomers.js';
 import { IEditorService } from '../../../../services/editor/common/editorService.js';
-import { IRuntimeSessionService } from '../../../../services/runtimeSession/common/runtimeSessionService.js';
+import { ILanguageRuntimeSession, IRuntimeSessionMetadata, IRuntimeSessionService } from '../../../../services/runtimeSession/common/runtimeSessionService.js';
 import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
 import { ILogService } from '../../../../../platform/log/common/log.js';
+import { EditorsOrder, IVisibleEditorPane } from '../../../../common/editor.js';
+import { EditorInput } from '../../../../common/editor/editorInput.js';
 import { IPositronNotebookService } from '../../../../contrib/positronNotebook/browser/positronNotebookService.js';
 import { IPositronNotebookInstance } from '../../../../contrib/positronNotebook/browser/IPositronNotebookInstance.js';
 import { IPositronNotebookCodeCell, NotebookCellOutputs } from '../../../../contrib/positronNotebook/browser/PositronNotebookCells/IPositronNotebookCell.js';
+import { SelectionState, SelectionStateMachine } from '../../../../contrib/positronNotebook/browser/selectionMachine.js';
+import { UNSUPPORTED_NOTEBOOK_EDITOR_MESSAGE } from '../../../../contrib/positronNotebook/browser/notebookUtils.js';
+import { POSITRON_NOTEBOOK_EDITOR_INPUT_ID } from '../../../../contrib/positronNotebook/common/positronNotebookCommon.js';
+import { NOTEBOOK_EDITOR_ID } from '../../../../contrib/notebook/common/notebookCommon.js';
 import { MainThreadNotebookFeatures } from '../../../browser/positron/mainThreadNotebookFeatures.js';
 
 const { mockRasterizeSvgToPng } = vi.hoisted(() => ({ mockRasterizeSvgToPng: vi.fn() }));
@@ -110,3 +118,184 @@ describe('MainThreadNotebookFeatures $getCellOutputs SVG handling', () => {
 		features.dispose();
 	});
 });
+
+describe('MainThreadNotebookFeatures $getActiveNotebookContext', () => {
+	createTestContainer().build();
+
+	const TEXT_FILE_EDITOR_PANE_ID = 'workbench.editors.files.textFileEditor';
+	const TEXT_FILE_EDITOR_INPUT_ID = 'workbench.editors.files.fileEditorInput';
+
+	/** A notebook instance stub with no cells, no kernel, and no selection. */
+	function createNotebookInstance(uriString: string): IPositronNotebookInstance {
+		return stubInterface<IPositronNotebookInstance>({
+			uri: URI.parse(uriString),
+			cells: constObservable([]),
+			kernel: constObservable(undefined),
+			selectionStateMachine: stubInterface<SelectionStateMachine>({
+				state: constObservable({ type: SelectionState.NoCells }),
+			}),
+		});
+	}
+
+	/** An editor input stub as it appears in the editor service's MRU list. */
+	function createEditorInput(typeId: string, uriString: string): EditorInput {
+		return stubInterface<EditorInput>({
+			typeId,
+			resource: URI.parse(uriString),
+		});
+	}
+
+	/** A foreground session stub attached to the given notebook. */
+	function createForegroundNotebookSession(uriString: string): ILanguageRuntimeSession {
+		return stubInterface<ILanguageRuntimeSession>({
+			metadata: stubInterface<IRuntimeSessionMetadata>({
+				notebookUri: URI.parse(uriString),
+			}),
+		});
+	}
+
+	/**
+	 * Builds a MainThreadNotebookFeatures against a stubbed editor state: the
+	 * given active pane, editors in most-recently-active order, the open
+	 * Positron notebook instances, and optionally the foreground session.
+	 */
+	function createContextFeatures(options: {
+		activeEditorPane: IVisibleEditorPane | undefined;
+		mruEditors: EditorInput[];
+		instances: IPositronNotebookInstance[];
+		foregroundSession?: ILanguageRuntimeSession;
+	}): MainThreadNotebookFeatures {
+		const editorService = stubInterface<IEditorService>({
+			activeEditorPane: options.activeEditorPane,
+			// Order-sensitive: only the most-recently-active query sees the
+			// editors, so a regression to another EditorsOrder fails here.
+			getEditors: (order: EditorsOrder) => order === EditorsOrder.MOST_RECENTLY_ACTIVE
+				? options.mruEditors.map((editor, groupId) => ({ groupId, editor }))
+				: [],
+		});
+		const notebookService = stubInterface<IPositronNotebookService>({
+			listInstances: (uri?: URI) => options.instances.filter(instance => !uri || isEqual(instance.uri, uri)),
+		});
+		return new MainThreadNotebookFeatures(
+			stubInterface<IExtHostContext>(),
+			editorService,
+			notebookService,
+			stubInterface<ILogService>(),
+			stubInterface<IConfigurationService>(),
+			stubInterface<IRuntimeSessionService>({
+				getNotebookSessionForNotebookUri: () => undefined,
+				foregroundSession: options.foregroundSession,
+			}),
+		);
+	}
+
+	it('resolves an open Positron notebook when it is not the active editor pane (#14762)', async () => {
+		const notebookUri = 'file:///test/notebook.ipynb';
+		const notebook = createNotebookInstance(notebookUri);
+		// Focus is elsewhere: the user is typing in another editor (chat
+		// editor, split group, another file tab) while the notebook stays
+		// open in its own tab.
+		const features = createContextFeatures({
+			activeEditorPane: stubInterface<IVisibleEditorPane>({ getId: () => TEXT_FILE_EDITOR_PANE_ID }),
+			mruEditors: [
+				createEditorInput(TEXT_FILE_EDITOR_INPUT_ID, 'file:///test/script.py'),
+				createEditorInput(POSITRON_NOTEBOOK_EDITOR_INPUT_ID, notebookUri),
+			],
+			instances: [notebook],
+		});
+
+		const context = await features.$getActiveNotebookContext();
+
+		expect(context?.uri).toBe(notebookUri);
+		features.dispose();
+	});
+
+	it('resolves the most recently active notebook when multiple are open and none is active', async () => {
+		const olderUri = 'file:///test/older.ipynb';
+		const recentUri = 'file:///test/recent.ipynb';
+		// Registration order (older first) deliberately differs from MRU
+		// order (recent first): the fallback must follow the editor
+		// service's most-recently-active order, not instance registration.
+		const features = createContextFeatures({
+			activeEditorPane: stubInterface<IVisibleEditorPane>({ getId: () => TEXT_FILE_EDITOR_PANE_ID }),
+			mruEditors: [
+				createEditorInput(TEXT_FILE_EDITOR_INPUT_ID, 'file:///test/script.py'),
+				createEditorInput(POSITRON_NOTEBOOK_EDITOR_INPUT_ID, recentUri),
+				createEditorInput(POSITRON_NOTEBOOK_EDITOR_INPUT_ID, olderUri),
+			],
+			instances: [createNotebookInstance(olderUri), createNotebookInstance(recentUri)],
+		});
+
+		const context = await features.$getActiveNotebookContext();
+
+		expect(context?.uri).toBe(recentUri);
+		features.dispose();
+	});
+
+	it('prefers the foreground session\'s notebook over a more recently active notebook editor', async () => {
+		const attachedUri = 'file:///test/attached.ipynb';
+		const recentUri = 'file:///test/recent.ipynb';
+		// The user's session (what the interpreter picker shows, and what the
+		// assistant's notebook mode is keyed on) is attached to one notebook
+		// while another notebook's editor was touched more recently: the
+		// attached notebook wins.
+		const features = createContextFeatures({
+			activeEditorPane: stubInterface<IVisibleEditorPane>({ getId: () => TEXT_FILE_EDITOR_PANE_ID }),
+			mruEditors: [
+				createEditorInput(POSITRON_NOTEBOOK_EDITOR_INPUT_ID, recentUri),
+				createEditorInput(POSITRON_NOTEBOOK_EDITOR_INPUT_ID, attachedUri),
+			],
+			instances: [createNotebookInstance(attachedUri), createNotebookInstance(recentUri)],
+			foregroundSession: createForegroundNotebookSession(attachedUri),
+		});
+
+		const context = await features.$getActiveNotebookContext();
+
+		expect(context?.uri).toBe(attachedUri);
+		features.dispose();
+	});
+
+	it('falls back to the most recently active notebook editor when the foreground notebook is closed', async () => {
+		const closedUri = 'file:///test/closed.ipynb';
+		const openUri = 'file:///test/open.ipynb';
+		// The foreground session's notebook editor was closed (its session may
+		// still be running); resolution must not target a closed notebook.
+		const features = createContextFeatures({
+			activeEditorPane: stubInterface<IVisibleEditorPane>({ getId: () => TEXT_FILE_EDITOR_PANE_ID }),
+			mruEditors: [createEditorInput(POSITRON_NOTEBOOK_EDITOR_INPUT_ID, openUri)],
+			instances: [createNotebookInstance(openUri)],
+			foregroundSession: createForegroundNotebookSession(closedUri),
+		});
+
+		const context = await features.$getActiveNotebookContext();
+
+		expect(context?.uri).toBe(openUri);
+		features.dispose();
+	});
+
+	it('resolves undefined when no Positron notebook is open anywhere', async () => {
+		const features = createContextFeatures({
+			activeEditorPane: stubInterface<IVisibleEditorPane>({ getId: () => TEXT_FILE_EDITOR_PANE_ID }),
+			mruEditors: [createEditorInput(TEXT_FILE_EDITOR_INPUT_ID, 'file:///test/script.py')],
+			instances: [],
+		});
+
+		expect(await features.$getActiveNotebookContext()).toBeUndefined();
+		features.dispose();
+	});
+
+	it('surfaces the unsupported-editor error when the built-in notebook editor is active, even if a Positron notebook is open elsewhere', async () => {
+		// The user is looking at a notebook in the built-in editor; falling
+		// back to a different open Positron notebook would make assistant
+		// tools operate on the wrong notebook.
+		const notebookUri = 'file:///test/notebook.ipynb';
+		const features = createContextFeatures({
+			activeEditorPane: stubInterface<IVisibleEditorPane>({ getId: () => NOTEBOOK_EDITOR_ID }),
+			mruEditors: [createEditorInput(POSITRON_NOTEBOOK_EDITOR_INPUT_ID, notebookUri)],
+			instances: [createNotebookInstance(notebookUri)],
+		});
+
+		await expect(features.$getActiveNotebookContext()).rejects.toThrow(UNSUPPORTED_NOTEBOOK_EDITOR_MESSAGE);
+		features.dispose();
+	});
+});