zag_agent/

builder.rs

//! High-level builder API for driving agents programmatically.
//!
//! Instead of shelling out to the `agent` CLI binary, Rust programs can
//! use `AgentBuilder` to configure and execute agent sessions directly.
//!
//! # Examples
//!
//! ```no_run
//! use zag_agent::builder::AgentBuilder;
//!
//! # async fn example() -> anyhow::Result<()> {
//! // Non-interactive exec — returns structured output
//! let output = AgentBuilder::new()
//!     .provider("claude")
//!     .model("sonnet")
//!     .auto_approve(true)
//!     .exec("write a hello world program")
//!     .await?;
//!
//! println!("{}", output.result.unwrap_or_default());
//!
//! // Interactive session
//! AgentBuilder::new()
//!     .provider("claude")
//!     .run(Some("initial prompt"))
//!     .await?;
//! # Ok(())
//! # }
//! ```

use crate::agent::Agent;
use crate::attachment::{self, Attachment};
use crate::config::Config;
use crate::factory::AgentFactory;
use crate::json_validation;
use crate::output::AgentOutput;
use crate::progress::{ProgressHandler, SilentProgress};
use crate::providers::claude::Claude;
use crate::providers::ollama::Ollama;
use crate::sandbox::SandboxConfig;
use crate::streaming::StreamingSession;
use crate::worktree;
use anyhow::{Result, bail};
use log::{debug, warn};
use std::time::Duration;

/// Format a Duration as a human-readable string (e.g., "5m", "1h30m").
fn format_duration(d: Duration) -> String {
    let total_secs = d.as_secs();
    let h = total_secs / 3600;
    let m = (total_secs % 3600) / 60;
    let s = total_secs % 60;
    let mut parts = Vec::new();
    if h > 0 {
        parts.push(format!("{h}h"));
    }
    if m > 0 {
        parts.push(format!("{m}m"));
    }
    if s > 0 || parts.is_empty() {
        parts.push(format!("{s}s"));
    }
    parts.join("")
}

/// Builder for configuring and running agent sessions.
///
/// Use the builder pattern to set options, then call a terminal method
/// (`exec`, `run`, `resume`, `continue_last`) to execute.
pub struct AgentBuilder {
    provider: Option<String>,
    /// Set to true when the caller explicitly pinned a provider via
    /// `.provider()`. When false (default), the fallback tier list is
    /// allowed to downgrade to the next provider on binary/probe failure.
    provider_explicit: bool,
    model: Option<String>,
    system_prompt: Option<String>,
    root: Option<String>,
    auto_approve: bool,
    add_dirs: Vec<String>,
    files: Vec<String>,
    env_vars: Vec<(String, String)>,
    worktree: Option<Option<String>>,
    sandbox: Option<Option<String>>,
    size: Option<String>,
    json_mode: bool,
    json_schema: Option<serde_json::Value>,
    session_id: Option<String>,
    output_format: Option<String>,
    input_format: Option<String>,
    replay_user_messages: bool,
    include_partial_messages: bool,
    verbose: bool,
    quiet: bool,
    show_usage: bool,
    max_turns: Option<u32>,
    timeout: Option<std::time::Duration>,
    mcp_config: Option<String>,
    progress: Box<dyn ProgressHandler>,
}

impl Default for AgentBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl AgentBuilder {
    /// Create a new builder with default settings.
    pub fn new() -> Self {
        Self {
            provider: None,
            provider_explicit: false,
            model: None,
            system_prompt: None,
            root: None,
            auto_approve: false,
            add_dirs: Vec::new(),
            files: Vec::new(),
            env_vars: Vec::new(),
            worktree: None,
            sandbox: None,
            size: None,
            json_mode: false,
            json_schema: None,
            session_id: None,
            output_format: None,
            input_format: None,
            replay_user_messages: false,
            include_partial_messages: false,
            verbose: false,
            quiet: false,
            show_usage: false,
            max_turns: None,
            timeout: None,
            mcp_config: None,
            progress: Box::new(SilentProgress),
        }
    }

    /// Set the provider (e.g., "claude", "codex", "gemini", "copilot", "ollama").
    ///
    /// Calling this method pins the provider — it will NOT be downgraded to
    /// another provider in the tier list if its binary is missing or the
    /// startup probe fails. Omit this call (or set `provider` via the config
    /// file) to allow automatic downgrading.
    pub fn provider(mut self, provider: &str) -> Self {
        self.provider = Some(provider.to_string());
        self.provider_explicit = true;
        self
    }

    /// Set the model (e.g., "sonnet", "opus", "small", "large").
    pub fn model(mut self, model: &str) -> Self {
        self.model = Some(model.to_string());
        self
    }

    /// Set a system prompt to configure agent behavior.
    pub fn system_prompt(mut self, prompt: &str) -> Self {
        self.system_prompt = Some(prompt.to_string());
        self
    }

    /// Set the root directory for the agent to operate in.
    pub fn root(mut self, root: &str) -> Self {
        self.root = Some(root.to_string());
        self
    }

    /// Enable auto-approve mode (skip permission prompts).
    pub fn auto_approve(mut self, approve: bool) -> Self {
        self.auto_approve = approve;
        self
    }

    /// Add an additional directory for the agent to include.
    pub fn add_dir(mut self, dir: &str) -> Self {
        self.add_dirs.push(dir.to_string());
        self
    }

    /// Attach a file to the prompt (text files ≤50 KB inlined, others referenced).
    pub fn file(mut self, path: &str) -> Self {
        self.files.push(path.to_string());
        self
    }

    /// Add an environment variable for the agent subprocess.
    pub fn env(mut self, key: &str, value: &str) -> Self {
        self.env_vars.push((key.to_string(), value.to_string()));
        self
    }

    /// Enable worktree mode with an optional name.
    pub fn worktree(mut self, name: Option<&str>) -> Self {
        self.worktree = Some(name.map(String::from));
        self
    }

    /// Enable sandbox mode with an optional name.
    pub fn sandbox(mut self, name: Option<&str>) -> Self {
        self.sandbox = Some(name.map(String::from));
        self
    }

    /// Set the Ollama parameter size (e.g., "2b", "9b", "35b").
    pub fn size(mut self, size: &str) -> Self {
        self.size = Some(size.to_string());
        self
    }

    /// Request JSON output from the agent.
    pub fn json(mut self) -> Self {
        self.json_mode = true;
        self
    }

    /// Set a JSON schema for structured output validation.
    /// Implies `json()`.
    pub fn json_schema(mut self, schema: serde_json::Value) -> Self {
        self.json_schema = Some(schema);
        self.json_mode = true;
        self
    }

    /// Set a specific session ID (UUID).
    pub fn session_id(mut self, id: &str) -> Self {
        self.session_id = Some(id.to_string());
        self
    }

    /// Set the output format (e.g., "text", "json", "json-pretty", "stream-json").
    pub fn output_format(mut self, format: &str) -> Self {
        self.output_format = Some(format.to_string());
        self
    }

    /// Set the input format (Claude only, e.g., "text", "stream-json").
    ///
    /// No-op for Codex, Gemini, Copilot, and Ollama. See `docs/providers.md`
    /// for the full per-provider support matrix.
    pub fn input_format(mut self, format: &str) -> Self {
        self.input_format = Some(format.to_string());
        self
    }

    /// Re-emit user messages from stdin on stdout (Claude only).
    ///
    /// Only works with `--input-format stream-json` and `--output-format stream-json`.
    /// [`exec_streaming`](Self::exec_streaming) auto-enables this flag, so most
    /// callers never need to set it manually. No-op for non-Claude providers.
    pub fn replay_user_messages(mut self, replay: bool) -> Self {
        self.replay_user_messages = replay;
        self
    }

    /// Include partial message chunks in streaming output (Claude only).
    ///
    /// Only works with `--output-format stream-json`. Defaults to `false`.
    ///
    /// When `false` (the default), streaming surfaces one `assistant_message`
    /// event per complete assistant turn. When `true`, the agent instead emits
    /// a stream of token-level partial `assistant_message` chunks as the model
    /// generates them — use this for responsive, token-by-token UIs over
    /// [`exec_streaming`](Self::exec_streaming). No-op for non-Claude providers.
    pub fn include_partial_messages(mut self, include: bool) -> Self {
        self.include_partial_messages = include;
        self
    }

    /// Enable verbose output.
    pub fn verbose(mut self, v: bool) -> Self {
        self.verbose = v;
        self
    }

    /// Enable quiet mode (suppress all non-essential output).
    pub fn quiet(mut self, q: bool) -> Self {
        self.quiet = q;
        self
    }

    /// Show token usage statistics.
    pub fn show_usage(mut self, show: bool) -> Self {
        self.show_usage = show;
        self
    }

    /// Set the maximum number of agentic turns.
    pub fn max_turns(mut self, turns: u32) -> Self {
        self.max_turns = Some(turns);
        self
    }

    /// Set a timeout for exec. If the agent doesn't complete within this
    /// duration, it will be killed and an error returned.
    pub fn timeout(mut self, duration: std::time::Duration) -> Self {
        self.timeout = Some(duration);
        self
    }

    /// Set MCP server config for this invocation (Claude only).
    ///
    /// Accepts either a JSON string (`{"mcpServers": {...}}`) or a path to a JSON file.
    /// No-op for Codex, Gemini, Copilot, and Ollama — those providers manage
    /// MCP configuration through their own CLIs or do not support it. See
    /// `docs/providers.md` for the full per-provider support matrix.
    pub fn mcp_config(mut self, config: &str) -> Self {
        self.mcp_config = Some(config.to_string());
        self
    }

    /// Set a custom progress handler for status reporting.
    pub fn on_progress(mut self, handler: Box<dyn ProgressHandler>) -> Self {
        self.progress = handler;
        self
    }

    /// Resolve file attachments and prepend them to a prompt.
    fn prepend_files(&self, prompt: &str) -> Result<String> {
        if self.files.is_empty() {
            return Ok(prompt.to_string());
        }
        let attachments: Vec<Attachment> = self
            .files
            .iter()
            .map(|f| Attachment::from_path(std::path::Path::new(f)))
            .collect::<Result<Vec<_>>>()?;
        let prefix = attachment::format_attachments_prefix(&attachments);
        Ok(format!("{prefix}{prompt}"))
    }

    /// Resolve the effective provider name.
    fn resolve_provider(&self) -> Result<String> {
        if let Some(ref p) = self.provider {
            let p = p.to_lowercase();
            if !Config::VALID_PROVIDERS.contains(&p.as_str()) {
                bail!(
                    "Invalid provider '{}'. Available: {}",
                    p,
                    Config::VALID_PROVIDERS.join(", ")
                );
            }
            return Ok(p);
        }
        let config = Config::load(self.root.as_deref()).unwrap_or_default();
        if let Some(p) = config.provider() {
            return Ok(p.to_string());
        }
        Ok("claude".to_string())
    }

    /// Create and configure the agent.
    ///
    /// Returns the constructed agent along with the provider name that
    /// actually succeeded. When `provider_explicit` is false, the factory
    /// may downgrade to another provider in the tier list, so the returned
    /// provider can differ from the one passed in.
    async fn create_agent(&self, provider: &str) -> Result<(Box<dyn Agent + Send + Sync>, String)> {
        // Apply system_prompt config fallback
        let base_system_prompt = self.system_prompt.clone().or_else(|| {
            Config::load(self.root.as_deref())
                .unwrap_or_default()
                .system_prompt()
                .map(String::from)
        });

        // Augment system prompt with JSON instructions for non-Claude agents
        let system_prompt = if self.json_mode && provider != "claude" {
            let mut prompt = base_system_prompt.unwrap_or_default();
            if let Some(ref schema) = self.json_schema {
                let schema_str = serde_json::to_string_pretty(schema).unwrap_or_default();
                prompt.push_str(&format!(
                    "\n\nYou MUST respond with valid JSON only. No markdown fences, no explanations. \
                     Your response must conform to this JSON schema:\n{schema_str}"
                ));
            } else {
                prompt.push_str(
                    "\n\nYou MUST respond with valid JSON only. No markdown fences, no explanations.",
                );
            }
            Some(prompt)
        } else {
            base_system_prompt
        };

        self.progress
            .on_spinner_start(&format!("Initializing {provider} agent"));

        let progress = &*self.progress;
        let mut on_downgrade = |from: &str, to: &str, reason: &str| {
            progress.on_warning(&format!("Downgrading provider: {from} → {to} ({reason})"));
        };
        let (mut agent, effective_provider) = AgentFactory::create_with_fallback(
            provider,
            self.provider_explicit,
            system_prompt,
            self.model.clone(),
            self.root.clone(),
            self.auto_approve,
            self.add_dirs.clone(),
            &mut on_downgrade,
        )
        .await?;
        let provider = effective_provider.as_str();

        // Apply max_turns: explicit > config > none
        let effective_max_turns = self.max_turns.or_else(|| {
            Config::load(self.root.as_deref())
                .unwrap_or_default()
                .max_turns()
        });
        if let Some(turns) = effective_max_turns {
            agent.set_max_turns(turns);
        }

        // Set output format
        let mut output_format = self.output_format.clone();
        if self.json_mode && output_format.is_none() {
            output_format = Some("json".to_string());
            if provider != "claude" {
                agent.set_capture_output(true);
            }
        }
        agent.set_output_format(output_format);

        // Configure Claude-specific options
        if provider == "claude"
            && let Some(claude_agent) = agent.as_any_mut().downcast_mut::<Claude>()
        {
            claude_agent.set_verbose(self.verbose);
            if let Some(ref session_id) = self.session_id {
                claude_agent.set_session_id(session_id.clone());
            }
            if let Some(ref input_fmt) = self.input_format {
                claude_agent.set_input_format(Some(input_fmt.clone()));
            }
            if self.replay_user_messages {
                claude_agent.set_replay_user_messages(true);
            }
            if self.include_partial_messages {
                claude_agent.set_include_partial_messages(true);
            }
            if self.json_mode
                && let Some(ref schema) = self.json_schema
            {
                let schema_str = serde_json::to_string(schema).unwrap_or_default();
                claude_agent.set_json_schema(Some(schema_str));
            }
            if self.mcp_config.is_some() {
                claude_agent.set_mcp_config(self.mcp_config.clone());
            }
        }

        // Configure Ollama-specific options
        if provider == "ollama"
            && let Some(ollama_agent) = agent.as_any_mut().downcast_mut::<Ollama>()
        {
            let config = Config::load(self.root.as_deref()).unwrap_or_default();
            if let Some(ref size) = self.size {
                let resolved = config.ollama_size_for(size);
                ollama_agent.set_size(resolved.to_string());
            }
        }

        // Configure sandbox
        if let Some(ref sandbox_opt) = self.sandbox {
            let sandbox_name = sandbox_opt
                .as_deref()
                .map(String::from)
                .unwrap_or_else(crate::sandbox::generate_name);
            let template = crate::sandbox::template_for_provider(provider);
            let workspace = self.root.clone().unwrap_or_else(|| ".".to_string());
            agent.set_sandbox(SandboxConfig {
                name: sandbox_name,
                template: template.to_string(),
                workspace,
            });
        }

        if !self.env_vars.is_empty() {
            agent.set_env_vars(self.env_vars.clone());
        }

        self.progress.on_spinner_finish();
        self.progress.on_success(&format!(
            "{} initialized with model {}",
            provider,
            agent.get_model()
        ));

        Ok((agent, effective_provider))
    }

    /// Run the agent non-interactively and return structured output.
    ///
    /// This is the primary entry point for programmatic use.
    pub async fn exec(self, prompt: &str) -> Result<AgentOutput> {
        let provider = self.resolve_provider()?;
        debug!("exec: provider={provider}");

        // Set up worktree if requested
        let effective_root = if let Some(ref wt_opt) = self.worktree {
            let wt_name = wt_opt
                .as_deref()
                .map(String::from)
                .unwrap_or_else(worktree::generate_name);
            let repo_root = worktree::git_repo_root(self.root.as_deref())?;
            let wt_path = worktree::create_worktree(&repo_root, &wt_name)?;
            self.progress
                .on_success(&format!("Worktree created at {}", wt_path.display()));
            Some(wt_path.to_string_lossy().to_string())
        } else {
            self.root.clone()
        };

        let mut builder = self;
        if effective_root.is_some() {
            builder.root = effective_root;
        }

        let (agent, provider) = builder.create_agent(&provider).await?;

        // Prepend file attachments
        let prompt_with_files = builder.prepend_files(prompt)?;

        // Handle JSON mode with prompt wrapping for non-Claude agents
        let effective_prompt = if builder.json_mode && provider != "claude" {
            format!(
                "IMPORTANT: You MUST respond with valid JSON only. No markdown, no explanation.\n\n{prompt_with_files}"
            )
        } else {
            prompt_with_files
        };

        let result = if let Some(timeout_dur) = builder.timeout {
            match tokio::time::timeout(timeout_dur, agent.run(Some(&effective_prompt))).await {
                Ok(r) => r?,
                Err(_) => {
                    agent.cleanup().await.ok();
                    bail!("Agent timed out after {}", format_duration(timeout_dur));
                }
            }
        } else {
            agent.run(Some(&effective_prompt)).await?
        };

        // Clean up
        agent.cleanup().await?;

        if let Some(output) = result {
            // Validate JSON output if schema is provided
            if let Some(ref schema) = builder.json_schema {
                if !builder.json_mode {
                    warn!(
                        "json_schema is set but json_mode is false — \
                         schema will not be sent to the agent, only used for output validation"
                    );
                }
                if let Some(ref result_text) = output.result {
                    debug!(
                        "exec: validating result ({} bytes): {:.300}",
                        result_text.len(),
                        result_text
                    );
                    if let Err(errors) = json_validation::validate_json_schema(result_text, schema)
                    {
                        let preview = if result_text.len() > 500 {
                            &result_text[..500]
                        } else {
                            result_text.as_str()
                        };
                        bail!(
                            "JSON schema validation failed: {}\nRaw agent output ({} bytes):\n{}",
                            errors.join("; "),
                            result_text.len(),
                            preview
                        );
                    }
                }
            }
            Ok(output)
        } else {
            // Agent returned no structured output — create a minimal one
            Ok(AgentOutput::from_text(&provider, ""))
        }
    }

    /// Run the agent with streaming input and output (Claude only).
    ///
    /// Returns a [`StreamingSession`] that allows sending NDJSON messages to
    /// the agent's stdin and reading events from stdout. Automatically
    /// configures `--input-format stream-json`, `--output-format stream-json`,
    /// and `--replay-user-messages`.
    ///
    /// # Default emission granularity
    ///
    /// By default `assistant_message` events are emitted **once per complete
    /// assistant turn** — you get one event when the model finishes speaking,
    /// not a stream of token chunks. For responsive, token-level UIs call
    /// [`include_partial_messages(true)`](Self::include_partial_messages)
    /// on the builder before `exec_streaming`; the session will then emit
    /// partial `assistant_message` chunks as the model generates them.
    ///
    /// The default is kept `false` so existing callers that render whole-turn
    /// bubbles are not broken. See `docs/providers.md` for the full
    /// per-provider flag support matrix.
    ///
    /// # Event lifecycle
    ///
    /// The session emits a unified
    /// [`Event::Result`](crate::output::Event::Result) at the **end of every
    /// agent turn** — not only at final session end. Use that event as the
    /// authoritative turn-boundary signal. After a `Result`, the session
    /// remains open and accepts another
    /// [`send_user_message`](StreamingSession::send_user_message) for the next
    /// turn. Call
    /// [`close_input`](StreamingSession::close_input) followed by
    /// [`wait`](StreamingSession::wait) to terminate the session cleanly.
    ///
    /// Do not depend on replayed `user_message` events to detect turn
    /// boundaries; those only appear while `--replay-user-messages` is set.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use zag_agent::builder::AgentBuilder;
    /// use zag_agent::output::Event;
    ///
    /// # async fn example() -> anyhow::Result<()> {
    /// let mut session = AgentBuilder::new()
    ///     .provider("claude")
    ///     .exec_streaming("initial prompt")
    ///     .await?;
    ///
    /// // Drain the first turn until Result.
    /// while let Some(event) = session.next_event().await? {
    ///     println!("{:?}", event);
    ///     if matches!(event, Event::Result { .. }) {
    ///         break;
    ///     }
    /// }
    ///
    /// // Follow-up turn.
    /// session.send_user_message("do something else").await?;
    /// while let Some(event) = session.next_event().await? {
    ///     if matches!(event, Event::Result { .. }) {
    ///         break;
    ///     }
    /// }
    ///
    /// session.close_input();
    /// session.wait().await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn exec_streaming(self, prompt: &str) -> Result<StreamingSession> {
        let provider = self.resolve_provider()?;
        debug!("exec_streaming: provider={provider}");

        if provider != "claude" {
            bail!("Streaming input is only supported by the Claude provider");
        }

        // Prepend file attachments
        let prompt_with_files = self.prepend_files(prompt)?;

        // Streaming only works on Claude — do not allow the fallback loop
        // to downgrade to a provider that can't stream.
        let mut builder = self;
        builder.provider_explicit = true;
        let (agent, _provider) = builder.create_agent(&provider).await?;

        // Downcast to Claude to call execute_streaming
        let claude_agent = agent
            .as_any_ref()
            .downcast_ref::<Claude>()
            .ok_or_else(|| anyhow::anyhow!("Failed to downcast agent to Claude"))?;

        claude_agent.execute_streaming(Some(&prompt_with_files))
    }

    /// Start an interactive agent session.
    ///
    /// This takes over stdin/stdout for the duration of the session.
    pub async fn run(self, prompt: Option<&str>) -> Result<()> {
        let provider = self.resolve_provider()?;
        debug!("run: provider={provider}");

        // Prepend file attachments
        let prompt_with_files = match prompt {
            Some(p) => Some(self.prepend_files(p)?),
            None if !self.files.is_empty() => {
                let attachments: Vec<Attachment> = self
                    .files
                    .iter()
                    .map(|f| Attachment::from_path(std::path::Path::new(f)))
                    .collect::<Result<Vec<_>>>()?;
                Some(attachment::format_attachments_prefix(&attachments))
            }
            None => None,
        };

        let (agent, _provider) = self.create_agent(&provider).await?;
        agent.run_interactive(prompt_with_files.as_deref()).await?;
        agent.cleanup().await?;
        Ok(())
    }

    /// Resume a previous session by ID.
    pub async fn resume(self, session_id: &str) -> Result<()> {
        let provider = self.resolve_provider()?;
        debug!("resume: provider={provider}, session={session_id}");

        // Resuming must stick with the recorded provider — no downgrade.
        let mut builder = self;
        builder.provider_explicit = true;
        let (agent, _provider) = builder.create_agent(&provider).await?;
        agent.run_resume(Some(session_id), false).await?;
        agent.cleanup().await?;
        Ok(())
    }

    /// Resume the most recent session.
    pub async fn continue_last(self) -> Result<()> {
        let provider = self.resolve_provider()?;
        debug!("continue_last: provider={provider}");

        // Resuming must stick with the recorded provider — no downgrade.
        let mut builder = self;
        builder.provider_explicit = true;
        let (agent, _provider) = builder.create_agent(&provider).await?;
        agent.run_resume(None, true).await?;
        agent.cleanup().await?;
        Ok(())
    }
}

#[cfg(test)]
#[path = "builder_tests.rs"]
mod tests;