notebooks create/dispose editors. this means controllers must be created eagerly (😢) and that notebooks need a custom way of plugging comparision keys for session. works unless creating another session for the same cell of a duplicated editor

Author: jrieken

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

@@ -11,11 +11,15 @@ import { IInteractiveEditorService, INTERACTIVE_EDITOR_ID } from 'vs/workbench/c
 import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/common/extensions';
 import { InteractiveEditorServiceImpl } from 'vs/workbench/contrib/interactiveEditor/common/interactiveEditorServiceImpl';
 import { IInteractiveEditorSessionService, InteractiveEditorSessionService } from 'vs/workbench/contrib/interactiveEditor/browser/interactiveEditorSession';
+import { Registry } from 'vs/platform/registry/common/platform';
+import { IWorkbenchContributionsRegistry, Extensions } from 'vs/workbench/common/contributions';
+import { LifecyclePhase } from 'vs/workbench/services/lifecycle/common/lifecycle';
+import { InteractiveEditorNotebookContribution } from 'vs/workbench/contrib/interactiveEditor/browser/interactiveEditorNotebook';
 
 registerSingleton(IInteractiveEditorService, InteractiveEditorServiceImpl, InstantiationType.Delayed);
 registerSingleton(IInteractiveEditorSessionService, InteractiveEditorSessionService, InstantiationType.Delayed);
 
-registerEditorContribution(INTERACTIVE_EDITOR_ID, InteractiveEditorController, EditorContributionInstantiation.Lazy);
+registerEditorContribution(INTERACTIVE_EDITOR_ID, InteractiveEditorController, EditorContributionInstantiation.Eager); // EAGER because of notebook dispose/create of editors
 
 registerAction2(interactiveEditorActions.StartSessionAction);
 registerAction2(interactiveEditorActions.UnstashSessionAction);
@@ -42,3 +46,7 @@ registerAction2(interactiveEditorActions.FeebackUnhelpfulCommand);
 registerAction2(interactiveEditorActions.ApplyPreviewEdits);
 
 registerAction2(interactiveEditorActions.CopyRecordings);
+
+
+Registry.as<IWorkbenchContributionsRegistry>(Extensions.Workbench)
+	.registerWorkbenchContribution(InteractiveEditorNotebookContribution, LifecyclePhase.Restored);

src/vs/workbench/contrib/interactiveEditor/browser/interactiveEditorNotebook.ts

@@ -0,0 +1,35 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { illegalState } from 'vs/base/common/errors';
+import { Schemas } from 'vs/base/common/network';
+import { isEqual } from 'vs/base/common/resources';
+import { IInteractiveEditorSessionService } from 'vs/workbench/contrib/interactiveEditor/browser/interactiveEditorSession';
+import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/services/notebookEditorService';
+import { CellUri } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+
+export class InteractiveEditorNotebookContribution {
+
+	constructor(
+		@IInteractiveEditorSessionService sessionService: IInteractiveEditorSessionService,
+		@INotebookEditorService notebookEditorService: INotebookEditorService,
+	) {
+
+		sessionService.registerSessionKeyComputer(Schemas.vscodeNotebookCell, {
+			getComparisonKey: (_editor, uri) => {
+				const data = CellUri.parse(uri);
+				if (!data) {
+					throw illegalState('Expected notebook');
+				}
+				for (const editor of notebookEditorService.listNotebookEditors()) {
+					if (isEqual(editor.textModel?.uri, data.notebook)) {
+						return `<notebook>${editor.getId()}#${uri}`;
+					}
+				}
+				throw illegalState('Expected notebook');
+			}
+		});
+	}
+}

src/vs/workbench/contrib/interactiveEditor/browser/interactiveEditorSession.ts

@@ -13,7 +13,6 @@ import { EditMode, IInteractiveEditorSessionProvider, IInteractiveEditorSession,
 import { IRange, Range } from 'vs/editor/common/core/range';
 import { IActiveCodeEditor, ICodeEditor } from 'vs/editor/browser/editorBrowser';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
-import { ResourceMap } from 'vs/base/common/map';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { IModelService } from 'vs/editor/common/services/model';
 import { ITextModelService } from 'vs/editor/common/services/resolverService';
@@ -266,6 +265,10 @@ export class EditResponse {
 	}
 }
 
+export interface ISessionKeyComputer {
+	getComparisonKey(editor: ICodeEditor, uri: URI): string;
+}
+
 export const IInteractiveEditorSessionService = createDecorator<IInteractiveEditorSessionService>('IInteractiveEditorSessionService');
 
 export interface IInteractiveEditorSessionService {
@@ -277,6 +280,8 @@ export interface IInteractiveEditorSessionService {
 
 	releaseSession(session: Session): void;
 
+	registerSessionKeyComputer(scheme: string, value: ISessionKeyComputer): IDisposable;
+
 	//
 
 	recordings(): readonly Recording[];
@@ -291,7 +296,8 @@ export class InteractiveEditorSessionService implements IInteractiveEditorSessio
 
 	declare _serviceBrand: undefined;
 
-	private readonly _sessions = new Map<ICodeEditor, ResourceMap<SessionData>>();
+	private readonly _sessions = new Map<string, SessionData>();
+	private readonly _keyComputers = new Map<string, ISessionKeyComputer>();
 	private _recordings: Recording[] = [];
 
 	constructor(
@@ -360,37 +366,27 @@ export class InteractiveEditorSessionService implements IInteractiveEditorSessio
 
 		const session = new Session(options.editMode, editor, textModel0, textModel, provider, raw, wholeRangeDecorationId);
 
-		// store: editor -> uri -> session
-		let map = this._sessions.get(editor);
-		if (!map) {
-			map = new ResourceMap<SessionData>();
-			this._sessions.set(editor, map);
+		// store: key -> session
+		const key = this._key(editor, textModel.uri);
+		if (this._sessions.has(key)) {
+			store.dispose();
+			throw new Error(`Session already stored for ${key}`);
 		}
-		if (map.has(textModel.uri)) {
-			throw new Error(`Session already stored for ${textModel.uri}`);
-		}
-		map.set(textModel.uri, { session, store });
+		this._sessions.set(key, { session, store });
 		return session;
 	}
 
 	releaseSession(session: Session): void {
 
-		const { editor, textModelN } = session;
+		const { editor } = session;
 
 		// cleanup
-		const map = this._sessions.get(editor);
-		if (map) {
-			const data = map.get(textModelN.uri);
-			if (data) {
-				data.store.dispose();
-				data.session.session.dispose?.();
-
-				map.delete(textModelN.uri);
-				this._logService.trace(`[IE] did RELEASE session for ${editor.getId()}, ${session.provider.debugName}`);
-
-			}
-			if (map.size === 0) {
-				this._sessions.delete(editor);
+		for (const [key, value] of this._sessions) {
+			if (value.session === session) {
+				value.store.dispose();
+				this._sessions.delete(key);
+				this._logService.trace(`[IE] did RELEASED session for ${editor.getId()}, ${session.provider.debugName}`);
+				break;
 			}
 		}
 
@@ -405,7 +401,21 @@ export class InteractiveEditorSessionService implements IInteractiveEditorSessio
 	}
 
 	getSession(editor: ICodeEditor, uri: URI): Session | undefined {
-		return this._sessions.get(editor)?.get(uri)?.session;
+		const key = this._key(editor, uri);
+		return this._sessions.get(key)?.session;
+	}
+
+	private _key(editor: ICodeEditor, uri: URI): string {
+		const item = this._keyComputers.get(uri.scheme);
+		return item
+			? item.getComparisonKey(editor, uri)
+			: `${editor.getId()}@${uri.toString()}`;
+
+	}
+
+	registerSessionKeyComputer(scheme: string, value: ISessionKeyComputer): IDisposable {
+		this._keyComputers.set(scheme, value);
+		return toDisposable(() => this._keyComputers.delete(scheme));
 	}
 
 	// --- debug