Fix errors arising from kernels running on a separate extension host from other extensions (#9221)

jmcphers · web-flow · commit 513ae0374bba · 2025-09-02T16:14:22.000-07:00
This change fixes some issues that arise when extensions that try to use R/Python sessions are running on a different extension host than the kernels themselves. This fundamentally does not work today because sessions live in the extension host, and many of their methods are accessible only from that extension specific host or the main thread. It is not possible for one extension host to talk directly to another extension host. Adding full support for this is more than I wanted to bite off, so this change is a bit of a stopgap. Here's how it works: - API methods no longer return full-fledged session objects. Instead they return a more basic `BaseLanguageRuntimeSession` object. - When an extension asks for a session object for a session that is not known to the current extension host, we generate a kind of proxy that lets us talk to the session via the main thread. - In some extensions, we really and truly need a full session object. In those cases, we play a little bit dirty: we use `package.json`'s `extensionDependencies` field to declare a dependency on the kernels and/or supervisor. When VS Code is deciding which extensions to run on which host, it ensures that extensions that depend on each other are on the same host -- so we're guaranteed that the sessions in question live on our own host. That allow us to safely cast the `BaseLanguageRuntimeSession` to a full `LanguageRuntimeSessiion`. This makes a for a lot of running around to get e.g. a code result: ```mermaid graph TD e1[Extension 1 - Consumer] -- execute code --> eh1 eh1[Ext Host 1] -- session proxy --> p[Positron] p -- execute code --> eh2[Ext Host 2] eh2 -- execute code --> e2[Extension 2 - Kernel] e2 -- result --> eh2 eh2 -- result --> p p -- result --> eh1 eh1 -- result --> e1 ``` But it's about the best we can do given the architecture today. Addresses #9166. ### Release Notes  #### New Features - N/A #### Bug Fixes - Fix errors when kernels are run on a separate extension host (#9166) ### QA Notes We don't actually know how Julia's system got into a state where it decided to put the kernels on a different extension host. But you can make yours do the same thing by adding this to your user settings JSON: ```json "extensions.experimental.affinity": { "positron.positron-r": 1, "positron.positron-supervisor": 1 } ``` It is possible that there are other problems in this configuration that I didn't fix. If adding the above setting gives you any trouble, check and see if it also repros on main. If it does, please file a separate issue; we think this config may be more common in the future.
diff --git a/extensions/ipynb/src/ipynbMain.ts b/extensions/ipynb/src/ipynbMain.ts
@@ -87,7 +87,7 @@ export function activate(context: vscode.ExtensionContext, serializer: vscode.No
 	// --- Start Positron ---
 	context.subscriptions.push(vscode.commands.registerCommand('ipynb.newUntitledIpynb', async (languageId?: string) => {
 		// Try to use Positron's foreground session's language, fall back to 'plaintext'.
-		const language = languageId ?? (await positron.runtime.getForegroundSession())?.runtimeMetadata?.languageId ?? 'plaintext';
+		const language = languageId ?? (await positron.runtime.getForegroundSession())?.runtimeMetadata.languageId ?? 'plaintext';
 		const cell = new vscode.NotebookCellData(vscode.NotebookCellKind.Code, '', language);
 		const data = new vscode.NotebookData([cell]);
 		data.metadata = {
diff --git a/extensions/positron-assistant/package.json b/extensions/positron-assistant/package.json
@@ -590,5 +590,8 @@
     "@anthropic-ai/sdk": "^0.57.0",
     "@github/copilot-language-server": "^1.335.0",
     "vscode-languageclient": "^9.0.1"
-  }
+  },
+  "extensionDependencies": [
+    "positron.positron-supervisor"
+  ]
 }
diff --git a/extensions/positron-assistant/src/tools.ts b/extensions/positron-assistant/src/tools.ts
@@ -315,13 +315,7 @@ export function registerAssistantTools(
 				]);
 			}
 
-			let session: positron.LanguageRuntimeSession | undefined;
-			const sessions = await positron.runtime.getActiveSessions();
-			if (sessions && sessions.length > 0) {
-				session = sessions.find(
-					(session) => session.metadata.sessionId === options.input.sessionIdentifier,
-				);
-			}
+			const session = await positron.runtime.getSession(options.input.sessionIdentifier);
 			if (!session) {
 				return new vscode.LanguageModelToolResult([
 					new vscode.LanguageModelTextPart('[[]]')
diff --git a/extensions/positron-connections/package.json b/extensions/positron-connections/package.json
@@ -141,6 +141,9 @@
       ]
     }
   },
+  "extensionDependencies": [
+    "positron.positron-supervisor"
+  ],
   "devDependencies": {
     "@types/glob": "^7.2.0",
     "@types/mocha": "^9.1.0",
diff --git a/extensions/positron-javascript/src/session.ts b/extensions/positron-javascript/src/session.ts
@@ -111,6 +111,10 @@ export class JavaScriptLanguageRuntimeSession implements positron.LanguageRuntim
 		return Promise.resolve(positron.RuntimeCodeFragmentStatus.Complete);
 	}
 
+	getDynState(): Thenable<positron.LanguageRuntimeDynState> {
+		return Promise.resolve(this.dynState);
+	}
+
 	createClient(id: string, type: positron.RuntimeClientType, params: any): Thenable<void> {
 		if (type === positron.RuntimeClientType.Variables) {
 			// The only client type we support right now is an environment.
diff --git a/extensions/positron-python/src/client/llm-tools/llm-tools.ts b/extensions/positron-python/src/client/llm-tools/llm-tools.ts
@@ -18,15 +18,9 @@ export function registerPythonLanguageModelTools(
         'getAttachedPythonPackages',
         {
             invoke: async (options, _token) => {
-                let session: positron.LanguageRuntimeSession | undefined;
+                let session: positron.BaseLanguageRuntimeSession | undefined;
                 if (options.input.sessionIdentifier) {
-                    const sessions = await positron.runtime.getActiveSessions();
-                    if (sessions && sessions.length > 0) {
-                        session = sessions.find(
-                            (session) => session.metadata.sessionId === options.input.sessionIdentifier,
-                        );
-                    }
-
+                    session = await positron.runtime.getSession(options.input.sessionIdentifier);
                     if (!session) {
                         session = await positron.runtime.getForegroundSession();
                     }
diff --git a/extensions/positron-python/src/client/positron/manager.ts b/extensions/positron-python/src/client/positron/manager.ts
@@ -435,7 +435,9 @@ export class PythonRuntimeManager implements IPythonRuntimeManager, Disposable {
             if (sessionsToShutdown.length > 0) {
                 traceInfo(`Shutting down ${sessionsToShutdown.length} sessions using Python runtime at ${pythonPath}`);
                 await Promise.all(
-                    sessionsToShutdown.map((session) => session.shutdown(positron.RuntimeExitReason.Shutdown)),
+                    sessionsToShutdown.map(async (session) => {
+                        session.shutdown(positron.RuntimeExitReason.Shutdown);
+                    }),
                 );
                 // Remove the runtime from our registry so we can recreate it
                 this.registeredPythonRuntimes.delete(pythonPath);
diff --git a/extensions/positron-python/src/client/positron/session.ts b/extensions/positron-python/src/client/positron/session.ts
@@ -118,7 +118,7 @@ export class PythonRuntimeSession implements positron.LanguageRuntimeSession, vs
     /** Information about the runtime that is only available after starting */
     private _runtimeInfo?: positron.LanguageRuntimeInfo;
 
-    dynState: positron.LanguageRuntimeDynState;
+    private dynState: positron.LanguageRuntimeDynState;
 
     onDidReceiveRuntimeMessage = this._messageEmitter.event;
 
@@ -168,6 +168,10 @@ export class PythonRuntimeSession implements positron.LanguageRuntimeSession, vs
         return this._runtimeInfo;
     }
 
+    getDynState(): Thenable<positron.LanguageRuntimeDynState> {
+        return Promise.resolve(this.dynState);
+    }
+
     async debug(request: positron.DebugProtocolRequest): Promise<positron.DebugProtocolResponse> {
         if (this._kernel) {
             return await this._kernel.debug(request);
diff --git a/extensions/positron-r/src/session.ts b/extensions/positron-r/src/session.ts
@@ -148,6 +148,10 @@ export class RSession implements positron.LanguageRuntimeSession, vscode.Disposa
 		return this._kernel?.runtimeInfo;
 	}
 
+	getDynState(): Thenable<positron.LanguageRuntimeDynState> {
+		return Promise.resolve(this.dynState);
+	}
+
 	/**
 	 * Opens a resource in the runtime.
 	 * @param resource The resource to open.
diff --git a/extensions/positron-reticulate/src/extension.ts b/extensions/positron-reticulate/src/extension.ts
@@ -179,7 +179,13 @@ export class ReticulateRuntimeManager implements positron.LanguageRuntimeManager
 				// We have a free R session, we can attach to it. First we need to figure out if there's
 				// a prefered one.
 				// TODO: maybe show a quick menu so the user can select the session they want to attach to?
-				return freeRSessions[0];
+				const session = freeRSessions[0];
+
+				// Resolve the session to a full session object. Note that this
+				// typecast is only safe because we ensure that this extension
+				// runs on the same extension host as the Python and R
+				// extensions, via declared dependencies in package.json.
+				return session as positron.LanguageRuntimeSession;
 			} else {
 				progress.report({ increment: 2, message: vscode.l10n.t('Starting a new R session') });
 				// We need to create a new R session.
@@ -297,8 +303,10 @@ export class ReticulateRuntimeManager implements positron.LanguageRuntimeManager
 		progress.report({ increment: 10, message: vscode.l10n.t('Finding the host R session') });
 		const rSession = await (async () => {
 			for (let attempt = 1; attempt <= 5; attempt++) {
-				const sessions = await positron.runtime.getActiveSessions();
-				const hostRSession = sessions.find((sess) => sess.metadata.sessionId === hostRSessionId);
+				// Note that this typecast is only safe because we ensure that
+				// this extension runs on the same extension host as the Python
+				// and R extensions, via declared dependencies in package.json.
+				const hostRSession = await positron.runtime.getSession(hostRSessionId) as positron.LanguageRuntimeSession | undefined;
 				if (hostRSession) {
 					return hostRSession;
 				}
@@ -652,6 +660,10 @@ class ReticulateRuntimeSession implements positron.LanguageRuntimeSession {
 		);
 	}
 
+	getDynState(): Thenable<positron.LanguageRuntimeDynState> {
+		return this.pythonSession.getDynState();
+	}
+
 	createPythonRuntimeSession(runtimeMetadata: positron.LanguageRuntimeMetadata, sessionMetadata: positron.RuntimeSessionMetadata, kernelSpec?: JupyterKernelSpec): positron.LanguageRuntimeSession {
 		const api = vscode.extensions.getExtension('ms-python.python');
 		if (!api) {
@@ -746,10 +758,6 @@ class ReticulateRuntimeSession implements positron.LanguageRuntimeSession {
 		return this.pythonSession.runtimeMetadata;
 	}
 
-	public get dynState() {
-		return this.pythonSession.dynState;
-	}
-
 	public get runtimeInfo() {
 		return this.pythonSession.runtimeInfo;
 	}
diff --git a/extensions/positron-runtime-debugger/package.json b/extensions/positron-runtime-debugger/package.json
@@ -49,5 +49,9 @@
   },
   "devDependencies": {
     "@types/node": "^24.1.0"
-  }
+  },
+  "extensionDependencies": [
+    "ms-python.python",
+    "positron.positron-r"
+  ]
 }
diff --git a/extensions/positron-runtime-debugger/src/notebookDebugService.ts b/extensions/positron-runtime-debugger/src/notebookDebugService.ts
@@ -76,9 +76,8 @@ function getActiveNotebookCell(): vscode.NotebookCell | undefined {
 
 /** Get the language runtime session for a notebook. */
 export async function getNotebookSession(notebookUri: vscode.Uri): Promise<positron.LanguageRuntimeSession | undefined> {
-	const runtimeSessions = await positron.runtime.getActiveSessions();
-	const runtimeSession = runtimeSessions.find(
-		(session) => session.metadata.notebookUri && isUriEqual(session.metadata.notebookUri, notebookUri)
-	);
+	// This cast is safe only because our package.json ensures that this
+	// extension runs in the same extension host as the notebook kernels.
+	const runtimeSession = await positron.runtime.getNotebookSession(notebookUri) as positron.LanguageRuntimeSession | undefined;
 	return runtimeSession;
 }
diff --git a/extensions/positron-supervisor/src/KallichoreAdapterApi.ts b/extensions/positron-supervisor/src/KallichoreAdapterApi.ts
@@ -1319,14 +1319,15 @@ export class KCApi implements PositronSupervisorApi {
 
 		// Find the session in our list
 		const kallichoreSession = this._sessions.find(s => s.metadata.sessionId === session.metadata.sessionId);
+		const sessionName = (await session.getDynState()).sessionName;
 		if (!kallichoreSession) {
-			vscode.window.showInformationMessage(vscode.l10n.t('Active session {0} not managed by the kernel supervisor', session.dynState.sessionName));
+			vscode.window.showInformationMessage(vscode.l10n.t('Active session {0} not managed by the kernel supervisor', sessionName));
 			return;
 		}
 
 		// Ensure the session is still active
 		if (kallichoreSession.runtimeState === positron.RuntimeState.Exited) {
-			vscode.window.showInformationMessage(vscode.l10n.t('Session {0} is not running', session.dynState.sessionName));
+			vscode.window.showInformationMessage(vscode.l10n.t('Session {0} is not running', sessionName));
 			return;
 		}
 
diff --git a/extensions/positron-supervisor/src/KallichoreSession.ts b/extensions/positron-supervisor/src/KallichoreSession.ts
@@ -1475,6 +1475,10 @@ export class KallichoreSession implements JupyterLanguageRuntimeSession {
 		}
 	}
 
+	getDynState(): Thenable<positron.LanguageRuntimeDynState> {
+		return Promise.resolve(this.dynState);
+	}
+
 	/**
 	 * Performs a restart of the kernel. Kallichore handles the mechanics of
 	 * stopping the process and starting a new one; we just need to listen for
diff --git a/extensions/positron-zed/src/positronZedLanguageRuntime.ts b/extensions/positron-zed/src/positronZedLanguageRuntime.ts
@@ -1192,6 +1192,10 @@ export class PositronZedRuntimeSession implements positron.LanguageRuntimeSessio
 		}
 	}
 
+	getDynState(): Thenable<positron.LanguageRuntimeDynState> {
+		return Promise.resolve(this.dynState);
+	}
+
 	/**
 	 * Restarts the runtime.
 	 */
diff --git a/extensions/vscode-api-tests/src/singlefolder-tests/positron/runtime.test.ts b/extensions/vscode-api-tests/src/singlefolder-tests/positron/runtime.test.ts
@@ -40,6 +40,10 @@ class TestLanguageRuntimeSession implements positron.LanguageRuntimeSession {
 		return this._runtimeInfo;
 	}
 
+	getDynState(): Promise<positron.LanguageRuntimeDynState> {
+		return Promise.resolve(this.dynState);
+	}
+
 	generateMessageId(): string {
 		return `msg-${TestLanguageRuntimeSession.messageId++}`;
 	}
diff --git a/src/positron-dts/positron.d.ts b/src/positron-dts/positron.d.ts
@@ -1036,12 +1036,11 @@ declare module 'positron' {
 	}
 
 	/**
-	 * LanguageRuntimeSession is an interface implemented by extensions that provide a
-	 * set of common tools for interacting with a language runtime, such as code
-	 * execution, LSP implementation, and plotting.
+	 * Basic metadata about an active language runtime session, including
+	 * immutable metadata about the session itself and metadata about the
+	 * runtime with which it is associated.
 	 */
-	export interface LanguageRuntimeSession extends vscode.Disposable {
-
+	export interface ActiveRuntimeSessionMetadata {
 		/** An object supplying immutable metadata about this specific session */
 		readonly metadata: RuntimeSessionMetadata;
 
@@ -1050,12 +1049,61 @@ declare module 'positron' {
 		 * session is associated.
 		 */
 		readonly runtimeMetadata: LanguageRuntimeMetadata;
+	}
 
-		/** Information about the runtime that is only available after starting. */
-		readonly runtimeInfo: LanguageRuntimeInfo | undefined;
+	/**
+	 * Base interface for a language runtime session.
+	 *
+	 * This is the version of a language runtime session that is returned by
+	 * Positron's API methods; it provides basic access to the session for use
+	 * by extensions other that the one that created the session.
+	 */
+	export interface BaseLanguageRuntimeSession extends ActiveRuntimeSessionMetadata {
 
 		/** The state of the runtime that changes during a user session */
-		dynState: LanguageRuntimeDynState;
+		getDynState(): Thenable<LanguageRuntimeDynState>;
+
+		/**
+		 * Calls a method in the runtime and returns the result.
+		 *
+		 * Throws a RuntimeMethodError if the method call fails.
+		 *
+		 * @param method The name of the method to call
+		 * @param args Arguments to pass to the method
+		 */
+		callMethod?(method: string, ...args: any[]): Thenable<any>;
+
+		/**
+		 * Execute code in the runtime
+		 *
+		 * @param code The code to execute
+		 * @param id The ID of the code
+		 * @param mode The code execution mode
+		 * @param errorBehavior The code execution error behavior
+		 * Note: The errorBehavior parameter is currently ignored by kernels
+		 */
+		execute(code: string,
+			id: string,
+			mode: RuntimeCodeExecutionMode,
+			errorBehavior: RuntimeErrorBehavior): void;
+
+		/**
+		 * Shut down the runtime; returns a Thenable that resolves when the
+		 * runtime shutdown sequence has been successfully started (not
+		 * necessarily when it has completed).
+		 */
+		shutdown(exitReason: RuntimeExitReason): Thenable<void>;
+	}
+
+	/**
+	 * LanguageRuntimeSession is the full interface implemented by extensions
+	 * that provide a set of common tools for interacting with a language
+	 * runtime, such as code execution, LSP implementation, and plotting.
+	 */
+	export interface LanguageRuntimeSession extends BaseLanguageRuntimeSession, vscode.Disposable {
+
+		/** Information about the runtime that is only available after starting. */
+		readonly runtimeInfo: LanguageRuntimeInfo | undefined;
 
 		/** An object that emits language runtime events */
 		onDidReceiveRuntimeMessage: vscode.Event<LanguageRuntimeMessage>;
@@ -1081,30 +1129,6 @@ declare module 'positron' {
 		 */
 		debug(request: DebugProtocolRequest): Thenable<DebugProtocolResponse>;
 
-		/**
-		 * Execute code in the runtime
-		 *
-		 * @param code The code to execute
-		 * @param id The ID of the code
-		 * @param mode The code execution mode
-		 * @param errorBehavior The code execution error behavior
-		 * Note: The errorBehavior parameter is currently ignored by kernels
-		 */
-		execute(code: string,
-			id: string,
-			mode: RuntimeCodeExecutionMode,
-			errorBehavior: RuntimeErrorBehavior): void;
-
-		/**
-		 * Calls a method in the runtime and returns the result.
-		 *
-		 * Throws a RuntimeMethodError if the method call fails.
-		 *
-		 * @param method The name of the method to call
-		 * @param args Arguments to pass to the method
-		 */
-		callMethod?(method: string, ...args: any[]): Thenable<any>;
-
 		/** Test a code fragment for completeness */
 		isCodeFragmentComplete(code: string): Thenable<RuntimeCodeFragmentStatus>;
 
@@ -1173,13 +1197,6 @@ declare module 'positron' {
 		 */
 		restart(workingDirectory?: string): Thenable<void>;
 
-		/**
-		 * Shut down the runtime; returns a Thenable that resolves when the
-		 * runtime shutdown sequence has been successfully started (not
-		 * necessarily when it has completed).
-		 */
-		shutdown(exitReason: RuntimeExitReason): Thenable<void>;
-
 		/**
 		 * Forcibly quits the runtime; returns a Thenable that resolves when the
 		 * runtime has been terminated. This may be called by Positron if the
@@ -1812,19 +1829,24 @@ declare module 'positron' {
 		/**
 		 * List all active sessions.
 		 */
-		export function getActiveSessions(): Thenable<LanguageRuntimeSession[]>;
+		export function getActiveSessions(): Thenable<BaseLanguageRuntimeSession[]>;
+
+		/**
+		 * Get a specific session by its ID.
+		 */
+		export function getSession(sessionId: string): Thenable<BaseLanguageRuntimeSession | undefined>;
 
 		/**
 		 * Get the active foreground session, if any.
 		 */
-		export function getForegroundSession(): Thenable<LanguageRuntimeSession | undefined>;
+		export function getForegroundSession(): Thenable<BaseLanguageRuntimeSession | 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>;
+		export function getNotebookSession(notebookUri: vscode.Uri): Thenable<BaseLanguageRuntimeSession | undefined>;
 
 		/**
 		 * Select and start a runtime previously registered with Positron. Any
diff --git a/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts b/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts
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
diff --git a/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts b/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts
diff --git a/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts b/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts
diff --git a/src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts b/src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts

Original file line number	Diff line number	Diff line change
`@@ -141,6 +141,9 @@`
`141`	`141`	`]`
`142`	`142`	`}`
`143`	`143`	`},`
	`144`	`+ "extensionDependencies": [`
	`145`	`+ "positron.positron-supervisor"`
	`146`	`+ ],`
`144`	`147`	`"devDependencies": {`
`145`	`148`	`"@types/glob": "^7.2.0",`
`146`	`149`	`"@types/mocha": "^9.1.0",`