diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index fb469b694f..08fd6ebb67 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -8,7 +8,7 @@ import * as fs from 'fs-extra'; import { inject, injectable, named } from 'inversify'; import * as os from 'os'; -import { CancellationToken, l10n, Uri } from 'vscode'; +import { CancellationToken, CancellationTokenSource, l10n, Uri } from 'vscode'; import { startServer, stopServer } from '@deepnote/runtime-core'; @@ -20,13 +20,15 @@ import { IAsyncDisposableRegistry, IDisposable, IOutputChannel } from '../../pla import { sleep } from '../../platform/common/utils/async'; import { generateUuid } from '../../platform/common/uuid'; import { DeepnoteServerStartupError } from '../../platform/errors/deepnoteKernelErrors'; +import { getCachedEnvironment } from '../../platform/interpreter/helpers'; +import { IInstaller, InstallerResponse, Product } from '../../platform/interpreter/installer/types'; import { logger } from '../../platform/logging'; import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { applyIntegrationEndpointEnv } from './deepnoteIntegrationEndpointEnv'; -import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; +import { DeepnoteServerInfo, IDeepnoteServerStarter } from './types'; const MAX_OUTPUT_TRACKING_LENGTH = 5000; const SERVER_STARTUP_TIMEOUT_MS = 120_000; @@ -49,7 +51,7 @@ type PendingOperation = }; interface ProjectContext { - environmentId: string; + interpreterId: string; serverInfo: DeepnoteServerInfo | null; } @@ -72,7 +74,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension constructor( @inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory, - @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller, + @inject(IInstaller) private readonly installer: IInstaller, @inject(DeepnoteAgentSkillsManager) private readonly agentSkillsManager: DeepnoteAgentSkillsManager, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry, @@ -98,14 +100,11 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension */ public async startServer( interpreter: PythonEnvironment, - venvPath: Uri, - managedVenv: boolean, - additionalPackages: string[], - environmentId: string, deepnoteFileUri: Uri, token?: CancellationToken ): Promise { const fileKey = deepnoteFileUri.fsPath; + const interpreterId = interpreter.id; let pendingOp = this.pendingOperations.get(fileKey); if (pendingOp) { @@ -119,12 +118,12 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension let existingContext = this.projectContexts.get(fileKey); if (existingContext != null) { - const { environmentId: existingEnvironmentId, serverInfo: existingServerInfo } = existingContext; + const { interpreterId: existingInterpreterId, serverInfo: existingServerInfo } = existingContext; - if (existingEnvironmentId === environmentId) { + if (existingInterpreterId === interpreterId) { if (existingServerInfo != null && (await this.isServerRunning(existingServerInfo))) { logger.info( - `Deepnote server already running at ${existingServerInfo.url} for ${fileKey} (environmentId ${environmentId})` + `Deepnote server already running at ${existingServerInfo.url} for ${fileKey} (interpreter ${interpreterId})` ); return existingServerInfo; } @@ -136,14 +135,15 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } else { logger.info( - `Stopping existing server for ${fileKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...` + `Stopping existing server for ${fileKey} with interpreter ${existingInterpreterId} to start new one with interpreter ${interpreterId}...` ); await this.stopServerForEnvironment(existingContext, deepnoteFileUri, token); - existingContext.environmentId = environmentId; + existingContext = { interpreterId, serverInfo: null }; + this.projectContexts.set(fileKey, existingContext); } } else { const newContext: ProjectContext = { - environmentId, + interpreterId, serverInfo: null }; @@ -153,16 +153,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const operation = { type: 'start' as const, - promise: this.startServerForEnvironment( - existingContext, - interpreter, - venvPath, - managedVenv, - additionalPackages, - environmentId, - deepnoteFileUri, - token - ) + promise: this.startServerForEnvironment(existingContext, interpreter, deepnoteFileUri, token) }; this.pendingOperations.set(fileKey, operation); @@ -223,7 +214,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension * Core server start using @deepnote/runtime-core's `startServer`. * * Extension-specific layers: - * - Toolkit/venv installation (before start) + * - Toolkit check/install via IInstaller (before start) * - Integration endpoint env var injection (via ServerOptions.env) — these point the toolkit at the * extension's loopback `userpod-api` endpoint, which is how it fetches SQL credentials at kernel init * - Lock file creation (after start, using returned PID) @@ -232,34 +223,49 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension private async startServerForEnvironment( projectContext: ProjectContext, interpreter: PythonEnvironment, - venvPath: Uri, - managedVenv: boolean, - additionalPackages: string[], - environmentId: string, deepnoteFileUri: Uri, token?: CancellationToken ): Promise { const fileKey = deepnoteFileUri.fsPath; + const interpreterId = interpreter.id; Cancellation.throwIfCanceled(token); - logger.info(`Ensuring deepnote-toolkit is installed in venv for environment ${environmentId}...`); - const { pythonInterpreter: venvInterpreter } = await this.toolkitInstaller.ensureVenvAndToolkit( - interpreter, - venvPath, - managedVenv, - token - ); + // Check if deepnote-toolkit is installed, and install if needed + logger.info(`Checking deepnote-toolkit installation for interpreter ${interpreterId}...`); + const isInstalled = await this.installer.isInstalled(Product.deepnoteToolkit, interpreter); - this.agentSkillsManager.ensureSkillsUpdated(environmentId, venvInterpreter); + if (!isInstalled) { + logger.info(`deepnote-toolkit not installed, installing via IInstaller...`); + const cts = new CancellationTokenSource(); + let cancellationListener: IDisposable | undefined; - Cancellation.throwIfCanceled(token); + try { + if (token) { + cancellationListener = token.onCancellationRequested(() => cts.cancel()); + } + + const result = await this.installer.install(Product.deepnoteToolkit, interpreter, cts); + + if (result === InstallerResponse.Cancelled) { + throw new Error('deepnote-toolkit installation was cancelled by the user'); + } else if (result !== InstallerResponse.Installed) { + throw new Error('Failed to install deepnote-toolkit. Check the Output panel for details.'); + } + } finally { + cancellationListener?.dispose(); + cts.dispose(); + } + } - await this.toolkitInstaller.installAdditionalPackages(venvPath, additionalPackages, token); + this.agentSkillsManager.ensureSkillsUpdated(interpreterId, interpreter); Cancellation.throwIfCanceled(token); - logger.info(`Starting deepnote-toolkit server for ${fileKey} (environmentId ${environmentId})`); + // Derive the environment path from the interpreter + const envPath = this.deriveEnvPath(interpreter); + + logger.info(`Starting deepnote-toolkit server for ${fileKey} (interpreter ${interpreterId})`); this.outputChannel.appendLine(l10n.t('Starting Deepnote server...')); const extraEnv: Record = {}; @@ -276,7 +282,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension let serverInfo: DeepnoteServerInfo | undefined; try { serverInfo = await startServer({ - pythonEnv: venvPath.fsPath, + pythonEnv: envPath, workingDirectory: path.dirname(deepnoteFileUri.fsPath), startupTimeoutMs: SERVER_STARTUP_TIMEOUT_MS, env: extraEnv @@ -287,7 +293,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension throw new DeepnoteServerStartupError( interpreter.uri.fsPath, - serverInfo?.jupyterPort ?? 0, + 0, 'unknown', capturedOutput?.stdout || '', capturedOutput?.stderr || '', @@ -314,6 +320,29 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension return serverInfo; } + /** + * Derive the environment path from a Python interpreter. + * Uses the cached environment info, or falls back to navigating up from the executable. + */ + private deriveEnvPath(interpreter: PythonEnvironment): string { + const cachedEnv = getCachedEnvironment(interpreter); + // eslint-disable-next-line local-rules/dont-use-fspath + const folderPath = cachedEnv?.environment?.folderUri?.fsPath; + + if (folderPath) { + return folderPath; + } + + const sysPrefix = cachedEnv?.executable?.sysPrefix; + + if (sysPrefix) { + return sysPrefix; + } + + // Fallback: go up from bin/python (or Scripts/python.exe on Windows) + return path.dirname(path.dirname(interpreter.uri.fsPath)); + } + /** * Stop the server using @deepnote/runtime-core's `stopServer` (SIGTERM -> wait -> SIGKILL). */ diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index ab4c6ff245..4775e16153 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -2,11 +2,14 @@ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; -import { Uri } from 'vscode'; +import { EventEmitter, Uri } from 'vscode'; import { serializeProjectFile } from '../../notebooks/deepnote/deepnoteTestHelpers'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; +import { PythonExtension } from '@vscode/python-extension'; +import { setPythonApi } from '../../platform/interpreter/helpers'; +import { IInstaller, InstallerResponse } from '../../platform/interpreter/installer/types'; import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { @@ -15,10 +18,10 @@ import { __resetRuntimeCoreMock } from '../../test/mocks/deepnoteRuntimeCore'; import { stubReadFile } from '../../test/mocks/vscodeFs'; +import { resolvableInstance } from '../../test/datascience/helpers'; import { resetVSCodeMocks } from '../../test/vscode-mock'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; -import { IDeepnoteToolkitInstaller } from './types'; /** * Unit tests for DeepnoteServerStarter. @@ -33,14 +36,13 @@ suite('DeepnoteServerStarter', () => { id: '/usr/bin/python3', uri: Uri.file('/usr/bin/python3') }; - const venvPath = Uri.file('/venvs/env1'); // Two notebooks in the SAME project directory but different files (sibling files). const uriA = Uri.file('/workspace/project/notebook-a.deepnote'); const uriB = Uri.file('/workspace/project/notebook-b.deepnote'); let serverStarter: DeepnoteServerStarter; let mockProcessServiceFactory: IProcessServiceFactory; - let mockToolkitInstaller: IDeepnoteToolkitInstaller; + let mockInstaller: IInstaller; let mockAgentSkillsManager: DeepnoteAgentSkillsManager; let mockOutputChannel: IOutputChannel; let mockAsyncRegistry: IAsyncDisposableRegistry; @@ -51,7 +53,7 @@ suite('DeepnoteServerStarter', () => { resetVSCodeMocks(); mockProcessServiceFactory = mock(); - mockToolkitInstaller = mock(); + mockInstaller = mock(); mockAgentSkillsManager = mock(); mockOutputChannel = mock(); mockAsyncRegistry = mock(); @@ -59,22 +61,26 @@ suite('DeepnoteServerStarter', () => { when(mockAsyncRegistry.push(anything())).thenReturn(); when(mockOutputChannel.appendLine(anything())).thenReturn(); + when(mockInstaller.isInstalled(anything(), anything())).thenResolve(true); + when(mockInstaller.install(anything(), anything(), anything())).thenResolve(InstallerResponse.Installed); + when(mockInstaller.onInstalled).thenReturn(new EventEmitter().event); when(mockUserpodApiEndpoints.ready).thenReturn(Promise.resolve()); when(mockUserpodApiEndpoints.baseUrl).thenReturn(undefined); - // The toolkit install step runs before runtime-core's startServer; stub it so the - // start path reaches startServer. (ts-mockito methods that are not stubbed return null.) - when(mockToolkitInstaller.ensureVenvAndToolkit(anything(), anything(), anything(), anything())).thenResolve({ - pythonInterpreter: interpreter, - toolkitVersion: '1.0.0' - }); - when(mockToolkitInstaller.installAdditionalPackages(anything(), anything(), anything())).thenResolve(); when(mockAgentSkillsManager.ensureSkillsUpdated(anything(), anything())).thenReturn(); + // startServer derives the env path via getCachedEnvironment, which needs the Python API. + const mockedApi = mock(); + sinon.stub(PythonExtension, 'api').resolves(resolvableInstance(mockedApi)); + const environments = mock(); + when(mockedApi.environments).thenReturn(instance(environments)); + when(environments.known).thenReturn([]); + setPythonApi(instance(mockedApi)); + serverStarter = new DeepnoteServerStarter( instance(mockProcessServiceFactory), - instance(mockToolkitInstaller), + instance(mockInstaller), instance(mockAgentSkillsManager), instance(mockOutputChannel), instance(mockAsyncRegistry), @@ -83,6 +89,7 @@ suite('DeepnoteServerStarter', () => { }); teardown(async () => { + setPythonApi(undefined as any); sinon.restore(); await serverStarter.dispose(); }); @@ -100,8 +107,8 @@ suite('DeepnoteServerStarter', () => { when(mockUserpodApiEndpoints.getAuthToken('project-a')).thenReturn('token-a'); when(mockUserpodApiEndpoints.getAuthToken('project-b')).thenReturn('token-b'); - await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriB); + await serverStarter.startServer(interpreter, uriA); + await serverStarter.startServer(interpreter, uriB); assert.deepStrictEqual( __getStartServerCalls().map((c) => c.env), @@ -128,8 +135,8 @@ suite('DeepnoteServerStarter', () => { suite('per-notebook keying (startServer/stopServer)', () => { test('starts SEPARATE servers for two different notebook URIs in the same dir (catches cross-sibling server reuse)', async () => { - const infoA = await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - const infoB = await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriB); + const infoA = await serverStarter.startServer(interpreter, uriA); + const infoB = await serverStarter.startServer(interpreter, uriB); // runtime-core startServer must be invoked once per notebook — NOT reused across siblings. const calls = __getStartServerCalls(); @@ -147,8 +154,8 @@ suite('DeepnoteServerStarter', () => { // Simulate a live server: the health probe (GET {url}/api) succeeds. const fetchStub = sinon.stub(globalThis, 'fetch').resolves(new Response()); - const first = await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - const second = await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); + const first = await serverStarter.startServer(interpreter, uriA); + const second = await serverStarter.startServer(interpreter, uriA); assert.strictEqual( __getStartServerCalls().length, @@ -166,16 +173,16 @@ suite('DeepnoteServerStarter', () => { // Health probes report "running" for everything; a stopped notebook must still respawn. sinon.stub(globalThis, 'fetch').resolves(new Response()); - await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); + await serverStarter.startServer(interpreter, uriA); await serverStarter.stopServer(uriA); - await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); + await serverStarter.startServer(interpreter, uriA); assert.strictEqual(__getStartServerCalls().length, 2, 'restart after stop must spawn a new server'); }); test('stopServer(uriA) tears down ONLY notebook A; B keeps running (catches cross-notebook teardown)', async () => { - const infoA = await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriB); + const infoA = await serverStarter.startServer(interpreter, uriA); + await serverStarter.startServer(interpreter, uriB); await serverStarter.stopServer(uriA); @@ -196,8 +203,8 @@ suite('DeepnoteServerStarter', () => { suite('dispose', () => { test('stops every running server; a second dispose is a no-op', async () => { - const infoA = await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - const infoB = await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriB); + const infoA = await serverStarter.startServer(interpreter, uriA); + const infoB = await serverStarter.startServer(interpreter, uriB); await serverStarter.dispose(); await serverStarter.dispose(); @@ -225,7 +232,7 @@ suite('DeepnoteServerStarter', () => { }) ); - const startPromise = serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); + const startPromise = serverStarter.startServer(interpreter, uriA); let disposeResolved = false; const disposePromise = serverStarter.dispose().then(() => { diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index f33eba3fe3..57b13f2e92 100644 --- a/src/kernels/deepnote/types.ts +++ b/src/kernels/deepnote/types.ts @@ -151,32 +151,24 @@ export interface IDeepnoteToolkitInstaller { export const IDeepnoteServerStarter = Symbol('IDeepnoteServerStarter'); export interface IDeepnoteServerStarter { /** - * Starts a deepnote-toolkit Jupyter server for a kernel environment. - * Environment-based method. + * Starts a deepnote-toolkit Jupyter server using the active Python interpreter. + * Handles checking/installing deepnote-toolkit via the IInstaller infrastructure. * @param interpreter The Python interpreter to use - * @param venvPath The path to the venv - * @param managedVenv Whether the venv is managed by this extension (created by us) - * @param environmentId The environment ID (for server management) * @param deepnoteFileUri The URI of the .deepnote file * @param token Cancellation token to cancel the operation * @returns Connection information (URL, port, etc.) */ startServer( interpreter: PythonEnvironment, - venvPath: vscode.Uri, - managedVenv: boolean, - additionalPackages: string[], - environmentId: string, deepnoteFileUri: vscode.Uri, token?: vscode.CancellationToken ): Promise; /** - * Stops the deepnote-toolkit server for a kernel environment. - * @param environmentId The environment ID + * Stops the deepnote-toolkit server for a .deepnote file. + * @param deepnoteFileUri The URI of the .deepnote file * @param token Cancellation token to cancel the operation */ - // stopServer(environmentId: string, token?: vscode.CancellationToken): Promise; stopServer(deepnoteFileUri: vscode.Uri, token?: vscode.CancellationToken): Promise; /** diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index e09236d914..b4ce53ad36 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -6,33 +6,24 @@ import * as fs from 'fs'; import { inject, injectable, named, optional } from 'inversify'; import { CancellationToken, - CancellationTokenSource, - NotebookController, NotebookControllerAffinity, NotebookDocument, NotebookEditor, ProgressLocation, - QuickPickItem, Uri, commands, env, l10n, - notebooks, window, workspace } from 'vscode'; -import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; import { DEEPNOTE_NOTEBOOK_TYPE, - DEEPNOTE_TOOLKIT_VERSION, DeepnoteKernelConnectionMetadata, - IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, IDeepnoteLspClientManager, - IDeepnoteNotebookEnvironmentMapper, IDeepnoteServerProvider, IDeepnoteServerStarter, - IDeepnoteToolkitInstaller, IServerHandleRegistry } from '../../kernels/deepnote/types'; import { createJupyterConnectionInfo } from '../../kernels/jupyter/jupyterUtils'; @@ -44,24 +35,23 @@ import { } from '../../kernels/jupyter/types'; import { IJupyterKernelSpec, IKernelProvider } from '../../kernels/types'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { ITelemetryService } from '../../platform/analytics/types'; import { IPythonExtensionChecker } from '../../platform/api/types'; import { Cancellation, isCancellationError } from '../../platform/common/cancellation'; import { JVSC_EXTENSION_ID, STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants'; import { getDisplayPath } from '../../platform/common/platform/fs-paths.node'; import { IConfigurationService, IDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { disposeAsync } from '../../platform/common/utils'; +import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { createDeepnoteServerConfigHandle } from '../../platform/deepnote/deepnoteServerUtils.node'; -import { DeepnoteKernelError, DeepnoteToolkitMissingError } from '../../platform/errors/deepnoteKernelErrors'; +import { DeepnoteKernelError } from '../../platform/errors/deepnoteKernelErrors'; +import { IInterpreterService } from '../../platform/interpreter/contracts'; import { logger } from '../../platform/logging'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; -import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; -// Constants for NotebookEditor retry logic const NOTEBOOK_EDITOR_RETRY_COUNT = 10; const NOTEBOOK_EDITOR_RETRY_DELAY_MS = 100; @@ -74,10 +64,8 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, private readonly notebookConnectionMetadata = new Map(); // Track registered controllers per NOTEBOOK (full URI with query) - one controller per notebook private readonly notebookControllers = new Map(); - // Track environment for each notebook - private readonly notebookEnvironmentsIds = new Map(); - // Track per-notebook placeholder controllers for notebooks without configured environments - private readonly placeholderControllers = new Map(); + // Track interpreter ID for each notebook + private readonly notebookInterpreterIds = new Map(); constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, @@ -93,14 +81,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, @inject(IDeepnoteRequirementsHelper) private readonly requirementsHelper: IDeepnoteRequirementsHelper, - @inject(IDeepnoteEnvironmentManager) private readonly environmentManager: IDeepnoteEnvironmentManager, @inject(IDeepnoteServerStarter) private readonly serverStarter: IDeepnoteServerStarter, - @inject(IDeepnoteNotebookEnvironmentMapper) - private readonly notebookEnvironmentMapper: IDeepnoteNotebookEnvironmentMapper, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, - @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller, - @inject(IServerHandleRegistry) private readonly serverHandleRegistry: IServerHandleRegistry, - @inject(ITelemetryService) private readonly analytics: ITelemetryService + @inject(IInterpreterService) private readonly interpreterService: IInterpreterService, + @inject(IServerHandleRegistry) private readonly serverHandleRegistry: IServerHandleRegistry ) {} public activate() { @@ -161,14 +145,11 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, (result) => { logger.info(`Auto-selecting Deepnote kernel for ${getDisplayPath(notebook.uri)} result: ${result}`); if (!result) { - logger.info(`No environment configured for ${getDisplayPath(notebook.uri)}, showing warning`); - this.showNoEnvironmentWarning(notebook).catch((error) => { - logger.error( - `Error showing no environment warning for ${getDisplayPath(notebook.uri)}`, - error - ); - void this.handleKernelSelectionError(error, notebook); - }); + logger.warn( + `No active Python interpreter found for ${getDisplayPath( + notebook.uri + )}, kernel not selected` + ); } }, (error) => { @@ -178,83 +159,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, ); } - private async showNoEnvironmentWarning(notebook: NotebookDocument): Promise { - logger.info(`Showing no environment warning for ${getDisplayPath(notebook.uri)}`); - const selectEnvironmentAction = l10n.t('Select Environment'); - const cancelAction = l10n.t('Cancel'); - - const selectedAction = await window.showWarningMessage( - l10n.t('No environment configured for this notebook. Please select an environment to continue.'), - { modal: false }, - selectEnvironmentAction, - cancelAction - ); - - logger.info(`Selected action: ${selectedAction}`); - if (selectedAction === selectEnvironmentAction) { - logger.info(`Executing command to pick environment for ${getDisplayPath(notebook.uri)}`); - void commands.executeCommand('deepnote.environments.selectForNotebook', { notebook }); - } - } - - public async pickEnvironment(notebookUri: Uri): Promise { - logger.info(`Picking environment for notebook ${getDisplayPath(notebookUri)}`); - - // Wait for environment manager to finish loading environments from storage - await this.environmentManager.waitForInitialization(); - - const environments = this.environmentManager.listEnvironments(); - const items: (QuickPickItem & { environment?: DeepnoteEnvironment })[] = environments.map((env) => { - return { - label: env.name, - description: getDisplayPath(env.pythonInterpreter.uri), - detail: env.packages?.length - ? l10n.t('Packages: {0}', env.packages.join(', ')) - : l10n.t('No additional packages'), - environment: env - }; - }); - - items.push({ - label: '$(add) Create New Environment', - description: 'Set up a new kernel environment', - alwaysShow: true - }); - - const selected = await window.showQuickPick(items, { - placeHolder: `Select an environment for ${getDisplayPath(notebookUri)}`, - matchOnDescription: true, - matchOnDetail: true - }); - - if (!selected) { - logger.info('User cancelled environment selection'); - return; // User cancelled - } - - if (!selected.environment) { - logger.info('User chose to create new environment - triggering create command'); - - await commands.executeCommand('deepnote.environments.create'); - - const newEnvironments = this.environmentManager.listEnvironments(); - - if (newEnvironments.length > environments.length) { - logger.info('Environment created, showing picker again'); - - return this.pickEnvironment(notebookUri); - } - - logger.info('No new environment created'); - - return; - } - - logger.info(`Selected environment "${selected.environment.name}" for notebook ${getDisplayPath(notebookUri)}`); - - return selected.environment; - } - private onControllerSelectionChanged(event: { notebook: NotebookDocument; controller: IVSCodeNotebookController; @@ -273,24 +177,16 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } private onDidCloseNotebook(notebook: NotebookDocument) { - logger.info(`Notebook closed: ${getDisplayPath(notebook.uri)}, with type: ${notebook.notebookType}`); - - // Only handle deepnote notebooks if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { return; } - logger.info(`Deepnote notebook closed: ${getDisplayPath(notebook.uri)}`); - - // Clean up placeholder controller if it exists const notebookKey = getNotebookKey(notebook.uri); - const placeholder = this.placeholderControllers.get(notebookKey); + this.notebookConnectionMetadata.delete(notebookKey); + this.notebookInterpreterIds.delete(notebookKey); + this.notebookControllers.delete(notebookKey); - if (placeholder) { - logger.info(`Disposing placeholder controller for closed notebook: ${getDisplayPath(notebook.uri)}`); - placeholder.dispose(); - this.placeholderControllers.delete(notebookKey); - } + logger.info(`Deepnote notebook closed, cleaned up: ${getDisplayPath(notebook.uri)}`); } /** @@ -332,19 +228,15 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // cause "command already exists" errors when trying to start new clients await this.lspClientManager.stopLspClients(notebook.uri, token); - // Update the controller with new environment's metadata - // Because we use notebook-based controller IDs, addOrUpdate will call updateConnection() - // on the existing controller instead of creating a new one - const environmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(notebook.uri); - const environment = environmentId ? this.environmentManager.getEnvironment(environmentId) : undefined; + // Get the active interpreter and re-setup the kernel + const interpreter = await this.interpreterService.getActiveInterpreter(notebook.uri); - if (environment == null) { - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(notebook.uri); - logger.error(`No environment found for notebook ${getDisplayPath(notebook.uri)}`); + if (!interpreter) { + logger.error(`No active Python interpreter found for ${getDisplayPath(notebook.uri)}`); return; } - await this.ensureKernelSelectedWithConfiguration(notebook, environment, notebookKey, progress, token); + await this.ensureKernelSelectedWithInterpreter(notebook, interpreter, notebookKey, progress, token); // Setup succeeded. If it registered a new server handle (full setup path), drop the old one. // The verified-controller early return reuses the existing handle, so nothing to clear then. @@ -366,47 +258,28 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // notebookKey uniquely identifies THIS NOTEBOOK - the same identity the controller/server use const notebookKey = getNotebookKey(notebook.uri); - const environmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(notebook.uri); - - if (environmentId == null) { - await this.selectPlaceholderController(notebook); + // Get the active Python interpreter + const interpreter = await this.interpreterService.getActiveInterpreter(notebook.uri); + if (!interpreter) { + logger.warn(`No active Python interpreter found for ${getDisplayPath(notebook.uri)}`); return false; } - const environment = environmentId ? this.environmentManager.getEnvironment(environmentId) : undefined; - - if (environment == null) { - logger.info(`No environment found for notebook ${getDisplayPath(notebook.uri)}`); - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(notebook.uri); - await this.selectPlaceholderController(notebook); - - return false; - } - - await this.ensureKernelSelectedWithConfiguration(notebook, environment, notebookKey, progress, token); + await this.ensureKernelSelectedWithInterpreter(notebook, interpreter, notebookKey, progress, token); return true; } - public async ensureKernelSelectedWithConfiguration( + public async ensureKernelSelectedWithInterpreter( notebook: NotebookDocument, - configuration: DeepnoteEnvironment, + interpreter: PythonEnvironment, notebookKey: string, progress: { report(value: { message?: string; increment?: number }): void }, progressToken: CancellationToken ): Promise { - // Dispose placeholder controller if it exists (real controller is taking over) - const placeholder = this.placeholderControllers.get(notebookKey); - - if (placeholder) { - logger.info(`Disposing placeholder controller for ${getDisplayPath(notebook.uri)}`); - placeholder.dispose(); - this.placeholderControllers.delete(notebookKey); - } - - logger.info(`Setting up kernel using configuration: ${configuration.name} (${configuration.id})`); - progress.report({ message: `Using ${configuration.name}...` }); + logger.info(`Setting up kernel using interpreter: ${interpreter.id}`); + progress.report({ message: `Using interpreter ${getDisplayPath(interpreter.uri)}...` }); // Check if Python extension is installed if (!this.pythonExtensionChecker.isPythonExtensionInstalled) { @@ -416,70 +289,38 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } const existingController = this.notebookControllers.get(notebookKey); - const existingEnvironmentId = this.notebookEnvironmentsIds.get(notebookKey); - - if (existingEnvironmentId != null && existingController != null && existingEnvironmentId === configuration.id) { - logger.info(`Existing controller found for notebook ${getDisplayPath(notebook.uri)}, verifying connection`); - - // Verify the controller's interpreter path matches the expected venv path - // This handles cases where notebooks were used in VS Code and now opened in Cursor - if (this.isControllerInterpreterValid(existingController, configuration.venvPath)) { - logger.info(`Controller verified, selecting it`); - await this.ensureControllerSelectedForNotebook(notebook, existingController, progressToken); - - return; - } - - const expectedInterpreter = this.getVenvInterpreterUri(configuration.venvPath); - logger.warn( - `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingController.connection.interpreter?.uri.fsPath}. Recreating controller.` - ); + const existingInterpreterId = this.notebookInterpreterIds.get(notebookKey); - // Dispose old controller and recreate it - existingController.dispose(); - this.notebookControllers.delete(notebookKey); + if (existingInterpreterId != null && existingController != null && existingInterpreterId === interpreter.id) { + logger.info(`Existing controller found for notebook ${getDisplayPath(notebook.uri)}, reusing`); + await this.ensureControllerSelectedForNotebook(notebook, existingController, progressToken); + return; } // Ensure server is running (startServer is idempotent - returns early if already running) - // Note: startServer() will create the venv if it doesn't exist - logger.info(`Ensuring server is running for configuration ${configuration.id}`); + // Server starter handles toolkit check/install via IInstaller internally + logger.info(`Ensuring server is running for interpreter ${interpreter.id}`); progress.report({ message: 'Starting Deepnote server...' }); - const serverInfo = await this.serverStarter.startServer( - configuration.pythonInterpreter, - configuration.venvPath, - configuration.managedVenv, - configuration.packages ?? [], - configuration.id, - notebook.uri, - progressToken - ); + const serverInfo = await this.serverStarter.startServer(interpreter, notebook.uri, progressToken); - this.notebookEnvironmentsIds.set(notebookKey, configuration.id); + this.notebookInterpreterIds.set(notebookKey, interpreter.id); logger.info(`Server running at ${serverInfo.url}`); - // Update last used timestamp - await this.environmentManager.updateLastUsed(configuration.id); - - // Create server provider handle + // Create server provider handle using interpreter ID const serverProviderHandle: JupyterServerProviderHandle = { extensionId: JVSC_EXTENSION_ID, id: 'deepnote-server', - handle: createDeepnoteServerConfigHandle(configuration.id, notebook.uri) + handle: createDeepnoteServerConfigHandle(interpreter.id, notebook.uri) }; + // Register the server with the provider (one server per PROJECT) this.serverProvider.registerServer(serverProviderHandle.handle, serverInfo); this.serverHandleRegistry.set(notebookKey, serverProviderHandle.handle); - const lspInterpreterUri = this.getVenvInterpreterUri(configuration.venvPath); - - const lspInterpreter: PythonEnvironment = { - uri: lspInterpreterUri, - id: lspInterpreterUri.fsPath - } as PythonEnvironment; - + // Use the active interpreter directly for LSP (it already has deepnote-toolkit installed) try { - await this.lspClientManager.startLspClients(serverInfo, notebook.uri, lspInterpreter, progressToken); + await this.lspClientManager.startLspClients(serverInfo, notebook.uri, interpreter, progressToken); logger.info(`✓ LSP clients started for ${notebookKey}`); } catch (error) { @@ -488,12 +329,14 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, progress.report({ message: 'Connecting to kernel...' }); + const displayName = `Deepnote: ${getDisplayPath(interpreter.uri)} (${notebookKey})`; + const connectionInfo = createJupyterConnectionInfo( serverProviderHandle, { baseUrl: serverInfo.url, token: serverInfo.token || '', - displayName: `Deepnote: ${configuration.name} (${notebookKey})`, + displayName, authorizationHeader: {} }, this.requestCreator, @@ -508,8 +351,8 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const kernelSpecs = await sessionManager.getKernelSpecs(); logger.info(`Available kernel specs on Deepnote server: ${kernelSpecs.map((s) => s.name).join(', ')}`); - // Use the extracted kernel selection logic - kernelSpec = this.selectKernelSpec(kernelSpecs, configuration.id); + // Select the default Python kernel (ipykernel-provided python3 spec) + kernelSpec = this.selectKernelSpec(kernelSpecs); logger.info(`✓ Using kernel spec: ${kernelSpec.name} (${kernelSpec.display_name})`); } finally { @@ -518,10 +361,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, progress.report({ message: 'Finalizing kernel setup...' }); - const venvInterpreter = this.getVenvInterpreterUri(configuration.venvPath); - - logger.info(`Using venv path: ${configuration.venvPath.fsPath}`); - logger.info(`Venv interpreter path: ${venvInterpreter.fsPath}`); + logger.info(`Using interpreter: ${interpreter.uri.fsPath}`); // CRITICAL: Use unique notebook-based ID (includes query with notebook ID) // This ensures each notebook gets its own controller/kernel, even within the same project. @@ -533,14 +373,14 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const projectTitle = notebook.metadata?.deepnoteProjectName || 'Untitled Project'; const newConnectionMetadata = DeepnoteKernelConnectionMetadata.create({ - interpreter: { uri: venvInterpreter, id: venvInterpreter.fsPath }, + interpreter, kernelSpec, baseUrl: serverInfo.url, id: controllerId, projectFilePath: getNotebookKey(notebook.uri), serverProviderHandle, serverInfo, - environmentName: configuration.name, + environmentName: getDisplayPath(interpreter.uri), projectName: projectTitle, notebookName: notebookKey }); @@ -564,7 +404,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Store the controller for reuse this.notebookControllers.set(notebookKey, controller); - // Prepare init notebook execution const projectId = notebook.metadata?.deepnoteProjectId; const notebookId = notebook.metadata?.deepnoteNotebookId; const project = @@ -610,7 +449,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Auto-select the controller await this.ensureControllerSelectedForNotebook(notebook, controller, progressToken); - logger.info(`Successfully set up kernel with configuration: ${configuration.name}`); + logger.info(`Successfully set up kernel with interpreter: ${interpreter.id}`); progress.report({ message: 'Kernel ready!' }); } @@ -653,35 +492,20 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } /** - * Select the appropriate kernel spec for an environment. + * Select the default Python kernel spec from the server. * Extracted for testability. * @param kernelSpecs Available kernel specs from the server - * @param environmentId The environment ID to find a kernel for * @returns The selected kernel spec * @throws Error if no suitable kernel spec is found */ - public selectKernelSpec(kernelSpecs: IJupyterKernelSpec[], environmentId: string): IJupyterKernelSpec { - // Look for environment-specific kernel first - const expectedKernelName = `deepnote-${environmentId}`; - logger.info(`Looking for environment-specific kernel: ${expectedKernelName}`); - - const kernelSpec = kernelSpecs.find((s) => s.name === expectedKernelName); + public selectKernelSpec(kernelSpecs: IJupyterKernelSpec[]): IJupyterKernelSpec { + const kernelSpec = + kernelSpecs.find((s) => s.language === 'python') || + kernelSpecs.find((s) => s.name === 'python3') || + kernelSpecs[0]; if (!kernelSpec) { - logger.warn( - `Environment-specific kernel '${expectedKernelName}' not found! Falling back to generic Python kernel.` - ); - // Fallback to any Python kernel - const fallbackKernel = - kernelSpecs.find((s) => s.language === 'python') || - kernelSpecs.find((s) => s.name === 'python3') || - kernelSpecs[0]; - - if (!fallbackKernel) { - throw new Error('No kernel specs available on Deepnote server'); - } - - return fallbackKernel; + throw new Error('No kernel specs available on Deepnote server'); } return kernelSpec; @@ -689,7 +513,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, /** * Ensure an environment is configured for the notebook before execution. - * If not configured, shows picker and sets up the kernel. + * Uses the active Python interpreter and the IInstaller infrastructure. * @returns true if environment is ready, false if user cancelled */ public async ensureEnvironmentConfiguredBeforeExecution( @@ -700,96 +524,21 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const notebookKey = getNotebookKey(notebook.uri); - const existingEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(notebook.uri); + const interpreter = await this.interpreterService.getActiveInterpreter(notebook.uri); - // No environment configured - need to pick one - if (!existingEnvironmentId) { - return this.pickAndSetupEnvironment(notebook, notebookKey, token); - } - - const environment = this.environmentManager.getEnvironment(existingEnvironmentId); - - // Environment no longer exists - remove stale mapping and pick a new one - if (!environment) { - logger.info(`Removing stale environment mapping for ${getDisplayPath(notebook.uri)}`); - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(notebook.uri); - - return this.pickAndSetupEnvironment(notebook, notebookKey, token); + if (!interpreter) { + logger.warn(`No active Python interpreter found for ${getDisplayPath(notebook.uri)}`); + return false; } const existingController = this.notebookControllers.get(notebookKey); + const existingInterpreterId = this.notebookInterpreterIds.get(notebookKey); - // Environment and controller already configured - but verify interpreter path still matches - if (existingController) { - if (!this.isControllerInterpreterValid(existingController, environment.venvPath)) { - const expectedInterpreter = this.getVenvInterpreterUri(environment.venvPath); - logger.warn( - `Controller interpreter path mismatch! Expected: ${expectedInterpreter.fsPath}, Got: ${existingController.connection.interpreter?.uri.fsPath}. Recreating controller.` - ); - - existingController.dispose(); - this.notebookControllers.delete(notebookKey); - - return this.setupKernelForEnvironment(notebook, environment, notebookKey, token); - } - - logger.info(`Environment "${environment.name}" already configured for ${getDisplayPath(notebook.uri)}`); - + if (existingController && existingInterpreterId === interpreter.id) { + logger.info(`Controller already configured for ${getDisplayPath(notebook.uri)}`); return true; } - // Environment exists but controller is missing - set it up - logger.info( - `Environment "${environment.name}" configured but controller missing for ${getDisplayPath( - notebook.uri - )}, triggering setup` - ); - - return this.setupKernelForEnvironment(notebook, environment, notebookKey, token); - } - - /** - * Pick an environment and set up the kernel for a notebook. - */ - private async pickAndSetupEnvironment( - notebook: NotebookDocument, - notebookKey: string, - token: CancellationToken - ): Promise { - Cancellation.throwIfCanceled(token); - - logger.info(`No environment configured for ${getDisplayPath(notebook.uri)}, showing picker`); - const selectedEnvironment = await this.pickEnvironment(notebook.uri); - - if (!selectedEnvironment) { - logger.info(`User cancelled environment selection for ${getDisplayPath(notebook.uri)}`); - - return false; - } - - Cancellation.throwIfCanceled(token); - - await this.notebookEnvironmentMapper.setEnvironmentForNotebook(notebook.uri, selectedEnvironment.id); - - const result = await this.setupKernelForEnvironment(notebook, selectedEnvironment, notebookKey, token); - - if (result) { - this.analytics.trackEvent({ eventName: 'select_environment' }); - logger.info(`Environment "${selectedEnvironment.name}" configured for ${getDisplayPath(notebook.uri)}`); - } - - return result; - } - - /** - * Set up the kernel for a given environment. - */ - private async setupKernelForEnvironment( - notebook: NotebookDocument, - environment: DeepnoteEnvironment, - notebookKey: string, - token: CancellationToken - ): Promise { try { await window.withProgress( { @@ -798,9 +547,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, cancellable: true }, async (progress, progressToken) => { - await this.ensureKernelSelectedWithConfiguration( + await this.ensureKernelSelectedWithInterpreter( notebook, - environment, + interpreter, notebookKey, progress, progressToken @@ -810,69 +559,50 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } catch (error) { if (token.isCancellationRequested || isCancellationError(error as Error)) { logger.info(`Kernel setup cancelled for ${getDisplayPath(notebook.uri)}`); - return false; } throw error; } - const createdController = this.notebookControllers.get(notebookKey); - - if (!createdController) { - logger.warn( - `Controller not created for "${environment.name}" on ${getDisplayPath(notebook.uri)} after setup` - ); - - return false; - } - - return true; + return !!this.notebookControllers.get(notebookKey); } /** - * Clear the controller selection for a notebook using a specific environment. - * This is used when deleting an environment to unselect its controller from any open notebooks. + * Clear the controller selection for a notebook if it was set up by this selector + * for the given environment. + * + * The caller passes an `environmentId` (UUID), but the auto-selector now tracks + * notebooks by interpreter.id. We match by comparing the notebook's tracked + * controller instance against the currently selected controller, so we only + * clear controllers we own — never an unrelated Deepnote kernel. */ public clearControllerForEnvironment(notebook: NotebookDocument, environmentId: string): void { - const selectedController = this.controllerRegistration.getSelected(notebook); - if (!selectedController || selectedController.connection.kind !== 'startUsingDeepnoteKernel') { - return; - } - - const expectedHandle = createDeepnoteServerConfigHandle(environmentId, notebook.uri); + const notebookKey = getNotebookKey(notebook.uri); + const trackedController = this.notebookControllers.get(notebookKey); - if (selectedController.connection.serverProviderHandle.handle === expectedHandle) { - // Unselect the controller by setting affinity to Default - selectedController.controller.updateNotebookAffinity(notebook, NotebookControllerAffinity.Default); - logger.info( - `Cleared controller for notebook ${getDisplayPath(notebook.uri)} (environment ${environmentId})` - ); + if (!trackedController) { + return; // We didn't set up a controller for this notebook } - } - private getVenvInterpreterUri(venvPath: Uri): Uri { - return process.platform === 'win32' - ? Uri.joinPath(venvPath, 'Scripts', 'python.exe') - : Uri.joinPath(venvPath, 'bin', 'python'); - } - - /** - * Check if a controller's interpreter path matches the expected venv path. - * Returns true when no interpreter is present (nothing to validate) or when paths match. - */ - private isControllerInterpreterValid( - controller: { connection: { interpreter?: { uri: Uri } } }, - venvPath: Uri - ): boolean { - const existingInterpreter = controller.connection.interpreter; + const selectedController = this.controllerRegistration.getSelected(notebook); + if (!selectedController || selectedController.id !== trackedController.id) { + return; // Selected controller isn't the one we own + } - if (!existingInterpreter) { - return true; + if (selectedController.connection.kind !== 'startUsingDeepnoteKernel') { + return; } - const expectedInterpreter = this.getVenvInterpreterUri(venvPath); + selectedController.controller.updateNotebookAffinity(notebook, NotebookControllerAffinity.Default); + + // Clean up our tracking state for this notebook + this.notebookControllers.delete(notebookKey); + this.notebookConnectionMetadata.delete(notebookKey); + this.notebookInterpreterIds.delete(notebookKey); - return existingInterpreter.uri.fsPath === expectedInterpreter.fsPath; + logger.info( + `Cleared Deepnote controller for notebook ${getDisplayPath(notebook.uri)} (environment ${environmentId})` + ); } /** @@ -881,7 +611,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, * Includes retry logic since the editor might not be visible immediately when the document opens. */ private async findNotebookEditor(notebook: NotebookDocument): Promise { - // Try to find immediately let editor = window.visibleNotebookEditors.find( (e) => getNotebookKey(e.notebook.uri) === getNotebookKey(notebook.uri) ); @@ -890,7 +619,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return editor; } - // If not found, wait briefly and retry (editor might not be visible yet) for (let i = 0; i < NOTEBOOK_EDITOR_RETRY_COUNT; i++) { await new Promise((resolve) => setTimeout(resolve, NOTEBOOK_EDITOR_RETRY_DELAY_MS)); @@ -906,28 +634,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - /** - * Create and select a placeholder controller for a notebook without a configured environment. - */ - private async selectPlaceholderController(notebook: NotebookDocument): Promise { - const placeholder = this.createPlaceholderController(notebook); - placeholder.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); - - const notebookEditor = await this.findNotebookEditor(notebook); - - if (notebookEditor) { - await commands.executeCommand('notebook.selectKernel', { - notebookEditor: notebookEditor, - id: placeholder.id, - extension: JVSC_EXTENSION_ID - }); - } else { - logger.warn( - `Could not find NotebookEditor for ${getDisplayPath(notebook.uri)}, kernel may not be selected` - ); - } - } - /** * Handle kernel selection errors with user-friendly messages and actions */ @@ -939,28 +645,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - if (error instanceof DeepnoteToolkitMissingError) { - const installAction = l10n.t('Install'); - const changeEnvironmentAction = l10n.t('Change Environment'); - const selectedAction = await window.showWarningMessage( - l10n.t( - 'Running Deepnote projects requires deepnote-toolkit[server]=={0} to be installed in the selected environment', - DEEPNOTE_TOOLKIT_VERSION - ), - { modal: true }, - installAction, - changeEnvironmentAction - ); - - if (selectedAction === installAction) { - await this.installToolkitAndNotify(error.venvPath, notebook); - } else if (selectedAction === changeEnvironmentAction) { - void commands.executeCommand('deepnote.environments.selectForNotebook', { notebook }); - } - - return; - } - // Handle DeepnoteKernelError types with specific guidance if (error instanceof DeepnoteKernelError) { // Log the technical details @@ -1014,41 +698,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } } - /** - * Install deepnote-toolkit in an existing venv and rebuild the controller. - */ - private async installToolkitAndNotify(venvPath: string, notebook: NotebookDocument): Promise { - try { - await window.withProgress( - { - location: ProgressLocation.Notification, - title: l10n.t('Installing deepnote-toolkit...'), - cancellable: true - }, - async (progress, token) => { - await this.toolkitInstaller.installToolkitInExistingVenv(Uri.file(venvPath), token); - - // After successful installation, rebuild the controller to use the new environment - progress.report({ message: l10n.t('Starting kernel...') }); - await this.rebuildController(notebook, progress, token); - } - ); - - void window.showInformationMessage(l10n.t('deepnote-toolkit installed successfully')); - } catch (installError) { - if (installError instanceof Error && isCancellationError(installError)) { - logger.info('deepnote-toolkit installation cancelled'); - - return; - } - - logger.error('Failed to install deepnote-toolkit', installError); - const errorMessage = installError instanceof Error ? installError.message : String(installError); - - void window.showErrorMessage(l10n.t('Failed to install deepnote-toolkit: {0}', errorMessage)); - } - } - /** * Read and hash the existing requirements.txt file if it exists. * Returns the same hash format as computeRequirementsHash for comparison. @@ -1084,78 +733,4 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return ''; } } - - /** - * Create a placeholder controller for a notebook without a configured environment. - * Each notebook gets its own placeholder with a unique ID. - * The placeholder's executeHandler shows the environment picker when user tries to run cells. - */ - private createPlaceholderController(notebook: NotebookDocument): NotebookController { - const notebookKey = getNotebookKey(notebook.uri); - - // Check if we already have one - const existing = this.placeholderControllers.get(notebookKey); - - if (existing) { - return existing; - } - - const controller = notebooks.createNotebookController( - `deepnote-placeholder-${notebookKey}`, - DEEPNOTE_NOTEBOOK_TYPE, - l10n.t('Deepnote: Select Environment') - ); - - controller.supportsExecutionOrder = true; - controller.supportedLanguages = ['python', 'sql', 'markdown', 'plaintext']; - - // Environment picker only; execution goes through the real controller on retry. - controller.executeHandler = async (cells, doc) => { - logger.info( - `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ - cells.length - } cells` - ); - - if (!workspace.isTrusted) { - logger.info(`Workspace is not trusted, skipping environment setup for ${getDisplayPath(doc.uri)}`); - - return; - } - - // Create a cancellation token that cancels when the notebook is closed - const cts = new CancellationTokenSource(); - const closeListener = workspace.onDidCloseNotebookDocument((closedDoc) => { - if (getNotebookKey(closedDoc.uri) === getNotebookKey(doc.uri)) { - logger.info(`Notebook closed during environment setup, cancelling operation`); - cts.cancel(); - } - }); - - try { - const hasEnvironment = await this.ensureEnvironmentConfiguredBeforeExecution(doc, cts.token); - - if (!hasEnvironment) { - logger.info(`User cancelled environment selection, not executing cells`); - - return; - } - - void window.showInformationMessage(l10n.t('Environment ready. Run the cells again to execute them.')); - } catch (error) { - if (isCancellationError(error)) { - logger.info(`Environment setup cancelled for ${getDisplayPath(doc.uri)}`); - } else { - logger.error(`Error in placeholder controller execute handler`, error); - } - } finally { - closeListener.dispose(); - cts.dispose(); - } - }; - - this.placeholderControllers.set(notebookKey, controller); - - return controller; - } } diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 7ae8452120..defb320fdd 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -1,20 +1,14 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; -import { anything, deepEqual, instance, mock, verify, when } from 'ts-mockito'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; import { DeepnoteKernelAutoSelector } from './deepnoteKernelAutoSelector.node'; -import { createMockChildProcess } from '../../kernels/deepnote/deepnoteTestHelpers.node'; -import { createMockCell } from './deepnoteTestHelpers'; import { ServerHandleRegistry } from '../../kernels/deepnote/deepnoteServerHandleRegistry.node'; import { - IDeepnoteEnvironmentManager, IDeepnoteLspClientManager, - IDeepnoteNotebookEnvironmentMapper, IDeepnoteServerProvider, - IDeepnoteServerStarter, - IDeepnoteToolkitInstaller + IDeepnoteServerStarter } from '../../kernels/deepnote/types'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; -import { ITelemetryService } from '../../platform/analytics/types'; import { IDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { IPythonExtensionChecker } from '../../platform/api/types'; import { IJupyterRequestCreator } from '../../kernels/jupyter/types'; @@ -24,23 +18,20 @@ import { IKernelProvider, IKernel, IJupyterKernelSpec, KernelConnectionMetadata import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; import { CancellationError, - EventEmitter, NotebookDocument, NotebookEditor, Uri, NotebookController, CancellationToken } from 'vscode'; -import { DeepnoteToolkitMissingError } from '../../platform/errors/deepnoteKernelErrors'; -import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; +import { IInterpreterService } from '../../platform/interpreter/contracts'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; suite('DeepnoteKernelAutoSelector - rebuildController', () => { let selector: DeepnoteKernelAutoSelector; - let registry: ServerHandleRegistry; let mockDisposableRegistry: IDisposableRegistry; let mockControllerRegistration: IControllerRegistration; let mockPythonExtensionChecker: IPythonExtensionChecker; @@ -51,12 +42,10 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { let mockNotebookManager: IDeepnoteNotebookManager; let mockKernelProvider: IKernelProvider; let mockRequirementsHelper: IDeepnoteRequirementsHelper; - let mockEnvironmentManager: IDeepnoteEnvironmentManager; let mockServerStarter: IDeepnoteServerStarter; - let mockNotebookEnvironmentMapper: IDeepnoteNotebookEnvironmentMapper; let mockOutputChannel: IOutputChannel; - let mockToolkitInstaller: IDeepnoteToolkitInstaller; - let mockTelemetryService: ITelemetryService; + let mockInterpreterService: IInterpreterService; + let registry: ServerHandleRegistry; let mockProgress: { report(value: { message?: string; increment?: number }): void }; let mockCancellationToken: CancellationToken; @@ -81,12 +70,10 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { mockNotebookManager = mock(); mockKernelProvider = mock(); mockRequirementsHelper = mock(); - mockEnvironmentManager = mock(); mockServerStarter = mock(); - mockToolkitInstaller = mock(); - mockNotebookEnvironmentMapper = mock(); mockOutputChannel = mock(); - mockTelemetryService = mock(); + mockInterpreterService = mock(); + registry = new ServerHandleRegistry(); mockProgress = { report: sandbox.stub() }; mockCancellationToken = mock(); @@ -134,8 +121,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { mockLoadingController ); - registry = new ServerHandleRegistry(); - // Create selector instance selector = new DeepnoteKernelAutoSelector( instance(mockDisposableRegistry), @@ -149,13 +134,10 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { instance(mockNotebookManager), instance(mockKernelProvider), instance(mockRequirementsHelper), - instance(mockEnvironmentManager), instance(mockServerStarter), - instance(mockNotebookEnvironmentMapper), instance(mockOutputChannel), - instance(mockToolkitInstaller), - registry, - instance(mockTelemetryService) + instance(mockInterpreterService), + registry ); }); @@ -174,31 +156,31 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { pendingCells: [{ index: 0 }, { index: 1 }] // 2 cells pending }; - // Create mock environment - const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); - when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); + // Mock interpreter service to return an active interpreter + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); when(mockKernelProvider.get(mockNotebook)).thenReturn(instance(mockKernel)); when(mockKernelProvider.getKernelExecution(instance(mockKernel))).thenReturn(mockExecution as any); - // Stub ensureKernelSelectedWithConfiguration to verify it's still called despite pending cells - const ensureKernelSelectedWithConfigurationStub = sandbox - .stub(selector, 'ensureKernelSelectedWithConfiguration') + // Stub ensureKernelSelectedWithInterpreter to verify it's still called despite pending cells + const ensureKernelSelectedWithInterpreterStub = sandbox + .stub(selector, 'ensureKernelSelectedWithInterpreter') .resolves(); // Act await selector.rebuildController(mockNotebook, mockProgress, instance(mockCancellationToken)); // Assert - should proceed despite pending cells assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.calledOnce, + ensureKernelSelectedWithInterpreterStub.calledOnce, true, 'ensureKernelSelected should be called even with pending cells' ); assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.firstCall.args[0], + ensureKernelSelectedWithInterpreterStub.firstCall.args[0], mockNotebook, 'ensureKernelSelected should be called with the notebook' ); @@ -211,16 +193,16 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Arrange when(mockKernelProvider.get(mockNotebook)).thenReturn(undefined); - // Create mock environment - const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); - when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); + // Mock interpreter service to return an active interpreter + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); - // Stub ensureKernelSelectedWithConfiguration to verify it's called - const ensureKernelSelectedWithConfigurationStub = sandbox - .stub(selector, 'ensureKernelSelectedWithConfiguration') + // Stub ensureKernelSelectedWithInterpreter to verify it's called + const ensureKernelSelectedWithInterpreterStub = sandbox + .stub(selector, 'ensureKernelSelectedWithInterpreter') .resolves(); // Act @@ -228,34 +210,34 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Assert - should proceed normally without a kernel assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.calledOnce, + ensureKernelSelectedWithInterpreterStub.calledOnce, true, 'ensureKernelSelected should be called even when no kernel exists' ); assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.firstCall.args[0], + ensureKernelSelectedWithInterpreterStub.firstCall.args[0], mockNotebook, 'ensureKernelSelected should be called with the notebook' ); }); - test('should complete successfully and delegate to ensureKernelSelectedWithConfiguration', async () => { - // This test verifies that ensureKernelSelectedWithConfiguration completes successfully + test('should complete successfully and delegate to ensureKernelSelectedWithInterpreter', async () => { + // This test verifies that ensureKernelSelectedWithInterpreter completes successfully // and delegates kernel setup to ensureKernelSelected // Arrange when(mockKernelProvider.get(mockNotebook)).thenReturn(undefined); - // Create mock environment - const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); - when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); + // Mock interpreter service to return an active interpreter + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); - // Stub ensureKernelSelectedWithConfiguration to verify delegation - const ensureKernelSelectedWithConfigurationStub = sandbox - .stub(selector, 'ensureKernelSelectedWithConfiguration') + // Stub ensureKernelSelectedWithInterpreter to verify delegation + const ensureKernelSelectedWithInterpreterStub = sandbox + .stub(selector, 'ensureKernelSelectedWithInterpreter') .resolves(); // Act @@ -263,30 +245,30 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Assert - method should complete without errors assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.calledOnce, + ensureKernelSelectedWithInterpreterStub.calledOnce, true, - 'ensureKernelSelectedWithConfiguration should be called to set up the new environment' + 'ensureKernelSelectedWithInterpreter should be called to set up the new environment' ); }); - test('should pass cancellation token to ensureKernelSelectedWithConfiguration', async () => { + test('should pass cancellation token to ensureKernelSelectedWithInterpreter', async () => { // This test verifies that rebuildController correctly passes the cancellation token - // to ensureKernelSelectedWithConfiguration, allowing the operation to be cancelled during execution + // to ensureKernelSelectedWithInterpreter, allowing the operation to be cancelled during execution // Arrange when(mockCancellationToken.isCancellationRequested).thenReturn(true); when(mockKernelProvider.get(mockNotebook)).thenReturn(undefined); - // Create mock environment - const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); - when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); + // Mock interpreter service to return an active interpreter + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); - // Stub ensureKernelSelectedWithConfiguration to verify it receives the token - const ensureKernelSelectedWithConfigurationStub = sandbox - .stub(selector, 'ensureKernelSelectedWithConfiguration') + // Stub ensureKernelSelectedWithInterpreter to verify it receives the token + const ensureKernelSelectedWithInterpreterStub = sandbox + .stub(selector, 'ensureKernelSelectedWithInterpreter') .resolves(); // Act @@ -294,67 +276,65 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Assert assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.calledOnce, + ensureKernelSelectedWithInterpreterStub.calledOnce, true, - 'ensureKernelSelectedWithConfiguration should be called once' + 'ensureKernelSelectedWithInterpreter should be called once' ); assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.firstCall.args[0], + ensureKernelSelectedWithInterpreterStub.firstCall.args[0], mockNotebook, 'ensureKernelSelected should be called with the notebook' ); assert.strictEqual( - ensureKernelSelectedWithConfigurationStub.firstCall.args[4], + ensureKernelSelectedWithInterpreterStub.firstCall.args[4], instance(mockCancellationToken), 'ensureKernelSelected should be called with the cancellation token' ); }); - test('should keep the old server handle registered when the environment switch fails', async () => { - // Arrange - old handle already tracked, no new handle registered because setup fails + test('should keep the old server handle registered when the interpreter switch fails', async () => { + // Old handle already tracked; setup fails so no new handle is ever registered. const notebookKey = getNotebookKey(mockNotebook.uri); const oldServerHandle = 'old-server-handle'; registry.set(notebookKey, oldServerHandle); - const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); - when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); - - sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').rejects(new Error('startServer failed')); + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; + when(mockKernelProvider.get(mockNotebook)).thenReturn(undefined); + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); + sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').rejects(new Error('startServer failed')); - // Act await assert.isRejected( selector.rebuildController(mockNotebook, mockProgress, instance(mockCancellationToken)) ); - // Assert - the old handle must remain registered so the selected controller keeps resolving + // The old handle must remain registered so the selected controller keeps resolving. verify(mockServerProvider.unregisterServer(oldServerHandle)).never(); }); - test('should unregister the old server handle after switching to a different environment', async () => { - // Arrange - old handle tracked; successful setup registers a different new handle + test('should unregister the old server handle after switching to a different interpreter', async () => { const notebookKey = getNotebookKey(mockNotebook.uri); const oldServerHandle = 'old-server-handle'; const newServerHandle = 'new-server-handle'; registry.set(notebookKey, oldServerHandle); - const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); - when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); - - // Real setup registers a new handle for the notebook - emulate that side effect - sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').callsFake(async () => { + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; + when(mockKernelProvider.get(mockNotebook)).thenReturn(undefined); + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); + // Real setup registers a new handle for the notebook - emulate that side effect. + sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').callsFake(async () => { registry.set(notebookKey, newServerHandle); }); - // Act await selector.rebuildController(mockNotebook, mockProgress, instance(mockCancellationToken)); - // Assert - only the stale old handle is unregistered; the new one stays verify(mockServerProvider.unregisterServer(oldServerHandle)).once(); verify(mockServerProvider.unregisterServer(newServerHandle)).never(); - - // Assert - the registry ends up tracking the new handle, not the stale one assert.strictEqual( registry.get(notebookKey), newServerHandle, @@ -363,133 +343,64 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); }); - suite('pickEnvironment', () => { - test('should return selected environment when user picks one', async () => { - // Arrange - const notebookUri = Uri.parse('file:///test/notebook.deepnote'); - const mockEnv1 = createMockEnvironment('env-1', 'Environment 1'); - const mockEnv2 = createMockEnvironment('env-2', 'Environment 2'); - const environments = [mockEnv1, mockEnv2]; - - // Mock environment manager - when(mockEnvironmentManager.waitForInitialization()).thenResolve(); - when(mockEnvironmentManager.listEnvironments()).thenReturn(environments); - - // Mock window.showQuickPick to simulate user selecting the first environment - when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenResolve({ - label: mockEnv1.name, - description: mockEnv1.pythonInterpreter.uri.fsPath, - environment: mockEnv1 - } as any); + suite('ensureControllerSelectedForNotebook', () => { + // Every Deepnote controller for one notebook carries the same id, derived from the notebook URI. + const CONTROLLER_ID = 'deepnote-notebook-/test/notebook.deepnote'; - // Act - const result = await selector.pickEnvironment(notebookUri); + function createController(): IVSCodeNotebookController { + const controller = mock(); + when(controller.id).thenReturn(CONTROLLER_ID); + when(controller.connection).thenReturn({ id: CONTROLLER_ID } as KernelConnectionMetadata); + when(controller.controller).thenReturn({ + updateNotebookAffinity: sandbox.stub() + } as unknown as NotebookController); - // Assert - assert.strictEqual(result, mockEnv1, 'Should return the selected environment'); - }); - }); + return instance(controller); + } - suite('ensureEnvironmentConfiguredBeforeExecution', () => { - // The first-run picker is a second environment-selection path alongside the environments view; - // both must report select_environment or the metric only counts users who switch. - function arrangeFirstRunPicker(picked: DeepnoteEnvironment | undefined) { + setup(() => { when(mockCancellationToken.isCancellationRequested).thenReturn(false); - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn(undefined); - when(mockNotebookEnvironmentMapper.setEnvironmentForNotebook(anything(), anything())).thenResolve(); - when(mockEnvironmentManager.waitForInitialization()).thenResolve(); - when(mockEnvironmentManager.listEnvironments()).thenReturn(picked ? [picked] : []); - when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenResolve( - picked ? ({ label: picked.name, environment: picked } as any) : undefined - ); - } + when(mockedVSCodeNamespaces.window.visibleNotebookEditors).thenReturn([ + { notebook: mockNotebook } as NotebookEditor + ]); + when(mockedVSCodeNamespaces.commands.executeCommand('notebook.selectKernel', anything())).thenResolve(); + }); - test('should track select_environment after the execution picker configures an environment', async () => { - const mockEnvironment = createMockEnvironment('env-1', 'Environment 1'); - arrangeFirstRunPicker(mockEnvironment); - const setupStub = sandbox.stub(selector as any, 'setupKernelForEnvironment').resolves(true); + test('selects the kernel when the recorded controller was replaced by one with the same id', async () => { + // The recorded controller was disposed and rebuilt; only object identity tells the two apart, + // and skipping selectKernel here leaves the notebook bound to the dead one. + when(mockControllerRegistration.getSelected(anything())).thenReturn(createController()); - const result = await selector.ensureEnvironmentConfiguredBeforeExecution( + await selector.ensureControllerSelectedForNotebook( mockNotebook, + createController(), instance(mockCancellationToken) ); - assert.isTrue(result, 'Environment should be reported as configured'); - assert.isTrue(setupStub.calledOnce, 'Kernel setup should have run'); - verify(mockTelemetryService.trackEvent(deepEqual({ eventName: 'select_environment' }))).once(); - }); - - test('should not track select_environment when the picker is cancelled or setup fails', async () => { - arrangeFirstRunPicker(undefined); - - assert.isFalse( - await selector.ensureEnvironmentConfiguredBeforeExecution( - mockNotebook, - instance(mockCancellationToken) - ), - 'Cancelling the picker should not configure an environment' - ); - - arrangeFirstRunPicker(createMockEnvironment('env-1', 'Environment 1')); - sandbox.stub(selector as any, 'setupKernelForEnvironment').resolves(false); - - assert.isFalse( - await selector.ensureEnvironmentConfiguredBeforeExecution( - mockNotebook, - instance(mockCancellationToken) - ), - 'Failed setup should not report success' - ); - verify(mockTelemetryService.trackEvent(anything())).never(); + verify(mockedVSCodeNamespaces.commands.executeCommand('notebook.selectKernel', anything())).once(); }); - }); - - suite('ensureKernelSelected', () => { - test('should return false when no environment ID is assigned to the notebook', async () => { - // Mock environment mapper to return null (no environment assigned) - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn(undefined); - - // Stub ensureKernelSelectedWithConfiguration to track if it gets called - const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').resolves(); - // Mock commands.executeCommand - when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + test('skips the kernel picker when the recorded controller is the one being selected', async () => { + const controller = createController(); + when(mockControllerRegistration.getSelected(anything())).thenReturn(controller); - // Act - const result = await selector.ensureKernelSelected( + await selector.ensureControllerSelectedForNotebook( mockNotebook, - mockProgress, + controller, instance(mockCancellationToken) ); - // Assert - assert.strictEqual(result, false, 'Should return false when no environment is assigned'); - assert.strictEqual( - ensureKernelSelectedStub.called, - false, - 'ensureKernelSelectedWithConfiguration should not be called' - ); - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).once(); + verify(mockedVSCodeNamespaces.commands.executeCommand('notebook.selectKernel', anything())).never(); }); + }); - test('should return false and remove mapping when environment is not found', async () => { - // Arrange - const environmentId = 'missing-env-id'; - - // Mock environment mapper to return an ID - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn(environmentId); - - // Mock environment manager to return null (environment not found) - when(mockEnvironmentManager.getEnvironment(environmentId)).thenReturn(undefined); - - // Mock remove environment mapping - when(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).thenResolve(); - - // Stub ensureKernelSelectedWithConfiguration to track if it gets called - const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').resolves(); + suite('ensureKernelSelected', () => { + test('should return false when no active interpreter is found', async () => { + // Mock interpreter service to return undefined (no active interpreter) + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(undefined); - // Mock commands.executeCommand - when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + // Stub ensureKernelSelectedWithInterpreter to track if it gets called + const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').resolves(); // Act const result = await selector.ensureKernelSelected( @@ -499,34 +410,27 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { ); // Assert - assert.strictEqual(result, false, 'Should return false when environment is not found'); + assert.strictEqual(result, false, 'Should return false when no active interpreter is found'); assert.strictEqual( ensureKernelSelectedStub.called, false, - 'ensureKernelSelectedWithConfiguration should not be called' + 'ensureKernelSelectedWithInterpreter should not be called' ); - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).once(); - verify(mockEnvironmentManager.getEnvironment(environmentId)).once(); - verify(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).once(); }); - test('should return true and call ensureKernelSelectedWithConfiguration when environment is found', async () => { + test('should return true and call ensureKernelSelectedWithInterpreter when interpreter is found', async () => { // Arrange const notebookKey = getNotebookKey(mockNotebook.uri); - const environmentId = 'test-env-id'; - const mockEnvironment = createMockEnvironment(environmentId, 'Test Environment'); - - // Mock environment mapper to return an ID - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn(environmentId); - - // Mock environment manager to return the environment - when(mockEnvironmentManager.getEnvironment(environmentId)).thenReturn(mockEnvironment); + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; - // Stub ensureKernelSelectedWithConfiguration to track calls - const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').resolves(); + // Mock interpreter service to return an active interpreter + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); - // Mock commands.executeCommand - when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + // Stub ensureKernelSelectedWithInterpreter to track calls + const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').resolves(); // Act const result = await selector.ensureKernelSelected( @@ -536,74 +440,96 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { ); // Assert - assert.strictEqual(result, true, 'Should return true when environment is found'); + assert.strictEqual(result, true, 'Should return true when interpreter is found'); assert.strictEqual( ensureKernelSelectedStub.calledOnce, true, - 'ensureKernelSelectedWithConfiguration should be called once' + 'ensureKernelSelectedWithInterpreter should be called once' ); // Verify it was called with correct arguments const callArgs = ensureKernelSelectedStub.firstCall.args; assert.strictEqual(callArgs[0], mockNotebook, 'First arg should be notebook'); - assert.strictEqual(callArgs[1], mockEnvironment, 'Second arg should be environment'); + assert.deepStrictEqual(callArgs[1], mockInterpreter, 'Second arg should be interpreter'); assert.strictEqual(callArgs[2], notebookKey, 'Third arg should be notebookKey'); assert.strictEqual(callArgs[3], mockProgress, 'Fourth arg should be progress'); assert.strictEqual(callArgs[4], instance(mockCancellationToken), 'Fifth arg should be token'); - - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).once(); - verify(mockEnvironmentManager.getEnvironment(environmentId)).once(); }); }); - suite('ensureControllerSelectedForNotebook', () => { - // Every Deepnote controller for one notebook carries the same id, derived from the notebook URI. - const CONTROLLER_ID = 'deepnote-notebook-/test/notebook.deepnote'; - - function createController(): IVSCodeNotebookController { - const controller = mock(); - when(controller.id).thenReturn(CONTROLLER_ID); - when(controller.connection).thenReturn({ id: CONTROLLER_ID } as KernelConnectionMetadata); - when(controller.controller).thenReturn({ - updateNotebookAffinity: sandbox.stub() - } as unknown as NotebookController); + suite('ensureEnvironmentConfiguredBeforeExecution', () => { + const nonCancelledToken: CancellationToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => {} }) as any + }; + + test('should reconfigure when active interpreter differs from cached interpreter', async () => { + const notebookKey = mockNotebook.uri.toString(); + const interpreterA: PythonEnvironment = { + id: '/usr/bin/python3.10', + uri: Uri.parse('/usr/bin/python3.10') + }; + const interpreterB: PythonEnvironment = { + id: '/usr/bin/python3.12', + uri: Uri.parse('/usr/bin/python3.12') + }; - return instance(controller); - } + // Prime the internal maps: controller exists for interpreter A + const selectorAny = selector as any; + selectorAny.notebookControllers.set(notebookKey, instance(mockController)); + selectorAny.notebookInterpreterIds.set(notebookKey, interpreterA.id); - setup(() => { - when(mockCancellationToken.isCancellationRequested).thenReturn(false); - when(mockedVSCodeNamespaces.window.visibleNotebookEditors).thenReturn([ - { notebook: mockNotebook } as NotebookEditor - ]); - when(mockedVSCodeNamespaces.commands.executeCommand('notebook.selectKernel', anything())).thenResolve(); - }); + // Active interpreter is now B + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(interpreterB); - test('selects the kernel when the recorded controller was replaced by one with the same id', async () => { - // The recorded controller was disposed and rebuilt; only object identity tells the two apart, - // and skipping selectKernel here leaves the notebook bound to the dead one. - when(mockControllerRegistration.getSelected(anything())).thenReturn(createController()); + // Stub ensureKernelSelectedWithInterpreter to track calls + const ensureStub = sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').resolves(); - await selector.ensureControllerSelectedForNotebook( - mockNotebook, - createController(), - instance(mockCancellationToken) + // withProgress must call through to the task callback + when(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).thenCall( + (_opts: any, task: any) => { + return task({ report: sandbox.stub() }, nonCancelledToken); + } ); - verify(mockedVSCodeNamespaces.commands.executeCommand('notebook.selectKernel', anything())).once(); - }); + // Put a controller in the map so the final check returns true + ensureStub.callsFake(async () => { + selectorAny.notebookControllers.set(notebookKey, instance(mockNewController)); + }); - test('skips the kernel picker when the recorded controller is the one being selected', async () => { - const controller = createController(); - when(mockControllerRegistration.getSelected(anything())).thenReturn(controller); + const result = await selector.ensureEnvironmentConfiguredBeforeExecution(mockNotebook, nonCancelledToken); - await selector.ensureControllerSelectedForNotebook( - mockNotebook, - controller, - instance(mockCancellationToken) + assert.strictEqual(result, true, 'Should return true after reconfiguring'); + assert.strictEqual(ensureStub.calledOnce, true, 'Should call ensureKernelSelectedWithInterpreter'); + assert.deepStrictEqual( + ensureStub.firstCall.args[1], + interpreterB, + 'Should reconfigure with the new interpreter' ); + }); - verify(mockedVSCodeNamespaces.commands.executeCommand('notebook.selectKernel', anything())).never(); + test('should return true immediately when controller exists for the same interpreter', async () => { + const notebookKey = mockNotebook.uri.toString(); + const interpreterA: PythonEnvironment = { + id: '/usr/bin/python3.10', + uri: Uri.parse('/usr/bin/python3.10') + }; + + // Prime the internal maps: controller exists for interpreter A + const selectorAny = selector as any; + selectorAny.notebookControllers.set(notebookKey, instance(mockController)); + selectorAny.notebookInterpreterIds.set(notebookKey, interpreterA.id); + + // Active interpreter is still A + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(interpreterA); + + // Stub ensureKernelSelectedWithInterpreter — should NOT be called + const ensureStub = sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').resolves(); + + const result = await selector.ensureEnvironmentConfiguredBeforeExecution(mockNotebook, nonCancelledToken); + + assert.strictEqual(result, true, 'Should return true (fast path)'); + assert.strictEqual(ensureStub.called, false, 'Should NOT call ensureKernelSelectedWithInterpreter'); }); }); @@ -736,8 +662,21 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // - Validates server cleanup during rebuild // - Ensures old server is unregistered from provider // - // rebuildController drops the old server handle only AFTER setup registers a - // replacement, so a failed/cancelled switch never strands the controller on a dead handle. + // THE ACTUAL IMPLEMENTATION at deepnoteKernelAutoSelector.node.ts:269-291: + // + // // Clear cached state + // this.notebookControllers.delete(notebookKey); + // this.notebookConnectionMetadata.delete(notebookKey); + // + // // Unregister old server + // const oldServerHandle = this.notebookServerHandles.get(notebookKey); + // if (oldServerHandle) { + // this.serverProvider.unregisterServer(oldServerHandle); + // this.notebookServerHandles.delete(notebookKey); + // } + // + // These operations happen BEFORE calling ensureKernelSelected() to create the new controller, + // ensuring clean state for the environment switch. assert.ok(true, 'UT-2 is validated by existing tests and implementation (INV-9)'); }); @@ -838,7 +777,7 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // // CURRENT IMPLEMENTATION BEHAVIOR: // - // 1. If startServer() fails, the error propagates from ensureKernelSelectedWithConfiguration() + // 1. If startServer() fails, the error propagates from ensureKernelSelectedWithInterpreter() // (deepnoteKernelAutoSelector.node.ts:450-467) // // 2. The error is caught and shown to user in the UI layer @@ -874,72 +813,53 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // REAL TDD Tests - These should FAIL if bugs exist suite('Bug Detection: Kernel Selection', () => { - test('BUG-1: Should prefer environment-specific kernel over .env kernel', () => { - // REAL TEST: This will FAIL if the wrong kernel is selected - // - // The selectKernelSpec method is now extracted and testable! + test('Should select the first Python kernel from available specs', () => { + // The selectKernelSpec method selects the first Python kernel available - const envId = 'env123'; const kernelSpecs: IJupyterKernelSpec[] = [ createMockKernelSpec('.env', '.env Python', 'python'), - createMockKernelSpec(`deepnote-${envId}`, 'Deepnote Environment', 'python'), createMockKernelSpec('python3', 'Python 3', 'python') ]; - const selected = selector.selectKernelSpec(kernelSpecs, envId); + const selected = selector.selectKernelSpec(kernelSpecs); - // CRITICAL ASSERTION: Should select environment-specific kernel, NOT .env - assert.strictEqual( - selected?.name, - `deepnote-${envId}`, - `BUG DETECTED: Selected "${selected?.name}" instead of "deepnote-${envId}"! This would use wrong environment.` - ); + // Should select the first Python kernel + assert.strictEqual(selected.language, 'python', 'Should select a Python kernel'); + assert.strictEqual(selected.name, '.env', 'Should select the first Python kernel'); }); - test('BUG-1b: Current implementation falls back to Python kernel (documents expected behavior)', () => { - // This test documents that the current implementation DOES have fallback logic - // - // EXPECTED BEHAVIOR (current): Fall back to generic Python kernel when env-specific kernel not found - // This is a design decision - we don't want to block users if the environment-specific kernel isn't ready yet + test('Should fall back to python3 named kernel when no python language kernel exists first', () => { + // Documents fallback behavior - finds python3 by name if no python language match - const envId = 'env123'; const kernelSpecs: IJupyterKernelSpec[] = [ - createMockKernelSpec('.env', '.env Python', 'python'), + createMockKernelSpec('javascript', 'JavaScript', 'javascript'), createMockKernelSpec('python3', 'Python 3', 'python') ]; - // Should fall back to a Python kernel (this is the current behavior) - const selected = selector.selectKernelSpec(kernelSpecs, envId); + const selected = selector.selectKernelSpec(kernelSpecs); - // Should have selected a fallback kernel (either .env or python3) - assert.ok(selected, 'Should select a fallback kernel'); - assert.strictEqual(selected.language, 'python', 'Fallback should be a Python kernel'); + assert.strictEqual(selected.name, 'python3', 'Should find python3 kernel'); }); - test('Kernel selection: Should find environment-specific kernel when it exists', () => { - const envId = 'my-env'; + test('Kernel selection: Should fall back to first available kernel when no Python kernel exists', () => { const kernelSpecs: IJupyterKernelSpec[] = [ - createMockKernelSpec('python3', 'Python 3', 'python'), - createMockKernelSpec(`deepnote-${envId}`, 'My Environment', 'python') + createMockKernelSpec('javascript', 'JavaScript', 'javascript'), + createMockKernelSpec('r', 'R', 'r') ]; - const selected = selector.selectKernelSpec(kernelSpecs, envId); + const selected = selector.selectKernelSpec(kernelSpecs); - assert.strictEqual(selected?.name, `deepnote-${envId}`); + assert.strictEqual(selected.name, 'javascript', 'Should fall back to first available kernel'); }); - test('Kernel selection: Should fall back to python3 when env kernel missing', () => { - // Documents current fallback behavior - falls back to python3 when env kernel missing - const envId = 'my-env'; - const kernelSpecs: IJupyterKernelSpec[] = [ - createMockKernelSpec('python3', 'Python 3', 'python'), - createMockKernelSpec('javascript', 'JavaScript', 'javascript') - ]; - - // Should fall back to python3 (current behavior) - const selected = selector.selectKernelSpec(kernelSpecs, envId); + test('Kernel selection: Should throw when no kernel specs are available', () => { + const kernelSpecs: IJupyterKernelSpec[] = []; - assert.strictEqual(selected.name, 'python3', 'Should fall back to python3'); + assert.throws( + () => selector.selectKernelSpec(kernelSpecs), + /No kernel specs available/, + 'Should throw when no kernel specs are available' + ); }); }); @@ -970,9 +890,14 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // REAL TEST: This will FAIL if disposal happens too early // // Setup: Create a scenario where we have an old controller and create a new one - const baseFileUri = mockNotebook.uri.with({ query: '', fragment: '' }); + // const baseFileUri = mockNotebook.uri.with({ query: '', fragment: '' }); // const notebookKey = baseFileUri.fsPath; - const newEnv = createMockEnvironment('env-new', 'New Environment', true); + // Mock interpreter service to return an active interpreter + const mockInterpreter: PythonEnvironment = { + id: '/usr/bin/python3', + uri: Uri.parse('/usr/bin/python3') + }; + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(mockInterpreter); // Track call order const callOrder: string[] = []; @@ -996,8 +921,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { when(newController.controller).thenReturn({} as any); // Setup mocks - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri)).thenReturn('env-new'); - when(mockEnvironmentManager.getEnvironment('env-new')).thenReturn(newEnv); when(mockPythonExtensionChecker.isPythonExtensionInstalled).thenReturn(true); // Mock controller registration to track when new controller is added @@ -1154,103 +1077,119 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); }); - suite('Placeholder controller execution', () => { - function createPlaceholder() { - const placeholder = { - supportsExecutionOrder: false, - supportedLanguages: [] as string[], - updateNotebookAffinity: sandbox.stub(), - dispose: sandbox.stub(), - createNotebookCellExecution: sandbox.stub() - } as unknown as NotebookController; - - when( - mockedVSCodeNamespaces.notebooks!.createNotebookController(anything(), anything(), anything()) - ).thenReturn(placeholder); + suite('clearControllerForEnvironment', () => { + test('should unselect and clean up when tracked controller matches selected controller', () => { + const notebookKey = mockNotebook.uri.toString(); - const onDidCloseNotebookDocument = new EventEmitter(); - when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn( - onDidCloseNotebookDocument.event - ); - - const internals = selector as unknown as { - createPlaceholderController(notebook: NotebookDocument): NotebookController; - }; - - internals.createPlaceholderController(mockNotebook); - - return placeholder; - } + // Set up a tracked controller in the internal map + const trackedController = mock(); + when(trackedController.id).thenReturn('deepnote-notebook-123'); + when(trackedController.connection).thenReturn({ + kind: 'startUsingDeepnoteKernel' + } as any); + const mockNativeController = { + updateNotebookAffinity: sandbox.stub() + } as unknown as NotebookController; + when(trackedController.controller).thenReturn(mockNativeController); - const agentCell = createMockCell({ index: 0, metadata: { __deepnotePocket: { type: 'agent' } } }); - const codeCell = createMockCell({ index: 1 }); + const selectorAny = selector as any; + selectorAny.notebookControllers.set(notebookKey, instance(trackedController)); + selectorAny.notebookInterpreterIds.set(notebookKey, '/usr/bin/python3'); + selectorAny.notebookConnectionMetadata.set(notebookKey, {} as any); - test('configures the environment and executes nothing', async () => { - when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); - const placeholder = createPlaceholder(); - const ensureEnvironment = sandbox - .stub(selector, 'ensureEnvironmentConfiguredBeforeExecution') - .resolves(true); + // Selected controller is the same one we tracked + when(mockControllerRegistration.getSelected(mockNotebook)).thenReturn(instance(trackedController)); - await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); + selector.clearControllerForEnvironment(mockNotebook, 'env-uuid-123'); - assert.isTrue(ensureEnvironment.calledOnce, 'should prompt for an environment'); assert.isTrue( - (placeholder.createNotebookCellExecution as sinon.SinonStub).notCalled, - 'placeholder must not create executions' + (mockNativeController.updateNotebookAffinity as sinon.SinonStub).calledOnce, + 'Should have called updateNotebookAffinity' + ); + // Verify tracking state is cleaned up + assert.isFalse(selectorAny.notebookControllers.has(notebookKey), 'Should remove from notebookControllers'); + assert.isFalse( + selectorAny.notebookInterpreterIds.has(notebookKey), + 'Should remove from notebookInterpreterIds' + ); + assert.isFalse( + selectorAny.notebookConnectionMetadata.has(notebookKey), + 'Should remove from notebookConnectionMetadata' ); - verify(mockKernelProvider.getOrCreate(anything(), anything())).never(); }); - test('does nothing at all in an untrusted workspace', async () => { - when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(false); - const placeholder = createPlaceholder(); - const ensureEnvironment = sandbox - .stub(selector, 'ensureEnvironmentConfiguredBeforeExecution') - .resolves(true); + test('should NOT unselect when notebook has no tracked controller', () => { + // notebookControllers map is empty — we didn't set up this notebook + const trackedController = mock(); + const mockNativeController = { + updateNotebookAffinity: sandbox.stub() + } as unknown as NotebookController; + when(trackedController.controller).thenReturn(mockNativeController); + when(mockControllerRegistration.getSelected(mockNotebook)).thenReturn(instance(trackedController)); - await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); + selector.clearControllerForEnvironment(mockNotebook, 'env-uuid-123'); - assert.isTrue(ensureEnvironment.notCalled, 'should not prompt in an untrusted workspace'); + assert.isFalse( + (mockNativeController.updateNotebookAffinity as sinon.SinonStub).called, + 'Should NOT have called updateNotebookAffinity when we have no tracked controller' + ); }); - }); - /** - * Every exec in the installer can now be killed, so cancellation surfaces here immediately - * instead of after pip has finished anyway. A user-initiated Stop is not a failure and - * must not raise the error UI - but a genuine failure still has to. - */ - suite('cancellation is not reported as a failure', () => { - const toolkitMissing = () => new DeepnoteToolkitMissingError('/usr/bin/python3', '/fake/venv'); + test('should NOT unselect when selected controller differs from tracked controller', () => { + const notebookKey = mockNotebook.uri.toString(); - function chooseInstall(): void { - when( - mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything(), anything()) - ).thenResolve('Install' as never); - } + // Track controller A + const controllerA = mock(); + when(controllerA.id).thenReturn('deepnote-notebook-A'); + const selectorAny = selector as any; + selectorAny.notebookControllers.set(notebookKey, instance(controllerA)); - test('cancelling the toolkit install does not show an error message', async () => { - chooseInstall(); - when(mockToolkitInstaller.installToolkitInExistingVenv(anything(), anything())).thenReject( - new CancellationError() - ); + // But VS Code has controller B selected (different id) + const controllerB = mock(); + when(controllerB.id).thenReturn('deepnote-notebook-B'); + const mockNativeController = { + updateNotebookAffinity: sandbox.stub() + } as unknown as NotebookController; + when(controllerB.controller).thenReturn(mockNativeController); + when(mockControllerRegistration.getSelected(mockNotebook)).thenReturn(instance(controllerB)); - await selector.handleKernelSelectionError(toolkitMissing(), mockNotebook); + selector.clearControllerForEnvironment(mockNotebook, 'env-uuid-123'); - verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).never(); + assert.isFalse( + (mockNativeController.updateNotebookAffinity as sinon.SinonStub).called, + 'Should NOT unselect a controller we do not own' + ); }); - test('a failed toolkit install still shows an error message', async () => { - chooseInstall(); - when(mockToolkitInstaller.installToolkitInExistingVenv(anything(), anything())).thenReject( - new Error('pip exited with code 1') - ); + test('should NOT unselect when selected controller is not a Deepnote kernel', () => { + const notebookKey = mockNotebook.uri.toString(); - await selector.handleKernelSelectionError(toolkitMissing(), mockNotebook); + // Track a controller + const trackedController = mock(); + when(trackedController.id).thenReturn('deepnote-notebook-123'); + when(trackedController.connection).thenReturn({ + kind: 'startUsingLocalKernelSpec' + } as any); + const mockNativeController = { + updateNotebookAffinity: sandbox.stub() + } as unknown as NotebookController; + when(trackedController.controller).thenReturn(mockNativeController); + + const selectorAny = selector as any; + selectorAny.notebookControllers.set(notebookKey, instance(trackedController)); - verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + when(mockControllerRegistration.getSelected(mockNotebook)).thenReturn(instance(trackedController)); + + selector.clearControllerForEnvironment(mockNotebook, 'env-uuid-123'); + + assert.isFalse( + (mockNativeController.updateNotebookAffinity as sinon.SinonStub).called, + 'Should NOT unselect a non-Deepnote kernel' + ); }); + }); + suite('cancellation is not reported as a failure', () => { test('a cancelled kernel selection does not show an error message', async () => { await selector.handleKernelSelectionError(new CancellationError(), mockNotebook); @@ -1265,37 +1204,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); }); -/** - * Helper function to create mock environments - */ -function createMockEnvironment(id: string, name: string, hasServer: boolean = false): DeepnoteEnvironment { - const mockPythonInterpreter: PythonEnvironment = { - id: `/usr/bin/python3`, - uri: Uri.parse(`/usr/bin/python3`) - }; - - return { - id, - name, - description: `Test environment ${name}`, - pythonInterpreter: mockPythonInterpreter, - venvPath: Uri.file(`/test/venvs/${id}`), - managedVenv: true, - packages: [], - createdAt: new Date(), - lastUsedAt: new Date(), - serverInfo: hasServer - ? { - url: `http://localhost:8888`, - jupyterPort: 8888, - lspPort: 8889, - token: 'test-token', - process: createMockChildProcess() - } - : undefined - }; -} - /** * Helper function to create mock kernel specs */ diff --git a/src/platform/api/pythonApi.ts b/src/platform/api/pythonApi.ts index 3e68cc60f6..18e620063f 100644 --- a/src/platform/api/pythonApi.ts +++ b/src/platform/api/pythonApi.ts @@ -49,6 +49,10 @@ import { trackInterpreterDiscovery, trackPythonExtensionActivation } from '../.. import { findPythonEnvBelongingToFolder } from '../../notebooks/controllers/preferredKernelConnectionService.node'; import { DisposableMap } from '../common/utils/lifecycle'; +// The Python extension completes its Jupyter handshake asynchronously; if the callback has not +// landed by then it never will (see failApiIfHandshakeNeverLands). +const PYTHON_API_HANDSHAKE_TIMEOUT = 5_000; + export function deserializePythonEnvironment( pythonVersion: Partial | undefined, pythonEnvId: string @@ -161,6 +165,25 @@ export class OldPythonApiProvider implements IPythonApiProvider { return extension?.exports; } + /** + * The Python extension ends the handshake by calling `registerPythonApi` on the extension it + * knows as `ms-toolsai.jupyter`. This fork ships under a different id, so that callback never + * arrives and every `getApi()` awaiter would hang forever. Fail the promise instead — callers + * already degrade to unactivated execution when activation variables are unavailable. + */ + private failApiIfHandshakeNeverLands() { + const timer = setTimeout(() => { + if (!this.api.resolved && !this.api.rejected) { + logger.warn('Python extension did not complete the Jupyter API handshake; continuing without it'); + this.api.reject(new PythonExtensionApiNotExportedError()); + } + }, PYTHON_API_HANDSHAKE_TIMEOUT); + + // Nothing may be awaiting the promise at rejection time. + this.api.promise.catch(noop); + this.disposables.push({ dispose: () => clearTimeout(timer) }); + } + public setApi(api: PythonApi): void { // Never allow accessing python API (we don't want to ever use the API and run code in untrusted API). // Don't assume Python API will always be disabled in untrusted workspaces. @@ -219,6 +242,7 @@ export class OldPythonApiProvider implements IPythonApiProvider { this.api.reject(new PythonExtensionApiNotExportedError()); } else { pythonExtension.exports.jupyter.registerHooks(); + this.failApiIfHandshakeNeverLands(); } this._pythonExtensionHooked.resolve(); } diff --git a/src/platform/interpreter/installer/pipInstaller.node.ts b/src/platform/interpreter/installer/pipInstaller.node.ts index ae53dde627..6548baf90e 100644 --- a/src/platform/interpreter/installer/pipInstaller.node.ts +++ b/src/platform/interpreter/installer/pipInstaller.node.ts @@ -15,6 +15,8 @@ import { Environment } from '@vscode/python-extension'; import { getEnvironmentType } from '../helpers'; import { workspace } from 'vscode'; +import { DEEPNOTE_TOOLKIT_VERSION } from '../../../kernels/deepnote/types'; + /** * Installer for pip. Default installer for most everything. */ @@ -85,8 +87,14 @@ export class PipInstaller extends ModuleInstaller { if (getEnvironmentType(interpreter) === EnvironmentType.Unknown) { args.push('--user'); } + // deepnote_toolkit's import name differs from the pip package name (deepnote-toolkit[server]) + const pipPackageName = + moduleName === translateProductToModule(Product.deepnoteToolkit) + ? `deepnote-toolkit[server]==${DEEPNOTE_TOOLKIT_VERSION}` + : moduleName; + return { - args: ['-m', 'pip', ...args, moduleName].concat(getPinnedPackages('pip', moduleName)) + args: ['-m', 'pip', ...args, pipPackageName].concat(getPinnedPackages('pip', moduleName)) }; } private isPipAvailable(interpreter: PythonEnvironment | Environment): Promise { diff --git a/src/platform/interpreter/installer/productInstaller.node.ts b/src/platform/interpreter/installer/productInstaller.node.ts index bf59cc22f3..e9e8dac97a 100644 --- a/src/platform/interpreter/installer/productInstaller.node.ts +++ b/src/platform/interpreter/installer/productInstaller.node.ts @@ -7,10 +7,12 @@ import { ProductNames } from './productNames'; import { IInstallationChannelManager, IInstaller, + IModuleInstaller, InstallerResponse, IProductPathService, IProductService, ModuleInstallFlags, + ModuleInstallerType, Product, ProductType } from './types'; @@ -93,8 +95,20 @@ export class DataScienceInstaller { installPipIfRequired?: boolean, silent?: boolean ): Promise { - const channels = this.serviceContainer.get(IInstallationChannelManager); - const installer = await channels.getInstallationChannel(product, interpreter); + let installer: IModuleInstaller | undefined; + + // deepnote-toolkit is PyPI-only with pip-specific [server] extras syntax, + // so always use PipInstaller regardless of environment type. + // We bypass getInstallationChannels() because it filters by isSupported(), + // and PipInstaller.isSupported() rejects Conda/Pipenv/Poetry interpreters. + if (product === Product.deepnoteToolkit) { + const allInstallers = this.serviceContainer.getAll(IModuleInstaller); + installer = allInstallers.find((i) => i.type === ModuleInstallerType.Pip); + } else { + const channels = this.serviceContainer.get(IInstallationChannelManager); + installer = await channels.getInstallationChannel(product, interpreter); + } + if (!installer) { return InstallerResponse.Ignore; } diff --git a/src/platform/interpreter/installer/productInstaller.unit.test.ts b/src/platform/interpreter/installer/productInstaller.unit.test.ts index dbe9dba1db..27daf48151 100644 --- a/src/platform/interpreter/installer/productInstaller.unit.test.ts +++ b/src/platform/interpreter/installer/productInstaller.unit.test.ts @@ -206,4 +206,63 @@ suite('DataScienceInstaller install', async () => { const result = await dataScienceInstaller.install(Product.ipykernel, testEnvironment, tokenSource); expect(result).to.equal(InstallerResponse.Installed, 'Should be Installed'); }); + + test('Will use pip for deepnoteToolkit even on Conda interpreter (bypasses isSupported filter)', async () => { + const testEnvironment: PythonEnvironment = { + id: interpreterPath.fsPath, + uri: interpreterPath + }; + + // Create a pip installer mock + const pipInstaller = TypeMoq.Mock.ofType(); + pipInstaller.setup((c) => c.type).returns(() => ModuleInstallerType.Pip); + pipInstaller + .setup((c) => + c.installModule( + TypeMoq.It.isValue(Product.deepnoteToolkit), + TypeMoq.It.isValue(testEnvironment), + TypeMoq.It.isAny(), + TypeMoq.It.isAny(), + TypeMoq.It.isAny() + ) + ) + .returns(() => Promise.resolve()); + pipInstaller.setup((p) => (p as any).then).returns(() => undefined); + + // Create a conda installer mock (would normally be selected for Conda envs) + const condaInstaller = TypeMoq.Mock.ofType(); + condaInstaller.setup((c) => c.type).returns(() => ModuleInstallerType.Conda); + + // serviceContainer.getAll returns both installers — the code must pick pip + serviceContainer + .setup((c) => c.getAll(TypeMoq.It.isValue(IModuleInstaller))) + .returns(() => [condaInstaller.object, pipInstaller.object]); + + const result = await dataScienceInstaller.install(Product.deepnoteToolkit, testEnvironment, tokenSource); + expect(result).to.equal(InstallerResponse.Installed, 'Should be Installed via pip'); + + // Verify pip was called, not conda + pipInstaller.verify( + (c) => + c.installModule( + TypeMoq.It.isValue(Product.deepnoteToolkit), + TypeMoq.It.isValue(testEnvironment), + TypeMoq.It.isAny(), + TypeMoq.It.isAny(), + TypeMoq.It.isAny() + ), + TypeMoq.Times.once() + ); + condaInstaller.verify( + (c) => + c.installModule( + TypeMoq.It.isAny(), + TypeMoq.It.isAny(), + TypeMoq.It.isAny(), + TypeMoq.It.isAny(), + TypeMoq.It.isAny() + ), + TypeMoq.Times.never() + ); + }); }); diff --git a/src/platform/interpreter/installer/productNames.ts b/src/platform/interpreter/installer/productNames.ts index 5c9eaa559b..246432faa1 100644 --- a/src/platform/interpreter/installer/productNames.ts +++ b/src/platform/interpreter/installer/productNames.ts @@ -12,3 +12,4 @@ ProductNames.set(Product.kernelspec, 'kernelspec'); ProductNames.set(Product.pandas, 'pandas'); ProductNames.set(Product.pip, 'pip'); ProductNames.set(Product.ensurepip, 'ensurepip'); +ProductNames.set(Product.deepnoteToolkit, 'deepnote-toolkit'); diff --git a/src/platform/interpreter/installer/productService.node.ts b/src/platform/interpreter/installer/productService.node.ts index 5059e2b2a0..8af4dfb549 100644 --- a/src/platform/interpreter/installer/productService.node.ts +++ b/src/platform/interpreter/installer/productService.node.ts @@ -20,6 +20,7 @@ export class ProductService implements IProductService { this.ProductTypes.set(Product.pandas, ProductType.DataScience); this.ProductTypes.set(Product.pip, ProductType.DataScience); this.ProductTypes.set(Product.ensurepip, ProductType.DataScience); + this.ProductTypes.set(Product.deepnoteToolkit, ProductType.DataScience); } public getProductType(product: Product): ProductType { return this.ProductTypes.get(product)!; diff --git a/src/platform/interpreter/installer/types.ts b/src/platform/interpreter/installer/types.ts index bdb5bfee39..550886232c 100644 --- a/src/platform/interpreter/installer/types.ts +++ b/src/platform/interpreter/installer/types.ts @@ -21,7 +21,8 @@ export enum Product { nbconvert = 22, pandas = 23, pip = 27, - ensurepip = 28 + ensurepip = 28, + deepnoteToolkit = 29 } export enum ProductInstallStatus { diff --git a/src/platform/interpreter/installer/utils.ts b/src/platform/interpreter/installer/utils.ts index b85bd38d76..2e80f6d481 100644 --- a/src/platform/interpreter/installer/utils.ts +++ b/src/platform/interpreter/installer/utils.ts @@ -23,6 +23,8 @@ export function translateProductToModule(product: Product): string { return 'pip'; case Product.ensurepip: return 'ensurepip'; + case Product.deepnoteToolkit: + return 'deepnote_toolkit'; default: { throw new WrappedError( `Product ${product} cannot be installed as a Python Module.`, diff --git a/test/e2e/fixtures/interpreter-kernel.deepnote b/test/e2e/fixtures/interpreter-kernel.deepnote new file mode 100644 index 0000000000..ecb1a47e88 --- /dev/null +++ b/test/e2e/fixtures/interpreter-kernel.deepnote @@ -0,0 +1,23 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-interpreter-kernel-project + name: E2E Interpreter Kernel + notebooks: + - id: e2e-interpreter-kernel-notebook + name: Interpreter Kernel + blocks: + - id: e2e-interpreter-kernel-block + blockGroup: e2e-interpreter-kernel-group + type: code + content: |- + import sys + print("interpreter-kernel-ok") + print(sys.prefix) + sortingKey: a0 + metadata: {} + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/helpers/notifications.ts b/test/e2e/helpers/notifications.ts index 7f5ccadded..b5d3ca23c4 100644 --- a/test/e2e/helpers/notifications.ts +++ b/test/e2e/helpers/notifications.ts @@ -60,3 +60,30 @@ export async function waitForNotification( return undefined; } } + +/** + * Waits until no visible notification matches `pattern` any more. A progress notification is + * removed when its operation settles, so this gates on "that work finished" rather than on a + * fixed sleep. + */ +export async function waitForNotificationToClear(pattern: RegExp, timeout: number): Promise { + await VSBrowser.instance.driver.wait( + async () => { + const notifications = await new Workbench().getNotifications().catch((error) => { + console.warn('[deepnote-e2e] get notifications:', error); + + return [] as Notification[]; + }); + for (const notification of notifications) { + const message = await notification.getMessage().catch(() => ''); + if (pattern.test(message)) { + return false; + } + } + + return true; + }, + timeout, + `timed out waiting for notifications matching ${pattern} to clear` + ); +} diff --git a/test/e2e/suite/interpreterKernel.e2e.test.ts b/test/e2e/suite/interpreterKernel.e2e.test.ts new file mode 100644 index 0000000000..156adb317b --- /dev/null +++ b/test/e2e/suite/interpreterKernel.e2e.test.ts @@ -0,0 +1,163 @@ +/** + * End-to-end UI test for kernel setup WITHOUT Deepnote environments. + * + * This is the flow the extension uses now that environments are gone: the kernel is built from the + * workspace's *active Python interpreter*, and deepnote-toolkit is installed into that interpreter + * through the Python extension's installer infrastructure (`IInstaller`) — the same mechanism the + * Jupyter extension uses for its own missing dependencies: + * 1. open a one-notebook `.deepnote` file against a workspace whose active interpreter is a bare venv + * 2. the auto-selector picks that interpreter — no environment is created and no picker appears + * 3. deepnote-toolkit is missing, so an "Installing deepnote-toolkit" progress notification shows + * 4. once it installs, the server starts and the kernel controller is bound + * 5. run the cell and assert the rendered stdout + * + * The cell prints `sys.prefix`, so the output proves the kernel really ran inside the venv this test + * created rather than in a Deepnote-managed environment. The suite also asserts the toolkit landed in + * that same venv, which is the load-bearing difference from the environment-based flow. + * + * Prerequisites: + * - The Python extension (`ms-python.python`) must be installed in the test instance + * (`npm run setup:e2e:deps`). + * - `python3` must be on PATH and able to create a venv (CI installs `python3.12-venv`). + * - Network access: the toolkit is installed from PyPI on first kernel start, which is slow. + */ + +import { expect } from 'chai'; +import { execFileSync } from 'child_process'; +import * as fs from 'fs'; +import * as path from 'path'; +import { EditorView, VSBrowser, WebView } from 'vscode-extension-tester'; + +import { + FIRST_RUN_OUTPUT_TIMEOUT, + KERNEL_CONNECT_TIMEOUT, + SUITE_TIMEOUT, + WORKBENCH_TIMEOUT, + copyFixtureToTempDir, + openFolderViaDialog, + openWorkspaceFile, + runOnceAndAwaitOutput, + waitForNotification, + waitForNotificationToClear +} from '../helpers'; + +const NOTEBOOK_FILE_NAME = 'interpreter-kernel.deepnote'; +const EXPECTED_OUTPUT = 'interpreter-kernel-ok'; + +// The install toast is observed opportunistically; the venv check below is the real gate, so a +// missed toast must not cost the suite a full kernel-connect timeout. +const TOOLKIT_NOTIFICATION_TIMEOUT = 90_000; + +/** Path to the interpreter inside a venv, for the platform the test is running on. */ +function venvPython(venvDir: string): string { + return process.platform === 'win32' + ? path.join(venvDir, 'Scripts', 'python.exe') + : path.join(venvDir, 'bin', 'python'); +} + +/** True when `deepnote_toolkit` imports in the given interpreter. */ +function isToolkitInstalled(python: string): boolean { + try { + execFileSync(python, ['-c', 'import deepnote_toolkit'], { stdio: 'ignore' }); + + return true; + } catch { + return false; + } +} + +describe('Deepnote E2E — run on the active interpreter (no Deepnote environment)', function () { + this.timeout(SUITE_TIMEOUT); + + let cleanupTempDir: (() => void) | undefined; + let venvDir: string; + let interpreter: string; + + before(async function () { + const { cleanup, tempDir } = copyFixtureToTempDir(NOTEBOOK_FILE_NAME); + cleanupTempDir = cleanup; + + // A throwaway venv is what makes this test deterministic: it is guaranteed not to have + // deepnote-toolkit, so the install path runs on every execution rather than only on a + // machine that happens to be missing the package. + venvDir = path.join(tempDir, '.venv'); + execFileSync('python3', ['-m', 'venv', venvDir], { stdio: 'inherit' }); + interpreter = venvPython(venvDir); + + expect(isToolkitInstalled(interpreter)).to.equal( + false, + 'precondition: the fresh venv must not already provide deepnote-toolkit' + ); + + // Pin the workspace's interpreter so the auto-selector resolves this venv and not whatever + // the Python extension would otherwise discover on the machine. + const vscodeDir = path.join(tempDir, '.vscode'); + fs.mkdirSync(vscodeDir, { recursive: true }); + fs.writeFileSync( + path.join(vscodeDir, 'settings.json'), + JSON.stringify({ 'python.defaultInterpreterPath': interpreter }, undefined, 4) + ); + + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + + // Opening the folder reloads the window, so the notebook is opened in the test body — that + // keeps the install notification, which fires during the open, inside the assertions. + await openFolderViaDialog(tempDir); + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + }); + + after(async function () { + await new WebView().switchBack().catch((error) => { + console.warn('[deepnote-e2e] switch back from webview during cleanup:', error); + }); + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[deepnote-e2e] close all editors during cleanup:', error); + }); + + try { + cleanupTempDir?.(); + } catch (error) { + console.warn('[deepnote-e2e] remove temp workspace dir during cleanup:', error); + } + }); + + it('installs deepnote-toolkit into the active interpreter, then runs the cell', async function () { + await openWorkspaceFile(NOTEBOOK_FILE_NAME); + + await VSBrowser.instance.driver.wait( + async () => (await new EditorView().getOpenEditorTitles()).some((t) => t.includes(NOTEBOOK_FILE_NAME)), + WORKBENCH_TIMEOUT, + 'Deepnote notebook editor did not open' + ); + + // Opening the notebook auto-selects the kernel, which finds the toolkit missing and installs + // it — no environment is created and the user is never asked to pick one. The progress toast + // is the visible half of that, but it is transient (and absent on a retry, where the toolkit + // is already installed), so it is observed best-effort; the durable assertions are below. + await waitForNotification(/Installing deepnote-toolkit/i, TOOLKIT_NOTIFICATION_TIMEOUT, false); + + // The load-bearing gate: the install landed in the active interpreter. Polling the venv is + // race-free, unlike matching a toast that may already have gone. + await VSBrowser.instance.driver.wait( + () => isToolkitInstalled(interpreter), + KERNEL_CONNECT_TIMEOUT, + 'deepnote-toolkit was never installed into the active interpreter' + ); + + // Waiting for the auto-select toast to be *gone* gates "Run All" on a bound kernel without + // depending on catching it while it is shown. + await waitForNotificationToClear(/Auto-selecting Deepnote kernel/i, KERNEL_CONNECT_TIMEOUT); + + const renderedOutput = await runOnceAndAwaitOutput( + NOTEBOOK_FILE_NAME, + EXPECTED_OUTPUT, + FIRST_RUN_OUTPUT_TIMEOUT + ); + + expect(renderedOutput).to.contain(EXPECTED_OUTPUT); + + // The cell printed sys.prefix: the kernel must be the venv this test created, which is what + // separates "active interpreter" from the old Deepnote-managed environment. + expect(renderedOutput).to.contain(venvDir); + }); +});