Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions docs/init-hooks.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand Down
2 changes: 1 addition & 1 deletion docs/runtime.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions src/common/utils/tools/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) */
Expand Down
25 changes: 9 additions & 16 deletions src/node/runtime/LocalBaseRuntime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down Expand Up @@ -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}`);

Expand All @@ -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,
Expand Down
7 changes: 4 additions & 3 deletions src/node/runtime/LocalRuntime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand Down Expand Up @@ -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);
Expand Down
19 changes: 9 additions & 10 deletions src/node/runtime/SSHRuntime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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}`);
Expand All @@ -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
Expand Down Expand Up @@ -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);
Expand Down
7 changes: 4 additions & 3 deletions src/node/runtime/WorktreeRuntime.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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);
Expand Down
24 changes: 20 additions & 4 deletions src/node/runtime/initHook.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
6 changes: 6 additions & 0 deletions src/node/services/aiService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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,
Expand Down
57 changes: 57 additions & 0 deletions src/node/services/tools/bash.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/node/services/tools/bash.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down