🤖 feat: add MUX_ env vars to all bash spawns

Set MUX_PROJECT_PATH, MUX_RUNTIME, and MUX_WORKSPACE_NAME environment
variables in both init hooks and bash tool executions for parity.

- Add getMuxEnv() function to centralize env var creation
- Add muxEnv field to ToolConfiguration, used by bash tool
- Refactor runInitHook to receive muxEnv from caller (no duplication)
- Add tests for muxEnv injection in bash tool

_Generated with mux_

Author: ammar-agent

src/common/utils/tools/tools.ts

@@ -23,6 +23,8 @@ export interface ToolConfiguration {
   runtime: Runtime;
   /** Environment secrets to inject (optional) */
   secrets?: Record<string, string>;
+  /** MUX_ environment variables (MUX_PROJECT_PATH, MUX_RUNTIME) - set from init hook env */
+  muxEnv?: Record<string, string>;
   /** Process niceness level (optional, -20 to 19, lower = higher priority) */
   niceness?: number;
   /** Temporary directory for tool outputs in runtime's context (local or remote) */

src/node/runtime/LocalBaseRuntime.ts

@@ -22,12 +22,7 @@ import { getBashPath } from "@/node/utils/main/bashPath";
 import { EXIT_CODE_ABORTED, EXIT_CODE_TIMEOUT } from "@/common/constants/exitCodes";
 import { DisposableProcess } from "@/node/utils/disposableExec";
 import { expandTilde } from "./tildeExpansion";
-import {
-  checkInitHookExists,
-  getInitHookPath,
-  createLineBufferedLoggers,
-  getInitHookEnv,
-} from "./initHook";
+import { getInitHookPath, createLineBufferedLoggers } from "./initHook";
 
 /**
  * Abstract base class for local runtimes (both WorktreeRuntime and LocalRuntime).
@@ -347,19 +342,17 @@ export abstract class LocalBaseRuntime implements Runtime {
   /**
    * Helper to run .mux/init hook if it exists and is executable.
    * Shared between WorktreeRuntime and LocalRuntime.
+   * @param workspacePath - Path to the workspace directory
+   * @param muxEnv - MUX_ environment variables (from getMuxEnv)
+   * @param initLogger - Logger for streaming output
    */
   protected async runInitHook(
-    projectPath: string,
     workspacePath: string,
-    initLogger: InitLogger,
-    runtimeType: "local" | "worktree"
+    muxEnv: Record<string, string>,
+    initLogger: InitLogger
   ): Promise<void> {
-    // Check if hook exists and is executable
-    const hookExists = await checkInitHookExists(projectPath);
-    if (!hookExists) {
-      return;
-    }
-
+    // Hook path is derived from MUX_PROJECT_PATH in muxEnv
+    const projectPath = muxEnv.MUX_PROJECT_PATH;
     const hookPath = getInitHookPath(projectPath);
     initLogger.logStep(`Running init hook: ${hookPath}`);
 
@@ -373,7 +366,7 @@ export abstract class LocalBaseRuntime implements Runtime {
         stdio: ["ignore", "pipe", "pipe"],
         env: {
           ...process.env,
-          ...getInitHookEnv(projectPath, runtimeType),
+          ...muxEnv,
         },
         // Prevent console window from appearing on Windows
         windowsHide: true,

src/node/runtime/LocalRuntime.ts

@@ -6,7 +6,7 @@ import type {
   WorkspaceForkParams,
   WorkspaceForkResult,
 } from "./Runtime";
-import { checkInitHookExists } from "./initHook";
+import { checkInitHookExists, getMuxEnv } from "./initHook";
 import { getErrorMessage } from "@/common/utils/errors";
 import { LocalBaseRuntime } from "./LocalBaseRuntime";
 
@@ -70,13 +70,14 @@ export class LocalRuntime extends LocalBaseRuntime {
   }
 
   async initWorkspace(params: WorkspaceInitParams): Promise<WorkspaceInitResult> {
-    const { projectPath, workspacePath, initLogger } = params;
+    const { projectPath, branchName, workspacePath, initLogger } = params;
 
     try {
       // Run .mux/init hook if it exists
       const hookExists = await checkInitHookExists(projectPath);
       if (hookExists) {
-        await this.runInitHook(projectPath, workspacePath, initLogger, "local");
+        const muxEnv = getMuxEnv(projectPath, "local", branchName);
+        await this.runInitHook(workspacePath, muxEnv, initLogger);
       } else {
         // No hook - signal completion immediately
         initLogger.logComplete(0);

src/node/runtime/SSHRuntime.ts

@@ -17,7 +17,7 @@ import type {
 import { RuntimeError as RuntimeErrorClass } from "./Runtime";
 import { EXIT_CODE_ABORTED, EXIT_CODE_TIMEOUT } from "@/common/constants/exitCodes";
 import { log } from "@/node/services/log";
-import { checkInitHookExists, createLineBufferedLoggers, getInitHookEnv } from "./initHook";
+import { checkInitHookExists, createLineBufferedLoggers, getMuxEnv } from "./initHook";
 import { NON_INTERACTIVE_ENV_VARS } from "@/common/constants/env";
 import { streamProcessToLogger } from "./streamProcess";
 import { expandTildeForSSH, cdCommandForSSH } from "./tildeExpansion";
@@ -732,19 +732,17 @@ export class SSHRuntime implements Runtime {
 
   /**
    * Run .mux/init hook on remote machine if it exists
+   * @param workspacePath - Path to the workspace directory on remote
+   * @param muxEnv - MUX_ environment variables (from getMuxEnv)
+   * @param initLogger - Logger for streaming output
+   * @param abortSignal - Optional abort signal
    */
   private async runInitHook(
-    projectPath: string,
     workspacePath: string,
+    muxEnv: Record<string, string>,
     initLogger: InitLogger,
     abortSignal?: AbortSignal
   ): Promise<void> {
-    // Check if hook exists locally (we synced the project, so local check is sufficient)
-    const hookExists = await checkInitHookExists(projectPath);
-    if (!hookExists) {
-      return;
-    }
-
     // Construct hook path - expand tilde if present
     const remoteHookPath = `${workspacePath}/.mux/init`;
     initLogger.logStep(`Running init hook: ${remoteHookPath}`);
@@ -759,7 +757,7 @@ export class SSHRuntime implements Runtime {
       cwd: workspacePath, // Run in the workspace directory
       timeout: 3600, // 1 hour - generous timeout for init hooks
       abortSignal,
-      env: getInitHookEnv(projectPath, "ssh"),
+      env: muxEnv,
     });
 
     // Create line-buffered loggers
@@ -921,7 +919,8 @@ export class SSHRuntime implements Runtime {
       // Note: runInitHook calls logComplete() internally if hook exists
       const hookExists = await checkInitHookExists(projectPath);
       if (hookExists) {
-        await this.runInitHook(projectPath, workspacePath, initLogger, abortSignal);
+        const muxEnv = getMuxEnv(projectPath, "ssh", branchName);
+        await this.runInitHook(workspacePath, muxEnv, initLogger, abortSignal);
       } else {
         // No hook - signal completion immediately
         initLogger.logComplete(0);

src/node/runtime/WorktreeRuntime.ts

@@ -10,7 +10,7 @@ import type {
   InitLogger,
 } from "./Runtime";
 import { listLocalBranches } from "@/node/git";
-import { checkInitHookExists } from "./initHook";
+import { checkInitHookExists, getMuxEnv } from "./initHook";
 import { execAsync } from "@/node/utils/disposableExec";
 import { getProjectName } from "@/node/utils/runtime/helpers";
 import { getErrorMessage } from "@/common/utils/errors";
@@ -139,14 +139,15 @@ export class WorktreeRuntime extends LocalBaseRuntime {
   }
 
   async initWorkspace(params: WorkspaceInitParams): Promise<WorkspaceInitResult> {
-    const { projectPath, workspacePath, initLogger } = params;
+    const { projectPath, branchName, workspacePath, initLogger } = params;
 
     try {
       // Run .mux/init hook if it exists
       // Note: runInitHook calls logComplete() internally if hook exists
       const hookExists = await checkInitHookExists(projectPath);
       if (hookExists) {
-        await this.runInitHook(projectPath, workspacePath, initLogger, "worktree");
+        const muxEnv = getMuxEnv(projectPath, "worktree", branchName);
+        await this.runInitHook(workspacePath, muxEnv, initLogger);
       } else {
         // No hook - signal completion immediately
         initLogger.logComplete(0);

src/node/runtime/initHook.ts

@@ -2,6 +2,8 @@ import * as fs from "fs";
 import * as fsPromises from "fs/promises";
 import * as path from "path";
 import type { InitLogger } from "./Runtime";
+import type { RuntimeConfig } from "@/common/types/runtime";
+import { isWorktreeRuntime, isSSHRuntime } from "@/common/types/runtime";
 
 /**
  * Check if .mux/init hook exists and is executable
@@ -27,21 +29,35 @@ export function getInitHookPath(projectPath: string): string {
 }
 
 /**
- * Get environment variables for init hook execution
- * Centralizes env var injection to avoid duplication across runtimes
+ * Get MUX_ environment variables for bash execution.
+ * Used by both init hook and regular bash tool calls.
  * @param projectPath - Path to project root (local path for LocalRuntime, remote path for SSHRuntime)
  * @param runtime - Runtime type: "local", "worktree", or "ssh"
+ * @param workspaceName - Name of the workspace (branch name or custom name)
  */
-export function getInitHookEnv(
+export function getMuxEnv(
   projectPath: string,
-  runtime: "local" | "worktree" | "ssh"
+  runtime: "local" | "worktree" | "ssh",
+  workspaceName: string
 ): Record<string, string> {
   return {
     MUX_PROJECT_PATH: projectPath,
     MUX_RUNTIME: runtime,
+    MUX_WORKSPACE_NAME: workspaceName,
   };
 }
 
+/**
+ * Get the effective runtime type from a RuntimeConfig.
+ * Handles legacy "local" with srcBaseDir → "worktree" mapping.
+ */
+export function getRuntimeType(config: RuntimeConfig | undefined): "local" | "worktree" | "ssh" {
+  if (!config) return "worktree"; // Default to worktree for undefined config
+  if (isSSHRuntime(config)) return "ssh";
+  if (isWorktreeRuntime(config)) return "worktree";
+  return "local";
+}
+
 /**
  * Line-buffered logger that splits stream output into lines and logs them
  * Handles incomplete lines by buffering until a newline is received

src/node/services/aiService.ts

@@ -19,6 +19,7 @@ import type { InitStateManager } from "./initStateManager";
 import type { SendMessageError } from "@/common/types/errors";
 import { getToolsForModel } from "@/common/utils/tools/tools";
 import { createRuntime } from "@/node/runtime/runtimeFactory";
+import { getMuxEnv, getRuntimeType } from "@/node/runtime/initHook";
 import { secretsToRecord } from "@/common/types/secrets";
 import type { MuxProviderOptions } from "@/common/types/providerOptions";
 import { log } from "./log";
@@ -1005,6 +1006,11 @@ export class AIService extends EventEmitter {
           cwd: workspacePath,
           runtime,
           secrets: secretsToRecord(projectSecrets),
+          muxEnv: getMuxEnv(
+            metadata.projectPath,
+            getRuntimeType(metadata.runtimeConfig),
+            metadata.name
+          ),
           runtimeTempDir,
         },
         workspaceId,

src/node/services/tools/bash.test.ts

@@ -1230,6 +1230,63 @@ fi
   });
 });
 
+describe("muxEnv environment variables", () => {
+  it("should inject MUX_ environment variables when muxEnv is provided", async () => {
+    using tempDir = new TestTempDir("test-mux-env");
+    const config = createTestToolConfig(process.cwd());
+    config.runtimeTempDir = tempDir.path;
+    config.muxEnv = {
+      MUX_PROJECT_PATH: "/test/project/path",
+      MUX_RUNTIME: "worktree",
+      MUX_WORKSPACE_NAME: "feature-branch",
+    };
+    const tool = createBashTool(config);
+
+    const args: BashToolArgs = {
+      script: 'echo "PROJECT:$MUX_PROJECT_PATH RUNTIME:$MUX_RUNTIME WORKSPACE:$MUX_WORKSPACE_NAME"',
+      timeout_secs: 5,
+    };
+
+    const result = (await tool.execute!(args, mockToolCallOptions)) as BashToolResult;
+
+    expect(result.success).toBe(true);
+    if (result.success) {
+      expect(result.output).toContain("PROJECT:/test/project/path");
+      expect(result.output).toContain("RUNTIME:worktree");
+      expect(result.output).toContain("WORKSPACE:feature-branch");
+    }
+  });
+
+  it("should allow secrets to override muxEnv", async () => {
+    using tempDir = new TestTempDir("test-mux-env-override");
+    const config = createTestToolConfig(process.cwd());
+    config.runtimeTempDir = tempDir.path;
+    config.muxEnv = {
+      MUX_PROJECT_PATH: "/mux/path",
+      CUSTOM_VAR: "from-mux",
+    };
+    config.secrets = {
+      CUSTOM_VAR: "from-secrets",
+    };
+    const tool = createBashTool(config);
+
+    const args: BashToolArgs = {
+      script: 'echo "MUX:$MUX_PROJECT_PATH CUSTOM:$CUSTOM_VAR"',
+      timeout_secs: 5,
+    };
+
+    const result = (await tool.execute!(args, mockToolCallOptions)) as BashToolResult;
+
+    expect(result.success).toBe(true);
+    if (result.success) {
+      // MUX_PROJECT_PATH from muxEnv should be present
+      expect(result.output).toContain("MUX:/mux/path");
+      // Secrets should override muxEnv when there's a conflict
+      expect(result.output).toContain("CUSTOM:from-secrets");
+    }
+  });
+});
+
 describe("SSH runtime redundant cd detection", () => {
   // Helper to create bash tool with SSH runtime configuration
   // Note: These tests check redundant cd detection logic only - they don't actually execute via SSH

src/node/services/tools/bash.ts

@@ -246,7 +246,7 @@ export const createBashTool: ToolFactory = (config: ToolConfiguration) => {
 ${script}`;
       const execStream = await config.runtime.exec(scriptWithClosedStdin, {
         cwd: config.cwd,
-        env: config.secrets,
+        env: { ...config.muxEnv, ...config.secrets },
         timeout: effectiveTimeout,
         niceness: config.niceness,
         abortSignal,