use anyhow::{Context, Result, anyhow};
use std::fs;
use std::path::{Path, PathBuf};
use crate::config::{MuxMode, WindowConfig};
use crate::multiplexer::{
CreateSessionParams, CreateWindowInSessionParams, CreateWindowParams, Multiplexer,
PaneSetupOptions,
};
use crate::{cmd, config, git, prompt::Prompt};
use tracing::{debug, info};
use super::file_ops::{handle_file_operations, symlink_claude_local_md};
use super::types::CreateResult;
#[allow(clippy::too_many_arguments)]
pub fn setup_environment(
mux: &dyn Multiplexer,
branch_name: &str,
handle: &str,
worktree_path: &Path,
config: &config::Config,
options: &super::types::SetupOptions,
agent: Option<&str>,
after_window: Option<String>,
) -> Result<CreateResult> {
let agent = agent.map(|a| {
config
.agents
.get(a)
.map(|e| e.command.as_str())
.unwrap_or(a)
});
debug!(
branch = branch_name,
handle = handle,
path = %worktree_path.display(),
run_hooks = options.run_hooks,
run_file_ops = options.run_file_ops,
"setup_environment:start"
);
let prefix = config.window_prefix();
let repo_root = git::get_main_worktree_root()?;
let effective_working_dir = options.working_dir.as_deref().unwrap_or(worktree_path);
let file_ops_source = options.config_root.as_deref().unwrap_or(&repo_root);
if options.run_file_ops {
handle_file_operations(file_ops_source, effective_working_dir, &config.files)
.context("Failed to perform file operations")?;
debug!(
branch = branch_name,
"setup_environment:file operations applied"
);
}
if options.run_file_ops {
symlink_claude_local_md(&repo_root, effective_working_dir)
.context("Failed to auto-symlink CLAUDE.local.md")?;
}
let mut hooks_run = 0;
if options.run_hooks
&& let Some(post_create) = &config.post_create
&& !post_create.is_empty()
{
hooks_run = post_create.len();
let abs_worktree_path = worktree_path
.canonicalize()
.unwrap_or_else(|_| worktree_path.to_path_buf());
let abs_project_root = repo_root
.canonicalize()
.unwrap_or_else(|_| repo_root.clone());
let abs_config_dir = effective_working_dir
.canonicalize()
.unwrap_or_else(|_| effective_working_dir.to_path_buf());
let worktree_path_str = abs_worktree_path.to_string_lossy();
let project_root_str = abs_project_root.to_string_lossy();
let config_dir_str = abs_config_dir.to_string_lossy();
let hook_env = [
("WORKMUX_HANDLE", handle),
("WM_HANDLE", handle),
("WM_WORKTREE_PATH", worktree_path_str.as_ref()),
("WM_PROJECT_ROOT", project_root_str.as_ref()),
("WM_CONFIG_DIR", config_dir_str.as_ref()),
];
for (idx, command) in post_create.iter().enumerate() {
info!(branch = branch_name, step = idx + 1, total = hooks_run, command = %command, "setup_environment:hook start");
info!(command = %command, "Running post-create hook {}/{}", idx + 1, hooks_run);
cmd::shell_command_with_env(command, effective_working_dir, &hook_env)
.with_context(|| format!("Failed to run post-create command: '{}'", command))?;
info!(branch = branch_name, step = idx + 1, total = hooks_run, command = %command, "setup_environment:hook complete");
}
info!(
branch = branch_name,
total = hooks_run,
"setup_environment:hooks complete"
);
}
let window_plans: Vec<WindowConfig> = if let Some(windows) = &config.windows {
windows.clone()
} else {
let panes = config.panes.clone();
vec![WindowConfig { name: None, panes }]
};
let all_panes: Vec<config::PaneConfig> = window_plans
.iter()
.flat_map(|w| w.panes.as_deref().unwrap_or(&[]).iter().cloned())
.collect();
let all_resolved_panes = resolve_pane_configuration(&all_panes, agent);
if options.prompt_file_path.is_some() {
validate_prompt_consumption(&all_resolved_panes, agent, config, options)?;
}
let lima_vm_name = pre_boot_lima_vm(
mux,
config,
&all_resolved_panes,
effective_working_dir,
worktree_path,
options,
agent,
)?;
let pane_setup_options = PaneSetupOptions {
run_commands: options.run_pane_commands,
prompt_file_path: options.prompt_file_path.as_deref(),
worktree_root: Some(worktree_path),
lima_vm_name: lima_vm_name.as_deref(),
resume_mode: options.resume_mode.clone(),
};
let mut focus_pane_id: Option<String> = None;
let mut zoom_pane_id: Option<String> = None;
match options.mode {
MuxMode::Window => {
let panes = window_plans[0].panes.as_deref().unwrap_or(&[]);
let resolved_panes = resolve_pane_configuration(panes, agent);
let last_wm_window =
after_window.or_else(|| mux.find_last_window_with_prefix(prefix).unwrap_or(None));
let initial_pane_id = mux
.create_window(CreateWindowParams {
prefix,
name: handle,
cwd: effective_working_dir,
after_window: last_wm_window.as_deref(),
})
.context("Failed to create window")?;
info!(
branch = branch_name,
handle = handle,
pane_id = %initial_pane_id,
"setup_environment:window created"
);
let result = mux
.setup_panes(
&initial_pane_id,
&resolved_panes,
effective_working_dir,
pane_setup_options,
config,
agent,
)
.context("Failed to setup panes")?;
focus_pane_id = Some(result.focus_pane_id);
zoom_pane_id = result.zoom_pane_id;
}
MuxMode::Session => {
let session_full_name = crate::multiplexer::util::prefixed(prefix, handle);
for (i, window_plan) in window_plans.iter().enumerate() {
let panes = window_plan.panes.as_deref().unwrap_or(&[]);
let resolved_panes = resolve_pane_configuration(panes, agent);
let initial_pane_id = if i == 0 {
let pane_id = mux
.create_session(CreateSessionParams {
prefix,
name: handle,
cwd: effective_working_dir,
initial_window_name: window_plan.name.as_deref(),
})
.context("Failed to create session")?;
info!(
branch = branch_name,
handle = handle,
window = ?window_plan.name,
pane_id = %pane_id,
"setup_environment:session created (window 0)"
);
pane_id
} else {
let pane_id = mux
.create_window_in_session(CreateWindowInSessionParams {
session_name: &session_full_name,
name: window_plan.name.as_deref(),
cwd: effective_working_dir,
})
.context("Failed to create window in session")?;
info!(
branch = branch_name,
handle = handle,
window = ?window_plan.name,
window_index = i,
pane_id = %pane_id,
"setup_environment:window created in session"
);
pane_id
};
let result = mux
.setup_panes(
&initial_pane_id,
&resolved_panes,
effective_working_dir,
pane_setup_options.clone(),
config,
agent,
)
.context("Failed to setup panes")?;
let has_explicit_focus = resolved_panes.iter().any(|p| p.focus || p.zoom);
if i == 0 || has_explicit_focus {
focus_pane_id = Some(result.focus_pane_id);
}
if result.zoom_pane_id.is_some() {
zoom_pane_id = result.zoom_pane_id;
}
}
}
}
let focus_pane_id = focus_pane_id.expect("at least one window must be created");
debug!(
branch = branch_name,
focus_id = %focus_pane_id,
"setup_environment:panes configured"
);
if options.focus_window {
match options.mode {
MuxMode::Window => {
mux.select_pane(&focus_pane_id)?;
mux.select_window(prefix, handle)?;
}
MuxMode::Session => {
mux.switch_to_pane(&focus_pane_id, None)?;
}
}
}
if let Some(ref zoom_id) = zoom_pane_id {
mux.zoom_pane(zoom_id)?;
}
Ok(CreateResult {
worktree_path: worktree_path.to_path_buf(),
branch_name: branch_name.to_string(),
post_create_hooks_run: hooks_run,
base_branch: None,
did_switch: false,
resolved_handle: handle.to_string(),
mode: options.mode,
})
}
#[allow(clippy::too_many_arguments)]
fn pre_boot_lima_vm(
mux: &dyn crate::multiplexer::Multiplexer,
config: &config::Config,
panes: &[config::PaneConfig],
working_dir: &Path,
worktree_path: &Path,
options: &super::types::SetupOptions,
agent: Option<&str>,
) -> Result<Option<String>> {
if !config.sandbox.is_enabled()
|| !matches!(
config.sandbox.backend(),
crate::config::SandboxBackend::Lima
)
{
return Ok(None);
}
let effective_agent = agent.or(config.agent.as_deref());
let shell = mux.get_default_shell()?;
let any_pane_needs_lima = panes.iter().any(|pane_config| {
let resolved = crate::multiplexer::util::resolve_pane_command(
pane_config.command.as_deref(),
options.run_pane_commands,
options.prompt_file_path.as_deref(),
working_dir,
effective_agent,
&shell,
config.agent_type.as_deref(),
);
if resolved.is_none() {
return false;
}
let is_agent_pane = pane_config.command.as_deref().is_some_and(|cmd| {
cmd == "<agent>"
|| crate::multiplexer::agent::is_known_agent(cmd)
|| effective_agent.is_some_and(|a| crate::config::is_agent_command(cmd, a))
});
match config.sandbox.target() {
crate::config::SandboxTarget::All => true,
crate::config::SandboxTarget::Agent => is_agent_pane,
}
});
if !any_pane_needs_lima {
return Ok(None);
}
info!("pre-booting Lima VM before window creation");
let vm_name = crate::sandbox::ensure_lima_vm(config, worktree_path)?;
Ok(Some(vm_name))
}
pub fn resolve_pane_configuration(
original_panes: &[config::PaneConfig],
agent: Option<&str>,
) -> Vec<config::PaneConfig> {
let Some(agent_cmd) = agent else {
return original_panes.to_vec();
};
if original_panes.iter().any(|pane| {
pane.command
.as_deref()
.is_some_and(|cmd| cmd == "<agent>" || crate::multiplexer::agent::is_known_agent(cmd))
}) {
return original_panes.to_vec();
}
let mut panes = original_panes.to_vec();
if let Some(focused) = panes.iter_mut().find(|pane| pane.focus) {
focused.command = Some(agent_cmd.to_string());
return panes;
}
if let Some(first) = panes.get_mut(0) {
first.command = Some(agent_cmd.to_string());
return panes;
}
vec![config::PaneConfig {
command: Some(agent_cmd.to_string()),
focus: true,
..Default::default()
}]
}
pub fn write_prompt_file(
working_dir: Option<&Path>,
branch_name: &str,
prompt: &Prompt,
) -> Result<PathBuf> {
let content = match prompt {
Prompt::Inline(text) => text.clone(),
Prompt::FromFile(path) => fs::read_to_string(path)
.with_context(|| format!("Failed to read prompt file '{}'", path.display()))?,
};
let safe_branch_name = branch_name.replace(['/', '\\', ':'], "-");
let prompt_path = if let Some(dir) = working_dir {
let workmux_dir = dir.join(".workmux");
fs::create_dir_all(&workmux_dir).with_context(|| {
format!("Failed to create .workmux directory in '{}'", dir.display())
})?;
if let Some(exclude_path) = resolve_git_exclude_path(dir)
&& exclude_path.exists()
&& let Ok(content) = fs::read_to_string(&exclude_path)
&& !content.lines().any(|line| line.trim() == ".workmux/")
&& let Ok(mut file) = fs::OpenOptions::new().append(true).open(&exclude_path)
{
use std::io::Write;
let _ = writeln!(file, "\n# workmux prompt files\n.workmux/");
}
let prompt_filename = format!("PROMPT-{}.md", safe_branch_name);
workmux_dir.join(prompt_filename)
} else {
let prompt_filename = format!("workmux-prompt-{}.md", safe_branch_name);
std::env::temp_dir().join(prompt_filename)
};
fs::write(&prompt_path, content)
.with_context(|| format!("Failed to write prompt file '{}'", prompt_path.display()))?;
Ok(prompt_path)
}
fn resolve_git_exclude_path(dir: &Path) -> Option<PathBuf> {
let git_path = dir.join(".git");
if git_path.is_dir() {
Some(git_path.join("info/exclude"))
} else if git_path.is_file() {
let content = fs::read_to_string(&git_path).ok()?;
let gitdir = content.strip_prefix("gitdir: ")?.trim();
let main_git = Path::new(gitdir).ancestors().nth(2)?;
Some(main_git.join("info/exclude"))
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_pane_configuration_no_agent_returns_original() {
let original_panes = vec![config::PaneConfig {
command: Some("vim".to_string()),
focus: true,
..Default::default()
}];
let result = resolve_pane_configuration(&original_panes, None);
assert_eq!(result.len(), 1);
assert_eq!(result[0].command, Some("vim".to_string()));
}
#[test]
fn resolve_pane_configuration_agent_placeholder_returns_original() {
let original_panes = vec![config::PaneConfig {
command: Some("<agent>".to_string()),
focus: true,
..Default::default()
}];
let result = resolve_pane_configuration(&original_panes, Some("claude"));
assert_eq!(result.len(), 1);
assert_eq!(result[0].command, Some("<agent>".to_string()));
}
#[test]
fn resolve_pane_configuration_agent_sets_focused_pane() {
let original_panes = vec![
config::PaneConfig {
command: Some("vim".to_string()),
..Default::default()
},
config::PaneConfig {
command: Some("npm run dev".to_string()),
focus: true,
..Default::default()
},
];
let result = resolve_pane_configuration(&original_panes, Some("claude"));
assert_eq!(result[0].command, Some("vim".to_string()));
assert_eq!(result[1].command, Some("claude".to_string()));
}
#[test]
fn resolve_pane_configuration_agent_sets_first_pane_when_no_focus() {
let original_panes = vec![config::PaneConfig {
command: Some("vim".to_string()),
..Default::default()
}];
let result = resolve_pane_configuration(&original_panes, Some("claude"));
assert_eq!(result[0].command, Some("claude".to_string()));
}
#[test]
fn resolve_pane_configuration_agent_creates_new_pane_when_empty() {
let result = resolve_pane_configuration(&[], Some("claude"));
assert_eq!(result.len(), 1);
assert_eq!(result[0].command, Some("claude".to_string()));
assert!(result[0].focus);
}
fn make_config_with_agent(agent: Option<&str>) -> config::Config {
config::Config {
agent: agent.map(|s| s.to_string()),
..Default::default()
}
}
fn make_options_with_prompt(run_pane_commands: bool) -> crate::workflow::types::SetupOptions {
crate::workflow::types::SetupOptions {
run_hooks: true,
run_file_ops: true,
run_pane_commands,
prompt_file_path: Some(std::path::PathBuf::from("/tmp/prompt.md")),
focus_window: true,
working_dir: None,
config_root: None,
open_if_exists: false,
mode: crate::config::MuxMode::default(),
resume_mode: crate::multiplexer::types::ResumeMode::default(),
}
}
#[test]
fn validate_prompt_errors_when_pane_commands_disabled() {
let panes = vec![config::PaneConfig {
command: Some("<agent>".to_string()),
focus: true,
..Default::default()
}];
let config = make_config_with_agent(Some("claude"));
let options = make_options_with_prompt(false);
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("pane commands are disabled")
);
}
#[test]
fn validate_prompt_errors_when_no_agent_configured() {
let panes = vec![config::PaneConfig {
command: Some("vim".to_string()),
focus: true,
..Default::default()
}];
let config = make_config_with_agent(None); let options = make_options_with_prompt(true);
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("no agent is configured")
);
}
#[test]
fn validate_prompt_errors_when_no_pane_runs_agent() {
let panes = vec![
config::PaneConfig {
focus: true,
..Default::default()
},
config::PaneConfig {
command: Some("clear".to_string()),
split: Some(config::SplitDirection::Horizontal),
..Default::default()
},
];
let config = make_config_with_agent(Some("claude"));
let options = make_options_with_prompt(true);
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_err());
let err_msg = result.unwrap_err().to_string();
assert!(err_msg.contains("no pane is configured to run the agent"));
assert!(err_msg.contains("claude"));
}
#[test]
fn validate_prompt_succeeds_with_agent_placeholder() {
let panes = vec![config::PaneConfig {
command: Some("<agent>".to_string()),
focus: true,
..Default::default()
}];
let config = make_config_with_agent(Some("claude"));
let options = make_options_with_prompt(true);
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_ok());
}
#[test]
fn validate_prompt_succeeds_with_matching_agent_command() {
let panes = vec![config::PaneConfig {
command: Some("claude".to_string()),
focus: true,
..Default::default()
}];
let config = make_config_with_agent(Some("claude"));
let options = make_options_with_prompt(true);
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_ok());
}
#[test]
fn validate_prompt_cli_agent_overrides_config() {
let panes = vec![config::PaneConfig {
command: Some("my-custom-agent".to_string()),
focus: true,
..Default::default()
}];
let config = make_config_with_agent(Some("claude")); let options = make_options_with_prompt(true);
let result =
super::validate_prompt_consumption(&panes, Some("my-custom-agent"), &config, &options);
assert!(result.is_ok());
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_err());
}
#[test]
fn validate_prompt_succeeds_when_any_pane_matches() {
let panes = vec![
config::PaneConfig {
command: Some("vim".to_string()), ..Default::default()
},
config::PaneConfig {
command: Some("claude --verbose".to_string()), focus: true,
split: Some(config::SplitDirection::Horizontal),
..Default::default()
},
];
let config = make_config_with_agent(Some("claude"));
let options = make_options_with_prompt(true);
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_ok());
}
#[test]
fn validate_prompt_succeeds_with_known_agent_command() {
let panes = vec![config::PaneConfig {
command: Some("codex --yolo".to_string()),
focus: true,
..Default::default()
}];
let config = make_config_with_agent(None); let options = make_options_with_prompt(true);
let result = super::validate_prompt_consumption(&panes, None, &config, &options);
assert!(result.is_ok());
}
#[test]
fn resolve_pane_configuration_known_agent_returns_original() {
let original_panes = vec![
config::PaneConfig {
command: Some("claude --dangerously-skip-permissions".to_string()),
focus: true,
..Default::default()
},
config::PaneConfig {
command: Some("codex --yolo".to_string()),
split: Some(config::SplitDirection::Vertical),
..Default::default()
},
];
let result = resolve_pane_configuration(&original_panes, Some("gemini"));
assert_eq!(
result[0].command.as_deref(),
Some("claude --dangerously-skip-permissions")
);
assert_eq!(result[1].command.as_deref(), Some("codex --yolo"));
}
#[test]
fn write_prompt_file_sanitizes_branch_with_slashes() {
use crate::prompt::Prompt;
let branch_name = "feature/nested/add-login";
let prompt = Prompt::Inline("test prompt content".to_string());
let path = super::write_prompt_file(None, branch_name, &prompt)
.expect("Should create prompt file");
let filename = path.file_name().unwrap().to_str().unwrap();
assert!(
filename.contains("feature-nested-add-login"),
"Expected sanitized branch name in filename, got: {}",
filename
);
assert!(
!filename.contains('/'),
"Filename should not contain slashes"
);
let content = std::fs::read_to_string(&path).unwrap();
assert_eq!(content, "test prompt content");
let _ = std::fs::remove_file(path);
}
#[test]
fn write_prompt_file_with_working_dir() {
use crate::prompt::Prompt;
use tempfile::TempDir;
let temp = TempDir::new().unwrap();
let branch_name = "feature/test";
let prompt = Prompt::Inline("test prompt".to_string());
let path = super::write_prompt_file(Some(temp.path()), branch_name, &prompt)
.expect("Should create prompt file");
assert!(path.starts_with(temp.path().join(".workmux")));
assert!(
path.file_name()
.unwrap()
.to_str()
.unwrap()
.contains("PROMPT-feature-test")
);
let content = std::fs::read_to_string(&path).unwrap();
assert_eq!(content, "test prompt");
}
}
fn validate_prompt_consumption(
panes: &[config::PaneConfig],
cli_agent: Option<&str>,
config: &config::Config,
options: &super::types::SetupOptions,
) -> Result<()> {
if !options.run_pane_commands {
return Err(anyhow!(
"Prompt provided (-p/-P/-e) but pane commands are disabled (--no-pane-cmds). \
The prompt would be ignored."
));
}
let has_self_identifying_agent = panes.iter().any(|pane| {
pane.command
.as_deref()
.is_some_and(crate::multiplexer::agent::is_known_agent)
});
if has_self_identifying_agent {
return Ok(());
}
let resolved_cli_agent = cli_agent.map(|a| {
config
.agents
.get(a)
.map(|e| e.command.as_str())
.unwrap_or(a)
});
let effective_agent = resolved_cli_agent.or(config.agent.as_deref());
let Some(agent_cmd) = effective_agent else {
return Err(anyhow!(
"Prompt provided but no agent is configured to consume it. \
Set 'agent' in config or use -a/--agent flag."
));
};
let consumes_prompt = panes.iter().any(|pane| {
pane.command
.as_deref()
.map(|cmd| config::is_agent_command(cmd, agent_cmd))
.unwrap_or(false)
});
if !consumes_prompt {
let commands: Vec<_> = panes
.iter()
.map(|p| p.command.as_deref().unwrap_or("<shell>"))
.collect();
return Err(anyhow!(
"Prompt provided, but no pane is configured to run the agent '{}'.\n\
Resolved pane commands: {:?}\n\
Ensure your panes config includes '<agent>' or runs the configured agent.",
agent_cmd,
commands
));
}
Ok(())
}