Fix memory management and add initial state detection to InstallPackagesTool

Co-authored-by: DonJayamanne <[email protected]>

Author: Copilot

src/standalone/chat/installPackageTool.node.ts

@@ -14,9 +14,10 @@ import { isPythonKernelConnection } from '../../kernels/helpers';
 import { RestartKernelTool } from './restartKernelTool.node';
 import { BaseTool, IBaseToolParams } from './helper';
 import { WrappedError } from '../../platform/errors/types';
+import { IDisposable } from '../../platform/common/types';
 
 
-export class InstallPackagesTool extends BaseTool<IInstallPackageParams> {
+export class InstallPackagesTool extends BaseTool<IInstallPackageParams> implements IDisposable {
     public static toolName = 'notebook_install_packages';
     private readonly executedNotebooks = new WeakSet<vscode.NotebookDocument>();
     private readonly disposables: vscode.Disposable[] = [];
@@ -26,6 +27,10 @@ export class InstallPackagesTool extends BaseTool<IInstallPackageParams> {
         private readonly installationManager: IInstallationChannelManager
     ) {
         super(InstallPackagesTool.toolName);
+        
+        // Check existing notebooks for already-executed cells
+        this.initializeExecutedNotebooks();
+        
         // Track cell execution for all notebooks
         this.disposables.push(
             vscode.workspace.onDidChangeNotebookDocument((e) => {
@@ -42,8 +47,29 @@ export class InstallPackagesTool extends BaseTool<IInstallPackageParams> {
         );
     }
 
+    dispose(): void {
+        this.disposables.forEach(d => d.dispose());
+        this.disposables.length = 0;
+    }
+
+    private initializeExecutedNotebooks(): void {
+        // Check all currently open notebooks for executed cells
+        for (const notebook of vscode.workspace.notebookDocuments) {
+            if (this.hasAnyExecutedCells(notebook)) {
+                this.executedNotebooks.add(notebook);
+            }
+        }
+    }
+
+    private hasAnyExecutedCells(notebook: vscode.NotebookDocument): boolean {
+        return notebook.getCells().some(cell => 
+            cell.kind === vscode.NotebookCellKind.Code &&
+            typeof cell.executionSummary?.executionOrder === 'number'
+        );
+    }
+
     private hasExecutedCells(notebook: vscode.NotebookDocument): boolean {
-        return this.executedNotebooks.has(notebook);
+        return this.executedNotebooks.has(notebook) || this.hasAnyExecutedCells(notebook);
     }
 
     async invokeImpl(

src/standalone/chat/installPackageTool.node.unit.test.ts

@@ -0,0 +1,188 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { assert } from 'chai';
+import * as sinon from 'sinon';
+import * as vscode from 'vscode';
+import { anything, instance, mock, when, verify, reset } from 'ts-mockito';
+import { InstallPackagesTool } from './installPackageTool.node';
+import { IKernelProvider } from '../../kernels/types';
+import { IControllerRegistration } from '../../notebooks/controllers/types';
+import { IInstallationChannelManager } from '../../platform/interpreter/installer/types';
+
+suite('InstallPackagesTool Unit Tests', () => {
+    let installPackagesTool: InstallPackagesTool;
+    let kernelProvider: IKernelProvider;
+    let controllerRegistration: IControllerRegistration;
+    let installationManager: IInstallationChannelManager;
+    
+    // Mock VS Code objects
+    let mockNotebook: vscode.NotebookDocument;
+    let mockWorkspace: typeof vscode.workspace;
+    let mockCells: vscode.NotebookCell[];
+
+    setup(() => {
+        kernelProvider = mock<IKernelProvider>();
+        controllerRegistration = mock<IControllerRegistration>();
+        installationManager = mock<IInstallationChannelManager>();
+        
+        // Create mock notebook and cells
+        mockNotebook = {
+            uri: vscode.Uri.file('/test/notebook.ipynb'),
+            getCells: sinon.stub()
+        } as any;
+        
+        mockCells = [];
+        (mockNotebook.getCells as sinon.SinonStub).returns(mockCells);
+        
+        // Mock vscode.workspace
+        mockWorkspace = {
+            notebookDocuments: [mockNotebook],
+            onDidChangeNotebookDocument: sinon.stub().returns({ dispose: sinon.stub() })
+        } as any;
+        
+        // Replace the actual workspace with our mock
+        sinon.stub(vscode, 'workspace').value(mockWorkspace);
+    });
+
+    teardown(() => {
+        if (installPackagesTool) {
+            installPackagesTool.dispose();
+        }
+        sinon.restore();
+        reset(kernelProvider);
+        reset(controllerRegistration);
+        reset(installationManager);
+    });
+
+    test('Should initialize without executed cells', () => {
+        // Create tool with no executed cells
+        installPackagesTool = new InstallPackagesTool(
+            instance(kernelProvider),
+            instance(controllerRegistration),
+            instance(installationManager)
+        );
+        
+        // hasExecutedCells should return false for notebook with no executed cells
+        const hasExecuted = (installPackagesTool as any).hasExecutedCells(mockNotebook);
+        assert.isFalse(hasExecuted, 'Should return false for notebook with no executed cells');
+    });
+
+    test('Should detect already executed cells during initialization', () => {
+        // Setup notebook with executed cells
+        const mockExecutedCell = {
+            kind: vscode.NotebookCellKind.Code,
+            executionSummary: { executionOrder: 1 }
+        } as vscode.NotebookCell;
+        
+        mockCells.push(mockExecutedCell);
+        
+        // Create tool - should detect the already executed cell
+        installPackagesTool = new InstallPackagesTool(
+            instance(kernelProvider),
+            instance(controllerRegistration),
+            instance(installationManager)
+        );
+        
+        // hasExecutedCells should return true for notebook with executed cells
+        const hasExecuted = (installPackagesTool as any).hasExecutedCells(mockNotebook);
+        assert.isTrue(hasExecuted, 'Should return true for notebook with executed cells');
+    });
+
+    test('Should track cell execution via event listener', () => {
+        let onDidChangeCallback: (e: vscode.NotebookDocumentChangeEvent) => void;
+        
+        // Capture the callback passed to onDidChangeNotebookDocument
+        (mockWorkspace.onDidChangeNotebookDocument as sinon.SinonStub).callsFake((callback) => {
+            onDidChangeCallback = callback;
+            return { dispose: sinon.stub() };
+        });
+        
+        installPackagesTool = new InstallPackagesTool(
+            instance(kernelProvider),
+            instance(controllerRegistration),
+            instance(installationManager)
+        );
+        
+        // Initially no executed cells
+        assert.isFalse((installPackagesTool as any).hasExecutedCells(mockNotebook));
+        
+        // Simulate cell execution event
+        const mockExecutedCell = {
+            kind: vscode.NotebookCellKind.Code,
+            executionSummary: { executionOrder: 1 }
+        } as vscode.NotebookCell;
+        
+        const changeEvent: vscode.NotebookDocumentChangeEvent = {
+            notebook: mockNotebook,
+            cellChanges: [{
+                cell: mockExecutedCell,
+                document: undefined,
+                executionSummary: undefined,
+                metadata: undefined,
+                outputs: undefined
+            }],
+            contentChanges: []
+        };
+        
+        // Trigger the event
+        onDidChangeCallback(changeEvent);
+        
+        // Now should detect executed cells
+        assert.isTrue((installPackagesTool as any).hasExecutedCells(mockNotebook));
+    });
+
+    test('Should ignore non-code cells in execution tracking', () => {
+        let onDidChangeCallback: (e: vscode.NotebookDocumentChangeEvent) => void;
+        
+        (mockWorkspace.onDidChangeNotebookDocument as sinon.SinonStub).callsFake((callback) => {
+            onDidChangeCallback = callback;
+            return { dispose: sinon.stub() };
+        });
+        
+        installPackagesTool = new InstallPackagesTool(
+            instance(kernelProvider),
+            instance(controllerRegistration),
+            instance(installationManager)
+        );
+        
+        // Simulate markdown cell change
+        const mockMarkdownCell = {
+            kind: vscode.NotebookCellKind.Markup,
+            executionSummary: { executionOrder: 1 }
+        } as vscode.NotebookCell;
+        
+        const changeEvent: vscode.NotebookDocumentChangeEvent = {
+            notebook: mockNotebook,
+            cellChanges: [{
+                cell: mockMarkdownCell,
+                document: undefined,
+                executionSummary: undefined,
+                metadata: undefined,
+                outputs: undefined
+            }],
+            contentChanges: []
+        };
+        
+        onDidChangeCallback(changeEvent);
+        
+        // Should still return false since only non-code cells were changed
+        assert.isFalse((installPackagesTool as any).hasExecutedCells(mockNotebook));
+    });
+
+    test('Should dispose of event listeners properly', () => {
+        const disposeSpy = sinon.spy();
+        (mockWorkspace.onDidChangeNotebookDocument as sinon.SinonStub).returns({ dispose: disposeSpy });
+        
+        installPackagesTool = new InstallPackagesTool(
+            instance(kernelProvider),
+            instance(controllerRegistration),
+            instance(installationManager)
+        );
+        
+        installPackagesTool.dispose();
+        
+        // Verify dispose was called
+        assert.isTrue(disposeSpy.calledOnce, 'Event listener dispose should be called');
+    });
+});