zesdex/src/app/subagent/engine.rs

use tokio::sync::mpsc;
use crate::dto::chat::message::ChatMessage;
use crate::dto::provider::request::ToolDef;
use crate::tool::{all_tools, tool_defs, tool_is_risky};
use super::context::SubagentContext;
use super::event::SubagentEvent;

pub const MAX_AGENT_STEPS: usize = usize::MAX;

/// Maps a subagent's allowed tool names to concrete Tool trait objects and
/// OpenAI-style tool definitions.  When `allowed_tools` is empty every tool is
/// available; otherwise only explicitly allowed ones are included.
fn build_subagent_tools(allowed_tools: &[String]) -> (Vec<Box<dyn crate::tool::Tool>>, Vec<ToolDef>) {
    let all = all_tools();
    let filtered: Vec<Box<dyn crate::tool::Tool>> = if allowed_tools.is_empty() {
        all
    } else {
        all.into_iter()
            .filter(|t| allowed_tools.contains(&t.name().to_string()))
            .collect()
    };
    let defs = tool_defs(&filtered);
    (filtered, defs)
}

/// Resolves the API key, model, and base URL from the persisted application
/// configuration rather than environment variables, matching how the main agent
/// resolves its credentials.
fn resolve_provider_config() -> (String, String, Option<String>) {
    let settings = crate::model::settings::Settings::load();
    let app_config = crate::model::app_config::AppConfig::load();

    let mut api_key = settings.api_keys.get(&settings.provider).cloned().unwrap_or_default();
    let model = settings.model.clone();
    let base_url = app_config.providers.get(&settings.provider)
        .map(|p| p.api_base.clone());

    if api_key.is_empty() {
        if let Some(provider_cfg) = app_config.providers.get(&settings.provider) {
            api_key = provider_cfg.api_key_env.as_ref()
                .and_then(|env| std::env::var(env).ok())
                .or_else(|| provider_cfg.default_api_key.clone())
                .unwrap_or_default();
        }
    }

    (api_key, model, base_url)
}

pub fn run_subagent(ctx: SubagentContext, tx: mpsc::Sender<SubagentEvent>) -> anyhow::Result<String> {
    let mut output = String::new();
    let mut messages: Vec<ChatMessage> = Vec::new();
    messages.push(ChatMessage::system(ctx.system_prompt.clone()));

    let tool_ctx = crate::tool::ToolCtx::builder()
        .session_dir(ctx.session_dir.clone())
        .origin(crate::app::state::types::Origin::SubAgent)
        .build();

    // Build tool list once before the loop
    let (tools, tdefs) = build_subagent_tools(&ctx.allowed_tools);
    let tdefs_opt: Option<Vec<ToolDef>> = if tdefs.is_empty() { None } else { Some(tdefs) };

    let max_steps = ctx.max_steps.min(MAX_AGENT_STEPS);
    for step in 0..max_steps {
        let (api_key, model, base_url) = resolve_provider_config();
        let client = crate::service::provider::LlmClient::new(api_key, model, base_url);

        // Use the structured tool-calling API so the LLM can request tools with
        // proper arguments, exactly like the main agent does.
        let (response, _usage) = match client.chat_with_tools_non_streaming(&messages, tdefs_opt.clone()) {
            Ok(result) => result,
            Err(e) => {
                let _ = tx.blocking_send(SubagentEvent::StepFailed {
                    _step: step,
                    _error: e.to_string(),
                });
                anyhow::bail!("subagent call failed at step {}: {}", step, e);
            }
        };

        let has_tool_calls = response.tool_calls.is_some()
            && response.tool_calls.as_ref().is_some_and(|tc| !tc.is_empty());

        let content = response.content.clone().unwrap_or_default();

        if has_tool_calls {
            let tool_calls = response.tool_calls.clone().unwrap_or_default();
            // Push the assistant message with tool_calls into the conversation
            messages.push(response);

            for tool_call in &tool_calls {
                let tool_name = &tool_call.function.name;
                let args = crate::dto::chat::tool::sanitize_tool_arguments(&tool_call.function.arguments);
                let explicitly_allowed = ctx.allowed_tools.contains(tool_name);
                let generally_allowed = ctx.allowed_tools.is_empty() || explicitly_allowed;

                let _ = tx.blocking_send(SubagentEvent::ToolCall {
                    _tool: tool_name.clone(),
                    _args: args.clone(),
                });

                if !generally_allowed {
                    let msg = format!("tool '{}' not allowed for this subagent", tool_name);
                    messages.push(ChatMessage::tool_result(tool_call.id.clone(), msg.clone()));
                    let _ = tx.blocking_send(SubagentEvent::ToolResult {
                        _tool: tool_name.clone(),
                        _output: msg,
                    });
                    continue;
                }

                if tool_is_risky(tool_name) && !explicitly_allowed {
                    let msg = format!("risky tool '{}' requires explicit permission; not allowed for this subagent", tool_name);
                    messages.push(ChatMessage::tool_result(tool_call.id.clone(), msg.clone()));
                    let _ = tx.blocking_send(SubagentEvent::ToolResult {
                        _tool: tool_name.clone(),
                        _output: msg,
                    });
                    continue;
                }

                let result = match tools.iter().find(|t| t.name() == tool_name.as_str()) {
                    Some(tool) => tool.run(&tool_ctx, &args),
                    None => Err(anyhow::anyhow!("tool '{}' not found", tool_name)),
                };

                match result {
                    Ok(output_text) => {
                        messages.push(ChatMessage::tool_result(tool_call.id.clone(), output_text.clone()));
                        let _ = tx.blocking_send(SubagentEvent::ToolResult {
                            _tool: tool_name.clone(),
                            _output: output_text,
                        });
                    }
                    Err(e) => {
                        let msg = format!("tool '{}' failed: {}", tool_name, e);
                        messages.push(ChatMessage::tool_result(tool_call.id.clone(), msg.clone()));
                        let _ = tx.blocking_send(SubagentEvent::ToolResult {
                            _tool: tool_name.clone(),
                            _output: msg,
                        });
                    }
                }
            }
        } else {
            // Text-only response — accumulate and finish
            if !content.is_empty() {
                output.push_str(&content);
                output.push('\n');
            }
            let _ = tx.blocking_send(SubagentEvent::StepCompleted {
                _step: step,
                _output: content.clone(),
            });
            // Break only when we got real content; empty means something went wrong
            if !content.is_empty() {
                break;
            }
        }
    }

    let _ = tx.blocking_send(SubagentEvent::Completed { _output: output.clone() });
    Ok(output)
}