Allow switching interpreter mid-native repl session

anthonykim1 · anthonykim1 · commit 75fc45150947 · 2026-02-04T15:49:07.000-08:00
diff --git a/src/client/repl/nativeRepl.ts b/src/client/repl/nativeRepl.ts
@@ -4,9 +4,10 @@
 // Native Repl class that holds instance of pythonServer and replController
 
 import { NotebookController, NotebookDocument, QuickPickItem, TextEditor, Uri, WorkspaceFolder } from 'vscode';
+import * as path from 'path';
 import { Disposable } from 'vscode-jsonrpc';
 import { PVSC_EXTENSION_ID } from '../common/constants';
-import { showQuickPick } from '../common/vscodeApis/windowApis';
+import { showNotebookDocument, showQuickPick } from '../common/vscodeApis/windowApis';
 import { getWorkspaceFolders, onDidCloseNotebookDocument } from '../common/vscodeApis/workspaceApis';
 import { PythonEnvironment } from '../pythonEnvironments/info';
 import { createPythonServer, PythonServer } from './pythonServer';
@@ -18,6 +19,8 @@ import { VariablesProvider } from './variables/variablesProvider';
 import { VariableRequester } from './variables/variableRequester';
 import { getTabNameForUri } from './replUtils';
 import { getWorkspaceStateValue, updateWorkspaceStateValue } from '../common/persistentState';
+import { onDidChangeEnvironmentEnvExt, useEnvExtension } from '../envExt/api.internal';
+import { getActiveInterpreterLegacy } from '../envExt/api.legacy';
 
 export const NATIVE_REPL_URI_MEMENTO = 'nativeReplUri';
 let nativeRepl: NativeRepl | undefined;
@@ -37,6 +40,10 @@ export class NativeRepl implements Disposable {
 
     public newReplSession: boolean | undefined = true;
 
+    private envChangeListenerRegistered = false;
+
+    private pendingInterpreterChangeTimer?: NodeJS.Timeout;
+
     // TODO: In the future, could also have attribute of URI for file specific REPL.
     private constructor() {
         this.watchNotebookClosed();
@@ -49,6 +56,7 @@ export class NativeRepl implements Disposable {
         await nativeRepl.setReplDirectory();
         nativeRepl.pythonServer = createPythonServer([interpreter.path as string], nativeRepl.cwd);
         nativeRepl.setReplController();
+        nativeRepl.registerInterpreterChangeHandler();
 
         return nativeRepl;
     }
@@ -116,8 +124,8 @@ export class NativeRepl implements Disposable {
     /**
      * Function that check if NotebookController for REPL exists, and returns it in Singleton manner.
      */
-    public setReplController(): NotebookController {
-        if (!this.replController) {
+    public setReplController(force: boolean = false): NotebookController {
+        if (!this.replController || force) {
             this.replController = createReplController(this.interpreter!.path, this.disposables, this.cwd);
             this.replController.variableProvider = new VariablesProvider(
                 new VariableRequester(this.pythonServer),
@@ -128,6 +136,64 @@ export class NativeRepl implements Disposable {
         return this.replController;
     }
 
+    private registerInterpreterChangeHandler(): void {
+        if (!useEnvExtension() || this.envChangeListenerRegistered) {
+            return;
+        }
+        this.envChangeListenerRegistered = true;
+        this.disposables.push(
+            onDidChangeEnvironmentEnvExt((event) => {
+                this.updateInterpreterForChange(event.uri).catch(() => undefined);
+            }),
+        );
+    }
+
+    private async updateInterpreterForChange(resource?: Uri): Promise<void> {
+        if (this.pythonServer?.isExecuting) {
+            this.scheduleInterpreterUpdate(resource);
+            return;
+        }
+        if (!this.shouldApplyInterpreterChange(resource)) {
+            return;
+        }
+        const scope = resource ?? (this.cwd ? Uri.file(this.cwd) : undefined);
+        const interpreter = await getActiveInterpreterLegacy(scope);
+        if (!interpreter || interpreter.path === this.interpreter?.path) {
+            return;
+        }
+
+        this.interpreter = interpreter;
+        this.pythonServer.dispose();
+        this.pythonServer = createPythonServer([interpreter.path as string], this.cwd);
+        if (this.replController) {
+            this.replController.dispose();
+        }
+        this.setReplController(true);
+
+        if (this.notebookDocument) {
+            const notebookEditor = await showNotebookDocument(this.notebookDocument, { preserveFocus: true });
+            await selectNotebookKernel(notebookEditor, this.replController.id, PVSC_EXTENSION_ID);
+        }
+    }
+
+    private scheduleInterpreterUpdate(resource?: Uri): void {
+        if (this.pendingInterpreterChangeTimer) {
+            clearTimeout(this.pendingInterpreterChangeTimer);
+        }
+        this.pendingInterpreterChangeTimer = setTimeout(() => {
+            this.pendingInterpreterChangeTimer = undefined;
+            this.updateInterpreterForChange(resource).catch(() => undefined);
+        }, 200);
+    }
+
+    private shouldApplyInterpreterChange(resource?: Uri): boolean {
+        if (!resource || !this.cwd) {
+            return true;
+        }
+        const relative = path.relative(this.cwd, resource.fsPath);
+        return relative === '' || (!relative.startsWith('..') && !path.isAbsolute(relative));
+    }
+
     /**
      * Function that checks if native REPL's text input box contains complete code.
      * @returns Promise<boolean> - True if complete/Valid code is present, False otherwise.
diff --git a/src/client/repl/pythonServer.ts b/src/client/repl/pythonServer.ts
@@ -16,6 +16,8 @@ export interface ExecutionResult {
 
 export interface PythonServer extends Disposable {
     onCodeExecuted: Event<void>;
+    readonly isExecuting: boolean;
+    readonly isDisposed: boolean;
     execute(code: string): Promise<ExecutionResult | undefined>;
     executeSilently(code: string): Promise<ExecutionResult | undefined>;
     interrupt(): void;
@@ -30,6 +32,18 @@ class PythonServerImpl implements PythonServer, Disposable {
 
     onCodeExecuted = this._onCodeExecuted.event;
 
+    private inFlightRequests = 0;
+
+    private disposed = false;
+
+    public get isExecuting(): boolean {
+        return this.inFlightRequests > 0;
+    }
+
+    public get isDisposed(): boolean {
+        return this.disposed;
+    }
+
     constructor(private connection: rpc.MessageConnection, private pythonServer: ch.ChildProcess) {
         this.initialize();
         this.input();
@@ -41,6 +55,14 @@ class PythonServerImpl implements PythonServer, Disposable {
                 traceLog('Log:', message);
             }),
         );
+        this.pythonServer.on('exit', (code) => {
+            traceError(`Python server exited with code ${code}`);
+            this.markDisposed();
+        });
+        this.pythonServer.on('error', (err) => {
+            traceError(err);
+            this.markDisposed();
+        });
         this.connection.listen();
     }
 
@@ -75,12 +97,15 @@ class PythonServerImpl implements PythonServer, Disposable {
     }
 
     private async executeCode(code: string): Promise<ExecutionResult | undefined> {
+        this.inFlightRequests += 1;
         try {
             const result = await this.connection.sendRequest('execute', code);
             return result as ExecutionResult;
         } catch (err) {
             const error = err as Error;
             traceError(`Error getting response from REPL server:`, error);
+        } finally {
+            this.inFlightRequests -= 1;
         }
         return undefined;
     }
@@ -93,39 +118,47 @@ class PythonServerImpl implements PythonServer, Disposable {
     }
 
     public async checkValidCommand(code: string): Promise<boolean> {
-        const completeCode: ExecutionResult = await this.connection.sendRequest('check_valid_command', code);
-        if (completeCode.output === 'True') {
-            return new Promise((resolve) => resolve(true));
+        this.inFlightRequests += 1;
+        try {
+            const completeCode: ExecutionResult = await this.connection.sendRequest('check_valid_command', code);
+            return completeCode.output === 'True';
+        } finally {
+            this.inFlightRequests -= 1;
         }
-        return new Promise((resolve) => resolve(false));
     }
 
     public dispose(): void {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
         this.connection.sendNotification('exit');
         this.disposables.forEach((d) => d.dispose());
         this.connection.dispose();
         serverInstance = undefined;
     }
+
+    private markDisposed(): void {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
+        this.connection.dispose();
+        serverInstance = undefined;
+    }
 }
 
 export function createPythonServer(interpreter: string[], cwd?: string): PythonServer {
-    if (serverInstance) {
+    if (serverInstance && !serverInstance.isDisposed) {
         return serverInstance;
     }
 
     const pythonServer = ch.spawn(interpreter[0], [...interpreter.slice(1), SERVER_PATH], {
         cwd, // Launch with correct workspace directory
     });
-
     pythonServer.stderr.on('data', (data) => {
         traceError(data.toString());
     });
-    pythonServer.on('exit', (code) => {
-        traceError(`Python server exited with code ${code}`);
-    });
-    pythonServer.on('error', (err) => {
-        traceError(err);
-    });
     const connection = rpc.createMessageConnection(
         new rpc.StreamMessageReader(pythonServer.stdout),
         new rpc.StreamMessageWriter(pythonServer.stdin),
diff --git a/src/test/repl/nativeRepl.test.ts b/src/test/repl/nativeRepl.test.ts
@@ -129,6 +129,8 @@ suite('REPL - Native REPL', () => {
             input: sinon.stub(),
             checkValidCommand: sinon.stub().resolves(true),
             dispose: sinon.stub(),
+            isExecuting: false,
+            isDisposed: false,
         };
 
         // Track the number of times createPythonServer was called
