Handle wheel installs correctly

extension/commands/run.ts

@@ -215,7 +215,7 @@ class ExecutionManager extends DisposableContext {
    * Get a terminal to use for the given file.
    */
   getTerminalForFile(doc: vscode.TextDocument, sdk: SDK): vscode.Terminal {
-    const fullId = `${doc.fileName} · ${sdk.homePath}`;
+    const fullId = `${doc.fileName} · ${sdk.mojoPath}`;
     // We have to keep the full terminal name short so that VS Code renders it nicely,
     // and we have to keep it unique among other files.
     const terminalName = `Mojo: ${path.basename(doc.fileName)} · ${md5(fullId).substring(0, 5)}`;

extension/debug/debug.ts

@@ -13,12 +13,9 @@ import { activatePickProcessToAttachCommand } from './attachQuickPick';
 import { initializeInlineLocalVariablesProvider } from './inlineVariables';
 import { MojoExtension } from '../extension';
 import { quote } from 'shell-quote';
-import * as util from 'util';
-import { execFile as execFileBase } from 'child_process';
 import { Optional } from '../types';
 import { PythonEnvironmentManager, SDK, SDKKind } from '../pyenv';
 import { Logger } from '../logging';
-const execFile = util.promisify(execFileBase);
 
 /**
  * Stricter version of vscode.DebugConfiguration intended to reduce the chances
@@ -70,6 +67,10 @@ type MojoCudaGdbDebugConfiguration = {
  */
 const DEBUG_TYPE: string = 'mojo-lldb';
 
+function envDictToList(dict: { [key: string]: string }): string[] {
+  return Object.entries(dict).map(([k, v]) => `${k}=${v}`);
+}
+
 /**
  * Some debug configurations come from an RPC call, which have an explicit SDK
  * to use. We should honor it when running the debug session.
@@ -117,31 +118,6 @@ class MojoDebugAdapterDescriptorFactory
       return undefined;
     }
     this.logger.info(`Using the SDK ${sdk.version} for the debug session`);
-    if (sdk.homePath.endsWith('.derived')) {
-      // Debug adapters from dev sdks tend to be corrupted because dependencies
-      // might need to be rebuilt, so we run a simple verification.
-      try {
-        await execFile(sdk.dapPath, ['--help']);
-      } catch (ex: any) {
-        const { stderr, stdout } = ex;
-        this.logger.main.outputChannel.appendLine(
-          '\n\n\n===== LLDB Debug Adapter verification =====',
-        );
-        this.logger.error('Unable to execute the LLDB Debug Adapter.', ex);
-        if (stdout) {
-          this.logger.info('stdout: ' + stdout);
-        }
-        if (stderr) {
-          this.logger.info('stderr: ' + stderr);
-        }
-        this.logger.main.outputChannel.show();
-
-        // this.envManager.showBazelwRunInstallPrompt(
-        //   'The LLDB Debug Adapter seems to be corrupted.',
-        //   sdk.homePath,
-        // );
-      }
-    }
 
     return new vscode.DebugAdapterExecutable(sdk.dapPath, [
       '--repl-mode',
@@ -204,6 +180,23 @@ class MojoDebugConfigurationResolver
     }
 
     if (debugConfiguration.mojoFile) {
+      // Wheel environments do not ship a modular.cfg file that we can use to
+      // invoke 'mojo run' directly. They instead use a wrapper, which can't be
+      // launched under lldb. It's technically possible to launch the underlying
+      // binary directly, but without a modular.cfg file we would need to set a
+      // number of environment variables. We choose not to do that and instead
+      // just disallow debugging a specific file when the SDK is installed via a
+      // Python wheel.
+      if (!sdk.supportsFileDebug) {
+        this.logger.error(
+          `Cannot launch debug session with mojoFile specified (was '${debugConfiguration.mojoFile}') because MAX was installed as a wheel.`,
+        );
+        vscode.window.showErrorMessage(
+          "Debugging a Mojo file using the 'mojoFile' option is not supported when the Mojo SDK is installed as a wheel.",
+        );
+        return undefined;
+      }
+
       if (
         !debugConfiguration.mojoFile.endsWith('.🔥') &&
         !debugConfiguration.mojoFile.endsWith('.mojo')
@@ -289,10 +282,9 @@ class MojoDebugConfigurationResolver
 
     const env = [
       `LLDB_VSCODE_RIT_TIMEOUT_IN_MS=${initializationTimeoutSec * 1000}`, // runInTerminal initialization timeout.
+      ...envDictToList(sdk.getProcessEnv()),
     ];
 
-    env.push(`MODULAR_HOME=${sdk.homePath}`);
-
     debugConfiguration.env = [...env, ...(debugConfiguration.env || [])];
     return debugConfiguration as vscode.DebugConfiguration;
   }
@@ -377,9 +369,7 @@ class MojoCudaGdbDebugConfigurationResolver
     debugConfig.args = quote(args || []);
     // cuda-gdb takes environment as a list of objects like:
     // [{"name": "HOME", "value": "/home/ubuntu"}]
-    let env = [];
-    env.push(`MODULAR_HOME=${sdk.homePath}`);
-    env = [...env, ...(debugConfigIn.env || [])];
+    const env = debugConfigIn.env || [];
     debugConfig.environment = env.map((envStr: string) => {
       const split = envStr.split('=');
       return { name: split[0], value: split.slice(1).join('=') };

extension/pyenv.ts

@@ -7,16 +7,20 @@
 import * as vscode from 'vscode';
 import * as ini from 'ini';
 import { DisposableContext } from './utils/disposableContext';
-import { PythonExtension } from '@vscode/python-extension';
+import { PythonExtension, ResolvedEnvironment } from '@vscode/python-extension';
 import assert from 'assert';
 import { Logger } from './logging';
 import path from 'path';
 import * as util from 'util';
-import { execFile as callbackExecFile } from 'child_process';
+import {
+  execFile as callbackExecFile,
+  exec as callbackExec,
+} from 'child_process';
 import { Memoize } from 'typescript-memoize';
 import { TelemetryReporter } from './telemetry';
-import { directoryExists } from './utils/files';
+import { fileExists } from './utils/files';
 const execFile = util.promisify(callbackExecFile);
+const exec = util.promisify(callbackExec);
 
 export enum SDKKind {
   Environment = 'environment',
@@ -26,14 +30,14 @@ export enum SDKKind {
 
 /// Represents a usable instance of the MAX SDK.
 export class SDK {
+  public readonly supportsFileDebug: boolean = false;
+
   constructor(
     private logger: Logger,
     /// What kind of SDK this is. Primarily used for logging and context hinting.
     readonly kind: SDKKind,
     /// The unparsed version string of the SDK.
     readonly version: string,
-    /// The home path of the SDK. This is always a directory containing a modular.cfg file.
-    readonly homePath: string,
     /// The path to the language server executable.
     readonly lspPath: string,
     /// The path to the mblack executable.
@@ -83,6 +87,43 @@ export class SDK {
 
   /// Gets an appropriate environment to spawn subprocesses from this SDK.
   public getProcessEnv(withTelemetry: boolean = true) {
+    return {
+      MODULAR_TELEMETRY_ENABLED: withTelemetry ? 'true' : 'false',
+    };
+  }
+}
+
+class HomeSDK extends SDK {
+  public override readonly supportsFileDebug: boolean = true;
+
+  constructor(
+    logger: Logger,
+    kind: SDKKind,
+    version: string,
+    private homePath: string,
+    lspPath: string,
+    mblackPath: string,
+    lldbPluginPath: string,
+    dapPath: string,
+    mojoPath: string,
+    visualizersPath: string,
+    lldbPath: string,
+  ) {
+    super(
+      logger,
+      kind,
+      version,
+      lspPath,
+      mblackPath,
+      lldbPluginPath,
+      dapPath,
+      mojoPath,
+      visualizersPath,
+      lldbPath,
+    );
+  }
+
+  public override getProcessEnv(withTelemetry: boolean = true) {
     return {
       MODULAR_HOME: this.homePath,
       MODULAR_TELEMETRY_ENABLED: withTelemetry ? 'true' : 'false',
@@ -155,20 +196,24 @@ export class PythonEnvironmentManager extends DisposableContext {
       return undefined;
     }
 
-    this.logger.info(`Found Python environment at ${envPath.path}`);
-
-    const homePath = path.join(env.executable.sysPrefix, 'share', 'max');
-    if (!(await directoryExists(homePath))) {
-      this.logger.error(
-        `SDK home path ${homePath} does not exist in the Python environment's system prefix. MAX is not installed.`,
+    // We cannot use the environment type information reported by the Python
+    // extension because it considers Conda and wheel-based installs to be the
+    // same, when we need to differentiate them.
+    this.logger.info(`Found Python environment at ${envPath.path}`, env);
+    if (await this.envHasModularCfg(env)) {
+      this.logger.info(
+        `Python environment '${envPath.path}' appears to be Conda-like; using modular.cfg method.`,
       );
-      await this.displaySDKError(
-        `MAX is not installed in Python environment located at ${envPath.path}. Please install the MAX SDK to proceed.`,
+      return this.createSDKFromHomePath(
+        SDKKind.Environment,
+        path.join(env.executable.sysPrefix, 'share', 'max'),
       );
-      return undefined;
+    } else {
+      this.logger.info(
+        `Python environment '${envPath.path}' does not have a modular.cfg file; assuming wheel installation.`,
+      );
+      return this.createSDKFromWheelEnv(env);
     }
-
-    return await this.createSDKFromHomePath(SDKKind.Environment, homePath);
   }
 
   private async displaySDKError(message: string) {
@@ -180,6 +225,85 @@ export class PythonEnvironmentManager extends DisposableContext {
     await vscode.window.showErrorMessage(message);
   }
 
+  private async envHasModularCfg(env: ResolvedEnvironment): Promise<boolean> {
+    return fileExists(
+      path.join(env.executable.sysPrefix, 'share', 'max', 'modular.cfg'),
+    );
+  }
+
+  private async createSDKFromWheelEnv(
+    env: ResolvedEnvironment,
+  ): Promise<SDK | undefined> {
+    const binPath = path.join(env.executable.sysPrefix, 'bin');
+    const libPath = path.join(
+      env.executable.sysPrefix,
+      'lib',
+      `python${env.version!.major}.${env.version!.minor}`,
+      'site-packages',
+      'modular',
+      'lib',
+    );
+    // helper to pull required files/folders out of the environment
+    const retrievePath = async (target: string) => {
+      this.logger.debug(`Retrieving tool path '${target}'.`);
+      try {
+        // stat-ing the path confirms it exists in some form; if an exception is thrown then it doesn't exist.
+        await vscode.workspace.fs.stat(vscode.Uri.file(target));
+        return target;
+      } catch {
+        this.logger.error(`Missing path ${target} in venv.`);
+        return undefined;
+      }
+    };
+
+    const libExt = process.platform == 'darwin' ? 'dylib' : 'so';
+
+    const mojoPath = await retrievePath(path.join(binPath, 'mojo'));
+    const lspPath = await retrievePath(path.join(binPath, 'mojo-lsp-server'));
+    const lldbPluginPath = await retrievePath(
+      path.join(libPath, `libMojoLLDB.${libExt}`),
+    );
+    const mblackPath = await retrievePath(path.join(binPath, 'mblack'));
+    const dapPath = await retrievePath(path.join(binPath, 'lldb-dap'));
+    const visualizerPath = await retrievePath(
+      path.join(libPath, 'lldb-visualizers'),
+    );
+    const lldbPath = await retrievePath(path.join(binPath, 'mojo-lldb'));
+    // The debugger requires that we avoid using the wrapped `mojo` entrypoint for specific scenarios.
+    const rawMojoPath = await retrievePath(
+      path.join(libPath, '..', 'bin', 'mojo'),
+    );
+
+    if (
+      !mojoPath ||
+      !lspPath ||
+      !lldbPluginPath ||
+      !rawMojoPath ||
+      !mblackPath ||
+      !lldbPluginPath ||
+      !dapPath ||
+      !visualizerPath ||
+      !lldbPath
+    ) {
+      return undefined;
+    }
+
+    // We don't know the version intrinsically so we need to invoke it ourselves.
+    const versionResult = await exec(`"${mojoPath}" --version`);
+    return new SDK(
+      this.logger,
+      SDKKind.Environment,
+      versionResult.stdout,
+      lspPath,
+      mblackPath,
+      lldbPluginPath,
+      dapPath,
+      mojoPath,
+      visualizerPath,
+      lldbPath,
+    );
+  }
+
   /// Attempts to create a SDK from a home path. Returns undefined if creation failed.
   public async createSDKFromHomePath(
     kind: SDKKind,
@@ -229,7 +353,7 @@ export class PythonEnvironmentManager extends DisposableContext {
         kind,
       });
 
-      return new SDK(
+      return new HomeSDK(
         this.logger,
         kind,
         version,