diff --git a/docs/init-hooks.md b/docs/init-hooks.md index dd84bea71d..01666048c8 100644 --- a/docs/init-hooks.md +++ b/docs/init-hooks.md @@ -32,13 +32,14 @@ The init script runs in the workspace directory with the workspace's environment ## Environment Variables -Init hooks receive the following environment variables: +Init hooks receive the following environment variables (also available in all agent bash tool executions): - `MUX_PROJECT_PATH` - Absolute path to the project root on the **local machine** - Always refers to your local project path, even on SSH workspaces - Useful for logging, debugging, or runtime-specific logic -- `MUX_RUNTIME` - Runtime type: `"local"` or `"ssh"` +- `MUX_RUNTIME` - Runtime type: `"local"`, `"worktree"`, or `"ssh"` - Use this to detect whether the hook is running locally or remotely +- `MUX_WORKSPACE_NAME` - Name of the workspace (typically the branch name) **Note for SSH workspaces:** Since the project is synced to the remote machine, files exist in both locations. The init hook runs in the workspace directory (`$PWD`), so use relative paths to reference project files: diff --git a/docs/runtime.md b/docs/runtime.md index 0a72f047cf..b1e597ccff 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -17,7 +17,7 @@ When creating a workspace, select the runtime from the dropdown in the workspace ## Init Hooks -[Init hooks](/init-hooks) can detect the runtime type via the `MUX_RUNTIME` environment variable: +[Init hooks](/init-hooks) and agent bash tool executions can detect the runtime type via the `MUX_RUNTIME` environment variable: - `local` — Local runtime - `worktree` — Worktree runtime diff --git a/src/common/utils/tools/tools.ts b/src/common/utils/tools/tools.ts index c837e9bb6b..31cf3aa466 100644 --- a/src/common/utils/tools/tools.ts +++ b/src/common/utils/tools/tools.ts @@ -23,6 +23,8 @@ export interface ToolConfiguration { runtime: Runtime; /** Environment secrets to inject (optional) */ secrets?: Record; + /** MUX_ environment variables (MUX_PROJECT_PATH, MUX_RUNTIME) - set from init hook env */ + muxEnv?: Record; /** Process niceness level (optional, -20 to 19, lower = higher priority) */ niceness?: number; /** Temporary directory for tool outputs in runtime's context (local or remote) */ diff --git a/src/node/runtime/LocalBaseRuntime.ts b/src/node/runtime/LocalBaseRuntime.ts index a4a57d10e3..eec2290089 100644 --- a/src/node/runtime/LocalBaseRuntime.ts +++ b/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, + initLogger: InitLogger ): Promise { - // 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, diff --git a/src/node/runtime/LocalRuntime.ts b/src/node/runtime/LocalRuntime.ts index 95a22b92ab..aaabea879f 100644 --- a/src/node/runtime/LocalRuntime.ts +++ b/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 { - 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); diff --git a/src/node/runtime/SSHRuntime.ts b/src/node/runtime/SSHRuntime.ts index 45d60e5dac..076c951987 100644 --- a/src/node/runtime/SSHRuntime.ts +++ b/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, initLogger: InitLogger, abortSignal?: AbortSignal ): Promise { - // 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); diff --git a/src/node/runtime/WorktreeRuntime.ts b/src/node/runtime/WorktreeRuntime.ts index 999aeb6bfe..90e4c500c7 100644 --- a/src/node/runtime/WorktreeRuntime.ts +++ b/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 { - 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); diff --git a/src/node/runtime/initHook.ts b/src/node/runtime/initHook.ts index e93addb889..f6bd9b6905 100644 --- a/src/node/runtime/initHook.ts +++ b/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 { 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 diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index e9862723d2..781a09f4c5 100644 --- a/src/node/services/aiService.ts +++ b/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, diff --git a/src/node/services/tools/bash.test.ts b/src/node/services/tools/bash.test.ts index 820c2bd532..4d4f4c9dde 100644 --- a/src/node/services/tools/bash.test.ts +++ b/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 diff --git a/src/node/services/tools/bash.ts b/src/node/services/tools/bash.ts index c0559a86d2..e5456084bd 100644 --- a/src/node/services/tools/bash.ts +++ b/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,