fix: forward cancellation token to toolkit installer subprocesses

cspell.json

@@ -45,6 +45,7 @@
         "Dremio",
         "duckdb",
         "ename",
+        "ensurepip",
         "esmock",
         "evalue",
         "findstr",

src/kernels/deepnote/deepnoteToolkitInstaller.node.ts

@@ -6,10 +6,10 @@ import { CancellationToken, l10n, Uri, workspace } from 'vscode';
 
 import { resolvePythonExecutable } from '@deepnote/runtime-core';
 
-import { Cancellation } from '../../platform/common/cancellation';
+import { Cancellation, isCancellationError } from '../../platform/common/cancellation';
 import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants';
 import { IFileSystem } from '../../platform/common/platform/types';
-import { IProcessServiceFactory } from '../../platform/common/process/types.node';
+import { ExecutionResult, IProcessServiceFactory } from '../../platform/common/process/types.node';
 import { IExtensionContext, IOutputChannel } from '../../platform/common/types';
 import {
     DeepnoteToolkitInstallError,
@@ -116,18 +116,11 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
         // Check if venv already exists with toolkit installed
         const existingVenv = await this.getVenvInterpreterByPath(venvPath);
         if (existingVenv) {
-            const toolkitVersion = await this.isToolkitInstalled(existingVenv);
+            const toolkitVersion = await this.isToolkitInstalled(existingVenv, token);
             if (toolkitVersion != null) {
                 logger.info(`deepnote-toolkit venv already exists at ${venvPath.fsPath}`);
 
-                // Ensure kernel spec is installed (may have been deleted or never installed)
-                try {
-                    Cancellation.throwIfCanceled(token);
-                    await this.installKernelSpec(existingVenv, venvPath, token);
-                } catch (ex) {
-                    logger.warn('Failed to ensure kernel spec installed', ex);
-                    // Don't fail - continue with existing venv
-                }
+                await this.tryInstallKernelSpec(existingVenv, venvPath, token);
 
                 logger.info(`Venv ready at ${venvPath.fsPath}`);
                 return { pythonInterpreter: existingVenv, toolkitVersion };
@@ -188,13 +181,10 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
         this.outputChannel.appendLine(l10n.t('Installing packages: {0}...', packages.join(', ')));
 
         try {
-            Cancellation.throwIfCanceled(token);
-
-            const venvProcessService = await this.processServiceFactory.create(undefined);
-            const installResult = await venvProcessService.exec(
+            const installResult = await this.runPython(
                 venvInterpreter.uri.fsPath,
                 ['-m', 'pip', 'install', '--upgrade', ...packages],
-                { throwOnStdErr: false }
+                token
             );
 
             if (installResult.stdout) {
@@ -207,6 +197,13 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
             logger.info('Additional packages installed successfully');
             this.outputChannel.appendLine(l10n.t('✓ Packages installed successfully'));
         } catch (ex) {
+            if (isCancellationError(ex)) {
+                logger.info('Package installation cancelled');
+                this.outputChannel.appendLine(l10n.t('Package installation cancelled'));
+
+                throw ex;
+            }
+
             logger.error('Failed to install additional packages', ex);
             this.outputChannel.appendLine(l10n.t('✗ Failed to install packages: {0}', ex));
             throw ex;
@@ -277,20 +274,13 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
                 await workspace.fs.delete(venvPath, { recursive: true });
             }
 
-            // Create new venv
-            // Use undefined as resource to get full system environment
-            const processService = await this.processServiceFactory.create(undefined);
-            const venvResult = await processService.exec(baseInterpreter.uri.fsPath, ['-m', 'venv', venvPath.fsPath], {
-                throwOnStdErr: false
-            });
+            const venvResult = await this.runPython(baseInterpreter.uri.fsPath, ['-m', 'venv', venvPath.fsPath], token);
 
             // Log any stderr output (warnings, etc.) but don't fail on it
             if (venvResult.stderr) {
                 logger.info('venv creation stderr', venvResult.stderr);
             }
 
-            Cancellation.throwIfCanceled(token);
-
             // Verify venv was created successfully by checking for the Python interpreter
             const venvInterpreter = await this.getVenvInterpreterByPath(venvPath);
             if (!venvInterpreter) {
@@ -310,6 +300,11 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
             // Use the shared helper method to install toolkit packages
             return await this.installToolkitPackages(venvInterpreter, venvPath, token);
         } catch (ex) {
+            // Unwrapped, so the caller's isCancellationError check suppresses the error UI
+            if (isCancellationError(ex)) {
+                throw ex;
+            }
+
             // If this is already a DeepnoteKernelError, rethrow it without wrapping
             if (ex instanceof DeepnoteVenvCreationError || ex instanceof DeepnoteToolkitInstallError) {
                 throw ex;
@@ -339,16 +334,14 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
         venvPath: Uri,
         token?: CancellationToken
     ): Promise<VenvAndToolkitInstallation> {
-        // Use undefined as resource to get full system environment (including git in PATH)
-        const venvProcessService = await this.processServiceFactory.create(undefined);
-
         // Upgrade pip in the venv to the latest version
         logger.info('Upgrading pip in venv to latest version...');
         this.outputChannel.appendLine(l10n.t('Upgrading pip...'));
-        const pipUpgradeResult = await venvProcessService.exec(
+
+        const pipUpgradeResult = await this.runPython(
             venvInterpreter.uri.fsPath,
             ['-m', 'pip', 'install', '--upgrade', 'pip'],
-            { throwOnStdErr: false }
+            token
         );
 
         if (pipUpgradeResult.stdout) {
@@ -358,8 +351,6 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
             logger.info('pip upgrade stderr', pipUpgradeResult.stderr);
         }
 
-        Cancellation.throwIfCanceled(token);
-
         // Install deepnote-toolkit, ipykernel, python-lsp-server, and deepnote-cli in venv
         logger.info(
             `Installing deepnote-toolkit (${DEEPNOTE_TOOLKIT_VERSION}), ipykernel, python-lsp-server, and deepnote-cli in venv from PyPI`
@@ -368,7 +359,7 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
             l10n.t('Installing deepnote-toolkit, ipykernel, python-lsp-server, and deepnote-cli...')
         );
 
-        const installResult = await venvProcessService.exec(
+        const installResult = await this.runPython(
             venvInterpreter.uri.fsPath,
             [
                 '-m',
@@ -380,11 +371,9 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
                 'python-lsp-server[all]',
                 'deepnote-cli'
             ],
-            { throwOnStdErr: false }
+            token
         );
 
-        Cancellation.throwIfCanceled(token);
-
         if (installResult.stdout) {
             this.outputChannel.appendLine(installResult.stdout);
         }
@@ -393,18 +382,12 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
         }
 
         // Verify installation
-        const installedToolkitVersion = await this.isToolkitInstalled(venvInterpreter);
+        const installedToolkitVersion = await this.isToolkitInstalled(venvInterpreter, token);
         if (installedToolkitVersion != null) {
             logger.info('deepnote-toolkit installed successfully in venv');
 
             // Install kernel spec so the kernel uses this venv's Python
-            try {
-                Cancellation.throwIfCanceled(token);
-                await this.installKernelSpec(venvInterpreter, venvPath, token);
-            } catch (ex) {
-                logger.warn('Failed to install kernel spec', ex);
-                // Don't fail the entire installation if kernel spec creation fails
-            }
+            await this.tryInstallKernelSpec(venvInterpreter, venvPath, token);
 
             this.outputChannel.appendLine(l10n.t('✓ Deepnote toolkit ready'));
             return { pythonInterpreter: venvInterpreter, toolkitVersion: installedToolkitVersion };
@@ -422,23 +405,79 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
         }
     }
 
-    private async isToolkitInstalled(interpreter: PythonEnvironment): Promise<string | undefined> {
+    private async isToolkitInstalled(
+        interpreter: PythonEnvironment,
+        token: CancellationToken | undefined
+    ): Promise<string | undefined> {
         try {
-            // Use undefined as resource to get full system environment
-            const processService = await this.processServiceFactory.create(undefined);
-            const result = await processService.exec(interpreter.uri.fsPath, [
-                '-c',
-                'import deepnote_toolkit; print(deepnote_toolkit.__version__)'
-            ]);
+            const result = await this.runPython(
+                interpreter.uri.fsPath,
+                ['-c', 'import deepnote_toolkit; print(deepnote_toolkit.__version__)'],
+                token
+            );
+
             logger.info(`isToolkitInstalled result: ${result.stdout}`);
             const version = result.stdout.trim();
+
             return version.length > 0 ? version : undefined;
         } catch (ex) {
+            if (isCancellationError(ex)) {
+                throw ex;
+            }
+
             logger.debug('deepnote-toolkit not found', ex);
+
             return undefined;
         }
     }
 
+    /**
+     * `ProcessService.exec` kills the subprocess when the token fires but still *resolves*, with the
+     * output captured up to that point. Hence the trailing check: without it a killed `pip install`
+     * reads as a successful one and a killed version probe reads as "toolkit not installed".
+     */
+    private async runPython(
+        pythonPath: string,
+        args: string[],
+        token: CancellationToken | undefined
+    ): Promise<ExecutionResult<string>> {
+        Cancellation.throwIfCanceled(token);
+
+        // Resource `undefined` gives the full system environment, which pip needs (e.g. git on PATH).
+        const processService = await this.processServiceFactory.create(undefined);
+        // Own process group: `python -m venv` runs ensurepip in a subprocess that would otherwise
+        // outlive cancellation and keep writing into the venv the retry is busy deleting.
+        const result = await processService.exec(pythonPath, args, {
+            throwOnStdErr: false,
+            token,
+            detached: true
+        });
+
+        Cancellation.throwIfCanceled(token);
+
+        return result;
+    }
+
+    /**
+     * Failure is tolerated because a venv without a kernel spec still runs and the next
+     * `ensureVenvAndToolkit` retries. Cancellation is not a failure, so it unwinds the caller.
+     */
+    private async tryInstallKernelSpec(
+        venvInterpreter: PythonEnvironment,
+        venvPath: Uri,
+        token: CancellationToken | undefined
+    ): Promise<void> {
+        try {
+            await this.installKernelSpec(venvInterpreter, venvPath, token);
+        } catch (ex) {
+            if (isCancellationError(ex)) {
+                throw ex;
+            }
+
+            logger.warn('Failed to install kernel spec', ex);
+        }
+    }
+
     /**
      * Generate a kernel spec name from a venv path.
      * This is used for both file-based and environment-based venvs.
@@ -480,22 +519,16 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
         const kernelSpecName = this.getKernelSpecName(venvPath);
         const kernelSpecPath = Uri.joinPath(venvPath, 'share', 'jupyter', 'kernels', kernelSpecName);
 
-        // Check if kernel spec already exists
-        if (await this.fs.exists(kernelSpecPath)) {
+        // Keyed on kernel.json, not the directory: a cancelled ipykernel install leaves the
+        // directory behind, and that must not short-circuit the reinstall.
+        if (await this.fs.exists(Uri.joinPath(kernelSpecPath, 'kernel.json'))) {
             logger.info(`Kernel spec already exists at ${kernelSpecPath.fsPath}`);
             return;
         }
 
-        Cancellation.throwIfCanceled(token);
-
         logger.info(`Installing kernel spec '${kernelSpecName}' for venv at ${venvPath.fsPath}...`);
-        const kernelDisplayName = this.getKernelDisplayName(venvPath);
-
-        const venvProcessService = await this.processServiceFactory.create(undefined);
-
-        Cancellation.throwIfCanceled(token);
 
-        await venvProcessService.exec(
+        await this.runPython(
             venvInterpreter.uri.fsPath,
             [
                 '-m',
@@ -506,9 +539,9 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller {
                 '--name',
                 kernelSpecName,
                 '--display-name',
-                kernelDisplayName
+                this.getKernelDisplayName(venvPath)
             ],
-            { throwOnStdErr: false }
+            token
         );
 
         logger.info(`Kernel spec installed successfully to ${kernelSpecPath.fsPath}`);