From af65545bff469e4bec3d44aebc4b84072f1fa516 Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Mon, 1 Apr 2024 17:31:43 +0200
Subject: [PATCH 01/10] use log output channel

---
 .../positron-notebook-controllers/src/extension.ts |  5 ++---
 .../positron-notebook-controllers/src/logging.ts   | 14 --------------
 .../src/notebookController.ts                      | 12 ++++++------
 .../src/notebookControllerManager.ts               |  6 +++---
 4 files changed, 11 insertions(+), 26 deletions(-)
 delete mode 100644 extensions/positron-notebook-controllers/src/logging.ts

diff --git a/extensions/positron-notebook-controllers/src/extension.ts b/extensions/positron-notebook-controllers/src/extension.ts
index 1c08cd948ca..2d08dcfe892 100644
--- a/extensions/positron-notebook-controllers/src/extension.ts
+++ b/extensions/positron-notebook-controllers/src/extension.ts
@@ -4,12 +4,11 @@
 
 import * as positron from 'positron';
 import * as vscode from 'vscode';
-import { initializeLogging } from './logging';
 import { NotebookControllerManager } from './notebookControllerManager';
 
-export async function activate(context: vscode.ExtensionContext): Promise<void> {
-	initializeLogging();
+export const log = vscode.window.createOutputChannel('Positron Notebook Controllers', { log: true });
 
+export async function activate(context: vscode.ExtensionContext): Promise<void> {
 	const manager = new NotebookControllerManager();
 	context.subscriptions.push(manager);
 
diff --git a/extensions/positron-notebook-controllers/src/logging.ts b/extensions/positron-notebook-controllers/src/logging.ts
deleted file mode 100644
index 85d1247f914..00000000000
--- a/extensions/positron-notebook-controllers/src/logging.ts
+++ /dev/null
@@ -1,14 +0,0 @@
-/*---------------------------------------------------------------------------------------------
- *  Copyright (C) 2023 Posit Software, PBC. All rights reserved.
- *--------------------------------------------------------------------------------------------*/
-
-import * as vscode from 'vscode';
-
-let channel: vscode.OutputChannel | undefined;
-export function initializeLogging() {
-	channel = vscode.window.createOutputChannel('Positron Notebook Controllers');
-}
-
-export function trace(message: string) {
-	channel?.appendLine(message);
-}
diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts
index dd4c6c8c91c..138d90766e0 100644
--- a/extensions/positron-notebook-controllers/src/notebookController.ts
+++ b/extensions/positron-notebook-controllers/src/notebookController.ts
@@ -3,7 +3,7 @@
  *--------------------------------------------------------------------------------------------*/
 import * as vscode from 'vscode';
 import * as positron from 'positron';
-import { trace } from './logging';
+import { log } from './extension';
 import { DeferredPromise, delay } from './util';
 
 /**
@@ -102,7 +102,7 @@ export class NotebookController implements vscode.Disposable {
 		}
 
 		if (!session) {
-			trace(`Tried to shutdown runtime for notebook without a running runtime: ${notebook.uri.path}`);
+			log.warn(`[${this.languageId}] Tried to shutdown runtime for notebook without a running runtime: ${notebook.uri.path}`);
 			return;
 		}
 
@@ -110,7 +110,7 @@ export class NotebookController implements vscode.Disposable {
 		session.dispose();
 		this._activeSessionsByNotebook.delete(notebook);
 		this._executionOrderBySessionId.delete(session.metadata.sessionId);
-		trace(`Shutdown runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebook.uri.path}`);
+		log.info(`Shutdown runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebook.uri.path}`);
 	}
 
 	/**
@@ -164,7 +164,7 @@ export class NotebookController implements vscode.Disposable {
 		try {
 			preferredRuntime = await positron.runtime.getPreferredRuntime(this.languageId);
 		} catch (err) {
-			trace(`Getting preferred runtime for language '${this.languageId}' failed. Reason: ${err}`);
+			log.error(`Getting preferred runtime for language '${this.languageId}' failed. Reason: ${err}`);
 			startPromise.error(err);
 			this._startingSessionsByNotebook.delete(notebook);
 
@@ -195,7 +195,7 @@ export class NotebookController implements vscode.Disposable {
 		this._executionOrderBySessionId.set(session.metadata.sessionId, 0);
 		this._startingSessionsByNotebook.delete(notebook);
 		startPromise.complete(session);
-		trace(`Session ${session.metadata.sessionId} is ready`);
+		log.info(`Session ${session.metadata.sessionId} is ready`);
 
 		return session;
 	}
@@ -239,7 +239,7 @@ export class NotebookController implements vscode.Disposable {
 		// Get the execution order for the session.
 		let executionOrder = this._executionOrderBySessionId.get(session.metadata.sessionId);
 		if (executionOrder === undefined) {
-			trace(`No execution order for session ${session.metadata.sessionId}, resetting to 0`);
+			log.error(`No execution order for session ${session.metadata.sessionId}, resetting to 0`);
 			executionOrder = 0;
 		}
 
diff --git a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
index f46fdb45351..f45d1900164 100644
--- a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
+++ b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
@@ -3,7 +3,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
-import { trace } from './logging';
+import { log } from './extension';
 import { NotebookController } from './notebookController';
 
 /**
@@ -24,7 +24,7 @@ export class NotebookControllerManager implements vscode.Disposable {
 		}
 		const controller = new NotebookController(languageId);
 		this.controllers.set(languageId, controller);
-		trace(`Registered notebook controller for language: ${languageId}`);
+		log.info(`Registered notebook controller for language: ${languageId}`);
 	}
 
 	/**
@@ -79,7 +79,7 @@ export class NotebookControllerManager implements vscode.Disposable {
 				? vscode.NotebookControllerAffinity.Preferred
 				: vscode.NotebookControllerAffinity.Default;
 			controller.controller.updateNotebookAffinity(notebook, affinity);
-			trace(`Updated notebook affinity for language: ${languageId}, notebook: ${notebook.uri.path}, affinity: ${affinity}`);
+			log.info(`Updated notebook affinity for language: ${languageId}, notebook: ${notebook.uri.path}, affinity: ${affinity}`);
 		}
 	}
 

From 53b55c99be69538358673fc9a69bb1dfa5cf6153 Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Tue, 2 Apr 2024 14:51:15 +0200
Subject: [PATCH 02/10] extract notebook session service; fix concurrent
 start/shutdown calls; restore sessions

---
 .../src/extension.ts                          |   4 +-
 .../positron-notebook-controllers/src/map.ts  | 123 ++++++++
 .../src/notebookController.ts                 | 133 ++-------
 .../src/notebookControllerManager.ts          |  11 +-
 .../src/notebookSessionService.ts             | 264 ++++++++++++++++++
 src/positron-dts/positron.d.ts                |   7 +
 .../positron/mainThreadLanguageRuntime.ts     |  10 +
 .../positron/extHost.positron.api.impl.ts     |   3 +
 .../positron/extHost.positron.protocol.ts     |   1 +
 .../common/positron/extHostLanguageRuntime.ts |  12 +
 .../runtimeSession/common/runtimeSession.ts   |  38 ++-
 .../common/runtimeSessionService.ts           |   6 +
 .../runtimeStartup/common/runtimeStartup.ts   |  10 +-
 13 files changed, 501 insertions(+), 121 deletions(-)
 create mode 100644 extensions/positron-notebook-controllers/src/map.ts
 create mode 100644 extensions/positron-notebook-controllers/src/notebookSessionService.ts

diff --git a/extensions/positron-notebook-controllers/src/extension.ts b/extensions/positron-notebook-controllers/src/extension.ts
index 2d08dcfe892..370db6008ed 100644
--- a/extensions/positron-notebook-controllers/src/extension.ts
+++ b/extensions/positron-notebook-controllers/src/extension.ts
@@ -5,11 +5,13 @@
 import * as positron from 'positron';
 import * as vscode from 'vscode';
 import { NotebookControllerManager } from './notebookControllerManager';
+import { NotebookSessionService } from './notebookSessionService';
 
 export const log = vscode.window.createOutputChannel('Positron Notebook Controllers', { log: true });
 
 export async function activate(context: vscode.ExtensionContext): Promise<void> {
-	const manager = new NotebookControllerManager();
+	const notebookSessionService = new NotebookSessionService();
+	const manager = new NotebookControllerManager(notebookSessionService);
 	context.subscriptions.push(manager);
 
 	// Register notebook controllers for newly registered runtimes.
diff --git a/extensions/positron-notebook-controllers/src/map.ts b/extensions/positron-notebook-controllers/src/map.ts
new file mode 100644
index 00000000000..791c3e89f86
--- /dev/null
+++ b/extensions/positron-notebook-controllers/src/map.ts
@@ -0,0 +1,123 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Uri as URI } from 'vscode';
+
+interface ResourceMapKeyFn {
+	(resource: URI): string;
+}
+
+class ResourceMapEntry<T> {
+	constructor(readonly uri: URI, readonly value: T) { }
+}
+
+function isEntries<T>(arg: ResourceMap<T> | ResourceMapKeyFn | readonly (readonly [URI, T])[] | undefined): arg is readonly (readonly [URI, T])[] {
+	return Array.isArray(arg);
+}
+
+export class ResourceMap<T> implements Map<URI, T> {
+
+	private static readonly defaultToKey = (resource: URI) => resource.toString();
+
+	readonly [Symbol.toStringTag] = 'ResourceMap';
+
+	private readonly map: Map<string, ResourceMapEntry<T>>;
+	private readonly toKey: ResourceMapKeyFn;
+
+	/**
+	 *
+	 * @param toKey Custom uri identity function, e.g use an existing `IExtUri#getComparison`-util
+	 */
+	constructor(toKey?: ResourceMapKeyFn);
+
+	/**
+	 *
+	 * @param other Another resource which this maps is created from
+	 * @param toKey Custom uri identity function, e.g use an existing `IExtUri#getComparison`-util
+	 */
+	constructor(other?: ResourceMap<T>, toKey?: ResourceMapKeyFn);
+
+	/**
+	 *
+	 * @param other Another resource which this maps is created from
+	 * @param toKey Custom uri identity function, e.g use an existing `IExtUri#getComparison`-util
+	 */
+	constructor(entries?: readonly (readonly [URI, T])[], toKey?: ResourceMapKeyFn);
+
+	constructor(arg?: ResourceMap<T> | ResourceMapKeyFn | readonly (readonly [URI, T])[], toKey?: ResourceMapKeyFn) {
+		if (arg instanceof ResourceMap) {
+			this.map = new Map(arg.map);
+			this.toKey = toKey ?? ResourceMap.defaultToKey;
+		} else if (isEntries(arg)) {
+			this.map = new Map();
+			this.toKey = toKey ?? ResourceMap.defaultToKey;
+
+			for (const [resource, value] of arg) {
+				this.set(resource, value);
+			}
+		} else {
+			this.map = new Map();
+			this.toKey = arg ?? ResourceMap.defaultToKey;
+		}
+	}
+
+	set(resource: URI, value: T): this {
+		this.map.set(this.toKey(resource), new ResourceMapEntry(resource, value));
+		return this;
+	}
+
+	get(resource: URI): T | undefined {
+		return this.map.get(this.toKey(resource))?.value;
+	}
+
+	has(resource: URI): boolean {
+		return this.map.has(this.toKey(resource));
+	}
+
+	get size(): number {
+		return this.map.size;
+	}
+
+	clear(): void {
+		this.map.clear();
+	}
+
+	delete(resource: URI): boolean {
+		return this.map.delete(this.toKey(resource));
+	}
+
+	forEach(clb: (value: T, key: URI, map: Map<URI, T>) => void, thisArg?: any): void {
+		if (typeof thisArg !== 'undefined') {
+			clb = clb.bind(thisArg);
+		}
+		for (const [_, entry] of this.map) {
+			clb(entry.value, entry.uri, <any>this);
+		}
+	}
+
+	*values(): IterableIterator<T> {
+		for (const entry of this.map.values()) {
+			yield entry.value;
+		}
+	}
+
+	*keys(): IterableIterator<URI> {
+		for (const entry of this.map.values()) {
+			yield entry.uri;
+		}
+	}
+
+	*entries(): IterableIterator<[URI, T]> {
+		for (const entry of this.map.values()) {
+			yield [entry.uri, entry.value];
+		}
+	}
+
+	*[Symbol.iterator](): IterableIterator<[URI, T]> {
+		for (const [, entry] of this.map) {
+			yield [entry.uri, entry.value];
+		}
+	}
+}
diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts
index 138d90766e0..07a728169b0 100644
--- a/extensions/positron-notebook-controllers/src/notebookController.ts
+++ b/extensions/positron-notebook-controllers/src/notebookController.ts
@@ -3,8 +3,7 @@
  *--------------------------------------------------------------------------------------------*/
 import * as vscode from 'vscode';
 import * as positron from 'positron';
-import { log } from './extension';
-import { DeferredPromise, delay } from './util';
+import { INotebookSessionService } from './notebookSessionService';
 
 /**
  * Wraps a vscode.NotebookController for a specific language, and manages a notebook runtime session
@@ -17,15 +16,6 @@ export class NotebookController implements vscode.Disposable {
 	/** The wrapped VSCode notebook controller. */
 	public readonly controller: vscode.NotebookController;
 
-	/**
-	 * A map of sessions currently starting, keyed by notebook. Values are promises that resolve
-	 * when the session has started and is ready to execute code.
-	 */
-	private readonly _startingSessionsByNotebook: Map<vscode.NotebookDocument, DeferredPromise<positron.LanguageRuntimeSession>> = new Map();
-
-	/** A map of the currently active sessions, keyed by notebook. */
-	private readonly _activeSessionsByNotebook: Map<vscode.NotebookDocument, positron.LanguageRuntimeSession> = new Map();
-
 	/** A map of the current execution order, keyed by session ID. */
 	private readonly _executionOrderBySessionId: Map<string, number> = new Map();
 
@@ -34,8 +24,12 @@ export class NotebookController implements vscode.Disposable {
 
 	/**
 	 * @param languageId The language ID for which this controller is responsible.
+	 * @param notebookSessionService The notebook session service.
 	 */
-	constructor(private readonly languageId: string) {
+	constructor(
+		private readonly languageId: string,
+		private readonly notebookSessionService: INotebookSessionService,
+	) {
 		// Create a VSCode notebook controller for this language.
 		this.controller = vscode.notebooks.createNotebookController(
 			`positron-${languageId}`,
@@ -57,15 +51,12 @@ export class NotebookController implements vscode.Disposable {
 		this._disposables.push(this.controller);
 
 		this._disposables.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => {
-			// Wait a few seconds before shutting down the runtime. If this was reached via a window
-			// reload, we want to give the runtime a chance to reconnect.
-			await delay(3000);
-			this.shutdownRuntimeSession(notebook);
+			if (notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
+				await notebookSessionService.shutdownRuntimeSession(notebook.uri);
+			}
 		}));
 
 		this._disposables.push(this.controller.onDidChangeSelectedNotebooks(async (e) => {
-			await this.shutdownRuntimeSession(e.notebook);
-
 			// Has this controller been selected for a notebook?
 			if (e.selected) {
 				// Note that this is also reached when a notebook is opened, if this controller was
@@ -75,6 +66,8 @@ export class NotebookController implements vscode.Disposable {
 					updateNotebookLanguage(e.notebook, this.languageId),
 					this.startNewRuntimeSession(e.notebook),
 				]);
+			} else {
+				await notebookSessionService.shutdownRuntimeSession(e.notebook.uri);
 			}
 		}));
 	}
@@ -84,35 +77,6 @@ export class NotebookController implements vscode.Disposable {
 		return `${this.languageId[0].toUpperCase()}${this.languageId.slice(1)}`;
 	}
 
-	/**
-	 * Shutdown the runtime session for a notebook.
-	 *
-	 * @param notebook Notebook whose runtime to shutdown.
-	 * @returns Promise that resolves when the runtime shutdown sequence has been started.
-	 */
-	private async shutdownRuntimeSession(notebook: vscode.NotebookDocument): Promise<void> {
-		// Get the notebook's session.
-		let session: positron.LanguageRuntimeSession | undefined;
-		const startingSessionPromise = this._startingSessionsByNotebook.get(notebook);
-		if (startingSessionPromise && !startingSessionPromise.isSettled) {
-			// If the runtime is still starting, wait for it to be ready.
-			session = await startingSessionPromise.p;
-		} else {
-			session = this._activeSessionsByNotebook.get(notebook);
-		}
-
-		if (!session) {
-			log.warn(`[${this.languageId}] Tried to shutdown runtime for notebook without a running runtime: ${notebook.uri.path}`);
-			return;
-		}
-
-		await session.shutdown(positron.RuntimeExitReason.Shutdown);
-		session.dispose();
-		this._activeSessionsByNotebook.delete(notebook);
-		this._executionOrderBySessionId.delete(session.metadata.sessionId);
-		log.info(`Shutdown runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebook.uri.path}`);
-	}
-
 	/**
 	 * Start a new runtime session for a notebook.
 	 *
@@ -121,9 +85,9 @@ export class NotebookController implements vscode.Disposable {
 	 */
 	public async startNewRuntimeSession(notebook: vscode.NotebookDocument): Promise<positron.LanguageRuntimeSession> {
 		try {
-			return await this.doStartNewRuntimeSession(notebook);
+			return await this.notebookSessionService.startRuntimeSession(notebook.uri, this.languageId);
 		} catch (err) {
-			const retry = vscode.l10n.t(`Retry`);
+			const retry = vscode.l10n.t("Retry");
 			const selection = await vscode.window.showErrorMessage(
 				vscode.l10n.t(
 					"Starting {0} interpreter for '{1}' failed. Reason: {2}",
@@ -141,65 +105,6 @@ export class NotebookController implements vscode.Disposable {
 		}
 	}
 
-	private async doStartNewRuntimeSession(notebook: vscode.NotebookDocument): Promise<positron.LanguageRuntimeSession> {
-		if (this._activeSessionsByNotebook.has(notebook)) {
-			throw new Error(`Tried to start a runtime for a notebook that already has one: ${notebook.uri.path}`);
-		}
-
-		const startingSessionPromise = this._startingSessionsByNotebook.get(notebook);
-		if (startingSessionPromise && !startingSessionPromise.isSettled) {
-			return startingSessionPromise.p;
-		}
-
-		// Update the starting sessions map. This needs to be set before any awaits.
-		// When a user executes code without a controller selected, they will be presented
-		// with a quickpick. Once they make a selection, this is event is fired, and
-		// the execute handler is called immediately after. We need to ensure that the map
-		// is updated before that happens.
-		const startPromise = new DeferredPromise<positron.LanguageRuntimeSession>();
-		this._startingSessionsByNotebook.set(notebook, startPromise);
-
-		// Get the preferred runtime for this language.
-		let preferredRuntime: positron.LanguageRuntimeMetadata;
-		try {
-			preferredRuntime = await positron.runtime.getPreferredRuntime(this.languageId);
-		} catch (err) {
-			log.error(`Getting preferred runtime for language '${this.languageId}' failed. Reason: ${err}`);
-			startPromise.error(err);
-			this._startingSessionsByNotebook.delete(notebook);
-
-			throw err;
-		}
-
-		// Start a session for the preferred runtime.
-		let session: positron.LanguageRuntimeSession;
-		try {
-			session = await positron.runtime.startLanguageRuntime(
-				preferredRuntime.runtimeId,
-				notebook.uri.path, // Use the notebook's path as the session name.
-				notebook.uri);
-			trace(
-				`Starting session for language runtime ${session.metadata.sessionId} `
-				+ `(language: ${this.label}, name: ${session.runtimeMetadata.runtimeName}, `
-				+ `version: ${session.runtimeMetadata.runtimeVersion}, notebook: ${notebook.uri.path})`
-			);
-		} catch (err) {
-			trace(`Starting session for language runtime ${preferredRuntime.runtimeName} failed. Reason: ${err}`);
-			startPromise.error(err);
-			this._startingSessionsByNotebook.delete(notebook);
-
-			throw err;
-		}
-
-		this._activeSessionsByNotebook.set(notebook, session);
-		this._executionOrderBySessionId.set(session.metadata.sessionId, 0);
-		this._startingSessionsByNotebook.delete(notebook);
-		startPromise.complete(session);
-		log.info(`Session ${session.metadata.sessionId} is ready`);
-
-		return session;
-	}
-
 	/**
 	 * Notebook controller execute handler.
 	 *
@@ -211,12 +116,12 @@ export class NotebookController implements vscode.Disposable {
 	private async executeCells(cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, _controller: vscode.NotebookController) {
 		// Get the notebook's session.
 		let session: positron.LanguageRuntimeSession | undefined;
-		const startingSessionPromise = this._startingSessionsByNotebook.get(notebook);
+		const startingSessionPromise = this.notebookSessionService.getStartingNotebookSessionPromise(notebook.uri);
 		if (startingSessionPromise && !startingSessionPromise.isSettled) {
 			// If the runtime is still starting, wait for it to be ready.
 			session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => startingSessionPromise.p);
 		} else {
-			session = this._activeSessionsByNotebook.get(notebook);
+			session = this.notebookSessionService.getNotebookSession(notebook.uri);
 		}
 
 		// No session has been started for this notebook, start one.
@@ -236,12 +141,8 @@ export class NotebookController implements vscode.Disposable {
 	 * @returns Promise that resolves when the runtime has finished executing the cell.
 	 */
 	private async executeCell(cell: vscode.NotebookCell, session: positron.LanguageRuntimeSession): Promise<void> {
-		// Get the execution order for the session.
-		let executionOrder = this._executionOrderBySessionId.get(session.metadata.sessionId);
-		if (executionOrder === undefined) {
-			log.error(`No execution order for session ${session.metadata.sessionId}, resetting to 0`);
-			executionOrder = 0;
-		}
+		// Get the execution order for the session, default to 0 for the first execution.
+		let executionOrder = this._executionOrderBySessionId.get(session.metadata.sessionId) ?? 0;
 
 		// Create a cell execution.
 		const currentExecution = this.controller.createNotebookCellExecution(cell);
diff --git a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
index f45d1900164..10621c92778 100644
--- a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
+++ b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
@@ -5,6 +5,7 @@
 import * as vscode from 'vscode';
 import { log } from './extension';
 import { NotebookController } from './notebookController';
+import { INotebookSessionService } from './notebookSessionService';
 
 /**
  * Manages notebook controllers.
@@ -13,6 +14,12 @@ export class NotebookControllerManager implements vscode.Disposable {
 	/** Notebook controllers keyed by languageId. */
 	public readonly controllers = new Map<string, NotebookController>();
 
+	/**
+	 *
+	 * @param notebookSessionService The notebook session service.
+	 */
+	constructor(private readonly notebookSessionService: INotebookSessionService) { }
+
 	/**
 	 * Create a notebook controller for a language.
 	 *
@@ -22,7 +29,7 @@ export class NotebookControllerManager implements vscode.Disposable {
 		if (this.controllers.has(languageId)) {
 			throw new Error(`Notebook controller already exists for language: ${languageId}`);
 		}
-		const controller = new NotebookController(languageId);
+		const controller = new NotebookController(languageId, this.notebookSessionService);
 		this.controllers.set(languageId, controller);
 		log.info(`Registered notebook controller for language: ${languageId}`);
 	}
@@ -79,7 +86,7 @@ export class NotebookControllerManager implements vscode.Disposable {
 				? vscode.NotebookControllerAffinity.Preferred
 				: vscode.NotebookControllerAffinity.Default;
 			controller.controller.updateNotebookAffinity(notebook, affinity);
-			log.info(`Updated notebook affinity for language: ${languageId}, notebook: ${notebook.uri.path}, affinity: ${affinity}`);
+			log.debug(`Updated notebook affinity for controller: ${controller.label}, notebook: ${notebook.uri.path}, affinity: ${affinity}`);
 		}
 	}
 
diff --git a/extensions/positron-notebook-controllers/src/notebookSessionService.ts b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
new file mode 100644
index 00000000000..5c2503b7f9b
--- /dev/null
+++ b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
@@ -0,0 +1,264 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (C) 2024 Posit Software, PBC. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as positron from 'positron';
+import { Uri } from 'vscode';
+import { log } from './extension';
+import { ResourceMap } from './map';
+import { DeferredPromise } from './util';
+
+/**
+ * The notebook session service is the main interface for interacting with
+ * runtime sessions; it manages the set of active sessions and provides
+ * facilities for starting, stopping, and interacting with them.
+ *
+ * TODO(seem): The intention is for functionality here to eventually be brought into Positron core.
+ */
+export interface INotebookSessionService {
+	/**
+	 * Checks for a starting notebook for the given notebook URI.
+	 *
+	 * @param notebookUri The notebook URI to check for.
+	 * @returns True if a starting notebook session exists for the given notebook URI.
+	 */
+	hasStartingNotebookSession(notebookUri: Uri): boolean;
+
+	/**
+	 * Checks for a starting or running notebook for the given notebook URI.
+	 *
+	 * @param notebookUri The notebook URI to check for.
+	 * @returns True if a starting or running notebook session exists for the given notebook URI.
+	 */
+	hasStartingOrRunningNotebookSession(notebookUri: Uri): boolean;
+
+	/**
+	 * Get the starting or running notebook session for the given notebook URI, if one exists.
+	 *
+	 * @param notebookUri The notebook URI of the session to retrieve.
+	 * @returns The starting or running notebook session for the given notebook URI, if one exists.
+	 */
+	getStartingNotebookSessionPromise(notebookUri: Uri): DeferredPromise<positron.LanguageRuntimeSession> | undefined;
+
+	/**
+	 * Get the running notebook session for the given notebook URI, if one exists.
+	 *
+	 * @param notebookUri The notebook URI of the session to retrieve.
+	 * @returns The running notebook session for the given notebook URI, if one exists.
+	 */
+	getNotebookSession(notebookUri: Uri): positron.LanguageRuntimeSession | undefined;
+
+	/**
+	 * Start a new runtime session for a notebook.
+	 *
+	 * @param notebookUri The notebook URI to start a runtime for.
+	 * @returns Promise that resolves when the runtime startup sequence has been started.
+	 */
+	startRuntimeSession(notebookUri: Uri, languageId: string): Promise<positron.LanguageRuntimeSession>;
+
+	/**
+	 * Shutdown the runtime session for a notebook.
+	 *
+	 * @param notebookUri The notebook URI whose runtime to shutdown.
+	 * @returns Promise that resolves when the runtime shutdown sequence has been started.
+	 */
+	shutdownRuntimeSession(notebookUri: Uri): Promise<void>;
+
+}
+
+/**
+ * The implementation of INotebookSessionService.
+ */
+export class NotebookSessionService implements INotebookSessionService {
+
+	/**
+	 * A map of sessions currently starting, keyed by notebook URI. Values are promises that resolve
+	 * when the session has started and is ready to execute code.
+	 */
+	private readonly _startingSessionsByNotebookUri = new ResourceMap<DeferredPromise<positron.LanguageRuntimeSession>>();
+
+	/**
+	 * A map of sessions currently shutting down, keyed by notebook URI. Values are promises that resolve
+	 * when the session has completed the shutdown sequence.
+	 */
+	private readonly _shuttingDownSessionsByNotebookUri = new ResourceMap<DeferredPromise<void>>();
+
+	/** A map of the currently active notebook sessions, keyed by notebook URI. */
+	private readonly _notebookSessionsByNotebookUri = new ResourceMap<positron.LanguageRuntimeSession>();
+
+	hasStartingNotebookSession(notebookUri: Uri): boolean {
+		return this._startingSessionsByNotebookUri.has(notebookUri);
+	}
+
+	hasStartingOrRunningNotebookSession(notebookUri: Uri): boolean {
+		return this._startingSessionsByNotebookUri.has(notebookUri) || this._notebookSessionsByNotebookUri.has(notebookUri);
+	}
+
+	getStartingNotebookSessionPromise(notebookUri: Uri): DeferredPromise<positron.LanguageRuntimeSession> | undefined {
+		return this._startingSessionsByNotebookUri.get(notebookUri);
+	}
+
+	getNotebookSession(notebookUri: Uri): positron.LanguageRuntimeSession | undefined {
+		return this._notebookSessionsByNotebookUri.get(notebookUri);
+	}
+
+	async startRuntimeSession(notebookUri: Uri, languageId: string): Promise<positron.LanguageRuntimeSession> {
+		// Return the existing promise, if there is one.
+		const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri);
+		if (startingSessionPromise && !startingSessionPromise.isSettled) {
+			return startingSessionPromise.p;
+		}
+
+		// Update the starting sessions map. This needs to be set before any awaits in case another
+		// caller tries to start a runtime or access the start promise concurrently.
+		const startPromise = new DeferredPromise<positron.LanguageRuntimeSession>();
+		this._startingSessionsByNotebookUri.set(notebookUri, startPromise);
+
+		// If the notebook has a session that is still shutting down, wait for it to finish.
+		const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri);
+		if (shuttingDownSessionPromise && !shuttingDownSessionPromise.isSettled) {
+			await shuttingDownSessionPromise.p;
+		}
+
+		// Ensure that we don't start a runtime for a notebook that already has one.
+		if (this._notebookSessionsByNotebookUri.has(notebookUri)) {
+			const err = new Error(`Tried to start a runtime for a notebook that already has one: ${notebookUri.path}`);
+			startPromise.error(err);
+			this._startingSessionsByNotebookUri.delete(notebookUri);
+			throw err;
+		}
+
+		// If there's already a session for this runtime e.g. one restored after a window reload, use it.
+		let session: positron.LanguageRuntimeSession | undefined;
+		try {
+			session = await positron.runtime.getNotebookSession(notebookUri);
+			if (session) {
+				log.info(
+					`Restored session for language runtime ${session.metadata.sessionId} `
+					+ `(language: ${session.runtimeMetadata.languageName}, name: ${session.runtimeMetadata.runtimeName}, `
+					+ `version: ${session.runtimeMetadata.runtimeVersion}, notebook: ${notebookUri.path})`
+				);
+			}
+		} catch (err) {
+			log.error(
+				`Getting existing session for notebook ${notebookUri.path}' failed. Reason: ${err}`
+			);
+			startPromise.error(err);
+			this._startingSessionsByNotebookUri.delete(notebookUri);
+			throw err;
+		}
+
+		// If we couldn't restore a session, start a new one.
+		if (!session) {
+			// Get the preferred runtime for this language.
+			let preferredRuntime: positron.LanguageRuntimeMetadata;
+			try {
+				preferredRuntime = await positron.runtime.getPreferredRuntime(languageId);
+			} catch (err) {
+				log.error(`Getting preferred runtime for language '${languageId}' failed. Reason: ${err}`);
+				startPromise.error(err);
+				this._startingSessionsByNotebookUri.delete(notebookUri);
+				throw err;
+			}
+
+			try {
+				session = await positron.runtime.startLanguageRuntime(
+					preferredRuntime.runtimeId,
+					notebookUri.path, // Use the notebook's path as the session name.
+					notebookUri);
+				log.info(
+					`Starting session for language runtime ${session.metadata.sessionId} `
+					+ `(language: ${session.runtimeMetadata.languageName}, name: ${session.runtimeMetadata.runtimeName}, `
+					+ `version: ${session.runtimeMetadata.runtimeVersion}, notebook: ${notebookUri.path})`
+				);
+			} catch (err) {
+				log.error(`Starting session for language runtime ${preferredRuntime.runtimeName} failed. Reason: ${err}`);
+				startPromise.error(err);
+				this._startingSessionsByNotebookUri.delete(notebookUri);
+				throw err;
+			}
+		}
+
+		// Complete the promise and update the session maps.
+		this._notebookSessionsByNotebookUri.set(notebookUri, session);
+		this._startingSessionsByNotebookUri.delete(notebookUri);
+		startPromise.complete(session);
+		log.info(`Session ${session.metadata.sessionId} is ready`);
+
+		return session;
+	}
+
+	async shutdownRuntimeSession(notebookUri: Uri): Promise<void> {
+		// Return the existing promise, if there is one.
+		const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri);
+		if (shuttingDownSessionPromise && !shuttingDownSessionPromise.isSettled) {
+			return shuttingDownSessionPromise.p;
+		}
+
+		// Update the shutting down sessions map. This needs to be set before any awaits in case
+		// another caller tries to shutdown a runtime or access the shutdown promise concurrently.
+		const shutDownPromise = new DeferredPromise<void>();
+		this._shuttingDownSessionsByNotebookUri.set(notebookUri, shutDownPromise);
+
+		// Get the notebook's session.
+		let session: positron.LanguageRuntimeSession | undefined;
+		const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri);
+		if (startingSessionPromise && !startingSessionPromise.isSettled) {
+			// If the runtime is still starting, wait for it to be ready.
+			session = await startingSessionPromise.p;
+		} else {
+			// Try to get an already running session.
+			session = this._notebookSessionsByNotebookUri.get(notebookUri);
+		}
+
+		// Ensure that we have a session.
+		if (!session) {
+			const err = new Error(`Tried to shutdown runtime for notebook without a running runtime: ${notebookUri.path}`);
+			this._shuttingDownSessionsByNotebookUri.delete(notebookUri);
+			shutDownPromise.error(err);
+			throw err;
+		}
+
+		// Start the shutdown sequence.
+		try {
+			await session.shutdown(positron.RuntimeExitReason.Shutdown);
+			log.info(`Shutting down runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebookUri.path}`);
+		} catch (err) {
+			log.error(`Shutting down runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebookUri.path} failed. Reason: ${err}`);
+			this._shuttingDownSessionsByNotebookUri.delete(notebookUri);
+			this._notebookSessionsByNotebookUri.delete(notebookUri);
+			shutDownPromise.error(err);
+			throw err;
+		}
+
+		// Wait for the session to end. This is necessary so that we know when to start the next
+		// session for the notebook, since at most one session can exist per notebook.
+		try {
+			const timeout = new Promise<void>((_, reject) => {
+				setTimeout(() => {
+					reject(new Error(`Shutting down runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebookUri.path} timed out`));
+				}, 5000);
+			});
+			const promise = new Promise<void>(resolve => {
+				const disposable = session.onDidEndSession(() => {
+					disposable.dispose();
+					resolve();
+				});
+			});
+			await Promise.race([promise, timeout]);
+		} catch (err) {
+			log.error(err);
+			this._shuttingDownSessionsByNotebookUri.delete(notebookUri);
+			this._notebookSessionsByNotebookUri.delete(notebookUri);
+			shutDownPromise.error(err);
+			throw err;
+		}
+
+		// Complete the promise and update the session maps.
+		this._shuttingDownSessionsByNotebookUri.delete(notebookUri);
+		this._notebookSessionsByNotebookUri.delete(notebookUri);
+		shutDownPromise.complete();
+		log.info(`Session ${session.metadata.sessionId} shutdown completed`);
+	}
+
+}
diff --git a/src/positron-dts/positron.d.ts b/src/positron-dts/positron.d.ts
index 010b9a6d774..960e3f89564 100644
--- a/src/positron-dts/positron.d.ts
+++ b/src/positron-dts/positron.d.ts
@@ -1155,6 +1155,13 @@ declare module 'positron' {
 		 */
 		export function getForegroundSession(): Thenable<LanguageRuntimeSession | undefined>;
 
+		/**
+		 * Get the session corresponding to a notebook, if any.
+		 *
+		 * @param notebookUri The URI of the notebook.
+		 */
+		export function getNotebookSession(notebookUri: vscode.Uri): Thenable<LanguageRuntimeSession | undefined>;
+
 		/**
 		 * Select and start a runtime previously registered with Positron. Any
 		 * previously active runtimes for the language will be shut down.
diff --git a/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts b/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts
index 56429f264b1..f97fa2cc2cd 100644
--- a/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts
+++ b/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts
@@ -1074,6 +1074,16 @@ export class MainThreadLanguageRuntime
 		return Promise.resolve(this._runtimeSessionService.foregroundSession?.sessionId);
 	}
 
+	$getNotebookSession(notebookUri: URI): Promise<string | undefined> {
+		// Revive the URI from the serialized form
+		const uri = URI.revive(notebookUri);
+
+		// Get the session for the notebook URI
+		const session = this._runtimeSessionService.getNotebookSessionForNotebookUri(uri);
+
+		return Promise.resolve(session?.sessionId);
+	}
+
 	// Called by the extension host to select a previously registered language runtime
 	$selectLanguageRuntime(runtimeId: string): Promise<void> {
 		return this._runtimeSessionService.selectRuntime(
diff --git a/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts b/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts
index c7be0e8b059..31425064f37 100644
--- a/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts
+++ b/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts
@@ -84,6 +84,9 @@ export function createPositronApiFactoryAndRegisterActors(accessor: ServicesAcce
 			getForegroundSession(): Thenable<positron.LanguageRuntimeSession | undefined> {
 				return extHostLanguageRuntime.getForegroundSession();
 			},
+			getNotebookSession(notebookUri: vscode.Uri): Thenable<positron.LanguageRuntimeSession | undefined> {
+				return extHostLanguageRuntime.getNotebookSession(notebookUri);
+			},
 			selectLanguageRuntime(runtimeId: string): Thenable<void> {
 				return extHostLanguageRuntime.selectLanguageRuntime(runtimeId);
 			},
diff --git a/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts b/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts
index 34e136b4363..099747dcd72 100644
--- a/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts
+++ b/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts
@@ -38,6 +38,7 @@ export interface MainThreadLanguageRuntimeShape extends IDisposable {
 	$executeCode(languageId: string, code: string, focus: boolean, allowIncomplete?: boolean): Promise<boolean>;
 	$getPreferredRuntime(languageId: string): Promise<ILanguageRuntimeMetadata>;
 	$getForegroundSession(): Promise<string | undefined>;
+	$getNotebookSession(notebookUri: URI): Promise<string | undefined>;
 	$restartSession(handle: number): Promise<void>;
 	$emitLanguageRuntimeMessage(handle: number, message: ILanguageRuntimeMessage): void;
 	$emitLanguageRuntimeState(handle: number, clock: number, state: RuntimeState): void;
diff --git a/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts b/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts
index 66b7aefbb4b..8b9728e4f16 100644
--- a/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts
+++ b/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts
@@ -545,6 +545,18 @@ export class ExtHostLanguageRuntime implements extHostProtocol.ExtHostLanguageRu
 		return session;
 	}
 
+	public async getNotebookSession(notebookUri: URI): Promise<positron.LanguageRuntimeSession | undefined> {
+		const sessionId = await this._proxy.$getNotebookSession(notebookUri);
+		if (!sessionId) {
+			return;
+		}
+		const session = this._runtimeSessions.find(session => session.metadata.sessionId === sessionId);
+		if (!session) {
+			throw new Error(`Session ID '${sessionId}' exists for notebook '${notebookUri.toString()}', but is not known to the extension host.`);
+		}
+		return session;
+	}
+
 	/**
 	 * Registers a new language runtime manager with the extension host.
 	 *
diff --git a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
index acc4775b12c..9b6192e659e 100644
--- a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
+++ b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
@@ -18,6 +18,7 @@ import { IModalDialogPromptInstance, IPositronModalDialogsService } from 'vs/wor
 import { IUiClientMessageInput, IUiClientMessageOutput, UiClientInstance } from 'vs/workbench/services/languageRuntime/common/languageRuntimeUiClient';
 import { UiFrontendEvent } from 'vs/workbench/services/languageRuntime/common/positronUiComm';
 import { ILanguageService } from 'vs/editor/common/languages/language';
+import { ResourceMap } from 'vs/base/common/map';
 
 /**
  * Utility class for tracking state changes in a language runtime session.
@@ -60,6 +61,10 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 	// (metadata.languageId) of the runtime owning the session.
 	private readonly _startingConsolesByLanguageId = new Map<string, ILanguageRuntimeMetadata>();
 
+	// A map of the starting notebooks. This is keyed by the notebook URI
+	// owning the session.
+	private readonly _startingNotebooksByNotebookUri = new ResourceMap<ILanguageRuntimeMetadata>();
+
 	// A map of runtimes currently starting to promises that resolve when the runtime
 	// is ready to use. This is keyed by the runtimeId (metadata.runtimeId) of the runtime.
 	private readonly _startingRuntimesByRuntimeId = new Map<string, DeferredPromise<string>>();
@@ -69,6 +74,10 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 	// languageId (metadata.languageId) of the session.
 	private readonly _consoleSessionsByLanguageId = new Map<string, ILanguageRuntimeSession>();
 
+	// A map of the currently active notebook sessions. This is keyed by the notebook URI
+	// owning the session.
+	private readonly _notebookSessionsByNotebookUri = new ResourceMap<ILanguageRuntimeSession>();
+
 	// The event emitter for the onWillStartRuntime event.
 	private readonly _onWillStartRuntimeEmitter =
 		this._register(new Emitter<IRuntimeSessionWillStartEvent>);
@@ -217,6 +226,17 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 		return this._consoleSessionsByLanguageId.get(runtimeId);
 	}
 
+	/**
+	 * Gets the notebook session for a notebook URI, if one exists.
+	 *
+	 * @param notebookUri The notebook URI of the session to retrieve.
+	 * @returns The notebook session with the given notebook URI, or undefined if
+	 *  no notebook session with the given notebook URI exists.
+	 */
+	getNotebookSessionForNotebookUri(notebookUri: URI): ILanguageRuntimeSession | undefined {
+		return this._notebookSessionsByNotebookUri.get(notebookUri);
+	}
+
 	/**
 	 * Selects and starts a new runtime session, after shutting down any currently active
 	 * sessions for the language.
@@ -448,7 +468,7 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 	/**
 	 * Restarts a runtime session.
 	 *
-	 * @param sessionId The session ID of the runtime to interrupt.
+	 * @param sessionId The session ID of the runtime to restart.
 	 * @param source The source of the request to restart the runtime.
 	 */
 	async restartSession(sessionId: string, source: string): Promise<void> {
@@ -740,6 +760,14 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 			if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) {
 				this._startingConsolesByLanguageId.delete(session.runtimeMetadata.languageId);
 				this._consoleSessionsByLanguageId.set(session.runtimeMetadata.languageId, session);
+			} else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) {
+				if (session.metadata.notebookUri) {
+					this._startingNotebooksByNotebookUri.delete(session.metadata.notebookUri);
+					this._notebookSessionsByNotebookUri.set(session.metadata.notebookUri, session);
+				} else {
+					this._logService.error(`Notebook session ${formatLanguageRuntimeSession(session)} ` +
+						`does not have a notebook URI.`);
+				}
 			}
 
 			// Fire the onDidStartRuntime event.
@@ -811,6 +839,14 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 					this._startingConsolesByLanguageId.delete(session.runtimeMetadata.languageId);
 					if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) {
 						this._consoleSessionsByLanguageId.delete(session.runtimeMetadata.languageId);
+					} else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) {
+						if (session.metadata.notebookUri) {
+							console.log(`WASIM removing notebook session from map`);
+							this._notebookSessionsByNotebookUri.delete(session.metadata.notebookUri);
+						} else {
+							this._logService.error(`Notebook session ${formatLanguageRuntimeSession(session)} ` +
+								`does not have a notebook URI.`);
+						}
 					}
 					break;
 			}
diff --git a/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts b/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts
index b8e7085a331..606d5b8d989 100644
--- a/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts
+++ b/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts
@@ -277,6 +277,12 @@ export interface IRuntimeSessionService {
 	 */
 	getConsoleSessionForLanguage(languageId: string): ILanguageRuntimeSession | undefined;
 
+	/**
+	 * Gets a specific notebook session by notebook URI. Currently, only one
+	 * notebook session can exist per notebook URI.
+	 */
+	getNotebookSessionForNotebookUri(notebookUri: URI): ILanguageRuntimeSession | undefined;
+
 	/**
 	 * Checks for a starting or running console for the given language ID.
 	 *
diff --git a/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts b/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts
index 4f72aceba15..1b3b8f893aa 100644
--- a/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts
+++ b/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts
@@ -20,6 +20,7 @@ import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
 import { ILifecycleService, ShutdownReason, StartupKind } from 'vs/workbench/services/lifecycle/common/lifecycle';
 import { INotificationService } from 'vs/platform/notification/common/notification';
 import { IWorkspaceTrustManagementService } from 'vs/platform/workspace/common/workspaceTrust';
+import { URI } from 'vs/base/common/uri';
 
 interface ILanguageRuntimeProviderMetadata {
 	languageId: string;
@@ -634,7 +635,14 @@ export class RuntimeStartupService extends Disposable implements IRuntimeStartup
 			StorageScope.WORKSPACE);
 		if (storedSessions) {
 			try {
-				const sessions = JSON.parse(storedSessions) as SerializedSessionMetadata[];
+				const serializedSessions = JSON.parse(storedSessions) as SerializedSessionMetadata[];
+				const sessions: SerializedSessionMetadata[] = serializedSessions.map(session => ({
+					...session,
+					metadata: {
+						...session.metadata,
+						notebookUri: URI.revive(session.metadata.notebookUri),
+					},
+				}));
 				if (sessions.length > 0) {
 					// If this workspace has sessions, attempt to reconnect to
 					// them.

From 88ea24a25ff1cef10e81ea7ebe9b9b31931e0824 Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Tue, 2 Apr 2024 17:56:50 +0200
Subject: [PATCH 03/10] rename

---
 .../src/notebookController.ts                          | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts
index 07a728169b0..dfc62a666c7 100644
--- a/extensions/positron-notebook-controllers/src/notebookController.ts
+++ b/extensions/positron-notebook-controllers/src/notebookController.ts
@@ -64,7 +64,7 @@ export class NotebookController implements vscode.Disposable {
 
 				await Promise.all([
 					updateNotebookLanguage(e.notebook, this.languageId),
-					this.startNewRuntimeSession(e.notebook),
+					this.startRuntimeSession(e.notebook),
 				]);
 			} else {
 				await notebookSessionService.shutdownRuntimeSession(e.notebook.uri);
@@ -78,12 +78,12 @@ export class NotebookController implements vscode.Disposable {
 	}
 
 	/**
-	 * Start a new runtime session for a notebook.
+	 * Start a runtime session for a notebook.
 	 *
 	 * @param notebook The notebook to start a runtime for.
 	 * @returns Promise that resolves when the runtime has started.
 	 */
-	public async startNewRuntimeSession(notebook: vscode.NotebookDocument): Promise<positron.LanguageRuntimeSession> {
+	public async startRuntimeSession(notebook: vscode.NotebookDocument): Promise<positron.LanguageRuntimeSession> {
 		try {
 			return await this.notebookSessionService.startRuntimeSession(notebook.uri, this.languageId);
 		} catch (err) {
@@ -98,7 +98,7 @@ export class NotebookController implements vscode.Disposable {
 				retry,
 			);
 			if (selection === retry) {
-				return vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startNewRuntimeSession(notebook));
+				return vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startRuntimeSession(notebook));
 			}
 
 			throw err;
@@ -126,7 +126,7 @@ export class NotebookController implements vscode.Disposable {
 
 		// No session has been started for this notebook, start one.
 		if (!session) {
-			session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startNewRuntimeSession(notebook));
+			session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startRuntimeSession(notebook));
 		}
 
 		for (const cell of cells) {

From ce944ceced000efafed5b953fec3d9571206f6b7 Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Tue, 2 Apr 2024 18:16:05 +0200
Subject: [PATCH 04/10] fix copyright

---
 extensions/positron-notebook-controllers/src/map.ts | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/extensions/positron-notebook-controllers/src/map.ts b/extensions/positron-notebook-controllers/src/map.ts
index 791c3e89f86..36a93ec1154 100644
--- a/extensions/positron-notebook-controllers/src/map.ts
+++ b/extensions/positron-notebook-controllers/src/map.ts
@@ -1,10 +1,11 @@
 /*---------------------------------------------------------------------------------------------
- *  Copyright (c) Microsoft Corporation. All rights reserved.
- *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *  Copyright (C) 2024 Posit Software, PBC. All rights reserved.
  *--------------------------------------------------------------------------------------------*/
 
 import { Uri as URI } from 'vscode';
 
+// The ResourceMap class and its dependencies are copied as is from the core project.
+
 interface ResourceMapKeyFn {
 	(resource: URI): string;
 }

From d77811c9629cf13088af491697a82585ff29bc4a Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Tue, 2 Apr 2024 18:20:09 +0200
Subject: [PATCH 05/10] remove debug log

---
 .../workbench/services/runtimeSession/common/runtimeSession.ts   | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
index 9b6192e659e..91de4d94fa3 100644
--- a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
+++ b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
@@ -841,7 +841,6 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
 						this._consoleSessionsByLanguageId.delete(session.runtimeMetadata.languageId);
 					} else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) {
 						if (session.metadata.notebookUri) {
-							console.log(`WASIM removing notebook session from map`);
 							this._notebookSessionsByNotebookUri.delete(session.metadata.notebookUri);
 						} else {
 							this._logService.error(`Notebook session ${formatLanguageRuntimeSession(session)} ` +

From a1f1e69ef45c1bbdd6e289089d99cda4eea4d5de Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Tue, 2 Apr 2024 18:20:48 +0200
Subject: [PATCH 06/10] docs

---
 .../workbench/services/runtimeStartup/common/runtimeStartup.ts | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts b/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts
index 1b3b8f893aa..d4ad895033f 100644
--- a/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts
+++ b/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts
@@ -636,6 +636,8 @@ export class RuntimeStartupService extends Disposable implements IRuntimeStartup
 		if (storedSessions) {
 			try {
 				const serializedSessions = JSON.parse(storedSessions) as SerializedSessionMetadata[];
+
+				// Revive the URIs in the session metadata.
 				const sessions: SerializedSessionMetadata[] = serializedSessions.map(session => ({
 					...session,
 					metadata: {
@@ -643,6 +645,7 @@ export class RuntimeStartupService extends Disposable implements IRuntimeStartup
 						notebookUri: URI.revive(session.metadata.notebookUri),
 					},
 				}));
+
 				if (sessions.length > 0) {
 					// If this workspace has sessions, attempt to reconnect to
 					// them.

From cadab9f4893de5ac35446c26ae5f3b253b838519 Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Wed, 3 Apr 2024 13:57:28 +0200
Subject: [PATCH 07/10] address review comments; renames

---
 .../src/notebookController.ts                 | 39 ++++----
 .../src/notebookControllerManager.ts          |  8 +-
 .../src/notebookSessionService.ts             | 95 ++++++++-----------
 3 files changed, 65 insertions(+), 77 deletions(-)

diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts
index dfc62a666c7..1dfa651dc7a 100644
--- a/extensions/positron-notebook-controllers/src/notebookController.ts
+++ b/extensions/positron-notebook-controllers/src/notebookController.ts
@@ -3,7 +3,7 @@
  *--------------------------------------------------------------------------------------------*/
 import * as vscode from 'vscode';
 import * as positron from 'positron';
-import { INotebookSessionService } from './notebookSessionService';
+import { NotebookSessionService } from './notebookSessionService';
 
 /**
  * Wraps a vscode.NotebookController for a specific language, and manages a notebook runtime session
@@ -23,16 +23,16 @@ export class NotebookController implements vscode.Disposable {
 	private static _CELL_COUNTER = 0;
 
 	/**
-	 * @param languageId The language ID for which this controller is responsible.
-	 * @param notebookSessionService The notebook session service.
+	 * @param _languageId The language ID for which this controller is responsible.
+	 * @param _notebookSessionService The notebook session service.
 	 */
 	constructor(
-		private readonly languageId: string,
-		private readonly notebookSessionService: INotebookSessionService,
+		private readonly _languageId: string,
+		private readonly _notebookSessionService: NotebookSessionService,
 	) {
 		// Create a VSCode notebook controller for this language.
 		this.controller = vscode.notebooks.createNotebookController(
-			`positron-${languageId}`,
+			`positron-${_languageId}`,
 			// The 'jupyter-notebook' notebook type is contributed via the built-in extension
 			// extensions/ipynb. Registering our notebook controllers with the same type ensures
 			// that they show up in the notebook UI's kernel picker for .ipynb files.
@@ -51,8 +51,8 @@ export class NotebookController implements vscode.Disposable {
 		this._disposables.push(this.controller);
 
 		this._disposables.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => {
-			if (notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
-				await notebookSessionService.shutdownRuntimeSession(notebook.uri);
+			if (this._notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
+				await this._notebookSessionService.shutdownRuntimeSession(notebook.uri);
 			}
 		}));
 
@@ -63,18 +63,18 @@ export class NotebookController implements vscode.Disposable {
 				// already selected.
 
 				await Promise.all([
-					updateNotebookLanguage(e.notebook, this.languageId),
+					updateNotebookLanguage(e.notebook, this._languageId),
 					this.startRuntimeSession(e.notebook),
 				]);
 			} else {
-				await notebookSessionService.shutdownRuntimeSession(e.notebook.uri);
+				await this._notebookSessionService.shutdownRuntimeSession(e.notebook.uri);
 			}
 		}));
 	}
 
 	/** The human-readable label of the controller. */
 	public get label(): string {
-		return `${this.languageId[0].toUpperCase()}${this.languageId.slice(1)}`;
+		return `${this._languageId[0].toUpperCase()}${this._languageId.slice(1)}`;
 	}
 
 	/**
@@ -85,7 +85,7 @@ export class NotebookController implements vscode.Disposable {
 	 */
 	public async startRuntimeSession(notebook: vscode.NotebookDocument): Promise<positron.LanguageRuntimeSession> {
 		try {
-			return await this.notebookSessionService.startRuntimeSession(notebook.uri, this.languageId);
+			return await this._notebookSessionService.startRuntimeSession(notebook.uri, this._languageId);
 		} catch (err) {
 			const retry = vscode.l10n.t("Retry");
 			const selection = await vscode.window.showErrorMessage(
@@ -115,13 +115,14 @@ export class NotebookController implements vscode.Disposable {
 	 */
 	private async executeCells(cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, _controller: vscode.NotebookController) {
 		// Get the notebook's session.
-		let session: positron.LanguageRuntimeSession | undefined;
-		const startingSessionPromise = this.notebookSessionService.getStartingNotebookSessionPromise(notebook.uri);
-		if (startingSessionPromise && !startingSessionPromise.isSettled) {
-			// If the runtime is still starting, wait for it to be ready.
-			session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => startingSessionPromise.p);
-		} else {
-			session = this.notebookSessionService.getNotebookSession(notebook.uri);
+		let session = this._notebookSessionService.getNotebookSession(notebook.uri);
+
+		// If the runtime is still starting, wait for it to be ready.
+		if (!session) {
+			const startingSessionPromise = this._notebookSessionService.getStartingNotebookSessionPromise(notebook.uri);
+			if (startingSessionPromise && !startingSessionPromise.isSettled) {
+				session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => startingSessionPromise.p);
+			}
 		}
 
 		// No session has been started for this notebook, start one.
diff --git a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
index 10621c92778..53548128373 100644
--- a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
+++ b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts
@@ -5,7 +5,7 @@
 import * as vscode from 'vscode';
 import { log } from './extension';
 import { NotebookController } from './notebookController';
-import { INotebookSessionService } from './notebookSessionService';
+import { NotebookSessionService } from './notebookSessionService';
 
 /**
  * Manages notebook controllers.
@@ -16,9 +16,9 @@ export class NotebookControllerManager implements vscode.Disposable {
 
 	/**
 	 *
-	 * @param notebookSessionService The notebook session service.
+	 * @param _notebookSessionService The notebook session service.
 	 */
-	constructor(private readonly notebookSessionService: INotebookSessionService) { }
+	constructor(private readonly _notebookSessionService: NotebookSessionService) { }
 
 	/**
 	 * Create a notebook controller for a language.
@@ -29,7 +29,7 @@ export class NotebookControllerManager implements vscode.Disposable {
 		if (this.controllers.has(languageId)) {
 			throw new Error(`Notebook controller already exists for language: ${languageId}`);
 		}
-		const controller = new NotebookController(languageId, this.notebookSessionService);
+		const controller = new NotebookController(languageId, this._notebookSessionService);
 		this.controllers.set(languageId, controller);
 		log.info(`Registered notebook controller for language: ${languageId}`);
 	}
diff --git a/extensions/positron-notebook-controllers/src/notebookSessionService.ts b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
index 5c2503b7f9b..e755325b3a4 100644
--- a/extensions/positron-notebook-controllers/src/notebookSessionService.ts
+++ b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
@@ -13,16 +13,36 @@ import { DeferredPromise } from './util';
  * runtime sessions; it manages the set of active sessions and provides
  * facilities for starting, stopping, and interacting with them.
  *
- * TODO(seem): The intention is for functionality here to eventually be brought into Positron core.
+ * TODO(seem): Most of this code is copied from the runtime session service. We should bring what's
+ * required into the runtime session service and expose what's needed via the Positron Extensions
+ * API.
  */
-export interface INotebookSessionService {
+export class NotebookSessionService {
+
+	/**
+	 * A map of sessions currently starting, keyed by notebook URI. Values are promises that resolve
+	 * when the session has started and is ready to execute code.
+	 */
+	private readonly _startingSessionsByNotebookUri = new ResourceMap<DeferredPromise<positron.LanguageRuntimeSession>>();
+
+	/**
+	 * A map of sessions currently shutting down, keyed by notebook URI. Values are promises that resolve
+	 * when the session has completed the shutdown sequence.
+	 */
+	private readonly _shuttingDownSessionsByNotebookUri = new ResourceMap<DeferredPromise<void>>();
+
+	/** A map of the currently active notebook sessions, keyed by notebook URI. */
+	private readonly _notebookSessionsByNotebookUri = new ResourceMap<positron.LanguageRuntimeSession>();
+
 	/**
 	 * Checks for a starting notebook for the given notebook URI.
 	 *
 	 * @param notebookUri The notebook URI to check for.
 	 * @returns True if a starting notebook session exists for the given notebook URI.
 	 */
-	hasStartingNotebookSession(notebookUri: Uri): boolean;
+	hasStartingNotebookSession(notebookUri: Uri): boolean {
+		return this._startingSessionsByNotebookUri.has(notebookUri);
+	}
 
 	/**
 	 * Checks for a starting or running notebook for the given notebook URI.
@@ -30,7 +50,9 @@ export interface INotebookSessionService {
 	 * @param notebookUri The notebook URI to check for.
 	 * @returns True if a starting or running notebook session exists for the given notebook URI.
 	 */
-	hasStartingOrRunningNotebookSession(notebookUri: Uri): boolean;
+	hasStartingOrRunningNotebookSession(notebookUri: Uri): boolean {
+		return this._startingSessionsByNotebookUri.has(notebookUri) || this._notebookSessionsByNotebookUri.has(notebookUri);
+	}
 
 	/**
 	 * Get the starting or running notebook session for the given notebook URI, if one exists.
@@ -38,7 +60,9 @@ export interface INotebookSessionService {
 	 * @param notebookUri The notebook URI of the session to retrieve.
 	 * @returns The starting or running notebook session for the given notebook URI, if one exists.
 	 */
-	getStartingNotebookSessionPromise(notebookUri: Uri): DeferredPromise<positron.LanguageRuntimeSession> | undefined;
+	getStartingNotebookSessionPromise(notebookUri: Uri): DeferredPromise<positron.LanguageRuntimeSession> | undefined {
+		return this._startingSessionsByNotebookUri.get(notebookUri);
+	}
 
 	/**
 	 * Get the running notebook session for the given notebook URI, if one exists.
@@ -46,7 +70,9 @@ export interface INotebookSessionService {
 	 * @param notebookUri The notebook URI of the session to retrieve.
 	 * @returns The running notebook session for the given notebook URI, if one exists.
 	 */
-	getNotebookSession(notebookUri: Uri): positron.LanguageRuntimeSession | undefined;
+	getNotebookSession(notebookUri: Uri): positron.LanguageRuntimeSession | undefined {
+		return this._notebookSessionsByNotebookUri.get(notebookUri);
+	}
 
 	/**
 	 * Start a new runtime session for a notebook.
@@ -54,54 +80,6 @@ export interface INotebookSessionService {
 	 * @param notebookUri The notebook URI to start a runtime for.
 	 * @returns Promise that resolves when the runtime startup sequence has been started.
 	 */
-	startRuntimeSession(notebookUri: Uri, languageId: string): Promise<positron.LanguageRuntimeSession>;
-
-	/**
-	 * Shutdown the runtime session for a notebook.
-	 *
-	 * @param notebookUri The notebook URI whose runtime to shutdown.
-	 * @returns Promise that resolves when the runtime shutdown sequence has been started.
-	 */
-	shutdownRuntimeSession(notebookUri: Uri): Promise<void>;
-
-}
-
-/**
- * The implementation of INotebookSessionService.
- */
-export class NotebookSessionService implements INotebookSessionService {
-
-	/**
-	 * A map of sessions currently starting, keyed by notebook URI. Values are promises that resolve
-	 * when the session has started and is ready to execute code.
-	 */
-	private readonly _startingSessionsByNotebookUri = new ResourceMap<DeferredPromise<positron.LanguageRuntimeSession>>();
-
-	/**
-	 * A map of sessions currently shutting down, keyed by notebook URI. Values are promises that resolve
-	 * when the session has completed the shutdown sequence.
-	 */
-	private readonly _shuttingDownSessionsByNotebookUri = new ResourceMap<DeferredPromise<void>>();
-
-	/** A map of the currently active notebook sessions, keyed by notebook URI. */
-	private readonly _notebookSessionsByNotebookUri = new ResourceMap<positron.LanguageRuntimeSession>();
-
-	hasStartingNotebookSession(notebookUri: Uri): boolean {
-		return this._startingSessionsByNotebookUri.has(notebookUri);
-	}
-
-	hasStartingOrRunningNotebookSession(notebookUri: Uri): boolean {
-		return this._startingSessionsByNotebookUri.has(notebookUri) || this._notebookSessionsByNotebookUri.has(notebookUri);
-	}
-
-	getStartingNotebookSessionPromise(notebookUri: Uri): DeferredPromise<positron.LanguageRuntimeSession> | undefined {
-		return this._startingSessionsByNotebookUri.get(notebookUri);
-	}
-
-	getNotebookSession(notebookUri: Uri): positron.LanguageRuntimeSession | undefined {
-		return this._notebookSessionsByNotebookUri.get(notebookUri);
-	}
-
 	async startRuntimeSession(notebookUri: Uri, languageId: string): Promise<positron.LanguageRuntimeSession> {
 		// Return the existing promise, if there is one.
 		const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri);
@@ -148,6 +126,9 @@ export class NotebookSessionService implements INotebookSessionService {
 			throw err;
 		}
 
+		// TODO: If it isn't running, log an error and start a new one.
+		// TODO: If it doesn't match the runtime ID, log an error, shut it down, and start a new one.
+
 		// If we couldn't restore a session, start a new one.
 		if (!session) {
 			// Get the preferred runtime for this language.
@@ -188,6 +169,12 @@ export class NotebookSessionService implements INotebookSessionService {
 		return session;
 	}
 
+	/**
+	 * Shutdown the runtime session for a notebook.
+	 *
+	 * @param notebookUri The notebook URI whose runtime to shutdown.
+	 * @returns Promise that resolves when the runtime shutdown sequence has been started.
+	 */
 	async shutdownRuntimeSession(notebookUri: Uri): Promise<void> {
 		// Return the existing promise, if there is one.
 		const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri);

From bca09a90804acfc32c4bb90d7d1ae5aaa5b67b66 Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Wed, 3 Apr 2024 14:56:45 +0200
Subject: [PATCH 08/10] simplify start on execute

---
 .../src/notebookController.ts                 |  8 --------
 .../src/notebookSessionService.ts             | 20 -------------------
 2 files changed, 28 deletions(-)

diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts
index 1dfa651dc7a..963fc227a1a 100644
--- a/extensions/positron-notebook-controllers/src/notebookController.ts
+++ b/extensions/positron-notebook-controllers/src/notebookController.ts
@@ -117,14 +117,6 @@ export class NotebookController implements vscode.Disposable {
 		// Get the notebook's session.
 		let session = this._notebookSessionService.getNotebookSession(notebook.uri);
 
-		// If the runtime is still starting, wait for it to be ready.
-		if (!session) {
-			const startingSessionPromise = this._notebookSessionService.getStartingNotebookSessionPromise(notebook.uri);
-			if (startingSessionPromise && !startingSessionPromise.isSettled) {
-				session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => startingSessionPromise.p);
-			}
-		}
-
 		// No session has been started for this notebook, start one.
 		if (!session) {
 			session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startRuntimeSession(notebook));
diff --git a/extensions/positron-notebook-controllers/src/notebookSessionService.ts b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
index e755325b3a4..54b47c59039 100644
--- a/extensions/positron-notebook-controllers/src/notebookSessionService.ts
+++ b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
@@ -34,16 +34,6 @@ export class NotebookSessionService {
 	/** A map of the currently active notebook sessions, keyed by notebook URI. */
 	private readonly _notebookSessionsByNotebookUri = new ResourceMap<positron.LanguageRuntimeSession>();
 
-	/**
-	 * Checks for a starting notebook for the given notebook URI.
-	 *
-	 * @param notebookUri The notebook URI to check for.
-	 * @returns True if a starting notebook session exists for the given notebook URI.
-	 */
-	hasStartingNotebookSession(notebookUri: Uri): boolean {
-		return this._startingSessionsByNotebookUri.has(notebookUri);
-	}
-
 	/**
 	 * Checks for a starting or running notebook for the given notebook URI.
 	 *
@@ -54,16 +44,6 @@ export class NotebookSessionService {
 		return this._startingSessionsByNotebookUri.has(notebookUri) || this._notebookSessionsByNotebookUri.has(notebookUri);
 	}
 
-	/**
-	 * Get the starting or running notebook session for the given notebook URI, if one exists.
-	 *
-	 * @param notebookUri The notebook URI of the session to retrieve.
-	 * @returns The starting or running notebook session for the given notebook URI, if one exists.
-	 */
-	getStartingNotebookSessionPromise(notebookUri: Uri): DeferredPromise<positron.LanguageRuntimeSession> | undefined {
-		return this._startingSessionsByNotebookUri.get(notebookUri);
-	}
-
 	/**
 	 * Get the running notebook session for the given notebook URI, if one exists.
 	 *

From e364067a8948814d774af0068fbec6326142936c Mon Sep 17 00:00:00 2001
From: seem <mwlorgat@gmail.com>
Date: Wed, 3 Apr 2024 16:23:07 +0200
Subject: [PATCH 09/10] register `onDidCloseNotebookDocument` listener at
 extension-level

---
 .../src/extension.ts                             |  8 ++++++++
 .../src/notebookController.ts                    |  9 +++------
 .../src/notebookSessionService.ts                | 16 ++++++++--------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/extensions/positron-notebook-controllers/src/extension.ts b/extensions/positron-notebook-controllers/src/extension.ts
index 370db6008ed..9fbfa8e5d87 100644
--- a/extensions/positron-notebook-controllers/src/extension.ts
+++ b/extensions/positron-notebook-controllers/src/extension.ts
@@ -11,6 +11,14 @@ export const log = vscode.window.createOutputChannel('Positron Notebook Controll
 
 export async function activate(context: vscode.ExtensionContext): Promise<void> {
 	const notebookSessionService = new NotebookSessionService();
+
+	// Shutdown any running sessions when a notebook is closed.
+	context.subscriptions.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => {
+		if (notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
+			await notebookSessionService.shutdownRuntimeSession(notebook.uri);
+		}
+	}));
+
 	const manager = new NotebookControllerManager(notebookSessionService);
 	context.subscriptions.push(manager);
 
diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts
index 963fc227a1a..87a91b0a4a6 100644
--- a/extensions/positron-notebook-controllers/src/notebookController.ts
+++ b/extensions/positron-notebook-controllers/src/notebookController.ts
@@ -4,6 +4,7 @@
 import * as vscode from 'vscode';
 import * as positron from 'positron';
 import { NotebookSessionService } from './notebookSessionService';
+import { log } from './extension';
 
 /**
  * Wraps a vscode.NotebookController for a specific language, and manages a notebook runtime session
@@ -50,13 +51,9 @@ export class NotebookController implements vscode.Disposable {
 
 		this._disposables.push(this.controller);
 
-		this._disposables.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => {
-			if (this._notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
-				await this._notebookSessionService.shutdownRuntimeSession(notebook.uri);
-			}
-		}));
-
 		this._disposables.push(this.controller.onDidChangeSelectedNotebooks(async (e) => {
+			log.debug(`Notebook ${e.notebook.uri}, controller ${this._languageId}, selected ${e.selected}`);
+
 			// Has this controller been selected for a notebook?
 			if (e.selected) {
 				// Note that this is also reached when a notebook is opened, if this controller was
diff --git a/extensions/positron-notebook-controllers/src/notebookSessionService.ts b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
index 54b47c59039..58a7bc352ee 100644
--- a/extensions/positron-notebook-controllers/src/notebookSessionService.ts
+++ b/extensions/positron-notebook-controllers/src/notebookSessionService.ts
@@ -168,14 +168,14 @@ export class NotebookSessionService {
 		this._shuttingDownSessionsByNotebookUri.set(notebookUri, shutDownPromise);
 
 		// Get the notebook's session.
-		let session: positron.LanguageRuntimeSession | undefined;
-		const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri);
-		if (startingSessionPromise && !startingSessionPromise.isSettled) {
-			// If the runtime is still starting, wait for it to be ready.
-			session = await startingSessionPromise.p;
-		} else {
-			// Try to get an already running session.
-			session = this._notebookSessionsByNotebookUri.get(notebookUri);
+		let session = this._notebookSessionsByNotebookUri.get(notebookUri);
+
+		// If the runtime is still starting, wait for it to be ready.
+		if (!session) {
+			const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri);
+			if (startingSessionPromise && !startingSessionPromise.isSettled) {
+				session = await startingSessionPromise.p;
+			}
 		}
 
 		// Ensure that we have a session.

From 2580968d4a56a9fb3efcd8e49156eae3224fcaa7 Mon Sep 17 00:00:00 2001
From: Wasim Lorgat <mwlorgat@gmail.com>
Date: Wed, 3 Apr 2024 16:24:44 +0200
Subject: [PATCH 10/10] Apply suggestions from code review

Co-authored-by: Nick Strayer <nick.strayer@posit.co>
Signed-off-by: Wasim Lorgat <mwlorgat@gmail.com>
---
 .../positron-notebook-controllers/src/notebookController.ts   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts
index 87a91b0a4a6..3794c359323 100644
--- a/extensions/positron-notebook-controllers/src/notebookController.ts
+++ b/extensions/positron-notebook-controllers/src/notebookController.ts
@@ -84,10 +84,10 @@ export class NotebookController implements vscode.Disposable {
 		try {
 			return await this._notebookSessionService.startRuntimeSession(notebook.uri, this._languageId);
 		} catch (err) {
-			const retry = vscode.l10n.t("Retry");
+			const retry = vscode.l10n.t('Retry');
 			const selection = await vscode.window.showErrorMessage(
 				vscode.l10n.t(
-					"Starting {0} interpreter for '{1}' failed. Reason: {2}",
+					'Starting {0} interpreter for "{1}" failed. Reason: {2}',
 					this.label,
 					notebook.uri.path,
 					err