diff --git a/extension/extension.ts b/extension/extension.ts index 0781464..4e4af03 100644 --- a/extension/extension.ts +++ b/extension/extension.ts @@ -93,7 +93,9 @@ Activating the Mojo Extension reason = 'no-python-extension' as const; } } - statusBar.update(sdk, reason); + const pickerDivergencePath = + this.pyenvManager!.getPickerEnvPathIfDivergent(); + statusBar.update(sdk, reason, pickerDivergencePath); }; this.pushSubscription( diff --git a/extension/lsp/lsp.ts b/extension/lsp/lsp.ts index c16756a..c4e1dcd 100644 --- a/extension/lsp/lsp.ts +++ b/extension/lsp/lsp.ts @@ -201,6 +201,10 @@ export class MojoLSPManager extends DisposableContext { const sdk = await this.envManager.getActiveSDK(); if (!sdk) { + // SDK detection failed for an actively-opened Mojo file. Start + // watching for newly-discovered pixi envs so a mid-session + // `pixi add mojo` is picked up automatically. + this.envManager.watchForEnvDiscoveryIfNeeded(); return; } diff --git a/extension/pyenv.ts b/extension/pyenv.ts index bae0f34..a137780 100644 --- a/extension/pyenv.ts +++ b/extension/pyenv.ts @@ -152,8 +152,16 @@ export class PythonEnvironmentManager extends DisposableContext { private displayedSDKError: boolean = false; private lastLoadedEnv: string | undefined = undefined; private activeSDK: SDK | undefined = undefined; + /// Filesystem path of the Python extension's active environment at the + /// time the active SDK was detected, recorded only when it differs from + /// the env that produced the SDK. `undefined` means no divergence. + private pickerEnvPathAtDetection: string | undefined = undefined; private overridePathState: OverridePathState = 'unset'; private sdkPathChangeTimer: NodeJS.Timeout | undefined = undefined; + /// Subscription to Python extension env-discovery changes. Created lazily + /// while we have no SDK and an open `.mojo` file (so we pick up a pixi env + /// that appears mid-session). Disposed once an SDK is found. + private envDiscoverySubscription: vscode.Disposable | undefined = undefined; constructor(logger: Logger, reporter: TelemetryReporter) { super(); @@ -171,18 +179,28 @@ export class PythonEnvironmentManager extends DisposableContext { vscode.extensions.onDidChange(() => this.handleExtensionChange()), ); // Debounce sdk.path edits so we don't thrash detection while the user is - // mid-typing in the Settings GUI. + // mid-typing in the Settings GUI. Also reacts to preferPixiEnv toggles, + // which can flip whether the workspace pixi env wins over the Python + // extension's active interpreter. this.pushSubscription( vscode.workspace.onDidChangeConfiguration((event) => { - if (!event.affectsConfiguration('mojo.sdk.path')) { + const sdkPathChanged = event.affectsConfiguration('mojo.sdk.path'); + const preferPixiChanged = + event.affectsConfiguration('mojo.preferPixiEnv'); + if (!sdkPathChanged && !preferPixiChanged) { return; } if (this.sdkPathChangeTimer) { clearTimeout(this.sdkPathChangeTimer); } + // Only sdk.path needs debouncing (user types into a text field); + // preferPixiEnv is a checkbox so a 0ms delay would also work, but + // the unified debounce simplifies the lifecycle. this.sdkPathChangeTimer = setTimeout(() => { this.sdkPathChangeTimer = undefined; - this.logger.info('mojo.sdk.path changed, refreshing SDK detection'); + this.logger.info( + 'SDK-related setting changed, refreshing SDK detection', + ); this.refresh(); }, 1500); }), @@ -201,6 +219,7 @@ export class PythonEnvironmentManager extends DisposableContext { /// `mojo.sdk.refresh` command surfaced on the SDK status bar. public refresh() { this.activeSDK = undefined; + this.pickerEnvPathAtDetection = undefined; this.displayedSDKError = false; this.envChangeEmitter.fire(); } @@ -227,6 +246,18 @@ export class PythonEnvironmentManager extends DisposableContext { this.handleEnvironmentChange(p.path), ), ); + // Ensure env discovery has populated `environments.known` before any SDK + // lookup runs against it. Without this, tryGetPixiSDK can miss a + // workspace pixi env on first call simply because the Python extension + // hasn't enumerated it yet. Per the API docs this triggers discovery + // only if it hasn't already happened in the session and returns the + // in-flight promise if discovery is already running, so the call is + // effectively free after the first time. + try { + await this.api.environments.refreshEnvironments(); + } catch (e) { + this.logger.warn('Failed to refresh Python environments:', e); + } } private async handleExtensionChange() { @@ -270,7 +301,8 @@ export class PythonEnvironmentManager extends DisposableContext { /// Finds the active SDK, in priority order: /// 1. `mojo.sdk.path` override (if set; fails loudly without falling back) /// 2. Monorepo `.derived/` SDK - /// 3. SDK from the active Python extension environment + /// 3. Workspace pixi env (gated on `mojo.preferPixiEnv`) + /// 4. SDK from the active Python extension environment public async findActiveSDK(): Promise { // 1. User-supplied override path beats every other source. If it's set // but doesn't resolve, do NOT fall back — that would silently violate @@ -300,6 +332,17 @@ export class PythonEnvironmentManager extends DisposableContext { return undefined; } + // 3. Workspace pixi env with `pixi add mojo`. Prefers a workspace-local + // pixi env over whatever the Python extension's interpreter picker has + // selected, because the picker's heuristics aren't pixi-aware and often + // pick a system or homebrew Python by default. + const pixiSDK = await this.tryGetPixiSDK(); + if (pixiSDK) { + this.logger.info('Using workspace pixi env for Mojo SDK detection.'); + return pixiSDK; + } + + // 4. Whatever the Python extension reports as the active environment. const envPath = this.api.environments.getActiveEnvironmentPath(); const env = await this.api.environments.resolveEnvironment(envPath); this.logger.info('Loading MAX SDK information from Python environment'); @@ -307,11 +350,10 @@ export class PythonEnvironmentManager extends DisposableContext { if (!env) { this.logger.error( - 'No Python enviroment could be retrieved from the Python extension.', - ); - await this.displaySDKError( - 'Unable to load a Python enviroment from the VS Code Python extension.', + 'No Python environment could be retrieved from the Python extension.', ); + // The SDK status bar surfaces "Mojo: No SDK" with diagnostic info in + // the tooltip; a transient toast on top of that is redundant noise. return undefined; } @@ -603,6 +645,155 @@ export class PythonEnvironmentManager extends DisposableContext { return undefined; } + /// Attempt to load a Mojo SDK from a workspace-local pixi environment + /// (`/.pixi/envs/*` containing `share/max/modular.cfg`). Returns + /// undefined if `mojo.preferPixiEnv` is disabled, no Python extension API + /// is available, no workspace folders are open, or no matching env contains + /// a valid SDK. Pixi envs are identified by path pattern rather than the + /// Python extension's `tools` tag — `KnownEnvironmentTools` does not + /// include 'Pixi'. + private async tryGetPixiSDK(): Promise { + const preferPixi = config.get( + 'preferPixiEnv', + /*workspaceFolder=*/ undefined, + true, + ); + if (!preferPixi) { + return undefined; + } + if (!this.api) { + return undefined; + } + const workspaceFolders = vscode.workspace.workspaceFolders; + if (!workspaceFolders || workspaceFolders.length === 0) { + return undefined; + } + + const pixiPathFragment = `${path.sep}.pixi${path.sep}envs${path.sep}`; + const workspacePrefixes = workspaceFolders.map((f) => f.uri.fsPath); + + const candidates = this.api.environments.known.filter((env) => { + const folderPath = env.environment?.folderUri.fsPath; + if (!folderPath || !folderPath.includes(pixiPathFragment)) { + return false; + } + return workspacePrefixes.some((prefix) => folderPath.startsWith(prefix)); + }); + + if (candidates.length === 0) { + return undefined; + } + + // Prefer the env named "default" (pixi's convention), then alphabetical + // by env folder name. Users with multi-env pixi setups who want a + // non-default env should use `mojo.sdk.path`. + candidates.sort((a, b) => { + const nameA = path.basename(a.environment!.folderUri.fsPath); + const nameB = path.basename(b.environment!.folderUri.fsPath); + + if (nameA === 'default' && nameB !== 'default') { + return -1; + } + + if (nameB === 'default' && nameA !== 'default') { + return 1; + } + return nameA.localeCompare(nameB); + }); + + for (const candidate of candidates) { + const resolved = + await this.api.environments.resolveEnvironment(candidate); + if (!resolved) { + continue; + } + if (!(await this.envHasModularCfg(resolved))) { + continue; + } + const sdk = await this.createSDKFromHomePath( + SDKKind.Environment, + path.join(resolved.executable.sysPrefix, 'share', 'max'), + resolved.executable.sysPrefix, + ); + if (!sdk) { + continue; + } + + // Record divergence (if any) between the pixi env we chose and the + // Python extension's active interpreter. The status bar surfaces this + // in the tooltip so the user understands why the picker selection + // isn't being honored for SDK lookup. + const pickerPath = this.api.environments.getActiveEnvironmentPath(); + const pickerResolved = + await this.api.environments.resolveEnvironment(pickerPath); + const pickerSysPrefix = pickerResolved?.executable.sysPrefix; + if ( + pickerSysPrefix && + pickerSysPrefix !== resolved.executable.sysPrefix + ) { + this.pickerEnvPathAtDetection = pickerSysPrefix; + } + + this.logger.info( + `Found workspace pixi env with Mojo SDK at ${ + resolved.environment?.folderUri.fsPath + }`, + ); + return sdk; + } + + return undefined; + } + + /// Returns the Python extension's active environment path if it differs + /// from the env that produced the currently active SDK, otherwise + /// undefined. Used by the SDK status bar to render a divergence note. + public getPickerEnvPathIfDivergent(): string | undefined { + return this.pickerEnvPathAtDetection; + } + + /// Subscribe to `onDidChangeEnvironments` only while there is no active + /// SDK. Called by `MojoLSPManager.tryStartLanguageClient` when a `.mojo` + /// file is opened but detection fails — that's the only situation where + /// we want to react to env-list changes (e.g., the user just ran + /// `pixi add mojo` mid-session and we want to pick it up). + public watchForEnvDiscoveryIfNeeded() { + if (this.activeSDK) { + // Already have an SDK; drop any subscription that was active. + this.envDiscoverySubscription?.dispose(); + this.envDiscoverySubscription = undefined; + return; + } + if (this.envDiscoverySubscription || !this.api) { + return; + } + + const pixiPathFragment = `${path.sep}.pixi${path.sep}envs${path.sep}`; + const workspacePrefixes = + vscode.workspace.workspaceFolders?.map((f) => f.uri.fsPath) ?? []; + + this.envDiscoverySubscription = + this.api.environments.onDidChangeEnvironments((event) => { + // Filter to env changes that affect a workspace pixi env. Avoids + // thrashing detection on irrelevant changes like system Python + // being indexed. + const folderPath = event.env.environment?.folderUri.fsPath; + if (!folderPath?.includes(pixiPathFragment)) { + return; + } + if ( + !workspacePrefixes.some((prefix) => folderPath.startsWith(prefix)) + ) { + return; + } + this.logger.info( + `Workspace pixi env ${event.type} at ${folderPath}, ` + + 'refreshing SDK detection', + ); + this.refresh(); + }); + } + /// Attempt to load a monorepo SDK from the currently open workspace folder. /// Resolves with the loaded SDK, or undefined if one doesn't exist. private async tryGetMonorepoSDK(): Promise { diff --git a/extension/statusBar.ts b/extension/statusBar.ts index a185c35..273b262 100644 --- a/extension/statusBar.ts +++ b/extension/statusBar.ts @@ -116,15 +116,26 @@ export class SDKStatusBar implements vscode.Disposable { this.statusBarItem.backgroundColor = undefined; } - update(sdk: SDK | undefined, reason?: SDKMissingReason) { + update( + sdk: SDK | undefined, + reason?: SDKMissingReason, + pickerDivergencePath?: string, + ) { if (sdk) { const version = sdk.version.replace(/^mojo\s*/i, '').trim(); const kindLabel = SDK_KIND_LABELS[sdk.kind]; this.statusBarItem.text = `$(check) Mojo ${version} (${kindLabel})`; - this.statusBarItem.tooltip = new vscode.MarkdownString( - `**Mojo SDK** (${kindLabel})\n\nVersion: ${version}\n\nPath: ${sdk.mojoPath}`, - ); + let tooltip = `**Mojo SDK** (${kindLabel})\n\nVersion: ${version}\n\nPath: ${sdk.mojoPath}`; + if (pickerDivergencePath) { + tooltip += + `\n\n*Note:* differs from the environment selected via ` + + `\`Python: Select Interpreter\` (\`${pickerDivergencePath}\`). ` + + 'Set `mojo.preferPixiEnv` to `false` if you prefer the picker as ' + + 'the authoritative source.'; + } + this.statusBarItem.tooltip = new vscode.MarkdownString(tooltip); this.statusBarItem.backgroundColor = undefined; + this.statusBarItem.color = undefined; this.statusBarItem.command = this.showOutputCommand; } else if (reason === 'no-python-extension') { this.statusBarItem.text = '$(warning) Mojo: Install Python extension'; @@ -132,9 +143,8 @@ export class SDKStatusBar implements vscode.Disposable { 'The Python extension (`ms-python.python`) is required to discover ' + 'Mojo SDKs in pixi or wheel environments.\n\nClick to open it in the marketplace.', ); - this.statusBarItem.backgroundColor = new vscode.ThemeColor( - 'statusBarItem.warningBackground', - ); + this.statusBarItem.backgroundColor = undefined; + this.statusBarItem.color = undefined; this.statusBarItem.command = { command: 'extension.open', arguments: ['ms-python.python'], @@ -152,19 +162,30 @@ export class SDKStatusBar implements vscode.Disposable { this.statusBarItem.backgroundColor = new vscode.ThemeColor( 'statusBarItem.errorBackground', ); + this.statusBarItem.color = undefined; this.statusBarItem.command = 'mojo.sdk.refresh'; } else { this.statusBarItem.text = '$(warning) Mojo: No SDK'; - this.statusBarItem.tooltip = 'No Mojo SDK detected. Click to view logs.'; + this.statusBarItem.tooltip = new vscode.MarkdownString( + "Couldn't find a Mojo SDK. The extension looked for, in order:\n\n" + + '1. The `mojo.sdk.path` override setting\n' + + '2. The monorepo `.derived/` SDK\n' + + '3. A workspace pixi env containing `share/max/modular.cfg` ' + + '(`pixi add mojo`)\n' + + '4. The active Python interpreter\n\n' + + 'Click to view logs.', + ); this.statusBarItem.backgroundColor = new vscode.ThemeColor( 'statusBarItem.warningBackground', ); + this.statusBarItem.color = new vscode.ThemeColor( + 'statusBarItem.warningForeground', + ); this.statusBarItem.command = this.showOutputCommand; } } updateLsp(state: vscodelc.State | undefined) { - const warningBg = new vscode.ThemeColor('statusBarItem.warningBackground'); const errorBg = new vscode.ThemeColor('statusBarItem.errorBackground'); switch (state) { @@ -188,7 +209,7 @@ export class SDKStatusBar implements vscode.Disposable { default: this.lspStatusBarItem.text = '$(circle-slash) Mojo LSP'; this.lspStatusBarItem.tooltip = 'Mojo language server has not started.'; - this.lspStatusBarItem.backgroundColor = warningBg; + this.lspStatusBarItem.backgroundColor = undefined; break; } } diff --git a/package.json b/package.json index 8571076..3e100f2 100644 --- a/package.json +++ b/package.json @@ -178,6 +178,12 @@ "type": "string", "default": "", "markdownDescription": "Absolute path to a Mojo SDK environment root. When set, this overrides all auto-detection (including the monorepo `.derived/` SDK and any environment selected via the Python extension) and will not fall back to auto-detection if invalid.\n\n**For pixi or conda installs**, point to the environment root that contains `share/max/modular.cfg` — for example, `/path/to/workspace/.pixi/envs/default`.\n\n**For wheel installs**, point to the environment root that contains `bin/mojo` and `lib/python*/site-packages/modular/` — for example, `/path/to/.venv`." + }, + "mojo.preferPixiEnv": { + "scope": "window", + "type": "boolean", + "default": true, + "markdownDescription": "When enabled, the extension prefers a workspace pixi environment (`.pixi/envs/*` containing `share/max/modular.cfg`) over the interpreter selected via `Python: Select Interpreter`. Disable this if you want the Python extension's active environment to be the authoritative source for Mojo SDK discovery." } } },