Skip to main content

zag_agent/
factory.rs

1use crate::agent::Agent;
2use crate::config::Config;
3use crate::providers::claude::Claude;
4use crate::providers::codex::Codex;
5use crate::providers::copilot::Copilot;
6use crate::providers::gemini::Gemini;
7#[cfg(test)]
8use crate::providers::mock::MockAgent;
9use crate::providers::ollama::Ollama;
10use anyhow::{Result, bail};
11use log::debug;
12
13pub struct AgentFactory;
14
15impl AgentFactory {
16    /// Create and configure an agent based on the provided parameters.
17    ///
18    /// This handles:
19    /// - Loading config from ~/.zag/projects/<id>/zag.toml
20    /// - Creating the appropriate agent implementation
21    /// - Resolving model size aliases (small/medium/large)
22    /// - Merging CLI flags with config file settings
23    /// - Configuring the agent with all settings
24    pub fn create(
25        agent_name: &str,
26        system_prompt: Option<String>,
27        model: Option<String>,
28        root: Option<String>,
29        auto_approve: bool,
30        add_dirs: Vec<String>,
31    ) -> Result<Box<dyn Agent + Send + Sync>> {
32        debug!("Creating agent: {}", agent_name);
33
34        // Skip pre-flight binary check for mock agent (test only)
35        #[cfg(test)]
36        let skip_preflight = agent_name == "mock";
37        #[cfg(not(test))]
38        let skip_preflight = false;
39
40        // Pre-flight: verify the agent CLI binary is available in PATH
41        if !skip_preflight {
42            crate::preflight::check_binary(agent_name)?;
43        }
44
45        // Initialize .agent directory and config on first run
46        let _ = Config::init(root.as_deref());
47
48        // Load config for defaults
49        let config = Config::load(root.as_deref()).unwrap_or_default();
50        debug!("Configuration loaded");
51
52        // Create the agent
53        let mut agent = Self::create_agent(agent_name)?;
54        debug!("Agent instance created");
55
56        // Configure system prompt
57        if let Some(ref sp) = system_prompt {
58            debug!("Setting system prompt (length: {})", sp.len());
59            agent.set_system_prompt(sp.clone());
60        }
61
62        // Configure model (CLI > config > agent default)
63        if let Some(model_input) = model {
64            let resolved = Self::resolve_model(agent_name, &model_input);
65            debug!("Model resolved from CLI: {} -> {}", model_input, resolved);
66            Self::validate_model(agent_name, &resolved)?;
67            agent.set_model(resolved);
68        } else if let Some(config_model) = config.get_model(agent_name) {
69            let resolved = Self::resolve_model(agent_name, config_model);
70            debug!(
71                "Model resolved from config: {} -> {}",
72                config_model, resolved
73            );
74            Self::validate_model(agent_name, &resolved)?;
75            agent.set_model(resolved);
76        } else {
77            debug!("Using default model for agent");
78        }
79
80        // Configure root directory
81        if let Some(root_dir) = root {
82            debug!("Setting root directory: {}", root_dir);
83            agent.set_root(root_dir);
84        }
85
86        // Configure permissions (CLI overrides config)
87        let skip = auto_approve || config.auto_approve();
88        agent.set_skip_permissions(skip);
89
90        // Configure additional directories
91        if !add_dirs.is_empty() {
92            agent.set_add_dirs(add_dirs);
93        }
94
95        Ok(agent)
96    }
97
98    /// Create the appropriate agent implementation based on name.
99    fn create_agent(agent_name: &str) -> Result<Box<dyn Agent + Send + Sync>> {
100        match agent_name.to_lowercase().as_str() {
101            "codex" => Ok(Box::new(Codex::new())),
102            "claude" => Ok(Box::new(Claude::new())),
103            "gemini" => Ok(Box::new(Gemini::new())),
104            "copilot" => Ok(Box::new(Copilot::new())),
105            "ollama" => Ok(Box::new(Ollama::new())),
106            #[cfg(test)]
107            "mock" => Ok(Box::new(MockAgent::new())),
108            _ => bail!("Unknown agent: {}", agent_name),
109        }
110    }
111
112    /// Resolve a model input (size alias or specific name) for a given agent.
113    fn resolve_model(agent_name: &str, model_input: &str) -> String {
114        match agent_name.to_lowercase().as_str() {
115            "claude" => Claude::resolve_model(model_input),
116            "codex" => Codex::resolve_model(model_input),
117            "gemini" => Gemini::resolve_model(model_input),
118            "copilot" => Copilot::resolve_model(model_input),
119            "ollama" => Ollama::resolve_model(model_input),
120            #[cfg(test)]
121            "mock" => MockAgent::resolve_model(model_input),
122            _ => model_input.to_string(), // Unknown agent, pass through
123        }
124    }
125
126    /// Validate a model for a given agent.
127    fn validate_model(agent_name: &str, model: &str) -> Result<()> {
128        match agent_name.to_lowercase().as_str() {
129            "claude" => Claude::validate_model(model, "Claude"),
130            "codex" => Codex::validate_model(model, "Codex"),
131            "gemini" => Gemini::validate_model(model, "Gemini"),
132            "copilot" => Copilot::validate_model(model, "Copilot"),
133            "ollama" => Ollama::validate_model(model, "Ollama"),
134            #[cfg(test)]
135            "mock" => MockAgent::validate_model(model, "Mock"),
136            _ => Ok(()), // Unknown agent, skip validation
137        }
138    }
139}
140
141#[cfg(test)]
142#[path = "factory_tests.rs"]
143mod tests;