Positron Notebooks: Inline Chat (#9147)

Addresses #8733.

### Release Notes

#### New Features

- N/A

#### Bug Fixes

- N/A

### QA Notes

@:notebooks @:assistant

Author: seeM

src/vs/workbench/api/browser/mainThreadNotebookEditors.ts

@@ -19,6 +19,7 @@ import { ExtHostContext, ExtHostNotebookEditorsShape, INotebookDocumentShowOptio
 // --- Start Positron ---
 import { IPositronNotebookService } from '../../services/positronNotebook/browser/positronNotebookService.js';
 import { PositronNotebookEditorInput } from '../../contrib/positronNotebook/browser/PositronNotebookEditorInput.js';
+import { isEqual } from '../../../base/common/resources.js';
 // --- End Positron ---
 
 class MainThreadNotebook {
@@ -107,16 +108,17 @@ export class MainThreadNotebookEditors implements MainThreadNotebookEditorsShape
 		// --- Start Positron ---
 		// Check if a Positron notebook is already open for this resource
 		const uri = URI.revive(resource);
-		const positronInstance = this._positronNotebookService.getInstance(uri);
-		if (positronInstance && positronInstance.connectedToEditor) {
-			// Find the editor pane containing this Positron notebook
-			for (const editorPane of this._editorService.visibleEditorPanes) {
-				const input = editorPane.input;
-				if (input instanceof PositronNotebookEditorInput && input.resource.toString() === uri.toString()) {
-					// Positron notebook is already open, just return a synthetic ID
-					// We can't return the actual notebook editor ID because Positron notebooks
-					// don't implement INotebookEditor interface
-					return `positron-notebook-${uri.toString()}`;
+		for (const positronInstance of this._positronNotebookService.listInstances(uri)) {
+			if (positronInstance.connectedToEditor) {
+				// Find the editor pane containing this Positron notebook
+				for (const editorPane of this._editorService.visibleEditorPanes) {
+					const input = editorPane.input;
+					if (input instanceof PositronNotebookEditorInput && isEqual(input.resource, uri)) {
+						// Positron notebook is already open, just return a synthetic ID
+						// We can't return the actual notebook editor ID because Positron notebooks
+						// don't implement INotebookEditor interface
+						return `positron-notebook-${uri.toString()}`;
+					}
 				}
 			}
 		}

src/vs/workbench/contrib/inlineChat/browser/inlineChatNotebook.ts

@@ -15,6 +15,10 @@ import { CellUri } from '../../notebook/common/notebookCommon.js';
 import { IEditorService } from '../../../services/editor/common/editorService.js';
 import { NotebookTextDiffEditor } from '../../notebook/browser/diff/notebookDiffEditor.js';
 import { NotebookMultiTextDiffEditor } from '../../notebook/browser/diff/notebookMultiDiffEditor.js';
+// --- Start Positron ---
+// Imports to support inline chat in Positron notebooks.
+import { IPositronNotebookService } from '../../../services/positronNotebook/browser/positronNotebookService.js';
+// --- End Positron ---
 
 export class InlineChatNotebookContribution {
 
@@ -24,6 +28,10 @@ export class InlineChatNotebookContribution {
 		@IInlineChatSessionService sessionService: IInlineChatSessionService,
 		@IEditorService editorService: IEditorService,
 		@INotebookEditorService notebookEditorService: INotebookEditorService,
+		// --- Start Positron ---
+		// Imports to support inline chat in Positron notebooks.
+		@IPositronNotebookService positronNotebookService: IPositronNotebookService,
+		// --- End Positron ---
 	) {
 
 		this._store.add(sessionService.registerSessionKeyComputer(Schemas.vscodeNotebookCell, {
@@ -57,6 +65,19 @@ export class InlineChatNotebookContribution {
 						// 	}
 					}
 				}
+				// --- Start Positron ---
+				// To support inline chat in Positron notebooks:
+				// construct a session comparison key from the corresponding notebook
+				for (const positronInstance of positronNotebookService.listInstances(data.notebook)) {
+					const candidate = `<positron-notebook>${positronInstance.id}#${uri}`;
+					if (!fallback) {
+						fallback = candidate;
+					}
+					if (positronInstance.hasCodeEditor(editor)) {
+						return candidate;
+					}
+				}
+				// --- End Positron ---
 
 				if (fallback) {
 					return fallback;
@@ -96,6 +117,20 @@ export class InlineChatNotebookContribution {
 					}
 				}
 			}
+			// --- Start Positron ---
+			// To support inline chat in Positron notebooks:
+			// cancel existing chat sessions when a new one is started.
+			for (const positronInstance of positronNotebookService.listInstances(candidate.notebook)) {
+				if (positronInstance.hasCodeEditor(newSessionEditor)) {
+					for (const { editor } of positronInstance.cells.get()) {
+						if (editor && editor !== newSessionEditor) {
+							InlineChatController.get(editor)?.acceptSession();
+						}
+					}
+					break;
+				}
+			}
+			// --- End Positron ---
 		}));
 	}
 

src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts

@@ -16,7 +16,7 @@ import { ILanguageSelection, ILanguageService } from '../../../../editor/common/
 import { ITextModelContentProvider, ITextModelService } from '../../../../editor/common/services/resolverService.js';
 import * as nls from '../../../../nls.js';
 // --- Start Positron ---
-// eslint-disable-next-line no-duplicate-imports
+/* eslint-disable no-duplicate-imports */
 import { ConfigurationScope } from '../../../../platform/configuration/common/configurationRegistry.js';
 // --- End Positron ---
 import { Extensions, IConfigurationPropertySchema, IConfigurationRegistry } from '../../../../platform/configuration/common/configurationRegistry.js';
@@ -807,11 +807,9 @@ class NotebookEditorManager implements IWorkbenchContribution {
 			if (model.isDirty() && !this._editorService.isOpened({ resource: model.resource, typeId: NotebookEditorInput.ID, editorId: model.viewType }) && extname(model.resource) !== '.interactive') {
 				// --- Start Positron ---
 				// Make sure that we dont try and open the same editor twice if we're using positron
-				// notebooks. This is a separate if-statement so we don't have to put the diff
-				// inside the inline conditional.
-				const positronNotebookInstance = this._positronNotebookService.getInstance(model.resource);
-				// Check to see if the instance is connected to the view.
-				if (positronNotebookInstance && positronNotebookInstance.connectedToEditor) {
+				// notebooks.
+				const positronInstances = this._positronNotebookService.listInstances(model.resource);
+				if (positronInstances.some(instance => instance.connectedToEditor)) {
 					continue;
 				}
 				// --- End Positron ---

src/vs/workbench/contrib/positronNotebook/browser/PositronNotebookCells/PositronNotebookCell.ts

@@ -14,6 +14,7 @@ import { CodeEditorWidget } from '../../../../../editor/browser/widget/codeEdito
 import { CellSelectionType } from '../../../../services/positronNotebook/browser/selectionMachine.js';
 import { PositronNotebookInstance } from '../PositronNotebookInstance.js';
 import { ISettableObservable, observableValue } from '../../../../../base/common/observable.js';
+import { ICodeEditor } from '../../../../../editor/browser/editorBrowser.js';
 
 export abstract class PositronNotebookCellGeneral extends Disposable implements IPositronNotebookCell {
 	kind!: CellKind;
@@ -30,6 +31,10 @@ export abstract class PositronNotebookCellGeneral extends Disposable implements
 		super();
 	}
 
+	get editor(): ICodeEditor | undefined {
+		return this._editor;
+	}
+
 	get uri(): URI {
 		return this.cellModel.uri;
 	}

src/vs/workbench/contrib/positronNotebook/browser/PositronNotebookEditor.tsx

@@ -12,7 +12,7 @@ import { PixelRatio } from '../../../../base/browser/pixelRatio.js';
 import { ISize, PositronReactRenderer } from '../../../../base/browser/positronReactRenderer.js';
 import { CancellationToken } from '../../../../base/common/cancellation.js';
 import { Emitter } from '../../../../base/common/event.js';
-import { DisposableStore } from '../../../../base/common/lifecycle.js';
+import { DisposableStore, MutableDisposable } from '../../../../base/common/lifecycle.js';
 import { FontMeasurements } from '../../../../editor/browser/config/fontMeasurements.js';
 import { IEditorOptions } from '../../../../editor/common/config/editorOptions.js';
 import { BareFontInfo, FontInfo } from '../../../../editor/common/config/fontInfo.js';
@@ -49,6 +49,7 @@ import { PositronNotebookEditorInput } from './PositronNotebookEditorInput.js';
 import { ILogService } from '../../../../platform/log/common/log.js';
 import { NotebookVisibilityProvider } from './NotebookVisibilityContext.js';
 import { observableValue } from '../../../../base/common/observable.js';
+import { PositronNotebookEditorControl } from './PositronNotebookEditorControl.js';
 
 
 /*
@@ -94,6 +95,10 @@ export class PositronNotebookEditor extends EditorPane {
 	 */
 	private readonly _editorMemento: IEditorMemento<INotebookEditorViewState>;
 
+	/**
+	 * The editor control, used by other features to access the code editor widget of the selected cell.
+	 */
+	private readonly _control = this._register(new MutableDisposable<PositronNotebookEditorControl>());
 
 	private _scopedContextKeyService?: IContextKeyService;
 	private _scopedInstantiationService?: IInstantiationService;
@@ -271,11 +276,22 @@ export class PositronNotebookEditor extends EditorPane {
 			);
 		}
 
+		if (input.notebookInstance === undefined) {
+			throw new Error(
+				'Notebook instance is undefined. This should have been created in the constructor.'
+			);
+		}
 
 		// We're setting the options on the input here so that the input can resolve the model
 		// without having to pass the options to the resolve method.
 		input.editorOptions = options;
 
+		// Update the editor control given the notebook instance.
+		// This has to be done before we `await super.setInput` since that fires events
+		// with listeners that call `this.getControl()` expecting an up-to-date control
+		// i.e. with `activeCodeEditor` being the editor of the selected cell in the notebook.
+		this._control.value = new PositronNotebookEditorControl(input.notebookInstance);
+
 		await super.setInput(input, options, context, token);
 
 		const model = await input.resolve(options);
@@ -299,12 +315,6 @@ export class PositronNotebookEditor extends EditorPane {
 			)
 		);
 
-		if (input.notebookInstance === undefined) {
-			throw new Error(
-				'Notebook instance is undefined. This should have been created in the constructor.'
-			);
-		}
-
 		this._renderReact();
 
 		input.notebookInstance.attachView(this._parentDiv);
@@ -326,6 +336,9 @@ export class PositronNotebookEditor extends EditorPane {
 		// Clear the input observable.
 		this._input = undefined;
 
+		// Clear the editor control.
+		this._control.clear();
+
 		this._disposeReactRenderer();
 
 		// Call the base class's method.
@@ -382,6 +395,10 @@ export class PositronNotebookEditor extends EditorPane {
 
 	}
 
+	override getControl() {
+		return this._control.value;
+	}
+
 	private _fontInfo: FontInfo | undefined;
 	private _dimension?: DOM.Dimension;
 	/**

src/vs/workbench/contrib/positronNotebook/browser/PositronNotebookEditorControl.ts

@@ -0,0 +1,53 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (C) 2025 Posit Software, PBC. All rights reserved.
+ *  Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
+ *--------------------------------------------------------------------------------------------*/
+import { Event } from '../../../../base/common/event.js';
+import { Disposable } from '../../../../base/common/lifecycle.js';
+import { ICompositeCodeEditor, IEditor } from '../../../../editor/common/editorCommon.js';
+import { SelectionState } from '../../../services/positronNotebook/browser/selectionMachine.js';
+import { PositronNotebookInstance } from './PositronNotebookInstance.js';
+
+/**
+ * The PositronNotebookEditorControl is used by features like inline chat, debugging, and outlines
+ * to access the code editor widget of the selected cell in a Positron notebook.
+ *
+ * TODO: Some notebook functionality (possibly debugging and outlines) require that the editor control
+ * also have a `notebookEditor: INotebookEditor` property. We'll need to investigate what that unlocks,
+ * whether to implement INotebookEditor, or find a different solution.
+ */
+export class PositronNotebookEditorControl extends Disposable implements ICompositeCodeEditor {
+	/**
+	 * Event that fires when the active cell, and therefore the active code editor, changes.
+	 */
+	public readonly onDidChangeActiveEditor = Event.None;
+
+	/**
+	 * The active cell's code editor.
+	 */
+	private _activeCodeEditor: IEditor | undefined;
+
+	constructor(
+		private readonly _notebookInstance: PositronNotebookInstance,
+	) {
+		super();
+
+		// Update the active code editor when the notebook selection state changes.
+		this._register(this._notebookInstance.selectionStateMachine.onNewState((state) => {
+			if (state.type === SelectionState.EditingSelection) {
+				this._activeCodeEditor = state.selectedCell.editor;
+			} else if (state.type === SelectionState.NoSelection) {
+				this._activeCodeEditor = undefined;
+			} else {
+				this._activeCodeEditor = state.selected[0]?.editor;
+			}
+		}));
+	}
+
+	/**
+	 * Gets the active cell's code editor.
+	 */
+	public get activeCodeEditor(): IEditor | undefined {
+		return this._activeCodeEditor;
+	}
+}

src/vs/workbench/contrib/positronNotebook/browser/PositronNotebookInstance.ts

@@ -35,6 +35,8 @@ import { ILanguageRuntimeSession, IRuntimeSessionService } from '../../../servic
 import { isEqual } from '../../../../base/common/resources.js';
 import { IPositronWebviewPreloadService } from '../../../services/positronWebviewPreloads/browser/positronWebviewPreloadService.js';
 import { ISettableObservable, observableValue } from '../../../../base/common/observable.js';
+import { ResourceMap } from '../../../../base/common/map.js';
+import { ICodeEditor } from '../../../../editor/browser/editorBrowser.js';
 
 interface IPositronNotebookInstanceRequiredTextModel extends IPositronNotebookInstance {
 	textModel: NotebookTextModel;
@@ -56,7 +58,7 @@ export class PositronNotebookInstance extends Disposable implements IPositronNot
 	// ===== Statics =====
 	// #region Statics
 	/** Map of all active notebook instances, keyed by notebook URI */
-	static _instanceMap: Map<string, PositronNotebookInstance> = new Map();
+	static _instanceMap = new ResourceMap<PositronNotebookInstance>();
 
 	/**
 	 * Either makes or retrieves an instance of a Positron Notebook based on the resource. This
@@ -72,8 +74,7 @@ export class PositronNotebookInstance extends Disposable implements IPositronNot
 		instantiationService: IInstantiationService,
 	): PositronNotebookInstance {
 
-		const pathOfNotebook = input.resource.toString();
-		const existingInstance = PositronNotebookInstance._instanceMap.get(pathOfNotebook);
+		const existingInstance = PositronNotebookInstance._instanceMap.get(input.resource);
 		if (existingInstance) {
 			// Update input
 			existingInstance._input = input;
@@ -84,7 +85,7 @@ export class PositronNotebookInstance extends Disposable implements IPositronNot
 		}
 
 		const instance = instantiationService.createInstance(PositronNotebookInstance, input, creationOptions);
-		PositronNotebookInstance._instanceMap.set(pathOfNotebook, instance);
+		PositronNotebookInstance._instanceMap.set(input.resource, instance);
 		return instance;
 	}
 
@@ -95,20 +96,17 @@ export class PositronNotebookInstance extends Disposable implements IPositronNot
 	 * @param newUri The new URI of the notebook
 	 */
 	static updateInstanceUri(oldUri: URI, newUri: URI): void {
-		const oldKey = oldUri.toString();
-		const newKey = newUri.toString();
-
-		if (oldKey === newKey) {
+		if (isEqual(oldUri, newUri)) {
 			return; // No change needed
 		}
 
-		const instance = PositronNotebookInstance._instanceMap.get(oldKey);
+		const instance = PositronNotebookInstance._instanceMap.get(oldUri);
 		if (instance) {
 			// Remove from old key
-			PositronNotebookInstance._instanceMap.delete(oldKey);
+			PositronNotebookInstance._instanceMap.delete(oldUri);
 
 			// Add to new key - the instance will be updated when getOrCreate is called with the new URI
-			PositronNotebookInstance._instanceMap.set(newKey, instance);
+			PositronNotebookInstance._instanceMap.set(newUri, instance);
 		}
 	}
 
@@ -447,7 +445,7 @@ export class PositronNotebookInstance extends Disposable implements IPositronNot
 		this._logService.info(this.id, 'dispose');
 		this._positronNotebookService.unregisterInstance(this);
 		// Remove from the instance map
-		PositronNotebookInstance._instanceMap.delete(this.uri.toString());
+		PositronNotebookInstance._instanceMap.delete(this.uri);
 
 		super.dispose();
 		this.detachView();
@@ -635,6 +633,15 @@ export class PositronNotebookInstance extends Disposable implements IPositronNot
 		this.selectionStateMachine.selectCell(cell, CellSelectionType.Edit);
 	}
 
+	hasCodeEditor(editor: ICodeEditor): boolean {
+		for (const cell of this._cells) {
+			if (cell.editor && cell.editor === editor) {
+				return true;
+			}
+		}
+		return false;
+	}
+
 	async attachView(container: HTMLElement) {
 		this.detachView();
 		this._container = container;

src/vs/workbench/services/positronNotebook/browser/IPositronNotebookCell.ts

@@ -7,6 +7,7 @@ import { VSBuffer } from '../../../../base/common/buffer.js';
 import { Disposable } from '../../../../base/common/lifecycle.js';
 import { ISettableObservable } from '../../../../base/common/observable.js';
 import { URI } from '../../../../base/common/uri.js';
+import { ICodeEditor } from '../../../../editor/browser/editorBrowser.js';
 import { CodeEditorWidget } from '../../../../editor/browser/widget/codeEditor/codeEditorWidget.js';
 import { NotebookPreloadOutputResults } from '../../positronWebviewPreloads/browser/positronWebviewPreloadService.js';
 
@@ -54,6 +55,11 @@ export interface IPositronNotebookCell extends Disposable {
 	 */
 	cellModel: PositronNotebookCellTextModel;
 
+	/**
+	 * The cell's code editor widget.
+	 */
+	editor: ICodeEditor | undefined;
+
 	/**
 	 * Get the handle number for cell from cell model
 	 */