diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 112f7f0e09..ce2d442b6b 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, optional } 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'; @@ -24,8 +24,10 @@ import { logger } from '../../platform/logging'; import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; -import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; +import { DeepnoteServerInfo, IDeepnoteServerStarter } from './types'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; +import { getCachedEnvironment } from '../../platform/interpreter/helpers'; +import { IInstaller, InstallerResponse, Product } from '../../platform/interpreter/installer/types'; const MAX_OUTPUT_TRACKING_LENGTH = 5000; const SERVER_STARTUP_TIMEOUT_MS = 120_000; @@ -48,7 +50,7 @@ type PendingOperation = }; interface ProjectContext { - environmentId: string; + interpreterId: string; serverInfo: DeepnoteServerInfo | null; } @@ -71,7 +73,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) * - SQL integration env var injection (via ServerOptions.env) * - Lock file creation (after start, using returned PID) * - Output channel logging (via process stdout/stderr streams) @@ -231,37 +222,52 @@ 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 = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); + const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, interpreterId, token); // Initialize output tracking for error reporting this.serverOutputByFile.set(fileKey, { stdout: '', stderr: '' }); @@ -269,7 +275,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 @@ -280,7 +286,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension throw new DeepnoteServerStartupError( interpreter.uri.fsPath, - serverInfo?.jupyterPort ?? 0, + 0, 'unknown', capturedOutput?.stdout || '', capturedOutput?.stderr || '', @@ -307,6 +313,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 f90e3a8ec1..2b7b37c2df 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -1,12 +1,13 @@ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; import { anything, instance, mock, when } from 'ts-mockito'; +import { EventEmitter } from 'vscode'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; -import { IDeepnoteToolkitInstaller } from './types'; +import { IInstaller, InstallerResponse } from '../../platform/interpreter/installer/types'; import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; /** @@ -19,7 +20,7 @@ import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnot suite('DeepnoteServerStarter', () => { let serverStarter: DeepnoteServerStarter; let mockProcessServiceFactory: IProcessServiceFactory; - let mockToolkitInstaller: IDeepnoteToolkitInstaller; + let mockInstaller: IInstaller; let mockAgentSkillsManager: DeepnoteAgentSkillsManager; let mockOutputChannel: IOutputChannel; let mockAsyncRegistry: IAsyncDisposableRegistry; @@ -32,7 +33,7 @@ suite('DeepnoteServerStarter', () => { setup(() => { mockProcessServiceFactory = mock(); - mockToolkitInstaller = mock(); + mockInstaller = mock(); mockAgentSkillsManager = mock(); mockOutputChannel = mock(); mockAsyncRegistry = mock(); @@ -40,10 +41,13 @@ 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); serverStarter = new DeepnoteServerStarter( instance(mockProcessServiceFactory), - instance(mockToolkitInstaller), + instance(mockInstaller), instance(mockAgentSkillsManager), instance(mockOutputChannel), instance(mockAsyncRegistry), @@ -62,7 +66,7 @@ suite('DeepnoteServerStarter', () => { // Create a starter without SQL provider const starterWithoutSql = new DeepnoteServerStarter( instance(mockProcessServiceFactory), - instance(mockToolkitInstaller), + instance(mockInstaller), instance(mockAgentSkillsManager), instance(mockOutputChannel), instance(mockAsyncRegistry) @@ -85,7 +89,7 @@ suite('DeepnoteServerStarter', () => { const starterWithCancelledSql = new DeepnoteServerStarter( instance(mockProcessServiceFactory), - instance(mockToolkitInstaller), + instance(mockInstaller), instance(mockAgentSkillsManager), instance(mockOutputChannel), instance(mockAsyncRegistry), diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index a0c17a31ba..444d3e6651 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 63e7129200..e7196633ca 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -8,32 +8,23 @@ import { CancellationToken, CancellationTokenSource, Disposable, - 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 + IDeepnoteServerStarter } from '../../kernels/deepnote/types'; import { createJupyterConnectionInfo } from '../../kernels/jupyter/jupyterUtils'; import { JupyterLabHelper } from '../../kernels/jupyter/session/jupyterLabHelper'; @@ -51,7 +42,8 @@ import { getDisplayPath } from '../../platform/common/platform/fs-paths.node'; import { IConfigurationService, IDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { disposeAsync } from '../../platform/common/utils'; 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'; @@ -60,10 +52,6 @@ import { IDeepnoteInitNotebookRunner } from './deepnoteInitNotebookRunner.node'; 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; - /** * Automatically selects and starts Deepnote kernel for .deepnote notebooks */ @@ -73,10 +61,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(); // Track server handles per PROJECT (baseFileUri) - one server per project private readonly projectServerHandles = new Map(); // Track projects where we need to run init notebook (set during controller setup) @@ -100,12 +86,9 @@ 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(IInterpreterService) private readonly interpreterService: IInterpreterService ) {} public activate() { @@ -170,14 +153,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) => { @@ -187,83 +167,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; @@ -300,24 +203,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 = notebook.uri.toString(); - 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)}`); } public async onKernelStarted(kernel: IKernel) { @@ -422,21 +317,17 @@ 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(baseFileUri); - 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(baseFileUri); - 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( + await this.ensureKernelSelectedWithInterpreter( notebook, - environment, + interpreter, baseFileUri, notebookKey, projectKey, @@ -459,27 +350,17 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // projectKey identifies the PROJECT for server tracking const projectKey = baseFileUri.fsPath; - const environmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); - - 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(baseFileUri); - await this.selectPlaceholderController(notebook); - - return false; - } - - await this.ensureKernelSelectedWithConfiguration( + await this.ensureKernelSelectedWithInterpreter( notebook, - environment, + interpreter, baseFileUri, notebookKey, projectKey, @@ -490,26 +371,17 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return true; } - public async ensureKernelSelectedWithConfiguration( + public async ensureKernelSelectedWithInterpreter( notebook: NotebookDocument, - configuration: DeepnoteEnvironment, + interpreter: PythonEnvironment, baseFileUri: Uri, notebookKey: string, projectKey: 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) { @@ -519,71 +391,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, - baseFileUri, - progressToken - ); + const serverInfo = await this.serverStarter.startServer(interpreter, baseFileUri, 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, baseFileUri) + handle: createDeepnoteServerConfigHandle(interpreter.id, baseFileUri) }; // Register the server with the provider (one server per PROJECT) this.serverProvider.registerServer(serverProviderHandle.handle, serverInfo); this.projectServerHandles.set(projectKey, 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) { @@ -592,12 +431,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, @@ -612,8 +453,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 { @@ -622,10 +463,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. @@ -637,14 +475,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: baseFileUri.toString(), serverProviderHandle, serverInfo, - environmentName: configuration.name, + environmentName: getDisplayPath(interpreter.uri), projectName: projectTitle, notebookName: notebookKey }); @@ -717,7 +555,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!' }); } @@ -746,35 +584,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; @@ -782,7 +605,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( @@ -795,113 +618,21 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const notebookKey = notebook.uri.toString(); const projectKey = baseFileUri.fsPath; - const existingEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); - - // No environment configured - need to pick one - if (!existingEnvironmentId) { - return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectKey, token); - } - - const environment = this.environmentManager.getEnvironment(existingEnvironmentId); + const interpreter = await this.interpreterService.getActiveInterpreter(notebook.uri); - // 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(baseFileUri); - - return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectKey, 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, - baseFileUri, - notebookKey, - projectKey, - 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, baseFileUri, notebookKey, projectKey, token); - } - - /** - * Pick an environment and set up the kernel for a notebook. - */ - private async pickAndSetupEnvironment( - notebook: NotebookDocument, - baseFileUri: Uri, - notebookKey: string, - projectKey: 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(baseFileUri, selectedEnvironment.id); - - const result = await this.setupKernelForEnvironment( - notebook, - selectedEnvironment, - baseFileUri, - notebookKey, - projectKey, - token - ); - - if (result) { - 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, - baseFileUri: Uri, - notebookKey: string, - projectKey: string, - token: CancellationToken - ): Promise { try { await window.withProgress( { @@ -910,9 +641,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, cancellable: true }, async (progress, progressToken) => { - await this.ensureKernelSelectedWithConfiguration( + await this.ensureKernelSelectedWithInterpreter( notebook, - environment, + interpreter, baseFileUri, notebookKey, projectKey, @@ -924,146 +655,56 @@ 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); - - 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})` - ); - } - } - - 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 notebookKey = notebook.uri.toString(); + const trackedController = this.notebookControllers.get(notebookKey); - if (!existingInterpreter) { - return true; + if (!trackedController) { + return; // We didn't set up a controller for this notebook } - const expectedInterpreter = this.getVenvInterpreterUri(venvPath); - - return existingInterpreter.uri.fsPath === expectedInterpreter.fsPath; - } - - /** - * Find the NotebookEditor for a given NotebookDocument. - * Required for properly selecting a kernel with the notebook.selectKernel command. - * 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) => e.notebook.uri.toString() === notebook.uri.toString()); - - if (editor) { - return editor; + const selectedController = this.controllerRegistration.getSelected(notebook); + if (!selectedController || selectedController.id !== trackedController.id) { + return; // Selected controller isn't the one we own } - // 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)); - - editor = window.visibleNotebookEditors.find((e) => e.notebook.uri.toString() === notebook.uri.toString()); - - if (editor) { - return editor; - } + if (selectedController.connection.kind !== 'startUsingDeepnoteKernel') { + return; } - 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); + selectedController.controller.updateNotebookAffinity(notebook, NotebookControllerAffinity.Default); - const notebookEditor = await this.findNotebookEditor(notebook); + // Clean up our tracking state for this notebook + this.notebookControllers.delete(notebookKey); + this.notebookConnectionMetadata.delete(notebookKey); + this.notebookInterpreterIds.delete(notebookKey); - 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` - ); - } + logger.info( + `Cleared Deepnote controller for notebook ${getDisplayPath(notebook.uri)} (environment ${environmentId})` + ); } /** * Handle kernel selection errors with user-friendly messages and actions */ - public async handleKernelSelectionError(error: unknown, notebook: NotebookDocument): Promise { - 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; - } - + public async handleKernelSelectionError(error: unknown, _notebook: NotebookDocument): Promise { // Handle DeepnoteKernelError types with specific guidance if (error instanceof DeepnoteKernelError) { // Log the technical details @@ -1117,35 +758,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) { - 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. @@ -1181,103 +793,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 = notebook.uri.toString(); - - // 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']; - - // Execution handler that shows environment picker when user tries to run without an environment - controller.executeHandler = async (cells, doc) => { - logger.info( - `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ - cells.length - } cells` - ); - - // Create a cancellation token that cancels when the notebook is closed - const cts = new CancellationTokenSource(); - const closeListener = workspace.onDidCloseNotebookDocument((closedDoc) => { - if (closedDoc.uri.toString() === doc.uri.toString()) { - 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; - } - - // Environment is now configured, execute the cells through the kernel - const docNotebookKey = doc.uri.toString(); - const realController = this.notebookControllers.get(docNotebookKey); - - if (!realController) { - logger.error(`No controller found after environment configuration for ${docNotebookKey}`); - - return; - } - - logger.info(`Executing ${cells.length} cells through kernel after environment configuration`); - - // Get or create a kernel for this notebook with the new connection - const kernel = this.kernelProvider.getOrCreate(doc, { - metadata: realController.connection, - controller: realController.controller, - resourceUri: doc.uri - }); - - // Execute cells through the kernel - const kernelExecution = this.kernelProvider.getKernelExecution(kernel); - - for (const cell of cells) { - try { - await kernelExecution.executeCell(cell); - } catch (cellError) { - logger.error(`Error executing cell ${cell.index}`, cellError); - // Continue with remaining cells - } - } - - logger.info(`Finished executing ${cells.length} cells`); - } 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 824635343c..0f986b7aaa 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -2,14 +2,10 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, verify, when } from 'ts-mockito'; import { DeepnoteKernelAutoSelector } from './deepnoteKernelAutoSelector.node'; -import { createMockChildProcess } from '../../kernels/deepnote/deepnoteTestHelpers.node'; import { - IDeepnoteEnvironmentManager, IDeepnoteLspClientManager, - IDeepnoteNotebookEnvironmentMapper, IDeepnoteServerProvider, - IDeepnoteServerStarter, - IDeepnoteToolkitInstaller + IDeepnoteServerStarter } from '../../kernels/deepnote/types'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { IDisposableRegistry, IOutputChannel } from '../../platform/common/types'; @@ -21,8 +17,8 @@ import { IDeepnoteNotebookManager } from '../types'; import { IKernelProvider, IKernel, IJupyterKernelSpec } from '../../kernels/types'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; import { NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; -import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; +import { IInterpreterService } from '../../platform/interpreter/contracts'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; @@ -39,11 +35,9 @@ 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 mockInterpreterService: IInterpreterService; let mockProgress: { report(value: { message?: string; increment?: number }): void }; let mockCancellationToken: CancellationToken; @@ -69,11 +63,9 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { mockNotebookManager = mock(); mockKernelProvider = mock(); mockRequirementsHelper = mock(); - mockEnvironmentManager = mock(); mockServerStarter = mock(); - mockToolkitInstaller = mock(); - mockNotebookEnvironmentMapper = mock(); mockOutputChannel = mock(); + mockInterpreterService = mock(); mockProgress = { report: sandbox.stub() }; mockCancellationToken = mock(); @@ -135,11 +127,9 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { instance(mockNotebookManager), instance(mockKernelProvider), instance(mockRequirementsHelper), - instance(mockEnvironmentManager), instance(mockServerStarter), - instance(mockNotebookEnvironmentMapper), instance(mockOutputChannel), - instance(mockToolkitInstaller) + instance(mockInterpreterService) ); }); @@ -158,31 +148,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' ); @@ -195,16 +185,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 @@ -212,34 +202,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 @@ -247,30 +237,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 @@ -278,50 +268,23 @@ 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[6], + ensureKernelSelectedWithInterpreterStub.firstCall.args[6], instance(mockCancellationToken), 'ensureKernelSelected should be called with the cancellation token' ); }); }); - 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); - - // Act - const result = await selector.pickEnvironment(notebookUri); - - // Assert - assert.strictEqual(result, mockEnv1, 'Should return the selected environment'); - }); - }); - suite('onKernelStarted', () => { test('should return early and not call initNotebookRunner for non-deepnote notebooks', async () => { // Arrange @@ -343,51 +306,12 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); 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(); - - // Act - const result = await selector.ensureKernelSelected( - mockNotebook, - mockProgress, - 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(); - }); - - 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(); + 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); - // 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(); + // Stub ensureKernelSelectedWithInterpreter to track if it gets called + const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').resolves(); // Act const result = await selector.ensureKernelSelected( @@ -397,36 +321,29 @@ 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 baseFileUri = mockNotebook.uri.with({ query: '', fragment: '' }); const notebookKey = mockNotebook.uri.toString(); const projectKey = baseFileUri.fsPath; - 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( @@ -436,25 +353,98 @@ 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].toString(), baseFileUri.toString(), 'Third arg should be baseFileUri'); assert.strictEqual(callArgs[3], notebookKey, 'Fourth arg should be notebookKey'); assert.strictEqual(callArgs[4], projectKey, 'Fifth arg should be projectKey'); assert.strictEqual(callArgs[5], mockProgress, 'Sixth arg should be progress'); assert.strictEqual(callArgs[6], instance(mockCancellationToken), 'Seventh arg should be token'); + }); + }); + + suite('ensureEnvironmentConfiguredBeforeExecution', () => { + const nonCancelledToken: CancellationToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => {} }) as any + }; - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).once(); - verify(mockEnvironmentManager.getEnvironment(environmentId)).once(); + 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') + }; + + // 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 now B + when(mockInterpreterService.getActiveInterpreter(anything())).thenResolve(interpreterB); + + // Stub ensureKernelSelectedWithInterpreter to track calls + const ensureStub = sandbox.stub(selector, 'ensureKernelSelectedWithInterpreter').resolves(); + + // 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); + } + ); + + // Put a controller in the map so the final check returns true + ensureStub.callsFake(async () => { + selectorAny.notebookControllers.set(notebookKey, instance(mockNewController)); + }); + + const result = await selector.ensureEnvironmentConfiguredBeforeExecution(mockNotebook, nonCancelledToken); + + 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' + ); + }); + + 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'); }); }); @@ -702,7 +692,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 @@ -738,72 +728,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' + ); }); }); @@ -834,9 +805,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[] = []; @@ -860,8 +836,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 @@ -1017,38 +991,119 @@ 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`) - }; + suite('clearControllerForEnvironment', () => { + test('should unselect and clean up when tracked controller matches selected controller', () => { + const notebookKey = mockNotebook.uri.toString(); - 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 - }; -} + // 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 selectorAny = selector as any; + selectorAny.notebookControllers.set(notebookKey, instance(trackedController)); + selectorAny.notebookInterpreterIds.set(notebookKey, '/usr/bin/python3'); + selectorAny.notebookConnectionMetadata.set(notebookKey, {} as any); + + // Selected controller is the same one we tracked + when(mockControllerRegistration.getSelected(mockNotebook)).thenReturn(instance(trackedController)); + + selector.clearControllerForEnvironment(mockNotebook, 'env-uuid-123'); + + assert.isTrue( + (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' + ); + }); + + 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)); + + selector.clearControllerForEnvironment(mockNotebook, 'env-uuid-123'); + + assert.isFalse( + (mockNativeController.updateNotebookAffinity as sinon.SinonStub).called, + 'Should NOT have called updateNotebookAffinity when we have no tracked controller' + ); + }); + + test('should NOT unselect when selected controller differs from tracked controller', () => { + const notebookKey = mockNotebook.uri.toString(); + + // Track controller A + const controllerA = mock(); + when(controllerA.id).thenReturn('deepnote-notebook-A'); + const selectorAny = selector as any; + selectorAny.notebookControllers.set(notebookKey, instance(controllerA)); + + // 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)); + + selector.clearControllerForEnvironment(mockNotebook, 'env-uuid-123'); + + assert.isFalse( + (mockNativeController.updateNotebookAffinity as sinon.SinonStub).called, + 'Should NOT unselect a controller we do not own' + ); + }); + + test('should NOT unselect when selected controller is not a Deepnote kernel', () => { + const notebookKey = mockNotebook.uri.toString(); + + // 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)); + + 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' + ); + }); + }); +}); /** * Helper function to create mock kernel specs 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.`,