Fix native repl not using env extension

anthonykim1 · anthonykim1 · commit 15f7e0b39eb9 · 2026-02-04T12:31:54.000-08:00
diff --git a/src/client/repl/nativeRepl.ts b/src/client/repl/nativeRepl.ts
@@ -9,7 +9,7 @@ import { PVSC_EXTENSION_ID } from '../common/constants';
 import { showQuickPick } from '../common/vscodeApis/windowApis';
 import { getWorkspaceFolders, onDidCloseNotebookDocument } from '../common/vscodeApis/workspaceApis';
 import { PythonEnvironment } from '../pythonEnvironments/info';
-import { createPythonServer, PythonServer } from './pythonServer';
+import { createPythonServer, createPythonServerEnvExt, PythonServer } from './pythonServer';
 import { executeNotebookCell, openInteractiveREPL, selectNotebookKernel } from './replCommandHandler';
 import { createReplController } from './replController';
 import { EventName } from '../telemetry/constants';
@@ -18,6 +18,8 @@ import { VariablesProvider } from './variables/variablesProvider';
 import { VariableRequester } from './variables/variableRequester';
 import { getTabNameForUri } from './replUtils';
 import { getWorkspaceStateValue, updateWorkspaceStateValue } from '../common/persistentState';
+import { getEnvironment, useEnvExtension } from '../envExt/api.internal';
+import { traceError, traceVerbose } from '../logging';
 
 export const NATIVE_REPL_URI_MEMENTO = 'nativeReplUri';
 let nativeRepl: NativeRepl | undefined;
@@ -29,6 +31,8 @@ export class NativeRepl implements Disposable {
 
     private interpreter!: PythonEnvironment;
 
+    private resourceUri: Uri | undefined;
+
     private disposables: Disposable[] = [];
 
     private replController!: NotebookController;
@@ -42,12 +46,38 @@ export class NativeRepl implements Disposable {
         this.watchNotebookClosed();
     }
 
+    /**
+     * Get the interpreter path associated with this REPL instance.
+     * Used to detect when interpreter has changed and REPL needs to be recreated.
+     */
+    public get interpreterPath(): string {
+        return this.interpreter.path;
+    }
+
     // Static async factory method to handle asynchronous initialization
-    public static async create(interpreter: PythonEnvironment): Promise<NativeRepl> {
+    public static async create(interpreter: PythonEnvironment, resource?: Uri): Promise<NativeRepl> {
         const nativeRepl = new NativeRepl();
         nativeRepl.interpreter = interpreter;
+        nativeRepl.resourceUri = resource;
         await nativeRepl.setReplDirectory();
-        nativeRepl.pythonServer = createPythonServer([interpreter.path as string], nativeRepl.cwd);
+
+        // Use env extension's runInBackground if available for proper environment handling
+        if (useEnvExtension()) {
+            const pythonEnv = await getEnvironment(resource);
+            if (pythonEnv) {
+                traceVerbose(
+                    `Creating REPL server using environment extension for: ${pythonEnv.execInfo.run.executable}`,
+                );
+                nativeRepl.pythonServer = await createPythonServerEnvExt(pythonEnv, nativeRepl.cwd);
+            } else {
+                traceError('Failed to get environment from env extension, falling back to direct spawn');
+                nativeRepl.pythonServer = createPythonServer([interpreter.path as string], nativeRepl.cwd);
+            }
+        } else {
+            traceVerbose(`Creating REPL server using direct spawn for: ${interpreter.path}`);
+            nativeRepl.pythonServer = createPythonServer([interpreter.path as string], nativeRepl.cwd);
+        }
+
         nativeRepl.setReplController();
 
         return nativeRepl;
@@ -69,7 +99,19 @@ export class NativeRepl implements Disposable {
                     this.newReplSession = true;
                     await updateWorkspaceStateValue<string | undefined>(NATIVE_REPL_URI_MEMENTO, undefined);
                     this.pythonServer.dispose();
-                    this.pythonServer = createPythonServer([this.interpreter.path as string], this.cwd);
+
+                    // Recreate Python server - use env extension if available
+                    if (useEnvExtension()) {
+                        const pythonEnv = await getEnvironment(this.resourceUri);
+                        if (pythonEnv) {
+                            this.pythonServer = await createPythonServerEnvExt(pythonEnv, this.cwd);
+                        } else {
+                            this.pythonServer = createPythonServer([this.interpreter.path as string], this.cwd);
+                        }
+                    } else {
+                        this.pythonServer = createPythonServer([this.interpreter.path as string], this.cwd);
+                    }
+
                     this.disposables.push(this.pythonServer);
                     if (this.replController) {
                         this.replController.dispose();
@@ -118,7 +160,12 @@ export class NativeRepl implements Disposable {
      */
     public setReplController(): NotebookController {
         if (!this.replController) {
-            this.replController = createReplController(this.interpreter!.path, this.disposables, this.cwd);
+            this.replController = createReplController(
+                this.interpreter!.path,
+                this.disposables,
+                this.cwd,
+                this.pythonServer,
+            );
             this.replController.variableProvider = new VariablesProvider(
                 new VariableRequester(this.pythonServer),
                 () => this.notebookDocument,
@@ -183,13 +230,28 @@ export class NativeRepl implements Disposable {
 }
 
 /**
- * Get Singleton Native REPL Instance
+ * Get Singleton Native REPL Instance.
+ * Recreates the REPL if the interpreter has changed (e.g., when using environment extension).
  * @param interpreter
+ * @param disposables
+ * @param resource - Optional resource URI used to get the environment from the env extension
  * @returns Native REPL instance
  */
-export async function getNativeRepl(interpreter: PythonEnvironment, disposables: Disposable[]): Promise<NativeRepl> {
+export async function getNativeRepl(
+    interpreter: PythonEnvironment,
+    disposables: Disposable[],
+    resource?: Uri,
+): Promise<NativeRepl> {
+    // Check if interpreter has changed - if so, dispose and recreate the REPL
+    // This serves as a fallback in case the event listener didn't catch the change
+    if (nativeRepl && nativeRepl.interpreterPath !== interpreter.path) {
+        traceVerbose(`Interpreter changed to ${interpreter.path}, disposing Native REPL`);
+        nativeRepl.dispose();
+        nativeRepl = undefined;
+    }
+
     if (!nativeRepl) {
-        nativeRepl = await NativeRepl.create(interpreter);
+        nativeRepl = await NativeRepl.create(interpreter, resource);
         disposables.push(nativeRepl);
     }
     if (nativeRepl && nativeRepl.newReplSession) {
diff --git a/src/client/repl/pythonServer.ts b/src/client/repl/pythonServer.ts
@@ -6,9 +6,30 @@ import { EXTENSION_ROOT_DIR } from '../constants';
 import { traceError, traceLog } from '../logging';
 import { captureTelemetry } from '../telemetry';
 import { EventName } from '../telemetry/constants';
+import { PythonEnvironment as PythonEnvironmentEnvExt, PythonProcess } from '../envExt/types';
+import { runInBackground } from '../envExt/api.internal';
 
 const SERVER_PATH = path.join(EXTENSION_ROOT_DIR, 'python_files', 'python_server.py');
+
+/**
+ * Wrapper interface to unify ch.ChildProcess and PythonProcess from env extension
+ */
+interface PythonProcessWrapper {
+    readonly stdout: NodeJS.ReadableStream;
+    readonly stdin: NodeJS.WritableStream;
+    readonly stderr: NodeJS.ReadableStream;
+    kill(signal?: string): boolean;
+}
+
 let serverInstance: PythonServer | undefined;
+
+/**
+ * Clear the server instance. Used when the environment changes and we need to recreate the server.
+ */
+export function clearServerInstance(): void {
+    serverInstance = undefined;
+}
+
 export interface ExecutionResult {
     status: boolean;
     output: string;
@@ -30,7 +51,7 @@ class PythonServerImpl implements PythonServer, Disposable {
 
     onCodeExecuted = this._onCodeExecuted.event;
 
-    constructor(private connection: rpc.MessageConnection, private pythonServer: ch.ChildProcess) {
+    constructor(private connection: rpc.MessageConnection, private pythonServer: PythonProcessWrapper) {
         this.initialize();
         this.input();
     }
@@ -108,6 +129,33 @@ class PythonServerImpl implements PythonServer, Disposable {
     }
 }
 
+/**
+ * Wrap a ch.ChildProcess to match PythonProcessWrapper interface
+ */
+function wrapChildProcess(proc: ch.ChildProcess): PythonProcessWrapper {
+    return {
+        stdout: proc.stdout!,
+        stdin: proc.stdin!,
+        stderr: proc.stderr!,
+        kill: (signal?: string) => proc.kill(signal as NodeJS.Signals),
+    };
+}
+
+/**
+ * Wrap a PythonProcess from env extension to match PythonProcessWrapper interface
+ */
+function wrapEnvExtProcess(proc: PythonProcess): PythonProcessWrapper {
+    return {
+        stdout: proc.stdout,
+        stdin: proc.stdin as NodeJS.WritableStream,
+        stderr: proc.stderr,
+        kill: () => {
+            proc.kill();
+            return true;
+        },
+    };
+}
+
 export function createPythonServer(interpreter: string[], cwd?: string): PythonServer {
     if (serverInstance) {
         return serverInstance;
@@ -127,9 +175,41 @@ export function createPythonServer(interpreter: string[], cwd?: string): PythonS
         traceError(err);
     });
     const connection = rpc.createMessageConnection(
-        new rpc.StreamMessageReader(pythonServer.stdout),
-        new rpc.StreamMessageWriter(pythonServer.stdin),
+        new rpc.StreamMessageReader(pythonServer.stdout!),
+        new rpc.StreamMessageWriter(pythonServer.stdin!),
+    );
+    serverInstance = new PythonServerImpl(connection, wrapChildProcess(pythonServer));
+    return serverInstance;
+}
+
+/**
+ * Create a Python server using the environment extension's runInBackground API.
+ * This properly handles environment activation and variables.
+ */
+export async function createPythonServerEnvExt(
+    pythonEnv: PythonEnvironmentEnvExt,
+    cwd?: string,
+): Promise<PythonServer> {
+    if (serverInstance) {
+        return serverInstance;
+    }
+
+    const pythonProcess = await runInBackground(pythonEnv, {
+        args: [SERVER_PATH],
+        cwd,
+    });
+
+    pythonProcess.stderr.on('data', (data: Buffer) => {
+        traceError(data.toString());
+    });
+    pythonProcess.onExit((code) => {
+        traceError(`Python server exited with code ${code}`);
+    });
+
+    const connection = rpc.createMessageConnection(
+        new rpc.StreamMessageReader(pythonProcess.stdout),
+        new rpc.StreamMessageWriter(pythonProcess.stdin as NodeJS.WritableStream),
     );
-    serverInstance = new PythonServerImpl(connection, pythonServer);
+    serverInstance = new PythonServerImpl(connection, wrapEnvExtProcess(pythonProcess));
     return serverInstance;
 }
diff --git a/src/client/repl/replCommands.ts b/src/client/repl/replCommands.ts
@@ -19,6 +19,17 @@ import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { ReplType } from './types';
 
+/**
+ * Get the resource URI, falling back to active editor if not provided.
+ */
+function getResourceUri(uri?: Uri): Uri | undefined {
+    if (uri) {
+        return uri;
+    }
+    // Fallback to active editor's document URI
+    return window.activeTextEditor?.document.uri;
+}
+
 /**
  * Register Start Native REPL command in the command palette
  */
@@ -29,9 +40,10 @@ export async function registerStartNativeReplCommand(
     disposables.push(
         registerCommand(Commands.Start_Native_REPL, async (uri: Uri) => {
             sendTelemetryEvent(EventName.REPL, undefined, { replType: 'Native' });
-            const interpreter = await getActiveInterpreter(uri, interpreterService);
+            const resource = getResourceUri(uri);
+            const interpreter = await getActiveInterpreter(resource, interpreterService);
             if (interpreter) {
-                const nativeRepl = await getNativeRepl(interpreter, disposables);
+                const nativeRepl = await getNativeRepl(interpreter, disposables, resource);
                 await nativeRepl.sendToNativeRepl(undefined, false);
             }
         }),
@@ -55,10 +67,11 @@ export async function registerReplCommands(
                 await executeInTerminal();
                 return;
             }
-            const interpreter = await getActiveInterpreter(uri, interpreterService);
+            const resource = getResourceUri(uri);
+            const interpreter = await getActiveInterpreter(resource, interpreterService);
 
             if (interpreter) {
-                const nativeRepl = await getNativeRepl(interpreter, disposables);
+                const nativeRepl = await getNativeRepl(interpreter, disposables, resource);
                 const activeEditor = window.activeTextEditor;
                 if (activeEditor) {
                     const code = await getSelectedTextToExecute(activeEditor);
@@ -107,13 +120,14 @@ async function onInputEnter(
     interpreterService: IInterpreterService,
     disposables: Disposable[],
 ): Promise<void> {
-    const interpreter = await interpreterService.getActiveInterpreter(uri);
+    const resource = getResourceUri(uri);
+    const interpreter = await interpreterService.getActiveInterpreter(resource);
     if (!interpreter) {
-        commands.executeCommand(Commands.TriggerEnvironmentSelection, uri).then(noop, noop);
+        commands.executeCommand(Commands.TriggerEnvironmentSelection, resource).then(noop, noop);
         return;
     }
 
-    const nativeRepl = await getNativeRepl(interpreter, disposables);
+    const nativeRepl = await getNativeRepl(interpreter, disposables, resource);
     const completeCode = await nativeRepl?.checkUserInputCompleteCode(window.activeTextEditor);
     const editor = window.activeTextEditor;
 
diff --git a/src/client/repl/replController.ts b/src/client/repl/replController.ts
@@ -1,14 +1,12 @@
 import * as vscode from 'vscode';
-import { createPythonServer } from './pythonServer';
+import { PythonServer } from './pythonServer';
 
 export function createReplController(
     interpreterPath: string,
     disposables: vscode.Disposable[],
-    cwd?: string,
+    cwd: string | undefined,
+    server: PythonServer,
 ): vscode.NotebookController {
-    const server = createPythonServer([interpreterPath], cwd);
-    disposables.push(server);
-
     const controller = vscode.notebooks.createNotebookController('pythonREPL', 'jupyter-notebook', 'Python REPL');
     controller.supportedLanguages = ['python'];
 
diff --git a/src/client/repl/replUtils.ts b/src/client/repl/replUtils.ts
@@ -66,7 +66,7 @@ export function isMultiLineText(textEditor: TextEditor): boolean {
  * Function will also return undefined or active interpreter
  */
 export async function getActiveInterpreter(
-    uri: Uri,
+    uri: Uri | undefined,
     interpreterService: IInterpreterService,
 ): Promise<PythonEnvironment | undefined> {
     const interpreter = await interpreterService.getActiveInterpreter(uri);
