ai_agent/

agent.rs

// Source: /data/home/swei/claudecode/openclaudecode/src/utils/model/agent.ts
use crate::engine::{QueryEngine, QueryEngineConfig};
use crate::env::EnvConfig;
use crate::error::AgentError;
use crate::tools::bash::BashTool;
use crate::tools::edit::FileEditTool;
use crate::tools::glob::GlobTool;
use crate::tools::grep::GrepTool;
use crate::tools::read::FileReadTool as ReadTool;
use crate::tools::write::FileWriteTool as WriteTool;
use crate::types::*;

/// Register all built-in tool executors
fn register_all_tool_executors(engine: &mut QueryEngine) {
    type BoxFuture<T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send>>;

    // Bash tool - clone tool and ctx into async block
    let bash_executor = move |input: serde_json::Value,
                              ctx: &ToolContext|
          -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_clone = BashTool::new();
        let cwd = ctx.cwd.clone();
        Box::pin(async move {
            let ctx2 = ToolContext {
                cwd,
                abort_signal: None,
            };
            tool_clone.execute(input, &ctx2).await
        })
    };
    engine.register_tool("Bash".to_string(), bash_executor);

    // FileRead tool
    let read_executor = move |input: serde_json::Value,
                              ctx: &ToolContext|
          -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_clone = ReadTool::new();
        let cwd = ctx.cwd.clone();
        Box::pin(async move {
            let ctx2 = ToolContext {
                cwd,
                abort_signal: None,
            };
            tool_clone.execute(input, &ctx2).await
        })
    };
    engine.register_tool("FileRead".to_string(), read_executor);

    // FileWrite tool
    let write_executor = move |input: serde_json::Value,
                               ctx: &ToolContext|
          -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_clone = WriteTool::new();
        let cwd = ctx.cwd.clone();
        Box::pin(async move {
            let ctx2 = ToolContext {
                cwd,
                abort_signal: None,
            };
            tool_clone.execute(input, &ctx2).await
        })
    };
    engine.register_tool("FileWrite".to_string(), write_executor);

    // Glob tool
    let glob_executor = move |input: serde_json::Value,
                              ctx: &ToolContext|
          -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_clone = GlobTool::new();
        let cwd = ctx.cwd.clone();
        Box::pin(async move {
            let ctx2 = ToolContext {
                cwd,
                abort_signal: None,
            };
            tool_clone.execute(input, &ctx2).await
        })
    };
    engine.register_tool("Glob".to_string(), glob_executor);

    // Grep tool
    let grep_executor = move |input: serde_json::Value,
                              ctx: &ToolContext|
          -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_clone = GrepTool::new();
        let cwd = ctx.cwd.clone();
        Box::pin(async move {
            let ctx2 = ToolContext {
                cwd,
                abort_signal: None,
            };
            tool_clone.execute(input, &ctx2).await
        })
    };
    engine.register_tool("Grep".to_string(), grep_executor);

    // FileEdit tool
    let edit_executor = move |input: serde_json::Value,
                              ctx: &ToolContext|
          -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_clone = FileEditTool::new();
        let cwd = ctx.cwd.clone();
        Box::pin(async move {
            let ctx2 = ToolContext {
                cwd,
                abort_signal: None,
            };
            tool_clone.execute(input, &ctx2).await
        })
    };
    engine.register_tool("FileEdit".to_string(), edit_executor);

    // Skill tool
    use crate::tools::skill::register_skills_from_dir;
    use crate::tools::skill::SkillTool;
    use std::path::Path;

    // Register skills from examples/skills directory
    register_skills_from_dir(Path::new("examples/skills"));

    let skill_executor = move |input: serde_json::Value,
                               ctx: &ToolContext|
          -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_clone = SkillTool::new();
        let cwd = ctx.cwd.clone();
        Box::pin(async move {
            let ctx2 = ToolContext {
                cwd,
                abort_signal: None,
            };
            tool_clone.execute(input, &ctx2).await
        })
    };
    engine.register_tool("Skill".to_string(), skill_executor);

    // Add stub executors for other tools (they have definitions but no full implementation)
    let stub_executor = |input: serde_json::Value,
                         _ctx: &ToolContext|
     -> BoxFuture<Result<ToolResult, AgentError>> {
        let tool_name = input
            .get("name")
            .and_then(|n| n.as_str())
            .unwrap_or("unknown")
            .to_string();
        Box::pin(async move {
            Ok(ToolResult {
                result_type: "text".to_string(),
                tool_use_id: tool_name.clone(),
                content: format!("Tool '{}' is not fully implemented yet", tool_name),
                is_error: Some(false),
            })
        })
    };

    // Register stub executors for tools without full implementations
    for tool_name in &[
        "TaskCreate",
        "TaskList",
        "TaskUpdate",
        "TaskGet",
        "TeamCreate",
        "TeamDelete",
        "SendMessage",
        "EnterWorktree",
        "ExitWorktree",
        "EnterPlanMode",
        "ExitPlanMode",
        "AskUserQuestion",
        "ToolSearch",
        "CronCreate",
        "CronDelete",
        "CronList",
        "Config",
        "TodoWrite",
        "NotebookEdit",
        "WebFetch",
        "WebSearch",
        "Agent",
    ] {
        engine.register_tool(tool_name.to_string(), stub_executor);
    }
}

pub struct Agent {
    config: AgentOptions,
    model: String,
    api_key: Option<String>,
    base_url: Option<String>,
    tool_pool: Vec<ToolDefinition>,
    messages: Vec<Message>,
    session_id: String,
}

impl From<AgentOptions> for Agent {
    fn from(options: AgentOptions) -> Self {
        Agent::create(options)
    }
}

impl Agent {
    /// Create a new agent with model name and max turns
    pub fn new(model: &str, max_turns: u32) -> Self {
        Self::create(AgentOptions {
            model: Some(model.to_string()),
            max_turns: Some(max_turns),
            ..Default::default()
        })
    }

    /// Create a new agent with model, max turns, and event callback for streaming
    pub fn with_event_callback<F>(model: &str, max_turns: u32, on_event: F) -> Self
    where
        F: Fn(AgentEvent) + Send + Sync + 'static,
    {
        let mut agent = Self::new(model, max_turns);
        agent.config.on_event = Some(std::sync::Arc::new(on_event));
        agent
    }

    /// Create agent from AgentOptions
    pub fn create(options: AgentOptions) -> Self {
        // Load env config for defaults
        let env_config = EnvConfig::load();

        // Use env value, then options value, then default
        let model = env_config
            .model
            .clone()
            .or_else(|| options.model.clone())
            .unwrap_or_else(|| "claude-sonnet-4-6".to_string());

        let api_key = env_config
            .auth_token
            .clone()
            .or_else(|| options.api_key.clone());

        let base_url = env_config
            .base_url
            .clone()
            .or_else(|| options.base_url.clone());

        let session_id = uuid::Uuid::new_v4().to_string();

        Self {
            config: options.clone(),
            model,
            api_key,
            base_url,
            tool_pool: options.tools.clone(),
            messages: vec![],
            session_id,
        }
    }

    pub fn get_model(&self) -> &str {
        &self.model
    }

    pub fn get_session_id(&self) -> &str {
        &self.session_id
    }

    /// Get all messages in the conversation history
    pub fn get_messages(&self) -> &[Message] {
        &self.messages
    }

    /// Get all tools available to the agent
    pub fn get_tools(&self) -> &[ToolDefinition] {
        &self.tool_pool
    }

    /// Set system prompt for the agent
    pub fn set_system_prompt(&mut self, prompt: &str) {
        self.config.system_prompt = Some(prompt.to_string());
    }

    /// Set the working directory for the agent
    pub fn set_cwd(&mut self, cwd: &str) {
        self.config.cwd = Some(cwd.to_string());
    }

    /// Set the event callback for agent events (tool start/complete/error, thinking, done)
    /// Note: This must be called BEFORE query() - it sets the callback on the engine
    pub fn set_event_callback<F>(&mut self, callback: F)
    where
        F: Fn(AgentEvent) + Send + Sync + 'static,
    {
        self.config.on_event = Some(std::sync::Arc::new(callback));
    }

    /// Execute a tool directly (for testing/demo purposes)
    pub async fn execute_tool(
        &mut self,
        name: &str,
        input: serde_json::Value,
    ) -> Result<ToolResult, AgentError> {
        // Create a temporary engine to execute the tool
        let cwd = self.config.cwd.clone().unwrap_or_else(|| {
            std::env::current_dir()
                .map(|p| p.to_string_lossy().to_string())
                .unwrap_or_else(|_| ".".to_string())
        });
        let model = self.model.clone();
        let api_key = self.api_key.clone();
        let base_url = self.base_url.clone();

        let mut engine = QueryEngine::new(QueryEngineConfig {
            cwd: cwd.clone(),
            model: model.clone(),
            api_key: api_key.clone(),
            base_url: base_url.clone(),
            tools: vec![],
            system_prompt: None,
            max_turns: 10,
            max_budget_usd: None,
            max_tokens: 16384,
            can_use_tool: None,
            on_event: None,
        });

        // Register all tool executors (including Bash, Read, Write, etc.)
        register_all_tool_executors(&mut engine);

        // Register Agent tool executor with full parameter support
        let agent_tool_executor = move |input: serde_json::Value,
                                        _ctx: &ToolContext|
              -> std::pin::Pin<
            Box<dyn std::future::Future<Output = Result<ToolResult, AgentError>> + Send>,
        > {
            let cwd = cwd.clone();
            let api_key = api_key.clone();
            let base_url = base_url.clone();
            let model = model.clone();

            Box::pin(async move {
                // Extract ALL parameters from input
                let description = input["description"].as_str().unwrap_or("subagent");
                let subagent_prompt = input["prompt"].as_str().unwrap_or("");
                let subagent_model = input["model"]
                    .as_str()
                    .map(|s| s.to_string())
                    .unwrap_or_else(|| model.clone());
                let max_turns = input["max_turns"]
                    .as_u64()
                    .or_else(|| input["maxTurns"].as_u64()) // Support camelCase too
                    .unwrap_or(10) as u32;

                // NEW: Extract subagent_type
                let subagent_type = input["subagent_type"]
                    .as_str()
                    .or_else(|| input["subagentType"].as_str())
                    .map(|s| s.to_string());

                // NEW: Extract run_in_background (ignored for now, async not supported)
                let _run_in_background = input["run_in_background"]
                    .as_bool()
                    .or_else(|| input["runInBackground"].as_bool())
                    .unwrap_or(false);

                // NEW: Extract name
                let agent_name = input["name"].as_str().map(|s| s.to_string());

                // NEW: Extract team_name
                let _team_name = input["team_name"]
                    .as_str()
                    .or_else(|| input["teamName"].as_str())
                    .map(|s| s.to_string());

                // NEW: Extract mode (permission mode - ignored for now)
                let _mode = input["mode"].as_str().map(|s| s.to_string());

                // NEW: Extract cwd (working directory override)
                let subagent_cwd = input["cwd"]
                    .as_str()
                    .map(|s| s.to_string())
                    .unwrap_or_else(|| cwd.clone());

                // NEW: Extract isolation
                let _isolation = input["isolation"].as_str().map(|s| s.to_string());

                // Build system prompt for subagent
                let system_prompt = build_agent_system_prompt(description, subagent_type.as_deref());

                // Create sub-agent engine with proper system prompt
                let mut sub_engine = QueryEngine::new(QueryEngineConfig {
                    cwd: subagent_cwd,
                    model: subagent_model.to_string(),
                    api_key,
                    base_url,
                    tools: vec![],
                    system_prompt: Some(system_prompt),
                    max_turns,
                    max_budget_usd: None,
                    max_tokens: 16384,
                    can_use_tool: None,
                    on_event: None,
                });

                match sub_engine.submit_message(subagent_prompt).await {
                    Ok((result_text, _)) => {
                        let mut content = format!("[Subagent: {}]", description);
                        if let Some(ref name) = agent_name {
                            content = format!("[Subagent: {} ({})]", description, name);
                        }
                        content = format!("{}\n\n{}", content, result_text);
                        Ok(ToolResult {
                            result_type: "text".to_string(),
                            tool_use_id: "agent_tool".to_string(),
                            content,
                            is_error: Some(false),
                        })
                    }
                    Err(e) => Ok(ToolResult {
                        result_type: "text".to_string(),
                        tool_use_id: "agent_tool".to_string(),
                        content: format!("[Subagent: {}] Error: {}", description, e),
                        is_error: Some(true),
                    }),
                }
            })
        };

        engine.register_tool("Agent".to_string(), agent_tool_executor);
        engine.execute_tool(name, input).await
    }

    /// Simple blocking prompt method - sends a prompt and returns the result.
    /// This matches the TypeScript SDK's agent.prompt() API.
    pub async fn prompt(&mut self, prompt: &str) -> Result<QueryResult, AgentError> {
        self.query(prompt).await
    }

    pub async fn query(&mut self, prompt: &str) -> Result<QueryResult, AgentError> {
        use crate::ai_md::load_ai_md;
        use crate::memory::load_memory_prompt;
        use crate::prompts::build_system_prompt;
        use crate::tools::get_all_base_tools;

        let cwd = self.config.cwd.clone().unwrap_or_else(|| {
            std::env::current_dir()
                .map(|p| p.to_string_lossy().to_string())
                .unwrap_or_else(|_| ".".to_string())
        });
        let cwd_path = std::path::Path::new(&cwd);
        let model = self.model.clone();
        let api_key = self.api_key.clone();
        let base_url = self.base_url.clone();

        // Build system prompt: AI.md + memory prompt + custom system prompt
        let ai_md_prompt = load_ai_md(cwd_path).ok().flatten();
        let memory_prompt = load_memory_prompt();

        // Use the full system prompt from prompts module (matches TypeScript)
        let base_system_prompt = build_system_prompt();

        // Combine: AI.md (highest priority) -> memory -> base prompt -> custom
        let system_prompt = match (&ai_md_prompt, &memory_prompt, &self.config.system_prompt) {
            (Some(ai_md), Some(mem), Some(custom)) => Some(format!(
                "{}\n\n{}\n\n{}\n\n{}",
                ai_md, mem, base_system_prompt, custom
            )),
            (Some(ai_md), Some(mem), None) => {
                Some(format!("{}\n\n{}\n\n{}", ai_md, mem, base_system_prompt))
            }
            (Some(ai_md), None, Some(custom)) => {
                Some(format!("{}\n\n{}\n\n{}", ai_md, base_system_prompt, custom))
            }
            (Some(ai_md), None, None) => Some(format!("{}\n\n{}", ai_md, base_system_prompt)),
            (None, Some(mem), Some(custom)) => {
                Some(format!("{}\n\n{}\n\n{}", mem, base_system_prompt, custom))
            }
            (None, Some(mem), None) => Some(format!("{}\n\n{}", mem, base_system_prompt)),
            (None, None, Some(custom)) => Some(format!("{}\n\n{}", base_system_prompt, custom)),
            (None, None, None) => Some(base_system_prompt),
        };

        // Use base tools if tool_pool is empty
        let tools = if self.tool_pool.is_empty() {
            get_all_base_tools()
        } else {
            self.tool_pool.clone()
        };

        let on_event = self.config.on_event.clone();
        let mut engine = QueryEngine::new(QueryEngineConfig {
            cwd: cwd.clone(),
            model: model.clone(),
            api_key: api_key.clone(),
            base_url: base_url.clone(),
            tools,
            system_prompt,
            max_turns: self.config.max_turns.unwrap_or(10),
            max_budget_usd: self.config.max_budget_usd,
            max_tokens: self.config.max_tokens.unwrap_or(16384),
            can_use_tool: None,
            on_event,
        });

        // Register all tool executors on the engine so they can be called
        register_all_tool_executors(&mut engine);

        // Clone tool_pool before the closure to avoid capturing self
        let tool_pool = self.tool_pool.clone();

        // Register the Agent tool executor

        // Register the Agent tool executor to spawn sub-agents with full parameter support
        let agent_tool_executor = move |input: serde_json::Value,
                                        _ctx: &ToolContext|
              -> std::pin::Pin<
            Box<dyn std::future::Future<Output = Result<ToolResult, AgentError>> + Send>,
        > {
            let cwd = cwd.clone();
            let api_key = api_key.clone();
            let base_url = base_url.clone();
            let model = model.clone();
            let tool_pool = tool_pool.clone();

            Box::pin(async move {
                // Extract ALL parameters from input
                let description = input["description"].as_str().unwrap_or("subagent");
                let subagent_prompt = input["prompt"].as_str().unwrap_or("");
                let subagent_model = input["model"]
                    .as_str()
                    .map(|s| s.to_string())
                    .unwrap_or_else(|| model.clone());
                let max_turns = input["max_turns"]
                    .as_u64()
                    .or_else(|| input["maxTurns"].as_u64()) // Support camelCase too
                    .unwrap_or(10) as u32;

                // NEW: Extract subagent_type
                let subagent_type = input["subagent_type"]
                    .as_str()
                    .or_else(|| input["subagentType"].as_str())
                    .map(|s| s.to_string());

                // NEW: Extract run_in_background (ignored for now)
                let _run_in_background = input["run_in_background"]
                    .as_bool()
                    .or_else(|| input["runInBackground"].as_bool())
                    .unwrap_or(false);

                // NEW: Extract name
                let agent_name = input["name"].as_str().map(|s| s.to_string());

                // NEW: Extract team_name
                let _team_name = input["team_name"]
                    .as_str()
                    .or_else(|| input["teamName"].as_str())
                    .map(|s| s.to_string());

                // NEW: Extract mode
                let _mode = input["mode"].as_str().map(|s| s.to_string());

                // NEW: Extract cwd (working directory override)
                let subagent_cwd = input["cwd"]
                    .as_str()
                    .map(|s| s.to_string())
                    .unwrap_or_else(|| cwd.clone());

                // NEW: Extract isolation
                let _isolation = input["isolation"].as_str().map(|s| s.to_string());

                // Build system prompt for subagent based on agent type
                let system_prompt = build_agent_system_prompt(description, subagent_type.as_deref());

                // Use parent agent's tool pool for the subagent
                let parent_tools = tool_pool;

                // Create a new engine for the subagent
                let mut sub_engine = QueryEngine::new(QueryEngineConfig {
                    cwd: subagent_cwd,
                    model: subagent_model.to_string(),
                    api_key,
                    base_url,
                    tools: parent_tools,
                    system_prompt: Some(system_prompt),
                    max_turns,
                    max_budget_usd: None,
                    max_tokens: 16384,
                    can_use_tool: None,
                    on_event: None,
                });

                // Run the subagent
                match sub_engine.submit_message(subagent_prompt).await {
                    Ok((result_text, _)) => {
                        let mut content = format!("[Subagent: {}]", description);
                        if let Some(ref name) = agent_name {
                            content = format!("[Subagent: {} ({})]", description, name);
                        }
                        content = format!("{}\n\n{}", content, result_text);
                        Ok(ToolResult {
                            result_type: "text".to_string(),
                            tool_use_id: "agent_tool".to_string(),
                            content,
                            is_error: Some(false),
                        })
                    }
                    Err(e) => Ok(ToolResult {
                        result_type: "text".to_string(),
                        tool_use_id: "agent_tool".to_string(),
                        content: format!("[Subagent: {}] Error: {}", description, e),
                        is_error: Some(true),
                    }),
                }
            })
        };

        // Register all tool executors
        register_all_tool_executors(&mut engine);
        engine.register_tool("Agent".to_string(), agent_tool_executor);

        // Pass existing messages to engine for continuing conversation
        engine.set_messages(self.messages.clone());

        let start = std::time::Instant::now();
        let (response_text, exit_reason) = engine.submit_message(prompt).await?;
        let messages = engine.get_messages();

        // Get actual usage from engine
        let engine_usage = engine.get_usage();
        let usage = TokenUsage {
            input_tokens: engine_usage.input_tokens,
            output_tokens: engine_usage.output_tokens,
            cache_creation_input_tokens: engine_usage.cache_creation_input_tokens,
            cache_read_input_tokens: engine_usage.cache_read_input_tokens,
        };

        // Store messages in agent
        self.messages = messages;

        Ok(QueryResult {
            text: response_text,
            usage,
            num_turns: engine.get_turn_count(),
            duration_ms: start.elapsed().as_millis() as u64,
            exit_reason,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::{QueryEngine, QueryEngineConfig};
    use crate::types::ToolContext;
    use std::sync::Arc;

    /// Test that Agent tool correctly extracts all parameters from input
    #[tokio::test]
    async fn test_agent_tool_parses_all_parameters() {
        // Test parameter extraction from various input formats
        // This verifies all parameters are now properly parsed

        // Test 1: subagent_type parameter (snake_case)
        let input1 = serde_json::json!({
            "description": "explore-agent",
            "prompt": "Explore the codebase",
            "subagent_type": "Explore"
        });
        assert_eq!(input1["subagent_type"].as_str(), Some("Explore"));
        assert_eq!(input1["subagentType"].as_str(), None); // snake_case works

        // Test 2: subagent_type parameter (camelCase)
        let input2 = serde_json::json!({
            "description": "explore-agent",
            "prompt": "Explore the codebase",
            "subagentType": "Plan"
        });
        assert_eq!(input2["subagentType"].as_str(), Some("Plan"));

        // Test 3: run_in_background (snake_case)
        let input3 = serde_json::json!({
            "description": "background-agent",
            "prompt": "Run in background",
            "run_in_background": true
        });
        assert_eq!(input3["run_in_background"].as_bool(), Some(true));

        // Test 4: runInBackground (camelCase)
        let input4 = serde_json::json!({
            "description": "background-agent",
            "runInBackground": true
        });
        assert_eq!(input4["runInBackground"].as_bool(), Some(true));

        // Test 5: max_turns (snake_case)
        let input5 = serde_json::json!({
            "description": "test",
            "max_turns": 5
        });
        assert_eq!(input5["max_turns"].as_u64(), Some(5));

        // Test 6: maxTurns (camelCase)
        let input6 = serde_json::json!({
            "description": "test",
            "maxTurns": 10
        });
        assert_eq!(input6["maxTurns"].as_u64(), Some(10));

        // Test 7: team_name (snake_case)
        let input7 = serde_json::json!({
            "description": "team-agent",
            "team_name": "my-team"
        });
        assert_eq!(input7["team_name"].as_str(), Some("my-team"));

        // Test 8: teamName (camelCase)
        let input8 = serde_json::json!({
            "description": "team-agent",
            "teamName": "my-team"
        });
        assert_eq!(input8["teamName"].as_str(), Some("my-team"));

        // Test 9: cwd parameter
        let input9 = serde_json::json!({
            "description": "custom-cwd",
            "cwd": "/custom/path"
        });
        assert_eq!(input9["cwd"].as_str(), Some("/custom/path"));

        // Test 10: name parameter
        let input10 = serde_json::json!({
            "name": "my-agent",
            "description": "named-agent"
        });
        assert_eq!(input10["name"].as_str(), Some("my-agent"));

        // Test 11: mode parameter
        let input11 = serde_json::json!({
            "description": "plan-mode",
            "mode": "plan"
        });
        assert_eq!(input11["mode"].as_str(), Some("plan"));

        // Test 12: isolation parameter
        let input12 = serde_json::json!({
            "description": "isolated",
            "isolation": "worktree"
        });
        assert_eq!(input12["isolation"].as_str(), Some("worktree"));

        // Verify all expected keys are now handled
        // The agent tool executor should handle all these parameters
    }

    /// Test that Agent tool creates subagent with proper system prompt based on agent type
    #[tokio::test]
    async fn test_agent_tool_system_prompt_by_type() {
        // Test system prompt generation for different agent types
        let explore_prompt = build_agent_system_prompt("Explore task", Some("Explore"));
        assert!(explore_prompt.contains("Explore agent"));

        let plan_prompt = build_agent_system_prompt("Plan task", Some("Plan"));
        assert!(plan_prompt.contains("Plan agent"));

        let review_prompt = build_agent_system_prompt("Review task", Some("Review"));
        assert!(review_prompt.contains("Review agent"));

        let general_prompt = build_agent_system_prompt("General task", None);
        assert!(general_prompt.contains("Task description: General task"));
    }

    /// Test that Agent tool creates subagent with proper system prompt
    #[tokio::test]
    async fn test_agent_tool_creates_subagent_with_system_prompt() {
        let mut engine = QueryEngine::new(QueryEngineConfig {
            cwd: "/tmp".to_string(),
            model: "test-model".to_string(),
            api_key: Some("test-key".to_string()),
            base_url: Some("http://localhost:8080".to_string()),
            tools: vec![],
            system_prompt: Some("Parent system prompt".to_string()),
            max_turns: 10,
            max_budget_usd: None,
            max_tokens: 4096,
            can_use_tool: None,
            on_event: None,
        });

        // Register Agent tool
        let agent_tool_executor = create_agent_tool_executor(
            "/tmp".to_string(),
            Some("test-key".to_string()),
            Some("http://localhost:8080".to_string()),
            "test-model".to_string(),
            vec![],
        );
        engine.register_tool("Agent".to_string(), agent_tool_executor);

        let input = serde_json::json!({
            "description": "test-subagent",
            "prompt": "What is 1+1?"
        });

        let result = engine.execute_tool("Agent", input).await;
        // The subagent should have a system prompt (not None)
        // We verify by checking the tool executes without error
        assert!(result.is_ok(), "Agent tool should execute with system prompt");
    }

    /// Test Agent creation with options
    #[tokio::test]
    async fn test_create_agent() {
        let agent = Agent::create(AgentOptions {
            model: Some("claude-sonnet-4-6".to_string()),
            ..Default::default()
        });
        assert!(!agent.get_model().is_empty());
    }

    /// Check if required environment variables are present for real API tests
    /// Returns true if AI_BASE_URL, AI_MODEL, and AI_AUTH_TOKEN can be loaded from .env
    fn has_required_env_vars() -> bool {
        let config = EnvConfig::load();
        config.base_url.is_some() && config.model.is_some() && config.auth_token.is_some()
    }

    /// Test Agent tool calling with real .env config
    /// This test makes an actual API call using the configured model
    #[tokio::test]
    async fn test_agent_tool_calling_with_real_env_config() {
        // Only run if required env vars are set
        if !tests::has_required_env_vars() {
            eprintln!("Skipping test: AI_BASE_URL, AI_MODEL, or AI_AUTH_TOKEN not set");
            return;
        }

        // Load config from .env file
        let config = EnvConfig::load();

        // Verify config is loaded
        assert!(config.base_url.is_some(), "Base URL should be configured");
        assert!(config.auth_token.is_some(), "Auth token should be configured");
        assert!(config.model.is_some(), "Model should be configured");

        // Create agent with real config
        let agent = Agent::create(AgentOptions {
            model: config.model.clone(),
            tools: vec![],
            ..Default::default()
        });

        // Verify agent was created with the configured model
        let model = agent.get_model();
        assert!(!model.is_empty(), "Agent should have a model set");
        println!("Using model: {}", model);
    }

    /// Test agent prompt with real API call using .env config
    /// This is an integration test that exercises the full agent flow
    #[tokio::test]
    async fn test_agent_prompt_with_real_api() {
        // Only run if required env vars are set
        if !tests::has_required_env_vars() {
            eprintln!("Skipping test: AI_BASE_URL, AI_MODEL, or AI_AUTH_TOKEN not set");
            return;
        }

        // Load config from .env file
        let config = EnvConfig::load();

        // Skip test if no API configured
        if config.base_url.is_none() || config.auth_token.is_none() {
            eprintln!("Skipping test: no API config found");
            return;
        }

        // Create agent with all tools and real config
        use crate::get_all_tools;
        let tools = get_all_tools();

        let mut agent = Agent::create(AgentOptions {
            model: config.model.clone(),
            max_turns: Some(3),
            tools,
            ..Default::default()
        });

        // Make a simple prompt that should trigger tool use
        let result = agent.prompt("What is 2 + 2? Just give me the answer.").await;

        // Verify we got a response
        assert!(result.is_ok(), "Agent should respond successfully");
        let response = result.unwrap();
        assert!(!response.text.is_empty(), "Response should not be empty");
        println!("Agent response: {}", response.text);
    }

    /// Test agent tool calling with multiple tools from .env config
    /// This tests that the agent can use tools when configured via .env
    #[tokio::test]
    async fn test_agent_with_multiple_tools_real_config() {
        // Only run if required env vars are set
        if !tests::has_required_env_vars() {
            eprintln!("Skipping test: AI_BASE_URL, AI_MODEL, or AI_AUTH_TOKEN not set");
            return;
        }

        // Load config from .env file
        let config = EnvConfig::load();

        // Skip if no API configured
        if config.base_url.is_none() || config.auth_token.is_none() {
            eprintln!("Skipping test: no API config found");
            return;
        }

        // Get all available tools
        use crate::get_all_tools;
        let tools = get_all_tools();

        // Verify we have tools available
        assert!(!tools.is_empty(), "Should have tools available");

        let mut agent = Agent::create(AgentOptions {
            model: config.model.clone(),
            max_turns: Some(3),
            tools,
            ..Default::default()
        });

        // Prompt that might use tools
        let result = agent.prompt("List all Rust files in the current directory using glob").await;

        // Should get a response (may or may not use tools depending on model)
        assert!(result.is_ok(), "Agent should respond");
        let response = result.unwrap();
        assert!(!response.text.is_empty(), "Response should not be empty");
        println!("Agent response: {}", response.text);
    }

    /// Test that tool executors are registered and can be invoked
    /// This verifies the fix for tool calling not working in TUI
    #[tokio::test]
    async fn test_tool_executors_registered() {
        // Only run if required env vars are set
        if !tests::has_required_env_vars() {
            eprintln!("Skipping test: AI_BASE_URL, AI_MODEL, or AI_AUTH_TOKEN not set");
            return;
        }

        // Load config from .env file
        let config = EnvConfig::load();

        // Skip if no API configured
        if config.base_url.is_none() || config.auth_token.is_none() {
            eprintln!("Skipping test: no API config found");
            return;
        }

        // Get all available tools
        use crate::get_all_tools;
        let tools = get_all_tools();

        // Verify tools are available
        let tool_names: Vec<&str> = tools.iter().map(|t| t.name.as_str()).collect();
        assert!(tool_names.contains(&"Bash"), "Should have Bash tool");
        assert!(tool_names.contains(&"FileRead"), "Should have FileRead tool");
        assert!(tool_names.contains(&"Glob"), "Should have Glob tool");
        println!("Available tools: {:?}", tool_names);

        // Create agent - this will call register_all_tool_executors internally
        let mut agent = Agent::create(AgentOptions {
            model: config.model.clone(),
            max_turns: Some(3),
            tools,
            ..Default::default()
        });

        // Prompt that should definitely use the Bash tool
        let result = agent
            .prompt("Run this command: echo 'hello from tool test'")
            .await;

        // Verify we got a response
        assert!(result.is_ok(), "Agent should respond successfully");
        let response = result.unwrap();
        assert!(!response.text.is_empty(), "Response should not be empty");

        // Check that the tool was actually used (response should contain output)
        let text_lower = response.text.to_lowercase();
        let tool_was_used =
            text_lower.contains("hello from tool test") || text_lower.contains("tool");
        println!(
            "Tool calling test - Response: {} (tool_used: {})",
            response.text, tool_was_used
        );
    }

    /// Test Glob tool directly via agent
    #[tokio::test]
    async fn test_glob_tool_via_agent() {
        // Only run if required env vars are set
        if !tests::has_required_env_vars() {
            eprintln!("Skipping test: AI_BASE_URL, AI_MODEL, or AI_AUTH_TOKEN not set");
            return;
        }

        // Load config from .env file
        let config = EnvConfig::load();

        // Skip if no API configured
        if config.base_url.is_none() || config.auth_token.is_none() {
            eprintln!("Skipping test: no API config found");
            return;
        }

        // Get all available tools
        use crate::get_all_tools;
        let tools = get_all_tools();

        let mut agent = Agent::create(AgentOptions {
            model: config.model.clone(),
            max_turns: Some(3),
            tools,
            ..Default::default()
        });

        // Prompt that should use Glob tool
        let result = agent
            .prompt("List all .rs files in the src directory using the Glob tool")
            .await;

        assert!(result.is_ok(), "Agent should respond");
        let response = result.unwrap();
        assert!(!response.text.is_empty(), "Response should not be empty");
        println!("Glob tool test response: {}", response.text);
    }

    /// Test FileRead tool directly via agent
    #[tokio::test]
    async fn test_fileread_tool_via_agent() {
        // Only run if required env vars are set
        if !tests::has_required_env_vars() {
            eprintln!("Skipping test: AI_BASE_URL, AI_MODEL, or AI_AUTH_TOKEN not set");
            return;
        }

        // Load config from .env file
        let config = EnvConfig::load();

        // Skip if no API configured
        if config.base_url.is_none() || config.auth_token.is_none() {
            eprintln!("Skipping test: no API config found");
            return;
        }

        // Get all available tools
        use crate::get_all_tools;
        let tools = get_all_tools();

        let mut agent = Agent::create(AgentOptions {
            model: config.model.clone(),
            max_turns: Some(3),
            tools,
            ..Default::default()
        });

        // Prompt that should use FileRead tool
        let result = agent
            .prompt("Read the Cargo.toml file from the current directory")
            .await;

        assert!(result.is_ok(), "Agent should respond");
        let response = result.unwrap();
        assert!(!response.text.is_empty(), "Response should not be empty");
        // The response should contain something from Cargo.toml
        println!("FileRead tool test response: {}", response.text);
    }

    /// Test multiple tool calls in one turn
    #[tokio::test]
    async fn test_multiple_tool_calls() {
        // Only run if required env vars are set
        if !tests::has_required_env_vars() {
            eprintln!("Skipping test: AI_BASE_URL, AI_MODEL, or AI_AUTH_TOKEN not set");
            return;
        }

        // Load config from .env file
        let config = EnvConfig::load();

        // Skip if no API configured
        if config.base_url.is_none() || config.auth_token.is_none() {
            eprintln!("Skipping test: no API config found");
            return;
        }

        // Get all available tools
        use crate::get_all_tools;
        let tools = get_all_tools();

        let mut agent = Agent::create(AgentOptions {
            model: config.model.clone(),
            max_turns: Some(5),
            tools,
            ..Default::default()
        });

        // Prompt that should use multiple tools
        let result = agent
            .prompt("First list all files in the current directory, then read the README.md file if it exists")
            .await;

        assert!(result.is_ok(), "Agent should respond");
        let response = result.unwrap();
        assert!(!response.text.is_empty(), "Response should not be empty");
        println!("Multiple tool calls test response: {}", response.text);
    }
}

/// Helper function to create the Agent tool executor with all parameters
fn create_agent_tool_executor(
    cwd: String,
    api_key: Option<String>,
    base_url: Option<String>,
    model: String,
    tool_pool: Vec<crate::tools::ToolDefinition>,
) -> impl Fn(serde_json::Value, &ToolContext) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<ToolResult, AgentError>> + Send>> + Send + 'static {
    move |input: serde_json::Value,
          _ctx: &ToolContext|
          -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<ToolResult, AgentError>> + Send>,
    > {
        let cwd = cwd.clone();
        let api_key = api_key.clone();
        let base_url = base_url.clone();
        let model = model.clone();
        let tool_pool = tool_pool.clone();

        Box::pin(async move {
            // Extract ALL parameters from input (not just description, prompt, model, max_turns)
            let description = input["description"].as_str().unwrap_or("subagent");
            let subagent_prompt = input["prompt"].as_str().unwrap_or("");
            let subagent_model = input["model"]
                .as_str()
                .map(|s| s.to_string())
                .unwrap_or_else(|| model.clone());
            let max_turns = input["max_turns"]
                .as_u64()
                .or_else(|| input["maxTurns"].as_u64()) // Support camelCase too
                .unwrap_or(10) as u32;

            // NEW: Extract subagent_type
            let subagent_type = input["subagent_type"]
                .as_str()
                .or_else(|| input["subagentType"].as_str())
                .map(|s| s.to_string());

            // NEW: Extract run_in_background
            let run_in_background = input["run_in_background"]
                .as_bool()
                .or_else(|| input["runInBackground"].as_bool())
                .unwrap_or(false);

            // NEW: Extract name
            let agent_name = input["name"]
                .as_str()
                .map(|s| s.to_string());

            // NEW: Extract team_name
            let team_name = input["team_name"]
                .as_str()
                .or_else(|| input["teamName"].as_str())
                .map(|s| s.to_string());

            // NEW: Extract mode (permission mode)
            let mode = input["mode"]
                .as_str()
                .map(|s| s.to_string());

            // NEW: Extract cwd (working directory override)
            let subagent_cwd = input["cwd"]
                .as_str()
                .map(|s| s.to_string())
                .unwrap_or_else(|| cwd.clone());

            // NEW: Extract isolation
            let isolation = input["isolation"]
                .as_str()
                .map(|s| s.to_string());

            // Log all parameters for debugging
            let params_log = format!(
                "Agent tool params: description={}, subagent_type={:?}, run_in_background={}, name={:?}, team_name={:?}, mode={:?}, cwd={}, isolation={:?}",
                description,
                subagent_type,
                run_in_background,
                agent_name,
                team_name,
                mode,
                subagent_cwd,
                isolation
            );
            eprintln!("{}", params_log);

            // Use the correct cwd for the subagent
            let actual_cwd = subagent_cwd.clone();

            // Build system prompt for subagent (matching TypeScript getAgentSystemPrompt)
            let system_prompt = build_agent_system_prompt(description, subagent_type.as_deref());

            // Use parent agent's tool pool for the subagent
            let tools = tool_pool;

            // Create a new engine for the subagent
            let mut sub_engine = QueryEngine::new(QueryEngineConfig {
                cwd: actual_cwd,
                model: subagent_model.to_string(),
                api_key,
                base_url,
                tools,
                system_prompt: Some(system_prompt),
                max_turns,
                max_budget_usd: None,
                max_tokens: 16384,
                can_use_tool: None,
                on_event: None,
            });

            // Run the subagent
            match sub_engine.submit_message(subagent_prompt).await {
                Ok((result_text, _)) => {
                    let mut content = format!("[Subagent: {}]", description);
                    if let Some(ref name) = agent_name {
                        content = format!("[Subagent: {} ({}))]", description, name);
                    }
                    content = format!("{}\n\n{}", content, result_text);
                    Ok(ToolResult {
                        result_type: "text".to_string(),
                        tool_use_id: "agent_tool".to_string(),
                        content,
                        is_error: Some(false),
                    })
                }
                Err(e) => Ok(ToolResult {
                    result_type: "text".to_string(),
                    tool_use_id: "agent_tool".to_string(),
                    content: format!("[Subagent: {}] Error: {}", description, e),
                    is_error: Some(true),
                }),
            }
        })
    }
}

/// Build system prompt for subagent based on agent type
fn build_agent_system_prompt(agent_description: &str, agent_type: Option<&str>) -> String {
    let base_prompt = "You are an agent that helps users with software engineering tasks. Use the tools available to you to assist the user.\n\nComplete the task fully—don't gold-plate, but don't leave it half-done. When you complete the task, respond with a concise report covering what was done and any key findings.";

    match agent_type {
        Some("Explore") => {
            format!(
                "{}\n\nYou are an Explore agent. Your goal is to explore and understand the codebase thoroughly. Use search and read tools to investigate. Report your findings in detail.",
                base_prompt
            )
        }
        Some("Plan") => {
            format!(
                "{}\n\nYou are a Plan agent. Your goal is to plan and analyze tasks before execution. Break down complex tasks into steps. Provide a detailed plan.",
                base_prompt
            )
        }
        Some("Review") => {
            format!(
                "{}\n\nYou are a Review agent. Your goal is to review code and provide constructive feedback. Be thorough and focus on best practices.",
                base_prompt
            )
        }
        _ => {
            // General purpose agent
            format!(
                "{}\n\nTask description: {}",
                base_prompt, agent_description
            )
        }
    }
}