diff --git a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts index 2c7ebdadb..89ead5f3e 100644 --- a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts +++ b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts @@ -6,7 +6,7 @@ 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'; @@ -116,7 +116,7 @@ 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}`); @@ -125,6 +125,9 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { Cancellation.throwIfCanceled(token); await this.installKernelSpec(existingVenv, venvPath, token); } catch (ex) { + if (isCancellationError(ex as Error)) { + throw ex; + } logger.warn('Failed to ensure kernel spec installed', ex); // Don't fail - continue with existing venv } @@ -194,9 +197,13 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { const installResult = await venvProcessService.exec( venvInterpreter.uri.fsPath, ['-m', 'pip', 'install', '--upgrade', ...packages], - { throwOnStdErr: false } + { throwOnStdErr: false, token } ); + // exec resolves with partial output when the token kills pip, + // so re-check before reporting success + Cancellation.throwIfCanceled(token); + if (installResult.stdout) { this.outputChannel.appendLine(installResult.stdout); } @@ -207,6 +214,11 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { logger.info('Additional packages installed successfully'); this.outputChannel.appendLine(l10n.t('✓ Packages installed successfully')); } catch (ex) { + if (isCancellationError(ex as Error)) { + 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; @@ -281,7 +293,8 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { // 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 + throwOnStdErr: false, + token }); // Log any stderr output (warnings, etc.) but don't fail on it @@ -310,6 +323,12 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { // Use the shared helper method to install toolkit packages return await this.installToolkitPackages(venvInterpreter, venvPath, token); } catch (ex) { + // Rethrow cancellation unwrapped so upstream isCancellationError checks + // can suppress the error UI instead of reporting an install failure + if (isCancellationError(ex as Error)) { + throw ex; + } + // If this is already a DeepnoteKernelError, rethrow it without wrapping if (ex instanceof DeepnoteVenvCreationError || ex instanceof DeepnoteToolkitInstallError) { throw ex; @@ -348,7 +367,7 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { const pipUpgradeResult = await venvProcessService.exec( venvInterpreter.uri.fsPath, ['-m', 'pip', 'install', '--upgrade', 'pip'], - { throwOnStdErr: false } + { throwOnStdErr: false, token } ); if (pipUpgradeResult.stdout) { @@ -380,7 +399,7 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { 'python-lsp-server[all]', 'deepnote-cli' ], - { throwOnStdErr: false } + { throwOnStdErr: false, token } ); Cancellation.throwIfCanceled(token); @@ -393,7 +412,7 @@ 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'); @@ -402,6 +421,9 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { Cancellation.throwIfCanceled(token); await this.installKernelSpec(venvInterpreter, venvPath, token); } catch (ex) { + if (isCancellationError(ex as Error)) { + throw ex; + } logger.warn('Failed to install kernel spec', ex); // Don't fail the entire installation if kernel spec creation fails } @@ -422,18 +444,28 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { } } - private async isToolkitInstalled(interpreter: PythonEnvironment): Promise { + private async isToolkitInstalled( + interpreter: PythonEnvironment, + token: CancellationToken | undefined + ): Promise { 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 processService.exec( + interpreter.uri.fsPath, + ['-c', 'import deepnote_toolkit; print(deepnote_toolkit.__version__)'], + { token } + ); + // exec resolves with partial output when the token kills the process, + // so a cancelled probe must not be reported as "toolkit missing" + Cancellation.throwIfCanceled(token); logger.info(`isToolkitInstalled result: ${result.stdout}`); const version = result.stdout.trim(); return version.length > 0 ? version : undefined; } catch (ex) { + if (isCancellationError(ex as Error)) { + throw ex; + } logger.debug('deepnote-toolkit not found', ex); return undefined; } @@ -480,8 +512,10 @@ 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)) { + // Check if kernel spec already exists. Check for kernel.json rather than the + // directory: a cancelled ipykernel install can leave a partially written + // directory, which 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; } @@ -508,9 +542,13 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { '--display-name', kernelDisplayName ], - { throwOnStdErr: false } + { throwOnStdErr: false, token } ); + // exec resolves even when the token killed the process mid-write, + // so re-check before declaring the kernel spec installed + Cancellation.throwIfCanceled(token); + logger.info(`Kernel spec installed successfully to ${kernelSpecPath.fsPath}`); } diff --git a/src/kernels/deepnote/deepnoteToolkitInstaller.unit.test.ts b/src/kernels/deepnote/deepnoteToolkitInstaller.unit.test.ts new file mode 100644 index 000000000..a11a8ccd2 --- /dev/null +++ b/src/kernels/deepnote/deepnoteToolkitInstaller.unit.test.ts @@ -0,0 +1,125 @@ +import { assert } from 'chai'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { CancellationTokenSource, Uri } from 'vscode'; + +import { DeepnoteToolkitInstaller } from './deepnoteToolkitInstaller.node'; +import { IFileSystem } from '../../platform/common/platform/types'; +import { IProcessService, IProcessServiceFactory } from '../../platform/common/process/types.node'; +import { IExtensionContext, IOutputChannel } from '../../platform/common/types'; + +/** + * Regression tests for SAL-105: "Hanging kernel can't be cancelled". + * + * Every processService.exec(...) in the toolkit installer must forward the + * CancellationToken it was given. The token is what wires VS Code's Stop / + * Cancel button to ProcessService.kill(pid) (see proc.node.ts), so omitting it + * makes long-running pip installs uninterruptible. + */ +suite('DeepnoteToolkitInstaller - cancellation token propagation (SAL-105)', () => { + let installer: DeepnoteToolkitInstaller; + let mockProcessService: IProcessService; + let mockProcessServiceFactory: IProcessServiceFactory; + let mockOutputChannel: IOutputChannel; + let mockContext: IExtensionContext; + let mockFs: IFileSystem; + + const venvPath = Uri.file('/fake/venv'); + const fakePython = Uri.file('/fake/venv/bin/python'); + + setup(() => { + mockProcessService = mock(); + mockProcessServiceFactory = mock(); + mockOutputChannel = mock(); + mockContext = mock(); + mockFs = mock(); + + const processService = instance(mockProcessService); + // Prevent the ts-mockito instance from being treated as a thenable when + // resolved through a promise (see kernelProcess.node.unit.test.ts). + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (processService as any).then = undefined; + when(mockProcessServiceFactory.create(anything())).thenResolve(processService); + when(mockProcessService.exec(anything(), anything(), anything())).thenResolve({ stdout: '', stderr: '' }); + + installer = new DeepnoteToolkitInstaller( + instance(mockProcessServiceFactory), + instance(mockOutputChannel), + instance(mockContext), + instance(mockFs) + ); + + // Seed the interpreter cache so getVenvInterpreterByPath() resolves + // without touching the real filesystem / resolvePythonExecutable. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (installer as any).venvPythonPaths.set(venvPath.fsPath, fakePython); + }); + + test('installAdditionalPackages forwards the cancellation token to processService.exec', async () => { + const cts = new CancellationTokenSource(); + + try { + await installer.installAdditionalPackages(venvPath, ['some-package'], cts.token); + + verify(mockProcessService.exec(anything(), anything(), anything())).once(); + const [file, args, options] = capture(mockProcessService.exec).first(); + assert.deepStrictEqual( + { file, args, options }, + { + file: fakePython.fsPath, + args: ['-m', 'pip', 'install', '--upgrade', 'some-package'], + options: { throwOnStdErr: false, token: cts.token } + }, + 'the cancellation token must be forwarded to exec so Stop can kill the process' + ); + } finally { + cts.dispose(); + } + }); + + test('installAdditionalPackages does not call exec when no packages are requested', async () => { + const cts = new CancellationTokenSource(); + + try { + await installer.installAdditionalPackages(venvPath, [], cts.token); + + verify(mockProcessService.exec(anything(), anything(), anything())).never(); + } finally { + cts.dispose(); + } + }); + + test('ensureVenvAndToolkit forwards the cancellation token to the toolkit version probe', async () => { + const cts = new CancellationTokenSource(); + + try { + when(mockProcessService.exec(anything(), anything(), anything())).thenResolve({ + stdout: '1.2.3\n', + stderr: '' + }); + // Kernel spec already installed, so the fast path runs the probe exec only + when(mockFs.exists(anything())).thenResolve(true); + + const result = await installer.ensureVenvAndToolkit( + { uri: fakePython, id: fakePython.fsPath }, + venvPath, + false, + cts.token + ); + + assert.strictEqual(result.toolkitVersion, '1.2.3'); + verify(mockProcessService.exec(anything(), anything(), anything())).once(); + const [file, args, options] = capture(mockProcessService.exec).first(); + assert.deepStrictEqual( + { file, args, options }, + { + file: fakePython.fsPath, + args: ['-c', 'import deepnote_toolkit; print(deepnote_toolkit.__version__)'], + options: { token: cts.token } + }, + 'the cancellation token must be forwarded to the isToolkitInstalled probe' + ); + } finally { + cts.dispose(); + } + }); +});