Skip to content

Commit 3275c34

Browse files
prevent native REPL from caching state between sessions (#24857)
Resolves: #24359 Definitely warrants scrutiny / input as I've never written typescript before. Solved with trial and error + LLMs. --------- Co-authored-by: Anthony Kim <anthonykim@microsoft.com> Co-authored-by: Anthony Kim <62267334+anthonykim1@users.noreply.github.com>
1 parent 6a20c9c commit 3275c34

File tree

5 files changed

+121
-18
lines changed

5 files changed

+121
-18
lines changed

.github/actions/smoke-tests/action.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,13 @@ runs:
1313
using: 'composite'
1414
steps:
1515
- name: Install Node
16-
uses: actions/setup-node@v2
16+
uses: actions/setup-node@v4
1717
with:
1818
node-version: ${{ inputs.node_version }}
1919
cache: 'npm'
2020

2121
- name: Install Python
22-
uses: actions/setup-python@v2
22+
uses: actions/setup-python@v5
2323
with:
2424
python-version: '3.x'
2525
cache: 'pip'

src/client/common/vscodeApis/workspaceApis.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,10 @@ export function onDidChangeConfiguration(handler: (e: vscode.ConfigurationChange
6060
return vscode.workspace.onDidChangeConfiguration(handler);
6161
}
6262

63+
export function onDidCloseNotebookDocument(handler: (e: vscode.NotebookDocument) => void): vscode.Disposable {
64+
return vscode.workspace.onDidCloseNotebookDocument(handler);
65+
}
66+
6367
export function createFileSystemWatcher(
6468
globPattern: vscode.GlobPattern,
6569
ignoreCreateEvents?: boolean,

src/client/repl/nativeRepl.ts

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
14
// Native Repl class that holds instance of pythonServer and replController
25

36
import {
@@ -7,13 +10,12 @@ import {
710
QuickPickItem,
811
TextEditor,
912
Uri,
10-
workspace,
1113
WorkspaceFolder,
1214
} from 'vscode';
1315
import { Disposable } from 'vscode-jsonrpc';
1416
import { PVSC_EXTENSION_ID } from '../common/constants';
1517
import { showQuickPick } from '../common/vscodeApis/windowApis';
16-
import { getWorkspaceFolders } from '../common/vscodeApis/workspaceApis';
18+
import { getWorkspaceFolders, onDidCloseNotebookDocument } from '../common/vscodeApis/workspaceApis';
1719
import { PythonEnvironment } from '../pythonEnvironments/info';
1820
import { createPythonServer, PythonServer } from './pythonServer';
1921
import { executeNotebookCell, openInteractiveREPL, selectNotebookKernel } from './replCommandHandler';
@@ -69,11 +71,18 @@ export class NativeRepl implements Disposable {
6971
*/
7072
private watchNotebookClosed(): void {
7173
this.disposables.push(
72-
workspace.onDidCloseNotebookDocument(async (nb) => {
74+
onDidCloseNotebookDocument(async (nb) => {
7375
if (this.notebookDocument && nb.uri.toString() === this.notebookDocument.uri.toString()) {
7476
this.notebookDocument = undefined;
7577
this.newReplSession = true;
7678
await updateWorkspaceStateValue<string | undefined>(NATIVE_REPL_URI_MEMENTO, undefined);
79+
this.pythonServer.dispose();
80+
this.pythonServer = createPythonServer([this.interpreter.path as string], this.cwd);
81+
this.disposables.push(this.pythonServer);
82+
if (this.replController) {
83+
this.replController.dispose();
84+
}
85+
nativeRepl = undefined;
7786
}
7887
}),
7988
);

src/client/repl/pythonServer.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ class PythonServerImpl implements PythonServer, Disposable {
104104
this.connection.sendNotification('exit');
105105
this.disposables.forEach((d) => d.dispose());
106106
this.connection.dispose();
107+
serverInstance = undefined;
107108
}
108109
}
109110

src/test/repl/nativeRepl.test.ts

Lines changed: 102 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,21 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
14
/* eslint-disable no-unused-expressions */
25
/* eslint-disable @typescript-eslint/no-explicit-any */
36
import * as TypeMoq from 'typemoq';
47
import * as sinon from 'sinon';
5-
import { Disposable } from 'vscode';
8+
import { Disposable, EventEmitter, NotebookDocument, Uri } from 'vscode';
69
import { expect } from 'chai';
710

811
import { IInterpreterService } from '../../client/interpreter/contracts';
912
import { PythonEnvironment } from '../../client/pythonEnvironments/info';
10-
import { getNativeRepl, NativeRepl } from '../../client/repl/nativeRepl';
13+
import * as NativeReplModule from '../../client/repl/nativeRepl';
1114
import * as persistentState from '../../client/common/persistentState';
15+
import * as PythonServer from '../../client/repl/pythonServer';
16+
import * as vscodeWorkspaceApis from '../../client/common/vscodeApis/workspaceApis';
17+
import * as replController from '../../client/repl/replController';
18+
import { executeCommand } from '../../client/common/vscodeApis/commandApis';
1219

1320
suite('REPL - Native REPL', () => {
1421
let interpreterService: TypeMoq.IMock<IInterpreterService>;
@@ -19,38 +26,51 @@ suite('REPL - Native REPL', () => {
1926
let setReplControllerSpy: sinon.SinonSpy;
2027
let getWorkspaceStateValueStub: sinon.SinonStub;
2128
let updateWorkspaceStateValueStub: sinon.SinonStub;
29+
let createReplControllerStub: sinon.SinonStub;
30+
let mockNotebookController: any;
2231

2332
setup(() => {
33+
(NativeReplModule as any).nativeRepl = undefined;
34+
35+
mockNotebookController = {
36+
id: 'mockController',
37+
dispose: sinon.stub(),
38+
updateNotebookAffinity: sinon.stub(),
39+
createNotebookCellExecution: sinon.stub(),
40+
variableProvider: null,
41+
};
42+
2443
interpreterService = TypeMoq.Mock.ofType<IInterpreterService>();
2544
interpreterService
2645
.setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
2746
.returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
2847
disposable = TypeMoq.Mock.ofType<Disposable>();
2948
disposableArray = [disposable.object];
3049

31-
setReplDirectoryStub = sinon.stub(NativeRepl.prototype as any, 'setReplDirectory').resolves(); // Stubbing private method
32-
// Use a spy instead of a stub for setReplController
33-
setReplControllerSpy = sinon.spy(NativeRepl.prototype, 'setReplController');
50+
createReplControllerStub = sinon.stub(replController, 'createReplController').returns(mockNotebookController);
51+
setReplDirectoryStub = sinon.stub(NativeReplModule.NativeRepl.prototype as any, 'setReplDirectory').resolves();
52+
setReplControllerSpy = sinon.spy(NativeReplModule.NativeRepl.prototype, 'setReplController');
3453
updateWorkspaceStateValueStub = sinon.stub(persistentState, 'updateWorkspaceStateValue').resolves();
3554
});
3655

37-
teardown(() => {
56+
teardown(async () => {
3857
disposableArray.forEach((d) => {
3958
if (d) {
4059
d.dispose();
4160
}
4261
});
4362
disposableArray = [];
4463
sinon.restore();
64+
executeCommand('workbench.action.closeActiveEditor');
4565
});
4666

4767
test('getNativeRepl should call create constructor', async () => {
48-
const createMethodStub = sinon.stub(NativeRepl, 'create');
68+
const createMethodStub = sinon.stub(NativeReplModule.NativeRepl, 'create');
4969
interpreterService
5070
.setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
5171
.returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
5272
const interpreter = await interpreterService.object.getActiveInterpreter();
53-
await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
73+
await NativeReplModule.getNativeRepl(interpreter as PythonEnvironment, disposableArray);
5474

5575
expect(createMethodStub.calledOnce).to.be.true;
5676
});
@@ -61,7 +81,7 @@ suite('REPL - Native REPL', () => {
6181
.setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
6282
.returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
6383
const interpreter = await interpreterService.object.getActiveInterpreter();
64-
const nativeRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
84+
const nativeRepl = await NativeReplModule.getNativeRepl(interpreter as PythonEnvironment, disposableArray);
6585

6686
nativeRepl.sendToNativeRepl(undefined, false);
6787

@@ -74,7 +94,7 @@ suite('REPL - Native REPL', () => {
7494
.setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
7595
.returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
7696
const interpreter = await interpreterService.object.getActiveInterpreter();
77-
const nativeRepl = await getNativeRepl(interpreter as PythonEnvironment, disposableArray);
97+
const nativeRepl = await NativeReplModule.getNativeRepl(interpreter as PythonEnvironment, disposableArray);
7898

7999
nativeRepl.sendToNativeRepl(undefined, false);
80100

@@ -87,12 +107,81 @@ suite('REPL - Native REPL', () => {
87107
.setup((i) => i.getActiveInterpreter(TypeMoq.It.isAny()))
88108
.returns(() => Promise.resolve(({ path: 'ps' } as unknown) as PythonEnvironment));
89109

90-
await NativeRepl.create(interpreter as PythonEnvironment);
110+
await NativeReplModule.NativeRepl.create(interpreter as PythonEnvironment);
91111

92112
expect(setReplDirectoryStub.calledOnce).to.be.true;
93113
expect(setReplControllerSpy.calledOnce).to.be.true;
114+
expect(createReplControllerStub.calledOnce).to.be.true;
115+
});
116+
117+
test('watchNotebookClosed should clean up resources when notebook is closed', async () => {
118+
const notebookCloseEmitter = new EventEmitter<NotebookDocument>();
119+
sinon.stub(vscodeWorkspaceApis, 'onDidCloseNotebookDocument').callsFake((handler) => {
120+
const disposable = notebookCloseEmitter.event(handler);
121+
return disposable;
122+
});
123+
124+
const mockPythonServer = {
125+
onCodeExecuted: new EventEmitter<void>().event,
126+
execute: sinon.stub().resolves({ status: true, output: 'test output' }),
127+
executeSilently: sinon.stub().resolves({ status: true, output: 'test output' }),
128+
interrupt: sinon.stub(),
129+
input: sinon.stub(),
130+
checkValidCommand: sinon.stub().resolves(true),
131+
dispose: sinon.stub(),
132+
};
133+
134+
// Track the number of times createPythonServer was called
135+
let createPythonServerCallCount = 0;
136+
sinon.stub(PythonServer, 'createPythonServer').callsFake(() => {
137+
// eslint-disable-next-line no-plusplus
138+
createPythonServerCallCount++;
139+
return mockPythonServer;
140+
});
141+
142+
const interpreter = await interpreterService.object.getActiveInterpreter();
94143

95-
setReplDirectoryStub.restore();
96-
setReplControllerSpy.restore();
144+
// Create NativeRepl directly to have more control over its state, go around private constructor.
145+
const nativeRepl = new (NativeReplModule.NativeRepl as any)();
146+
nativeRepl.interpreter = interpreter as PythonEnvironment;
147+
nativeRepl.cwd = '/helloJustMockedCwd/cwd';
148+
nativeRepl.pythonServer = mockPythonServer;
149+
nativeRepl.replController = mockNotebookController;
150+
nativeRepl.disposables = [];
151+
152+
// Make the singleton point to our instance for testing
153+
// Otherwise, it gets mixed with Native Repl from .create from test above.
154+
(NativeReplModule as any).nativeRepl = nativeRepl;
155+
156+
// Reset call count after initial setup
157+
createPythonServerCallCount = 0;
158+
159+
// Set notebookDocument to a mock document
160+
const mockReplUri = Uri.parse('untitled:Untitled-999.ipynb?jupyter-notebook');
161+
const mockNotebookDocument = ({
162+
uri: mockReplUri,
163+
toString: () => mockReplUri.toString(),
164+
} as unknown) as NotebookDocument;
165+
166+
nativeRepl.notebookDocument = mockNotebookDocument;
167+
168+
// Create a mock notebook document for closing event with same URI
169+
const closingNotebookDocument = ({
170+
uri: mockReplUri,
171+
toString: () => mockReplUri.toString(),
172+
} as unknown) as NotebookDocument;
173+
174+
notebookCloseEmitter.fire(closingNotebookDocument);
175+
await new Promise((resolve) => setTimeout(resolve, 50));
176+
177+
expect(
178+
updateWorkspaceStateValueStub.calledWith(NativeReplModule.NATIVE_REPL_URI_MEMENTO, undefined),
179+
'updateWorkspaceStateValue should be called with NATIVE_REPL_URI_MEMENTO and undefined',
180+
).to.be.true;
181+
expect(mockPythonServer.dispose.calledOnce, 'pythonServer.dispose() should be called once').to.be.true;
182+
expect(createPythonServerCallCount, 'createPythonServer should be called to create a new server').to.equal(1);
183+
expect(nativeRepl.notebookDocument, 'notebookDocument should be undefined after closing').to.be.undefined;
184+
expect(nativeRepl.newReplSession, 'newReplSession should be set to true after closing').to.be.true;
185+
expect(mockNotebookController.dispose.calledOnce, 'replController.dispose() should be called once').to.be.true;
97186
});
98187
});

0 commit comments

Comments
 (0)