Handle reloading of REPL Window

src/client/extensionActivation.ts

@@ -112,9 +112,9 @@ export function activateFeatures(ext: ExtensionState, _components: Components):
     const executionHelper = ext.legacyIOC.serviceContainer.get<ICodeExecutionHelper>(ICodeExecutionHelper);
     const commandManager = ext.legacyIOC.serviceContainer.get<ICommandManager>(ICommandManager);
     registerTriggerForTerminalREPL(ext.disposables);
-    registerStartNativeReplCommand(ext.disposables, interpreterService);
-    registerReplCommands(ext.disposables, interpreterService, executionHelper, commandManager);
-    registerReplExecuteOnEnter(ext.disposables, interpreterService, commandManager);
+    registerStartNativeReplCommand(ext.disposables, interpreterService, ext.context);
+    registerReplCommands(ext.disposables, interpreterService, executionHelper, commandManager, ext.context);
+    registerReplExecuteOnEnter(ext.disposables, interpreterService, commandManager, ext.context);
 }
 
 /// //////////////////////////

src/client/repl/nativeRepl.ts

@@ -1,11 +1,13 @@
 // Native Repl class that holds instance of pythonServer and replController
 
 import {
+    ExtensionContext,
     NotebookController,
     NotebookControllerAffinity,
     NotebookDocument,
     QuickPickItem,
     TextEditor,
+    Uri,
     workspace,
     WorkspaceFolder,
 } from 'vscode';
@@ -21,8 +23,10 @@ import { EventName } from '../telemetry/constants';
 import { sendTelemetryEvent } from '../telemetry';
 import { VariablesProvider } from './variables/variablesProvider';
 import { VariableRequester } from './variables/variableRequester';
+import { getTabNameForUri } from './replUtils';
 
-let nativeRepl: NativeRepl | undefined; // In multi REPL scenario, hashmap of URI to Repl.
+export const NATIVE_REPL_URI_MEMENTO = 'nativeReplUri';
+let nativeRepl: NativeRepl | undefined;
 export class NativeRepl implements Disposable {
     // Adding ! since it will get initialized in create method, not the constructor.
     private pythonServer!: PythonServer;
@@ -39,14 +43,17 @@ export class NativeRepl implements Disposable {
 
     public newReplSession: boolean | undefined = true;
 
+    private context: ExtensionContext;
+
     // TODO: In the future, could also have attribute of URI for file specific REPL.
-    private constructor() {
+    private constructor(context: ExtensionContext) {
         this.watchNotebookClosed();
+        this.context = context;
     }
 
     // Static async factory method to handle asynchronous initialization
-    public static async create(interpreter: PythonEnvironment): Promise<NativeRepl> {
-        const nativeRepl = new NativeRepl();
+    public static async create(interpreter: PythonEnvironment, context: ExtensionContext): Promise<NativeRepl> {
+        const nativeRepl = new NativeRepl(context);
         nativeRepl.interpreter = interpreter;
         await nativeRepl.setReplDirectory();
         nativeRepl.pythonServer = createPythonServer([interpreter.path as string], nativeRepl.cwd);
@@ -65,10 +72,11 @@ export class NativeRepl implements Disposable {
      */
     private watchNotebookClosed(): void {
         this.disposables.push(
-            workspace.onDidCloseNotebookDocument((nb) => {
+            workspace.onDidCloseNotebookDocument(async (nb) => {
                 if (this.notebookDocument && nb.uri.toString() === this.notebookDocument.uri.toString()) {
                     this.notebookDocument = undefined;
                     this.newReplSession = true;
+                    await this.context.globalState.update(NATIVE_REPL_URI_MEMENTO, undefined);
                 }
             }),
         );
@@ -145,9 +153,39 @@ export class NativeRepl implements Disposable {
     /**
      * Function that opens interactive repl, selects kernel, and send/execute code to the native repl.
      */
-    public async sendToNativeRepl(code?: string): Promise<void> {
-        const notebookEditor = await openInteractiveREPL(this.replController, this.notebookDocument);
+    public async sendToNativeRepl(code?: string | undefined, preserveFocus: boolean = true): Promise<void> {
+        const mementoValue = (await this.context.globalState.get(NATIVE_REPL_URI_MEMENTO)) as string | undefined;
+        let mementoUri = mementoValue ? Uri.parse(mementoValue) : undefined;
+        const openNotebookDocuments = workspace.notebookDocuments.map((doc) => doc.uri);
+
+        if (mementoUri) {
+            const replTabBeforeReload = openNotebookDocuments.find((uri) => uri.fsPath === mementoUri?.fsPath);
+            if (replTabBeforeReload) {
+                this.notebookDocument = workspace.notebookDocuments.find(
+                    (doc) => doc.uri.fsPath === replTabBeforeReload.fsPath,
+                );
+                await this.context.globalState.update(NATIVE_REPL_URI_MEMENTO, replTabBeforeReload.toString());
+
+                // If repl URI does not have tabLabel 'Python REPL', something has changed:
+                // e.g. creation of untitled notebook without Python extension knowing.
+                if (getTabNameForUri(replTabBeforeReload) !== 'Python REPL') {
+                    mementoUri = undefined;
+                    await this.cleanRepl();
+                }
+            }
+        } else {
+            await this.cleanRepl();
+        }
+
+        const notebookEditor = await openInteractiveREPL(
+            this.replController,
+            this.notebookDocument,
+            mementoUri,
+            preserveFocus,
+        );
+
         this.notebookDocument = notebookEditor.notebook;
+        await this.context.globalState.update(NATIVE_REPL_URI_MEMENTO, this.notebookDocument.uri.toString());
 
         if (this.notebookDocument) {
             this.replController.updateNotebookAffinity(this.notebookDocument, NotebookControllerAffinity.Default);
@@ -157,16 +195,29 @@ export class NativeRepl implements Disposable {
             }
         }
     }
+
+    /**
+     * Properly clean up notebook document stored inside Native REPL.
+     * Also remove the Native REPL URI from memento to prepare for brand new REPL creation.
+     */
+    private async cleanRepl(): Promise<void> {
+        this.notebookDocument = undefined;
+        await this.context.globalState.update(NATIVE_REPL_URI_MEMENTO, undefined);
+    }
 }
 
 /**
  * Get Singleton Native REPL Instance
  * @param interpreter
  * @returns Native REPL instance
  */
-export async function getNativeRepl(interpreter: PythonEnvironment, disposables: Disposable[]): Promise<NativeRepl> {
+export async function getNativeRepl(
+    interpreter: PythonEnvironment,
+    disposables: Disposable[],
+    context: ExtensionContext,
+): Promise<NativeRepl> {
     if (!nativeRepl) {
-        nativeRepl = await NativeRepl.create(interpreter);
+        nativeRepl = await NativeRepl.create(interpreter, context);
         disposables.push(nativeRepl);
     }
     if (nativeRepl && nativeRepl.newReplSession) {

src/client/repl/replCommandHandler.ts

@@ -10,6 +10,7 @@ import {
     NotebookEdit,
     WorkspaceEdit,
     workspace,
+    Uri,
 } from 'vscode';
 import { getExistingReplViewColumn } from './replUtils';
 import { PVSC_EXTENSION_ID } from '../common/constants';
@@ -20,21 +21,27 @@ import { PVSC_EXTENSION_ID } from '../common/constants';
 export async function openInteractiveREPL(
     notebookController: NotebookController,
     notebookDocument: NotebookDocument | undefined,
+    mementoValue: Uri | undefined,
+    preserveFocus: boolean = true,
 ): Promise<NotebookEditor> {
     let viewColumn = ViewColumn.Beside;
-
-    // Case where NotebookDocument (REPL document already exists in the tab)
-    if (notebookDocument) {
+    if (mementoValue) {
+        // also check if memento value URI tab has file name of Python REPL
+        // Cached NotebookDocument exists.
+        notebookDocument = await workspace.openNotebookDocument(mementoValue as Uri);
+    } else if (notebookDocument) {
+        // Case where NotebookDocument (REPL document already exists in the tab)
         const existingReplViewColumn = getExistingReplViewColumn(notebookDocument);
         viewColumn = existingReplViewColumn ?? viewColumn;
     } else if (!notebookDocument) {
-        // Case where NotebookDocument doesnt exist, create a blank one.
+        // Case where NotebookDocument doesnt exist, or
+        // became outdated (untitled.ipynb created without Python extension knowing, effectively taking over original Python REPL's URI)
         notebookDocument = await workspace.openNotebookDocument('jupyter-notebook');
     }
     const editor = window.showNotebookDocument(notebookDocument!, {
         viewColumn,
         asRepl: 'Python REPL',
-        preserveFocus: true,
+        preserveFocus,
     });
     await commands.executeCommand('notebook.selectKernel', {
         editor,

src/client/repl/replCommands.ts

@@ -1,4 +1,4 @@
-import { commands, Uri, window } from 'vscode';
+import { commands, ExtensionContext, Uri, window } from 'vscode';
 import { Disposable } from 'vscode-jsonrpc';
 import { ICommandManager } from '../common/application/types';
 import { Commands } from '../common/constants';
@@ -25,15 +25,16 @@ import { ReplType } from './types';
 export async function registerStartNativeReplCommand(
     disposables: Disposable[],
     interpreterService: IInterpreterService,
+    context: ExtensionContext,
 ): Promise<void> {
     disposables.push(
         registerCommand(Commands.Start_Native_REPL, async (uri: Uri) => {
             sendTelemetryEvent(EventName.REPL, undefined, { replType: 'Native' });
             const interpreter = await getActiveInterpreter(uri, interpreterService);
             if (interpreter) {
                 if (interpreter) {
-                    const nativeRepl = await getNativeRepl(interpreter, disposables);
-                    await nativeRepl.sendToNativeRepl();
+                    const nativeRepl = await getNativeRepl(interpreter, disposables, context);
+                    await nativeRepl.sendToNativeRepl(undefined, false);
                 }
             }
         }),
@@ -48,6 +49,7 @@ export async function registerReplCommands(
     interpreterService: IInterpreterService,
     executionHelper: ICodeExecutionHelper,
     commandManager: ICommandManager,
+    context: ExtensionContext,
 ): Promise<void> {
     disposables.push(
         commandManager.registerCommand(Commands.Exec_In_REPL, async (uri: Uri) => {
@@ -60,7 +62,7 @@ export async function registerReplCommands(
             const interpreter = await getActiveInterpreter(uri, interpreterService);
 
             if (interpreter) {
-                const nativeRepl = await getNativeRepl(interpreter, disposables);
+                const nativeRepl = await getNativeRepl(interpreter, disposables, context);
                 const activeEditor = window.activeTextEditor;
                 if (activeEditor) {
                     const code = await getSelectedTextToExecute(activeEditor);
@@ -90,15 +92,16 @@ export async function registerReplExecuteOnEnter(
     disposables: Disposable[],
     interpreterService: IInterpreterService,
     commandManager: ICommandManager,
+    context: ExtensionContext,
 ): Promise<void> {
     disposables.push(
         commandManager.registerCommand(Commands.Exec_In_REPL_Enter, async (uri: Uri) => {
-            await onInputEnter(uri, 'repl.execute', interpreterService, disposables);
+            await onInputEnter(uri, 'repl.execute', interpreterService, disposables, context);
         }),
     );
     disposables.push(
         commandManager.registerCommand(Commands.Exec_In_IW_Enter, async (uri: Uri) => {
-            await onInputEnter(uri, 'interactive.execute', interpreterService, disposables);
+            await onInputEnter(uri, 'interactive.execute', interpreterService, disposables, context);
         }),
     );
 }
@@ -108,14 +111,15 @@ async function onInputEnter(
     commandName: string,
     interpreterService: IInterpreterService,
     disposables: Disposable[],
+    context: ExtensionContext,
 ): Promise<void> {
     const interpreter = await interpreterService.getActiveInterpreter(uri);
     if (!interpreter) {
         commands.executeCommand(Commands.TriggerEnvironmentSelection, uri).then(noop, noop);
         return;
     }
 
-    const nativeRepl = await getNativeRepl(interpreter, disposables);
+    const nativeRepl = await getNativeRepl(interpreter, disposables, context);
     const completeCode = await nativeRepl?.checkUserInputCompleteCode(window.activeTextEditor);
     const editor = window.activeTextEditor;
 

src/client/repl/replUtils.ts

@@ -99,3 +99,22 @@ export function getExistingReplViewColumn(notebookDocument: NotebookDocument): V
     }
     return undefined;
 }
+
+/**
+ * Function that will return tab name for before reloading VS Code
+ * This is so we can make sure tab name is still 'Python REPL' after reloading VS Code,
+ * and make sure Python REPL does not get 'merged' into unaware untitled.ipynb tab.
+ */
+export function getTabNameForUri(uri: Uri): string | undefined {
+    const tabGroups = window.tabGroups.all;
+
+    for (const tabGroup of tabGroups) {
+        for (const tab of tabGroup.tabs) {
+            if (tab.input instanceof TabInputNotebook && tab.input.uri.toString() === uri.toString()) {
+                return tab.label;
+            }
+        }
+    }
+
+    return undefined;
+}

src/test/repl/nativeRepl.test.ts

@@ -2,33 +2,39 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import * as TypeMoq from 'typemoq';
 import * as sinon from 'sinon';
-import { Disposable } from 'vscode';
+import { Disposable, ExtensionContext } from 'vscode';
 import { expect } from 'chai';
 
 import { IInterpreterService } from '../../client/interpreter/contracts';
 import { PythonEnvironment } from '../../client/pythonEnvironments/info';
-import { getNativeRepl, NativeRepl } from '../../client/repl/nativeRepl';
+import { getNativeRepl, NATIVE_REPL_URI_MEMENTO, NativeRepl } from '../../client/repl/nativeRepl';
+import { IExtensionContext } from '../../client/common/types';
+import * as replUtils from '../../client/repl/replUtils';
 
 suite('REPL - Native REPL', () => {
     let interpreterService: TypeMoq.IMock<IInterpreterService>;
-
+    let extensionContext: TypeMoq.IMock<IExtensionContext>;
     let disposable: TypeMoq.IMock<Disposable>;
     let disposableArray: Disposable[] = [];
-
     let setReplDirectoryStub: sinon.SinonStub;
     let setReplControllerSpy: sinon.SinonSpy;
-
+    let memento: TypeMoq.IMock<ExtensionContext['globalState']>;
+    let getTabNameForUriStub: sinon.SinonStub;
     setup(() => {
         interpreterService = TypeMoq.Mock.ofType<IInterpreterService>();
         interpreterService
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
         disposable = TypeMoq.Mock.ofType<Disposable>();
         disposableArray = [disposable.object];
-
+        memento = TypeMoq.Mock.ofType<ExtensionContext['globalState']>();
         setReplDirectoryStub = sinon.stub(NativeRepl.prototype as any, 'setReplDirectory').resolves(); // Stubbing private method
         // Use a spy instead of a stub for setReplController
+        getTabNameForUriStub = sinon.stub(replUtils, 'getTabNameForUri').returns('tabName');
         setReplControllerSpy = sinon.spy(NativeRepl.prototype, 'setReplController');
+        extensionContext = TypeMoq.Mock.ofType<IExtensionContext>();
+        extensionContext.setup((c) => c.globalState).returns(() => memento.object);
+        memento.setup((m) => m.get(NATIVE_REPL_URI_MEMENTO)).returns(() => undefined);
     });
 
     teardown(() => {
@@ -37,9 +43,10 @@ suite('REPL - Native REPL', () => {
                 d.dispose();
             }
         });
-
         disposableArray = [];
         sinon.restore();
+        extensionContext?.reset();
+        memento?.reset();
     });
 
     test('getNativeRepl should call create constructor', async () => {
@@ -48,18 +55,36 @@ suite('REPL - Native REPL', () => {
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
         const interpreter = await interpreterService.object.getActiveInterpreter();
-        await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
+        await getNativeRepl(interpreter as PythonEnvironment, disposableArray, extensionContext.object);
 
         expect(createMethodStub.calledOnce).to.be.true;
     });
 
+    test('sendToNativeRepl with undefined URI should not try to reload', async () => {
+        memento.setup((m) => m.get(NATIVE_REPL_URI_MEMENTO)).returns(() => undefined);
+
+        interpreterService
+            .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
+            .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
+        const interpreter = await interpreterService.object.getActiveInterpreter();
+        const nativeRepl = await getNativeRepl(
+            interpreter as PythonEnvironment,
+            disposableArray,
+            extensionContext.object,
+        );
+
+        nativeRepl.sendToNativeRepl(undefined, false);
+
+        expect(getTabNameForUriStub.notCalled).to.be.true;
+    });
+
     test('create should call setReplDirectory, setReplController', async () => {
         const interpreter = await interpreterService.object.getActiveInterpreter();
         interpreterService
             .setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
             .returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
 
-        await NativeRepl.create(interpreter as PythonEnvironment);
+        await NativeRepl.create(interpreter as PythonEnvironment, extensionContext.object);
 
         expect(setReplDirectoryStub.calledOnce).to.be.true;
         expect(setReplControllerSpy.calledOnce).to.be.true;