Add API method to shutdown Jupyter kernels programmatically (Proposed API) (#16775)

* AI ready repo

* Updates

* Updates

* Initial plan

* Implement shutdown kernel command infrastructure

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

* Add shutdown kernel command to VS Code UI and add unit test

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

* Replace string URI comparison with isEqual function in shutdownKernel

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

* Address feedback: Remove kernelConnector changes, call kernel.shutdown() and dispose() directly, add integration test

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

* Fix integration test to properly call both shutdown() and dispose() methods

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

* Remove 'shutdown' from KernelAction type as requested

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

* Address feedback: Remove kernelConnector test, separate try-catch for dispose, support multiple argument types, move test to new file

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

* Updates

* Fix formatting

* Convert shutdown command to API-based approach

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

* Updated instructions

* Fix test assertion and address linting/formatting issues

- Replace assert.isNotEmpty() with assert.isDefined() in shutdown test
- Fix incorrect assertion that was causing test failure when apiKernel is undefined
- Run format-fix to ensure consistent code formatting

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

* Move shutdown test to api.vscode.test.ts and delete separate test file

- Moved "Can shutdown a kernel correctly via API" test from shutdown.vscode.test.ts to api.vscode.test.ts
- Added necessary imports for kernels API, test helpers, and execution functions
- Updated test to create separate notebook to avoid conflicts with existing tests
- Deleted standalone shutdown.vscode.test.ts file as requested
- Fixed TypeScript compilation by removing unused IKernel import

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

* Create api.proposed.shutdown.d.ts with authorization check

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

* Fix formatting

* Fix API test by setting disableUI: false for user-started kernel

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

* Fix tests

* Updates to allow list

---------

Co-authored-by: Don Jayamanne <[email protected]>
Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: DonJayamanne <[email protected]>

Author: Copilot

.github/chatmodes/bugfix.chatmode.md

@@ -0,0 +1,22 @@
+---
+description: Plan to fix a bug
+tools: ['codebase', 'editFiles', 'fetch', 'findTestFiles', 'problems', 'runTasks', 'runTests', 'search', 'terminalLastCommand', 'testFailure', 'usages', 'vscodeAPI', 'github', 'get_issue', 'get_issue_comments', 'get_me', 'copilotCodingAgent']
+---
+# Bug fixing mode instructions
+You are an expert (TypeScript and Python) software engineer tasked with fixing a bug in the codebase.
+Your goal is to prepare a detailed plan to fix the bug, for this you first need to:
+* Understand the context of the bug by reading the issue description and comments.
+* Understand the codebase by reading the relevant instruction files.
+* Identify the root cause of the bug, and explain this to the user.
+
+Based on your above understanding generate a plan to fix the bug.
+Ensure the plan consists of a Markdown document that has the following sections:
+
+* Overview: A brief description of the bug.
+* Root Cause: A detailed explanation of the root cause of the bug, including any relevant code snippets or references to the codebase.
+* Requirements: A list of requirements to resolve the bug.
+* Implementation Steps: A detailed list of steps to implement the bug fix.
+
+Finally prompt the user if they would like to proceed with the implementation of the bug fix.
+Remember, do not make any code edits, just generate a plan.
+When implementing the fix, ensure to add unit tests to verify the fix.

.github/copilot-instructions.md

@@ -87,9 +87,9 @@ src/                                # Source code for the extension
 └── test/                           # Integration, unit and end-to-end tests
 ```
 
-### Component Instructions
+### Component-Specific Instructions
 
-**IMPORTANT**: Before modifying code in directories below, you **MUST** read the corresponding instruction file(s).
+**IMPORTANT**: Before modifying any code in the directories listed below, you **MUST** first read the corresponding instruction file to understand the specific conventions, patterns, and architectural requirements for that component.
 
 Each component has detailed guidelines that cover:
 - Architecture patterns and design principles
@@ -99,17 +99,18 @@ Each component has detailed guidelines that cover:
 - Error handling approaches
 - Component-specific best practices
 
-
-| Directory | Instruction File | Purpose |
-|-----------|------------------|---------|
-| `src/platform/**` | `.github/instructions/platform.instructions.md` | Cross-platform abstractions |
-| `src/kernels/**` | `.github/instructions/kernel.instructions.md` | Kernel management |
-| `src/kernels/jupyter/**` | `.github/instructions/kernel-jupyter.instructions.md` | Jupyter protocol |
-| `src/notebooks/**` | `.github/instructions/notebooks.instructions.md` | Notebook controllers |
-| `src/interactive-window/**` | `.github/instructions/interactiveWindow.instructions.md` | REPL functionality |
-| `src/standalone/**` | `.github/instructions/standalone.instructions.md` | Standalone features |
-| `src/notebooks/controllers/ipywidgets/**` | `.github/instructions/ipywidgets.instructions.md` | IPython widgets (interactive Notebook outputs) |
-| `src/webviews/extension-side/ipywidgets/**` | `.github/instructions/ipywidgets.instructions.md` | IPython Widget (interactive Notebook outputs) communication |
-| `src/webviews/webview-side/ipywidgets/**` | `.github/instructions/ipywidgets.instructions.md` | IPython Widget (interactive Notebook outputs) rendering |
-
-**For AI/Copilot**: Always read the relevant instruction file(s) before modifying code in these directories to ensure adherence to component-specific patterns.
+#### Required Reading by Directory
+
+| Directory/Component | Instruction File | When to Read |
+|-------------------|------------------|--------------|
+| `src/platform/**` | `.github/instructions/platform.instructions.md` | Before working with cross-platform abstractions, utilities, or core services |
+| `src/kernels/**` | `.github/instructions/kernel.instructions.md` | Before modifying kernel management, execution, or communication logic |
+| `src/kernels/jupyter/**` | `.github/instructions/kernel-jupyter.instructions.md` | Before working with Jupyter protocol implementation or Jupyter-specific features |
+| `src/notebooks/**` | `.github/instructions/notebooks.instructions.md` | Before modifying notebook controllers, editing, or management features |
+| `src/interactive-window/**` | `.github/instructions/interactiveWindow.instructions.md` | Before working with Python Interactive window (REPL) functionality |
+| `src/standalone/**` | `.github/instructions/standalone.instructions.md` | Before modifying standalone mode or isolated execution features |
+| `src/notebooks/controllers/ipywidgets/**` | `.github/instructions/ipywidgets.instructions.md` | Before working with IPython widget support |
+| `src/webviews/extension-side/ipywidgets/**` | `.github/instructions/ipywidgets.instructions.md` | Before modifying extension-side widget communication |
+| `src/webviews/webview-side/ipywidgets/**` | `.github/instructions/ipywidgets.instructions.md` | Before working with frontend widget rendering |
+
+**For AI/Copilot**: Always use the `read_file` tool to read the relevant instruction file(s) before analyzing, modifying, or creating code in these directories. This ensures adherence to component-specific patterns and prevents architectural violations.

.github/instructions/typescript.instructions.md

@@ -35,3 +35,9 @@ mockInstance.then = undefined; // Ensure 'then' is undefined to prevent hanging
 - Use `npm run lint` to check for linter issues
 - Use `npm run format` to check code style
 - Use `npm run format-fix` to auto-fix formatting issues
+
+## PreCommit
+Always run the following scripts before committing changes:
+- `npm run format-fix` to ensure proper code formatting
+- `npm run lint` to check for linter issues and attempt to fix before committing.
+- `npm run test:unittests` to check for tests failures and attempt to fix before committing.

src/api.proposed.shutdown.d.ts

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+// Importing a type to ensure this is treated as a module
+import type {} from 'vscode';
+
+declare module './api' {
+    /**
+     * Represents a Jupyter Kernel.
+     */
+    export interface Kernel {
+        /**
+         * Shuts down the kernel and all its associated resources.
+         * This operation will terminate the kernel process and cannot be undone.
+         * After shutdown, the kernel becomes unusable and should be disposed of.
+         */
+        shutdown(): Promise<void>;
+    }
+}

src/kernels/kernel.ts

@@ -449,6 +449,41 @@ abstract class BaseKernel implements IBaseKernel {
             Promise.all(Array.from(this.hooks.get('restartCompleted') || new Set<Hook>()).map((h) => h())).catch(noop);
         }
     }
+
+    public async shutdown(): Promise<void> {
+        try {
+            logger.info(`Shutdown requested ${getDisplayPath(this.uri)}`);
+
+            // Cancel any pending starts
+            this.startCancellation.cancel();
+            this.startCancellation.dispose();
+
+            // Get the current session if it exists
+            const session = this._jupyterSessionPromise
+                ? await this._jupyterSessionPromise.catch(() => undefined)
+                : undefined;
+
+            if (session) {
+                logger.info('Shutting down kernel session');
+                try {
+                    await session.shutdown();
+                } catch (ex) {
+                    logger.warn(`Failed to shutdown kernel session ${getDisplayPath(this.uri)}`, ex);
+                    // Continue with disposal even if shutdown fails
+                }
+            }
+
+            // Clean up the session references
+            this._session = undefined;
+            this._jupyterSessionPromise = undefined;
+            this._postInitializedOnStartPromise = undefined;
+
+            logger.info(`Shutdown completed ${getDisplayPath(this.uri)}`);
+        } catch (ex) {
+            logger.error(`Failed to shutdown kernel ${getDisplayPath(this.uri)}`, ex);
+            throw ex;
+        }
+    }
     protected async startJupyterSession(options: IDisplayOptions = new DisplayOptions(false)): Promise<IKernelSession> {
         this._startedAtLeastOnce = true;
         if (!options.disableUI) {

src/kernels/types.ts

@@ -407,6 +407,7 @@ export interface IBaseKernel extends IAsyncDisposable {
     start(options?: IStartOptions): Promise<IKernelSession>;
     interrupt(): Promise<void>;
     restart(): Promise<void>;
+    shutdown(): Promise<void>;
     addHook(
         event: 'didStart',
         hook: (session: IKernelSession | undefined, token: CancellationToken) => Promise<void>,

src/standalone/api/kernels/kernel.ts

@@ -131,7 +131,17 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel {
 
                 return that.onDidReceiveDisplayUpdate.bind(this);
             },
-            executeCode: (code: string, token: CancellationToken) => this.executeCode(code, token)
+            executeCode: (code: string, token: CancellationToken) => this.executeCode(code, token),
+            get shutdown() {
+                if (
+                    ![JVSC_EXTENSION_ID, POWER_TOYS_EXTENSION_ID].includes(extensionId) &&
+                    !PROPOSED_API_ALLOWED_PUBLISHERS.includes(extensionId.split('.')[0])
+                ) {
+                    throw new Error(`Proposed API is not supported for extension ${extensionId}`);
+                }
+
+                return () => that.shutdown();
+            }
         });
     }
     static createApiKernel(
@@ -178,6 +188,26 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel {
         }
     }
 
+    async shutdown(): Promise<void> {
+        await this.checkAccess();
+        logger.debug(`Shutting down kernel ${this.kernel.id} via API for extension ${this.extensionId}`);
+        try {
+            // First shutdown the kernel
+            await this.kernel.shutdown();
+        } catch (ex) {
+            logger.error(`Failed to shutdown kernel ${this.kernel.id} for extension ${this.extensionId}`, ex);
+            throw ex;
+        }
+
+        try {
+            // Then dispose the kernel
+            await this.kernel.dispose();
+        } catch (ex) {
+            logger.error(`Failed to dispose kernel ${this.kernel.id} for extension ${this.extensionId}`, ex);
+            throw ex;
+        }
+    }
+
     async *executeCodeInternal(
         code: string,
         handlers: Record<string, (...data: any[]) => Promise<any>> = {},

src/standalone/api/kernels/kernel.unit.test.ts

@@ -22,6 +22,7 @@ import { dispose } from '../../../platform/common/utils/lifecycle';
 import { ServiceContainer } from '../../../platform/ioc/container';
 import {
     IKernelProvider,
+    IKernelSession,
     INotebookKernelExecution,
     KernelConnectionMetadata,
     type IKernel
@@ -30,6 +31,9 @@ import { createMockedNotebookDocument } from '../../../test/datascience/editor-i
 import { IControllerRegistration, IVSCodeNotebookController } from '../../../notebooks/controllers/types';
 import { createKernelApiForExtension } from './kernel';
 import { noop } from '../../../test/core';
+import { JVSC_EXTENSION_ID_FOR_TESTS } from '../../../test/constants';
+import { IKernelConnection } from '@jupyterlab/services/lib/kernel/kernel';
+import { NotebookCellOutput } from 'vscode';
 
 suite('Kernel Api', () => {
     let disposables: IDisposable[] = [];
@@ -80,9 +84,13 @@ suite('Kernel Api', () => {
         when(kernelConnection.kind).thenReturn('connectToLiveRemoteKernel');
         kernel = mock<IKernel>();
         when(kernel.kernelConnectionMetadata).thenReturn(instance(kernelConnection));
+        when(kernel.disposed).thenReturn(false);
         when(kernel.startedAtLeastOnce).thenReturn(true);
         notebook = createMockedNotebookDocument([]);
         when(kernel.notebook).thenReturn(notebook);
+        const kernelSession = mock<IKernelSession>();
+        when(kernel.session).thenReturn(instance(kernelSession));
+        when(kernelSession.kernel).thenReturn(instance(mock<IKernelConnection>()));
 
         const controllerRegistration = mock<IControllerRegistration>();
         when(serviceContainer.get<IControllerRegistration>(IControllerRegistration)).thenReturn(
@@ -92,6 +100,10 @@ suite('Kernel Api', () => {
         const controller = mock<NotebookController>();
         const execution = mock<INotebookKernelExecution>();
         const kernelProvider = mock<IKernelProvider>();
+        when(execution.executeCode(anything(), anything(), anything(), anything())).thenCall(async function* () {
+            // Yield a dummy NotebookCellOutput to match the expected type
+            yield new NotebookCellOutput([]);
+        });
         when(kernelProvider.getKernelExecution(instance(kernel))).thenReturn(instance(execution));
         when(serviceContainer.get<IKernelProvider>(IKernelProvider)).thenReturn(instance(kernelProvider));
         when(vscController.controller).thenReturn(instance(controller));
@@ -110,4 +122,17 @@ suite('Kernel Api', () => {
             assert.equal(ex.name, 'vscode.jupyter.apiAccessRevoked');
         }
     });
+    test('Verify Kernel Shutdown', async () => {
+        when(kernel.status).thenReturn('idle');
+        when(kernel.shutdown()).thenResolve();
+        when(kernel.dispose()).thenCall(() => when(kernel.status).thenReturn('dead'));
+
+        const { api } = createKernelApiForExtension(JVSC_EXTENSION_ID_FOR_TESTS, instance(kernel));
+        for await (const _ of api.executeCode('bogus', token)) {
+            //
+        }
+        assert.equal(api.status, 'idle');
+        await api.shutdown();
+        assert.equal(api.status, 'dead');
+    });
 });