Add Reticulate support to the kernel supervisor (#5854)

This change makes it possible to use Reticulate sessions with the kernel
supervisor.

<img width="885" alt="image"
src="https://github.com/user-attachments/assets/493a7c36-6d49-498f-86f1-886d4255d4e2"
/>

The supervisor has a new `adopt` endpoint that allows it to connect to a
session that is already running somewhere else. Now, when you start a
Reticulate session, the following things happen:

- An R session is started, if necessary.
- A Reticulate session is started inside the R session.
- The Python kernel is started inside the Reticulate session.
- The supervisor "adopts" the Python kernel, connecting to its ZeroMQ
sockets and proxying messages as it does for other kernels.

In addition to adding the orchestration to use the supervisor, there are
a couple of other small changes in this PR:

- There is a new _Restart the Kernel Supervisor_ command, which shuts
down the supervisor process and all sessions, regardless of what state
they're in. This is mostly intended as a debugging tool, and it was
added in order to make it possible to pick up e.g. debug log level
changes without restarting Positron, but could also be helpful as a last
resort if things get stuck.
- The progress shown when starting a Reticulate session is now much more
chatty; it appears as soon as you try to start the session and shows
more information along the way.

Progress towards #4579; this enables the use of Reticulate on Posit
Workbench, and is the last major piece of functionality that required
the Jupyter Adapter.

(Note that the main body of this change lives in the supervisor itself;
this is just the front end / UI bits.)

### QA Notes

There's a _lot_ of orchestration involved already in starting Reticulate
and this makes it even more complicated. The highest risk areas are
around lifecycle management: shutting down, restarting, reconnecting,
etc.

There is a known issue in which if you have an exited R session in your
Console and try to start a Reticulate session, it never happens. I
didn't try to fix that in this PR.

extensions/positron-python/src/client/positron/session.ts

@@ -468,7 +468,7 @@ export class PythonRuntimeSession implements positron.LanguageRuntimeSession, vs
 
     private async createKernel(): Promise<JupyterLanguageRuntimeSession> {
         const config = vscode.workspace.getConfiguration('kernelSupervisor');
-        if (config.get<boolean>('enable', true) && this.runtimeMetadata.runtimeId !== 'reticulate') {
+        if (config.get<boolean>('enable', true)) {
             // Use the Positron kernel supervisor if enabled
             const ext = vscode.extensions.getExtension('positron.positron-supervisor');
             if (!ext) {

extensions/positron-reticulate/src/async.ts

@@ -0,0 +1,74 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (C) 2024 Posit Software, PBC. All rights reserved.
+ *  Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+/**
+ * PromiseHandles is a class that represents a promise that can be resolved or
+ * rejected externally.
+ */
+export class PromiseHandles<T> {
+	resolve!: (value: T | Promise<T>) => void;
+
+	reject!: (error: unknown) => void;
+
+	promise: Promise<T>;
+
+	constructor() {
+		this.promise = new Promise((resolve, reject) => {
+			this.resolve = resolve;
+			this.reject = reject;
+		});
+	}
+}
+
+/**
+ * A barrier that is initially closed and then becomes opened permanently.
+ * Ported from VS Code's async.ts.
+ */
+
+export class Barrier {
+	private _isOpen: boolean;
+	private _promise: Promise<boolean>;
+	private _completePromise!: (v: boolean) => void;
+
+	constructor() {
+		this._isOpen = false;
+		this._promise = new Promise<boolean>((c, _e) => {
+			this._completePromise = c;
+		});
+	}
+
+	isOpen(): boolean {
+		return this._isOpen;
+	}
+
+	open(): void {
+		this._isOpen = true;
+		this._completePromise(true);
+	}
+
+	wait(): Promise<boolean> {
+		return this._promise;
+	}
+}
+
+/**
+ * Wraps a promise in a timeout that rejects the promise if it does not resolve
+ * within the given time.
+ *
+ * @param promise The promise to wrap
+ * @param timeout The timeout interval in milliseconds
+ * @param message The error message to use if the promise times out
+ *
+ * @returns The wrapped promise
+ */
+export function withTimeout<T>(promise: Promise<T>,
+	timeout: number,
+	message: string): Promise<T> {
+	return Promise.race([
+		promise,
+		new Promise<T>((_, reject) => setTimeout(() => reject(new Error(message)), timeout))
+	]);
+}
+

extensions/positron-reticulate/src/extension.ts

@@ -8,6 +8,7 @@ import * as positron from 'positron';
 import path = require('path');
 import fs = require('fs');
 import { JupyterKernelSpec, JupyterSession, JupyterKernel } from './jupyter-adapter.d';
+import { Barrier, PromiseHandles } from './async';
 
 export class ReticulateRuntimeManager implements positron.LanguageRuntimeManager {
 
@@ -166,53 +167,121 @@ class InitializationError extends Error {
 class ReticulateRuntimeSession implements positron.LanguageRuntimeSession {
 
 	private kernel: JupyterKernel | undefined;
+	public started = new Barrier();
 	private pythonSession: positron.LanguageRuntimeSession;
 
 	// To create a reticulate runtime session we need to first create a python
 	// runtime session using the exported interface from the positron-python
 	// extension.
 
-	// The PythonRuntimeSession object in the positron-python extensions, is created
-	// by passing 'runtimeMetadata', 'sessionMetadata' and something called 'kernelSpec'
-	// that's further passed to the JupyterAdapter extension in order to actually initialize
-	// the session.
-
-	// ReticulateRuntimeSession are only different from Python runtime sessions in the
-	// way the kernel spec is provided. In general, the kernel spec contains a runtime
-	// path and some arguments that are used start the kernel process. (The kernel is started
-	// by the Jupyter Adapter in a vscode terminal). In the reticulate case, the kernel isn't
-	// started that way. Instead, we need to call into the R console to start the python jupyter
-	// kernel (that's actually running in the same process as R), and only then, ask JupyterAdapter
-	// to connect to that kernel.
+	// The PythonRuntimeSession object in the positron-python extensions, is
+	// created by passing 'runtimeMetadata', 'sessionMetadata' and something
+	// called 'kernelSpec' that's further passed to the JupyterAdapter
+	// extension in order to actually initialize the session.
+
+	// ReticulateRuntimeSession are only different from Python runtime sessions
+	// in the way the kernel spec is provided. In general, the kernel spec
+	// contains a runtime path and some arguments that are used start the
+	// kernel process. (The kernel is started by the Jupyter Adapter in a
+	// vscode terminal). In the reticulate case, the kernel isn't started that
+	// way. Instead, we need to call into the R console to start the python
+	// jupyter kernel (that's actually running in the same process as R), and
+	// only then, ask JupyterAdapter to connect to that kernel.
 	static async create(
 		runtimeMetadata: positron.LanguageRuntimeMetadata,
 		sessionMetadata: positron.RuntimeSessionMetadata,
 	): Promise<ReticulateRuntimeSession> {
-		const rSession = await getRSession();
-		const config = await ReticulateRuntimeSession.checkRSession(rSession);
-		const metadata = await ReticulateRuntimeSession.fixInterpreterPath(runtimeMetadata, config.python);
 
-		return new ReticulateRuntimeSession(
-			rSession,
-			metadata,
-			sessionMetadata,
-			ReticulateRuntimeSessionType.Create
-		);
+		// A deferred promise that will resolve when the session is created.
+		const sessionPromise = new PromiseHandles<ReticulateRuntimeSession>();
+
+		// Show a progress notification while we create the session.
+		vscode.window.withProgress({
+			location: vscode.ProgressLocation.Notification,
+			title: 'Creating the Reticulate Python session',
+			cancellable: false
+		}, async (progress, _token) => {
+			let session: ReticulateRuntimeSession | undefined;
+			try {
+				// Get the R session that we'll use to start the reticulate session.
+				progress.report({ increment: 10, message: 'Initializing the host R session' });
+				const rSession = await getRSession(progress);
+
+				// Make sure the R session has the necessary packages installed.
+				progress.report({ increment: 10, message: 'Checking prerequisites' });
+				const config = await ReticulateRuntimeSession.checkRSession(rSession);
+				const metadata = await ReticulateRuntimeSession.fixInterpreterPath(runtimeMetadata, config.python);
+
+				// Create the session itself.
+				session = new ReticulateRuntimeSession(
+					rSession,
+					metadata,
+					sessionMetadata,
+					ReticulateRuntimeSessionType.Create,
+					progress
+				);
+				sessionPromise.resolve(session);
+			} catch (err) {
+				sessionPromise.reject(err);
+			}
+
+			// Wait for the session to start (or fail to start) before
+			// returning from this callback, so that the progress bar stays up
+			// while we wait.
+			if (session) {
+				progress.report({ increment: 10, message: 'Waiting to connect' });
+				await session.started.wait();
+			}
+		});
+
+		return sessionPromise.promise;
 	}
 
 	static async restore(
 		runtimeMetadata: positron.LanguageRuntimeMetadata,
 		sessionMetadata: positron.RuntimeSessionMetadata,
 	): Promise<ReticulateRuntimeSession> {
-		const rSession = await getRSession();
-		const config = await ReticulateRuntimeSession.checkRSession(rSession);
-		const metadata = await ReticulateRuntimeSession.fixInterpreterPath(runtimeMetadata, config.python);
-		return new ReticulateRuntimeSession(
-			rSession,
-			metadata,
-			sessionMetadata,
-			ReticulateRuntimeSessionType.Restore
-		);
+
+		const sessionPromise = new PromiseHandles<ReticulateRuntimeSession>();
+
+		vscode.window.withProgress({
+			location: vscode.ProgressLocation.Notification,
+			title: 'Restoring the Reticulate Python session',
+			cancellable: false
+		}, async (progress, _token) => {
+			let session: ReticulateRuntimeSession | undefined;
+			try {
+				// Find the R session that we'll use to restore the reticulate session.
+				progress.report({ increment: 10, message: 'Initializing the host R session' });
+				const rSession = await getRSession(progress);
+
+				// Make sure the R session has the necessary packages installed.
+				progress.report({ increment: 10, message: 'Checking prerequisites' });
+				const config = await ReticulateRuntimeSession.checkRSession(rSession);
+				const metadata = await ReticulateRuntimeSession.fixInterpreterPath(runtimeMetadata, config.python);
+
+				// Create the session itself.
+				session = new ReticulateRuntimeSession(
+					rSession,
+					metadata,
+					sessionMetadata,
+					ReticulateRuntimeSessionType.Restore,
+					progress
+				);
+				sessionPromise.resolve(session);
+			} catch (err) {
+				sessionPromise.reject(err);
+			}
+
+			// Wait for the session to resume (or fail to resume) before
+			// returning
+			if (session) {
+				progress.report({ increment: 10, message: 'Waiting to reconnect' });
+				await session.started.wait();
+			}
+		});
+
+		return sessionPromise.promise;
 	}
 
 	static async checkRSession(rSession: positron.LanguageRuntimeSession): Promise<{ python: string }> {
@@ -330,6 +399,7 @@ class ReticulateRuntimeSession implements positron.LanguageRuntimeSession {
 		runtimeMetadata: positron.LanguageRuntimeMetadata,
 		sessionMetadata: positron.RuntimeSessionMetadata,
 		sessionType: ReticulateRuntimeSessionType,
+		readonly progress: vscode.Progress<{ message?: string; increment?: number }>
 	) {
 		// When the kernelSpec is undefined, the PythonRuntimeSession
 		// will perform a restore session.
@@ -360,20 +430,29 @@ class ReticulateRuntimeSession implements positron.LanguageRuntimeSession {
 		if (!api) {
 			throw new Error('Failed to find the Positron Python extension API.');
 		}
+		this.progress.report({ increment: 10, message: 'Creating the Python session' });
 		this.pythonSession = api.exports.positron.createPythonRuntimeSession(
 			runtimeMetadata,
 			sessionMetadata,
 			kernelSpec
 		);
 
+		// Open the start barrier once the session is ready.
+		this.pythonSession.onDidChangeRuntimeState((state) => {
+			if (state === positron.RuntimeState.Ready || state === positron.RuntimeState.Idle) {
+				this.started.open();
+			}
+		});
+
 		this.onDidReceiveRuntimeMessage = this.pythonSession.onDidReceiveRuntimeMessage;
 		this.onDidChangeRuntimeState = this.pythonSession.onDidChangeRuntimeState;
 		this.onDidEndSession = this.pythonSession.onDidEndSession;
 	}
 
 	// A function that starts a kernel and then connects to it.
 	async startKernel(session: JupyterSession, kernel: JupyterKernel) {
-		kernel.log('Starting the reticulate session!');
+		kernel.log('Starting the Reticulate session!');
+		this.progress.report({ increment: 10, message: 'Starting the Reticulate session in R' });
 
 		// Store a reference to the kernel, so the session can log, reconnect, etc.
 		this.kernel = kernel;
@@ -407,11 +486,15 @@ class ReticulateRuntimeSession implements positron.LanguageRuntimeSession {
 			throw new Error(`Reticulate initialization failed: ${init_err}`);
 		}
 
+		this.progress.report({ increment: 10, message: 'Connecting to the Reticulate session' });
+
 		try {
 			await kernel.connectToSession(session);
 		} catch (err: any) {
 			kernel.log('Failed connecting to the Reticulate Python session');
 			throw err;
+		} finally {
+			this.started.open();
 		}
 	}
 
@@ -533,15 +616,15 @@ class ReticulateRuntimeSession implements positron.LanguageRuntimeSession {
 	}
 }
 
-async function getRSession(): Promise<positron.LanguageRuntimeSession> {
+async function getRSession(progress: vscode.Progress<{ message?: string; increment?: number }>): Promise<positron.LanguageRuntimeSession> {
 
 	// Retry logic to start an R session.
 	const maxRetries = 5;
 	let session;
 	let error;
 	for (let i = 0; i < maxRetries; i++) {
 		try {
-			session = await getRSession_();
+			session = await getRSession_(progress);
 		}
 		catch (err: any) {
 			error = err; // Keep the last error so we can display it
@@ -566,12 +649,12 @@ class RSessionError extends Error {
 	}
 }
 
-async function getRSession_(): Promise<positron.LanguageRuntimeSession> {
+async function getRSession_(progress: vscode.Progress<{ message?: string; increment?: number }>): Promise<positron.LanguageRuntimeSession> {
 	let session = await positron.runtime.getForegroundSession();
 
 	if (session) {
 		// Get foreground session will return a runtime session even if it has
-		// already exitted. We check that it's still there before proceeding.
+		// already exited. We check that it's still there before proceeding.
 		// TODO: it would be nice to have an API to check for the session state.
 		try {
 			await session.callMethod?.('is_installed', 'reticulate', '1.39');
@@ -581,25 +664,15 @@ async function getRSession_(): Promise<positron.LanguageRuntimeSession> {
 	}
 
 	if (!session || session.runtimeMetadata.languageId !== 'r') {
-		session = await vscode.window.withProgress({
-			location: vscode.ProgressLocation.Notification,
-			title: 'Starting R session for reticulate',
-			cancellable: true
-		}, async (progress, token) => {
-			token.onCancellationRequested(() => {
-				throw new RSessionError('User requested cancellation', true);
-			});
-			progress.report({ increment: 0, message: 'Looking for prefered runtime...' });
-			const runtime = await positron.runtime.getPreferredRuntime('r');
-
-			progress.report({ increment: 20, message: 'Starting R runtime...' });
-			await positron.runtime.selectLanguageRuntime(runtime.runtimeId);
-
-			progress.report({ increment: 70, message: 'Getting R session...' });
-			session = await positron.runtime.getForegroundSession();
-
-			return session;
-		});
+		progress.report({ increment: 10, message: 'Looking for prefered runtime...' });
+
+		const runtime = await positron.runtime.getPreferredRuntime('r');
+
+		progress.report({ increment: 10, message: 'Starting R runtime...' });
+		await positron.runtime.selectLanguageRuntime(runtime.runtimeId);
+
+		progress.report({ increment: 10, message: 'Getting R session...' });
+		session = await positron.runtime.getForegroundSession();
 	}
 
 	if (!session) {

extensions/positron-supervisor/package.json

@@ -90,6 +90,12 @@
         "title": "%command.reconnectSession.title%",
         "shortTitle": "%command.reconnectSession.title%",
         "enablement": "isDevelopment"
+      },
+      {
+        "command": "positron.supervisor.restartSupervisor",
+        "category": "%command.positron.supervisor.category%",
+        "title": "%command.restartSupervisor.title%",
+        "shortTitle": "%command.restartSupervisor.title%"
       }
     ]
   },
@@ -109,7 +115,7 @@
   },
   "positron": {
     "binaryDependencies": {
-      "kallichore": "0.1.22"
+      "kallichore": "0.1.26"
     }
   },
   "dependencies": {

extensions/positron-supervisor/package.nls.json

@@ -14,5 +14,6 @@
 	"configuration.sleepOnStartup.description": "Sleep for n seconds before starting up Jupyter kernel (when supported)",
 	"command.positron.supervisor.category": "Kernel Supervisor",
 	"command.showKernelSupervisorLog.title": "Show the Kernel Supervisor Log",
-	"command.reconnectSession.title": "Reconnect the current session"
+	"command.reconnectSession.title": "Reconnect the Current Session",
+	"command.restartSupervisor.title": "Restart the Kernel Supervisor"
 }