use anyhow::{Context, Result, bail};
use log::{debug, info};
use zag_agent::config::Config;
use zag_agent::factory::AgentFactory;
use zag_agent::{auto_selector, mcp, sandbox, session, skills, worktree};
use crate::cleanup::{
print_resume_hint, print_session_resume_hint, prompt_sandbox_cleanup, prompt_worktree_cleanup,
};
use crate::cli::Commands;
use crate::json_mode::{augment_system_prompt_for_json, handle_json_output, wrap_prompt_for_json};
use crate::output::print_agent_output;
use crate::resume::{
current_workspace, discover_provider_session_id, resolve_continue_target, resolve_resume_target,
};
use crate::session_setup::{
SessionMetadata, save_session_mapping, setup_plain_session, setup_sandbox, setup_worktree,
update_provider_session_id, update_session_log_metadata,
};
const EXIT_AGENT_FAILURE: i32 = 1;
const EXIT_PROVIDER_ERROR: i32 = 2;
pub(crate) struct AgentActionParams {
pub(crate) agent_name: String,
pub(crate) provider: String,
pub(crate) provider_explicit: bool,
pub(crate) action: Commands,
pub(crate) system_prompt: Option<String>,
pub(crate) model: Option<String>,
pub(crate) root: Option<String>,
pub(crate) auto_approve: bool,
pub(crate) add_dirs: Vec<String>,
pub(crate) show_usage: bool,
pub(crate) quiet: bool,
pub(crate) verbose: bool,
pub(crate) worktree: Option<Option<String>>,
pub(crate) sandbox: Option<Option<String>>,
pub(crate) size: Option<String>,
pub(crate) json_mode: bool,
pub(crate) json_schema: Option<serde_json::Value>,
pub(crate) json_stream: bool,
pub(crate) session: Option<String>,
pub(crate) max_turns: Option<u32>,
pub(crate) mcp_config: Option<String>,
pub(crate) timeout: Option<String>,
pub(crate) exit_on_failure: bool,
pub(crate) context_session: Option<String>,
pub(crate) plan_path: Option<String>,
pub(crate) env_vars: Vec<(String, String)>,
pub(crate) files: Vec<String>,
pub(crate) session_metadata: SessionMetadata,
}
pub(crate) fn run_resume_id(action: &Commands) -> Option<&str> {
match action {
Commands::Run { resume, .. } => resume.as_deref(),
_ => None,
}
}
fn run_continue_requested(action: &Commands) -> bool {
matches!(
action,
Commands::Run {
continue_session: true,
..
}
)
}
pub(crate) fn is_resume_run(action: &Commands) -> bool {
run_resume_id(action).is_some() || run_continue_requested(action)
}
fn run_prompt(action: &Commands) -> Option<&str> {
match action {
Commands::Run { prompt, .. } => prompt.as_deref(),
Commands::Exec { prompt, .. } => Some(prompt.as_str()),
_ => None,
}
}
fn is_new_interactive_run(action: &Commands, json_mode: bool) -> bool {
matches!(action, Commands::Run { .. })
&& !is_resume_run(action)
&& !(json_mode && run_prompt(action).is_some())
}
async fn resolve_auto_selection(params: &mut AgentActionParams) -> Result<()> {
let is_auto_provider = params.provider == "auto";
let is_auto_model = params.model.as_deref() == Some("auto");
if !is_auto_provider && !is_auto_model {
return Ok(());
}
let task_prompt = run_prompt(¶ms.action);
let task_prompt = task_prompt
.ok_or_else(|| anyhow::anyhow!("auto provider/model requires a prompt to analyze"))?;
let config = Config::load(params.root.as_deref()).unwrap_or_default();
let current_provider = if !is_auto_provider {
Some(params.provider.as_str())
} else {
None
};
let result = auto_selector::resolve(
task_prompt,
is_auto_provider,
is_auto_model,
current_provider,
&config,
params.root.as_deref(),
)
.await?;
if let Some(p) = result.provider {
params.provider = p;
}
if let Some(m) = result.model {
params.model = Some(m);
} else if is_auto_provider {
params.model = None;
}
params.agent_name = crate::capitalize(¶ms.provider);
let is_exec_action = matches!(params.action, Commands::Exec { .. });
let show_wrapper = !params.quiet && (!is_exec_action || params.verbose);
if show_wrapper {
let model_info = params
.model
.as_deref()
.map(|m| format!(" with model {}", m))
.unwrap_or_default();
println!(
"\x1b[32m✓\x1b[0m Auto-selected: {}{}",
params.agent_name, model_info
);
}
Ok(())
}
struct AgentSetupParams {
provider: String,
agent_name: String,
system_prompt: Option<String>,
model: Option<String>,
effective_root: Option<String>,
session_id: Option<String>,
auto_approve: bool,
add_dirs: Vec<String>,
output_format: Option<String>,
input_format: Option<String>,
replay_user_messages: bool,
include_partial_messages: bool,
verbose: bool,
json_mode: bool,
json_stream: bool,
max_turns: Option<u32>,
mcp_config: Option<String>,
env_vars: Vec<(String, String)>,
}
fn create_and_configure_agent(
p: AgentSetupParams,
json_schema: &Option<serde_json::Value>,
show_wrapper: bool,
) -> Result<(Box<dyn crate::agent::Agent + Send + Sync>, Option<String>)> {
let spinner = if show_wrapper {
crate::logging::spinner(format!("Initializing {} agent", p.agent_name))
} else {
let pb = indicatif::ProgressBar::new_spinner();
pb.set_draw_target(indicatif::ProgressDrawTarget::hidden());
pb
};
let mut agent = AgentFactory::create(
&p.provider,
p.system_prompt,
p.model,
p.effective_root,
p.auto_approve,
p.add_dirs,
)?;
let output_fmt_clone = p.output_format.clone();
agent.set_output_format(p.output_format);
if let Some(turns) = p.max_turns {
agent.set_max_turns(turns);
}
if !p.env_vars.is_empty() {
agent.set_env_vars(p.env_vars);
}
if p.provider == "claude"
&& let Some(claude_agent) = agent.as_any_mut().downcast_mut::<crate::claude::Claude>()
{
claude_agent.set_verbose(p.verbose);
if let Some(session_id) = p.session_id {
claude_agent.set_session_id(session_id);
}
if let Some(input_fmt) = p.input_format {
claude_agent.set_input_format(Some(input_fmt));
}
if p.replay_user_messages {
claude_agent.set_replay_user_messages(true);
}
if p.include_partial_messages {
claude_agent.set_include_partial_messages(true);
}
if p.json_mode
&& let Some(schema) = json_schema
{
let schema_str = serde_json::to_string(schema).unwrap_or_default();
claude_agent.set_json_schema(Some(schema_str));
}
if p.mcp_config.is_some() {
claude_agent.set_mcp_config(p.mcp_config);
}
let is_stream_json = p.json_stream || output_fmt_clone.as_deref() == Some("stream-json");
claude_agent.set_event_handler(Box::new(move |event, verbose| {
use crate::output::{ContentBlock, Event};
if is_stream_json {
if let Ok(json) = serde_json::to_string(event) {
println!("{}", json);
}
} else {
match event {
Event::Result { .. } => {
if !verbose {
use std::io::Write;
println!();
let _ = std::io::stdout().flush();
}
}
_ => {
if verbose {
if let Some(formatted) = crate::output::format_event_as_text(event) {
println!("{}", formatted);
}
} else if let Event::AssistantMessage { content, .. } = event {
for block in content {
if let ContentBlock::Text { text } = block {
print!("{}", text);
}
}
}
}
}
}
}));
}
let user_output_format = output_fmt_clone.clone();
if p.json_mode && user_output_format.is_none() {
agent.set_output_format(Some("json".to_string()));
if p.provider != "claude" {
agent.set_capture_output(true);
}
}
if p.json_stream && user_output_format.is_none() {
agent.set_output_format(Some("stream-json".to_string()));
}
crate::logging::finish_spinner_quiet(&spinner);
debug!("Agent configuration complete");
Ok((agent, output_fmt_clone))
}
struct ExecutionContext<'a> {
provider: &'a str,
json_mode: bool,
json_schema: &'a Option<serde_json::Value>,
output_fmt: Option<&'a str>,
show_usage: bool,
verbose: bool,
}
async fn execute_action(
action: Commands,
agent: &mut (dyn crate::agent::Agent + Send + Sync),
ctx: &ExecutionContext<'_>,
log_writer: Option<&crate::session_log::SessionLogWriter>,
) -> Result<bool> {
match action {
Commands::Run {
prompt,
resume,
continue_session,
..
} => {
if resume.is_some() || continue_session {
if let Some(ref session_id) = resume {
info!("Resuming session {}", session_id);
} else {
info!("Resuming latest session");
}
agent
.run_resume(resume.as_deref(), continue_session)
.await?;
} else if ctx.json_mode && prompt.is_some() {
info!("Starting non-interactive session (JSON mode)");
let wrapped = if ctx.provider != "claude" {
let w = prompt.as_deref().map(wrap_prompt_for_json);
if let Some(ref wp) = w {
debug!("JSON-wrapped run prompt: {}", wp);
}
w
} else {
debug!("Run prompt (JSON mode, Claude): {:?}", prompt);
None
};
let run_prompt = wrapped.as_deref().or(prompt.as_deref());
let agent_output = agent.run(run_prompt).await?;
if let (Some(writer), Some(agent_output)) = (log_writer, agent_output.as_ref()) {
crate::session_log::record_agent_output(writer, agent_output)?;
}
handle_json_output(
agent_output,
agent,
ctx.json_schema,
ctx.show_usage,
ctx.verbose,
)
.await?;
} else {
info!("Starting interactive session");
agent.run_interactive(prompt.as_deref()).await?;
}
}
Commands::Exec { prompt, .. } => {
info!("Starting non-interactive session");
let run_prompt = if ctx.json_mode && ctx.provider != "claude" {
let wrapped = wrap_prompt_for_json(&prompt);
debug!("JSON-wrapped prompt: {}", wrapped);
wrapped
} else {
debug!("Exec prompt: {}", prompt);
prompt.clone()
};
let agent_output = agent.run(Some(&run_prompt)).await?;
if let (Some(writer), Some(agent_output)) = (log_writer, agent_output.as_ref()) {
crate::session_log::record_agent_output(writer, agent_output)?;
}
let agent_success = agent_output.as_ref().map(|o| !o.is_error).unwrap_or(true);
if ctx.json_mode {
handle_json_output(
agent_output,
agent,
ctx.json_schema,
ctx.show_usage,
ctx.verbose,
)
.await?;
} else if let Some(agent_out) = agent_output {
print_agent_output(&agent_out, ctx.output_fmt, ctx.show_usage, ctx.verbose)?;
}
return Ok(agent_success);
}
_ => unreachable!(),
}
Ok(true)
}
fn log_config_details(params: &AgentActionParams) {
if let Some(ref m) = params.model {
debug!("Model specified: {}", m);
}
if let Some(ref r) = params.root {
debug!("Root directory: {}", r);
}
if params.auto_approve {
debug!("Auto-approve enabled");
}
if let Some(ref sp) = params.system_prompt {
debug!("System prompt: {}", sp);
}
if !params.add_dirs.is_empty() {
debug!("Additional directories: {:?}", params.add_dirs);
}
if params.worktree.is_some() {
debug!("Worktree mode enabled");
}
if params.sandbox.is_some() {
debug!("Sandbox mode enabled");
}
if params.json_mode {
debug!("JSON output mode enabled");
}
}
fn command_name(action: &Commands) -> &'static str {
match action {
Commands::Run { .. } => "run",
Commands::Exec { .. } => "exec",
Commands::Review { .. } => "review",
Commands::Config { .. } => "config",
Commands::Session { .. } => "session",
Commands::Capability { .. } => "capability",
Commands::Discover { .. } => "discover",
Commands::Listen { .. } => "listen",
Commands::Man { .. } => "man",
Commands::Skills { .. } => "skills",
Commands::Mcp { .. } => "mcp",
Commands::Ps { .. } => "ps",
Commands::Search { .. } => "search",
Commands::Input { .. } => "input",
Commands::Broadcast { .. } => "broadcast",
Commands::Whoami { .. } => "whoami",
Commands::Env { .. } => "env",
Commands::Collect { .. } => "collect",
Commands::Status { .. } => "status",
Commands::Wait { .. } => "wait",
Commands::Spawn { .. } => "spawn",
Commands::Pipe { .. } => "pipe",
Commands::Events { .. } => "events",
Commands::Cancel { .. } => "cancel",
Commands::Summary { .. } => "summary",
Commands::Watch { .. } => "watch",
Commands::Subscribe { .. } => "subscribe",
Commands::Log { .. } => "log",
Commands::Output { .. } => "output",
Commands::Retry { .. } => "retry",
Commands::Gc { .. } => "gc",
Commands::Serve { .. } => "serve",
Commands::Connect { .. } => "connect",
Commands::Disconnect => "disconnect",
Commands::Relay { .. } => "relay",
Commands::User { .. } => "user",
Commands::Plan { .. } => "plan",
}
}
fn action_prompt(action: &Commands) -> Option<&str> {
match action {
Commands::Run { prompt, .. } => prompt.as_deref(),
Commands::Exec { prompt, .. } => Some(prompt.as_str()),
_ => None,
}
}
fn should_enable_live_session_logs(action: &Commands, json_mode: bool) -> bool {
matches!(action, Commands::Run { .. }) && !json_mode
}
pub(crate) async fn run_agent_action(mut params: AgentActionParams) -> Result<()> {
resolve_auto_selection(&mut params).await?;
log_config_details(¶ms);
let session_metadata = std::mem::take(&mut params.session_metadata);
let AgentActionParams {
agent_name: _,
mut provider,
provider_explicit,
mut action,
system_prompt,
mut model,
root,
auto_approve,
add_dirs,
show_usage,
quiet,
verbose,
worktree: worktree_flag,
sandbox: sandbox_flag,
size,
json_mode,
json_schema,
json_stream,
session,
max_turns,
mcp_config,
timeout,
exit_on_failure,
context_session,
plan_path,
env_vars,
files,
session_metadata: _,
} = params;
let config = Config::load(root.as_deref()).unwrap_or_default();
let max_turns = max_turns.or(config.max_turns());
let system_prompt = system_prompt.or_else(|| config.system_prompt().map(String::from));
let is_exec = matches!(action, Commands::Exec { .. });
let show_wrapper = !quiet && (!is_exec || verbose);
let mut system_prompt =
augment_system_prompt_for_json(system_prompt, json_mode, &provider, &json_schema);
if let Err(e) = skills::setup_skills(&provider, &mut system_prompt) {
log::warn!("Failed to set up skills: {}", e);
}
if let Err(e) = mcp::setup_mcp(&provider, root.as_deref()) {
log::warn!("Failed to set up MCP servers: {}", e);
}
if let Some(ref sp) = system_prompt {
debug!("Effective system prompt: {}", sp);
}
let resume_target = if let Some(session_id) = run_resume_id(&action) {
resolve_resume_target(session_id, root.as_deref())
} else if run_continue_requested(&action) {
resolve_continue_target(root.as_deref())
} else {
None
};
if is_resume_run(&action) && resume_target.is_none() {
bail!("No matching session found to resume");
}
if let Some(target) = &resume_target {
if provider_explicit && provider != target.entry.provider {
bail!(
"Requested provider '{}' does not match the stored session provider '{}'",
provider,
target.entry.provider
);
}
provider = target.entry.provider.clone();
if !target.entry.model.is_empty() {
debug!(
"Restored model from session entry: '{}'",
target.entry.model
);
model = Some(target.entry.model.clone());
} else {
debug!("Session entry has empty model, will fall back to config/default");
}
}
if let Some(target) = &resume_target {
let native_id = target
.entry
.provider_session_id
.clone()
.unwrap_or_else(|| target.entry.session_id.clone());
if let Commands::Run {
resume,
continue_session,
..
} = &mut action
{
*resume = Some(native_id);
*continue_session = false;
}
}
let is_resume = is_resume_run(&action);
let plain = setup_plain_session(is_new_interactive_run(&action, json_mode), &root, &session);
let wrapper_session_id = plain.session_id.clone();
let log_session_id = wrapper_session_id
.clone()
.or_else(|| {
resume_target
.as_ref()
.map(|target| target.entry.session_id.clone())
})
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let wt = setup_worktree(
&worktree_flag,
is_resume,
&root,
show_wrapper,
wrapper_session_id.clone(),
)?;
let sb = setup_sandbox(&sandbox_flag, is_resume, &root, wrapper_session_id.clone())?;
let effective_root = if let Some(target) = &resume_target {
if target.entry.is_worktree {
let wt_path = std::path::Path::new(&target.entry.worktree_path);
if !wt_path.exists() && target.matched_by_wrapper_id {
log::warn!(
"Worktree no longer exists at {}, resuming without it",
target.entry.worktree_path
);
let mut store = session::SessionStore::load(root.as_deref()).unwrap_or_default();
store.remove(&target.entry.session_id);
let _ = store.save(root.as_deref());
Some(current_workspace(root.as_deref()))
} else {
Some(target.entry.worktree_path.clone())
}
} else {
Some(target.entry.worktree_path.clone())
}
} else {
wt.effective_root
.clone()
.or_else(|| plain.workspace_path.clone())
};
let (output_format, input_format, replay_user_messages, include_partial_messages) =
match &action {
Commands::Exec {
output,
input_format,
replay_user_messages,
include_partial_messages,
..
} => (
output.clone(),
input_format.clone(),
*replay_user_messages,
*include_partial_messages,
),
_ => (None, None, false, false),
};
if let Some(ref o) = output_format {
debug!("Output format: {}", o);
}
if let Some(ref i) = input_format {
debug!("Input format: {}", i);
}
let (mut agent, output_fmt_clone) = create_and_configure_agent(
AgentSetupParams {
provider: provider.clone(),
agent_name: crate::capitalize(&provider),
system_prompt,
model,
effective_root: effective_root.clone(),
session_id: wrapper_session_id.clone(),
auto_approve,
add_dirs,
output_format,
input_format,
replay_user_messages,
include_partial_messages,
verbose,
json_mode,
json_stream,
max_turns,
mcp_config,
env_vars,
},
&json_schema,
show_wrapper,
)?;
if sb.is_sandbox_session
&& let (Some(name), Some(workspace)) = (&sb.sandbox_name, &sb.workspace)
{
let config = sandbox::SandboxConfig {
name: name.clone(),
template: sandbox::template_for_provider(&provider).to_string(),
workspace: workspace.clone(),
};
agent.set_sandbox(config);
if show_wrapper {
println!("\x1b[32m✓\x1b[0m Sandbox configured: {}", name);
}
}
if let Some(target) = &resume_target
&& let Some(name) = &target.entry.sandbox_name
{
let config = sandbox::SandboxConfig {
name: name.clone(),
template: sandbox::template_for_provider(&provider).to_string(),
workspace: target.entry.worktree_path.clone(),
};
agent.set_sandbox(config);
if show_wrapper {
println!("\x1b[32m✓\x1b[0m Sandbox configured: {}", name);
}
}
if provider == "ollama" {
let config = Config::load(root.as_deref()).unwrap_or_default();
let current_model = agent.get_model().to_string();
let is_size_value = crate::ollama::AVAILABLE_SIZES.contains(¤t_model.as_str());
if is_size_value {
agent.set_model(config.ollama_model().to_string());
} else if current_model == crate::ollama::DEFAULT_MODEL {
agent.set_model(config.ollama_model().to_string());
}
if let Some(ollama_agent) = agent.as_any_mut().downcast_mut::<crate::ollama::Ollama>() {
if let Some(ref s) = size {
let resolved = config.ollama_size_for(s).to_string();
ollama_agent.set_size(resolved);
} else if is_size_value {
ollama_agent.set_size(current_model);
} else {
ollama_agent.set_size(config.ollama_size().to_string());
}
}
}
let model_display = if provider == "ollama" {
if let Some(ollama_agent) = agent.as_any_mut().downcast_mut::<crate::ollama::Ollama>() {
ollama_agent.display_model()
} else {
agent.get_model().to_string()
}
} else {
agent.get_model().to_string()
};
let persisted_model = agent.get_model().to_string();
let auto_approve_suffix = if auto_approve { " (auto approve)" } else { "" };
if show_wrapper {
println!(
"\x1b[32m✓\x1b[0m {} initialized with model {}{}",
crate::capitalize(&provider),
model_display,
auto_approve_suffix
);
}
save_session_mapping(
&plain,
&wt,
&sb,
&provider,
&persisted_model,
root.as_deref(),
&session_metadata,
);
let proc_id = uuid::Uuid::new_v4().to_string();
let proc_session_id = wt
.session_id
.clone()
.or_else(|| sb.session_id.clone())
.or_else(|| plain.session_id.clone());
let proc_prompt = action_prompt(&action).map(|p| p.chars().take(100).collect::<String>());
let proc_cmd = command_name(&action).to_string();
let parent_process_id = std::env::var("ZAG_PROCESS_ID").ok();
let parent_session_id = std::env::var("ZAG_SESSION_ID").ok();
if let Ok(mut pstore) = zag_agent::process_store::ProcessStore::load() {
pstore.add(zag_agent::process_store::ProcessEntry {
id: proc_id.clone(),
pid: std::process::id(),
session_id: proc_session_id.clone(),
provider: provider.clone(),
model: persisted_model.clone(),
command: proc_cmd,
prompt: proc_prompt,
started_at: chrono::Utc::now().to_rfc3339(),
status: "running".to_string(),
exit_code: None,
exited_at: None,
root: root.clone(),
parent_process_id,
parent_session_id,
});
let _ = pstore.save();
}
unsafe {
if let Some(ref sid) = proc_session_id {
std::env::set_var("ZAG_SESSION_ID", sid);
}
std::env::set_var("ZAG_PROCESS_ID", &proc_id);
std::env::set_var("ZAG_PROVIDER", &provider);
std::env::set_var("ZAG_MODEL", &persisted_model);
if let Some(ref r) = root {
std::env::set_var("ZAG_ROOT", r);
}
if let Some(ref name) = session_metadata.name {
std::env::set_var("ZAG_SESSION_NAME", name);
}
}
let lifecycle_session_id = wt
.session_id
.as_deref()
.or(sb.session_id.as_deref())
.or(plain.session_id.as_deref())
.unwrap_or(&log_session_id)
.to_string();
zag_orch::lifecycle::write_started_marker(&lifecycle_session_id);
zag_orch::lifecycle::prune_old_markers();
if show_wrapper {
let display_session_id = wt
.session_id
.as_deref()
.or(sb.session_id.as_deref())
.or(plain.session_id.as_deref())
.unwrap_or(&log_session_id);
println!("\x1b[33m>\x1b[0m Session: {}", display_session_id);
println!(
"\x1b[33m>\x1b[0m Listen: agent listen {}",
display_session_id
);
}
let initial_provider_session_id = if provider == "claude" {
wrapper_session_id.clone()
} else {
resume_target
.as_ref()
.and_then(|target| target.entry.provider_session_id.clone())
};
let log_metadata = crate::session_log::SessionLogMetadata {
provider: provider.clone(),
wrapper_session_id: log_session_id.clone(),
provider_session_id: initial_provider_session_id,
workspace_path: effective_root
.clone()
.or_else(|| plain.workspace_path.clone())
.or_else(|| wt.worktree_path.clone())
.or_else(|| sb.workspace.clone()),
command: command_name(&action).to_string(),
model: Some(persisted_model.clone()),
resumed: is_resume_run(&action),
backfilled: false,
};
let live_ctx = crate::session_log::LiveLogContext {
root: root.clone(),
provider_session_id: log_metadata.provider_session_id.clone(),
workspace_path: log_metadata.workspace_path.clone(),
started_at: chrono::Utc::now(),
is_worktree: wt.is_worktree_session,
};
let live_adapter = crate::session_log::live_adapter_for_provider(
&provider,
live_ctx,
should_enable_live_session_logs(&action, json_mode),
);
let log_coordinator = crate::session_log::SessionLogCoordinator::start(
&crate::session_log::logs_dir(root.as_deref()),
log_metadata,
live_adapter,
)?;
let _ = log_coordinator
.writer()
.set_global_index_dir(Config::global_base_dir());
crate::session_log::record_prompt(log_coordinator.writer(), action_prompt(&action))?;
if let Ok(log_path) = log_coordinator.writer().log_path() {
update_session_log_metadata(
wrapper_session_id
.as_deref()
.or(wt.session_id.as_deref())
.or(sb.session_id.as_deref()),
Some(log_path.to_string_lossy().to_string()),
"partial",
root.as_deref(),
);
}
if let Some(ref ctx_session_id) = context_session {
let context_text =
zag_orch::collect::extract_last_assistant_message(ctx_session_id, root.as_deref());
if let Some(context_text) = context_text {
let prefix = format!(
"Context from previous session ({}):\n\n{}\n\n---\n\n",
ctx_session_id, context_text
);
match &mut action {
Commands::Exec { prompt, .. } => {
*prompt = format!("{}{}", prefix, prompt);
}
Commands::Run {
prompt: Some(p), ..
} => {
*p = format!("{}{}", prefix, p);
}
_ => {}
}
}
}
if let Some(ref plan_file) = plan_path {
let plan_content = std::fs::read_to_string(plan_file)
.with_context(|| format!("Failed to read plan file: {}", plan_file))?;
let prefix = format!(
"Implementation plan:\n\n{}\n\n---\n\nFollow the plan above.\n\n",
plan_content
);
match &mut action {
Commands::Exec { prompt, .. } => {
*prompt = format!("{}{}", prefix, prompt);
}
Commands::Run {
prompt: Some(p), ..
} => {
*p = format!("{}{}", prefix, p);
}
Commands::Run { prompt, .. } => {
*prompt = Some(prefix);
}
_ => {}
}
}
if !files.is_empty() {
let attachments = files
.iter()
.map(|f| zag_agent::attachment::Attachment::from_path(std::path::Path::new(f)))
.collect::<Result<Vec<_>>>()?;
let prefix = zag_agent::attachment::format_attachments_prefix(&attachments);
match &mut action {
Commands::Exec { prompt, .. } => {
*prompt = format!("{}{}", prefix, prompt);
}
Commands::Run {
prompt: Some(p), ..
} => {
*p = format!("{}{}", prefix, p);
}
Commands::Run { prompt, .. } => {
*prompt = Some(prefix);
}
_ => {}
}
}
let is_worktree_session = wt.is_worktree_session;
let is_interactive_worktree = wt.is_worktree_session && matches!(action, Commands::Run { .. });
let is_interactive_sandbox = sb.is_sandbox_session && matches!(action, Commands::Run { .. });
let is_interactive_run = matches!(action, Commands::Run { .. });
let exec_ctx = ExecutionContext {
provider: &provider,
json_mode,
json_schema: &json_schema,
output_fmt: output_fmt_clone.as_deref(),
show_usage,
verbose,
};
let action_future = execute_action(
action,
&mut *agent,
&exec_ctx,
Some(log_coordinator.writer()),
);
let action_result = if let Some(ref timeout_str) = timeout {
let duration = zag_orch::duration::parse_duration(timeout_str)?;
match tokio::time::timeout(duration, action_future).await {
Ok(r) => r,
Err(_) => Err(anyhow::anyhow!("Agent timed out after {}", timeout_str)),
}
} else {
action_future.await
};
debug!("Cleaning up agent resources");
if let Err(cleanup_err) = agent.cleanup().await {
log::warn!("Agent cleanup failed: {}", cleanup_err);
}
let agent_success = match &action_result {
Err(err) => {
let process_err = err.downcast_ref::<zag_agent::process::ProcessError>();
let exit_code = process_err.and_then(|pe| pe.exit_code).unwrap_or(1);
if let Ok(mut pstore) = zag_agent::process_store::ProcessStore::load() {
pstore.update_status(&proc_id, "killed", Some(exit_code));
let _ = pstore.save();
}
if let Err(log_err) = log_coordinator.finish(false, Some(err.to_string())).await {
log::warn!("Failed to finish session log: {}", log_err);
}
zag_orch::lifecycle::write_ended_marker(&lifecycle_session_id, false, Some(exit_code));
eprintln!("\x1b[31merror\x1b[0m: {}", err);
eprintln!("\x1b[2mRun with --debug for full details\x1b[0m");
if process_err.is_some() {
std::process::exit(EXIT_PROVIDER_ERROR);
}
return Err(anyhow::anyhow!(err.to_string()));
}
Ok(success) => *success,
};
let wrapper_session_id = wt
.session_id
.as_deref()
.or(sb.session_id.as_deref())
.or(plain.session_id.as_deref());
let live_discovered_id = log_coordinator.writer().get_provider_session_id();
let discovery_workspace = effective_root
.as_deref()
.or(plain.workspace_path.as_deref());
let native_session_id = live_discovered_id
.filter(|id| wrapper_session_id.is_none_or(|wid| id != wid))
.or_else(|| {
discover_provider_session_id(
&provider,
wrapper_session_id,
root.as_deref(),
discovery_workspace,
)
});
if let Some(ref native_id) = native_session_id {
log_coordinator
.writer()
.set_provider_session_id(Some(native_id.clone()))?;
}
update_provider_session_id(wrapper_session_id, native_session_id, root.as_deref());
update_session_log_metadata(
wrapper_session_id,
log_coordinator
.writer()
.log_path()
.ok()
.map(|path| path.to_string_lossy().to_string()),
"partial",
root.as_deref(),
);
let final_exit_code = if agent_success { 0 } else { EXIT_AGENT_FAILURE };
log_coordinator.finish(agent_success, None).await?;
zag_orch::lifecycle::write_ended_marker(
&lifecycle_session_id,
agent_success,
Some(final_exit_code),
);
if let Ok(mut pstore) = zag_agent::process_store::ProcessStore::load() {
pstore.update_status(&proc_id, "exited", Some(final_exit_code));
let _ = pstore.save();
}
if !agent_success {
eprintln!("\x1b[31merror\x1b[0m: agent exited with failure");
eprintln!("\x1b[2mRun with --debug for full details\x1b[0m");
}
if show_wrapper {
println!("\x1b[32m✓\x1b[0m Session terminated");
}
if is_interactive_sandbox {
if let Some(ref name) = sb.sandbox_name {
prompt_sandbox_cleanup(
sb.session_id.as_deref().unwrap_or(""),
name,
root.as_deref(),
)?;
}
} else if let Some(target) = &resume_target
&& let Some(ref sandbox_name) = target.entry.sandbox_name
&& target.matched_by_wrapper_id
{
let sid = target.entry.session_id.as_str();
prompt_sandbox_cleanup(sid, sandbox_name, root.as_deref())?;
}
let cleanup_info = if is_worktree_session {
wt.session_id
.as_ref()
.zip(wt.worktree_path.as_ref())
.map(|(sid, wtp)| (sid.clone(), wtp.clone()))
} else if let Some(target) = &resume_target {
if target.entry.is_worktree && target.matched_by_wrapper_id {
Some((
target.entry.session_id.clone(),
target.entry.worktree_path.clone(),
))
} else {
None
}
} else {
None
};
if let Some((sid, wtp)) = cleanup_info {
let wt_path = std::path::Path::new(&wtp);
let has_changes = wt_path.exists()
&& (worktree::has_changes(wt_path).unwrap_or(true)
|| worktree::has_unpushed_commits(wt_path).unwrap_or(true));
if !has_changes {
if wt_path.exists() {
match worktree::remove_worktree(wt_path) {
Ok(()) => {
if show_wrapper {
println!("\x1b[32m✓\x1b[0m Worktree removed (no changes)");
}
}
Err(e) => {
log::warn!("Failed to remove worktree: {}", e);
}
}
}
let mut store = session::SessionStore::load(root.as_deref()).unwrap_or_default();
store.remove(&sid);
let _ = store.save(root.as_deref());
} else if is_interactive_worktree {
prompt_worktree_cleanup(&sid, &wtp, root.as_deref())?;
} else {
if show_wrapper {
let store = session::SessionStore::load(root.as_deref()).unwrap_or_default();
let provider_session_id = store
.find_by_session_id(&sid)
.and_then(|entry| entry.provider_session_id.as_deref());
print_resume_hint(&sid, provider_session_id, "Workspace");
}
}
} else if let Some(wrapper_session_id) = wrapper_session_id {
if is_interactive_run && show_wrapper {
let store = session::SessionStore::load(root.as_deref()).unwrap_or_default();
let provider_session_id = store
.find_by_session_id(wrapper_session_id)
.and_then(|entry| entry.provider_session_id.clone());
print_session_resume_hint(wrapper_session_id, provider_session_id.as_deref());
}
} else if is_interactive_run
&& show_wrapper
&& let Some(target) = &resume_target
&& !target.entry.is_worktree
&& target.entry.sandbox_name.is_none()
{
let sid = &target.entry.session_id;
let store = session::SessionStore::load(root.as_deref()).unwrap_or_default();
let provider_session_id = store
.find_by_session_id(sid)
.and_then(|entry| entry.provider_session_id.clone());
print_session_resume_hint(sid, provider_session_id.as_deref());
}
if exit_on_failure && !agent_success {
std::process::exit(EXIT_AGENT_FAILURE);
}
Ok(())
}