zesdex/src/app/subagent/engine.rs

use tokio::sync::mpsc;
use crate::dto::chat::message::ChatMessage;
use super::context::SubagentContext;
use super::event::SubagentEvent;

pub const MAX_AGENT_STEPS: usize = 25;

fn tool_call_from_response(response: &str) -> Vec<String> {
    let mut calls = Vec::new();
    for line in response.lines() {
        let trimmed = line.trim();
        if let Some(tool_call) = trimmed.strip_prefix("Tool: ") {
            calls.push(tool_call.to_string());
        }
    }
    calls
}

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 max_steps = ctx.max_steps.min(MAX_AGENT_STEPS);
    for step in 0..max_steps {
        let api_key = std::env::var("OPENROUTER_API_KEY").unwrap_or_default();
        let model = std::env::var("OPENROUTER_MODEL").unwrap_or_else(|_| "anthropic/claude-sonnet-5".to_string());

        let client = crate::service::openrouter::OpenRouterClient::new(api_key, model);
        let response = match client.chat(&messages) {
            Ok(r) => r,
            Err(e) => {
                let _ = tx.blocking_send(SubagentEvent::StepFailed {
                    step,
                    error: e.to_string(),
                });
                anyhow::bail!("subagent call failed at step {}: {}", step, e);
            }
        };

        let _ = tx.blocking_send(SubagentEvent::ToolCall {
            tool: "api".to_string(),
            args: serde_json::json!({"response": response}),
        });

        let tool_calls = tool_call_from_response(&response);
        if tool_calls.is_empty() {
            output.push_str(&response);
            output.push('\n');
            let _ = tx.blocking_send(SubagentEvent::StepCompleted {
                step,
                output: response.clone(),
            });
            if !response.contains("Tool:") {
                break;
            }
        } else {
            for tool_name in &tool_calls {
                if !ctx.allowed_tools.is_empty() && !ctx.allowed_tools.contains(tool_name) {
                    let msg = format!("tool '{}' not allowed for this subagent", tool_name);
                    messages.push(ChatMessage::tool_result("subagent".to_string(), msg));
                    continue;
                }
                let _ = tx.blocking_send(SubagentEvent::ToolResult {
                    tool: tool_name.clone(),
                    output: format!("{} executed", tool_name),
                });
            }
            let _ = tx.blocking_send(SubagentEvent::StepCompleted {
                step,
                output: response.clone(),
            });
        }

        let assistant_msg = ChatMessage::assistant(Some(response.clone()));
        messages.push(assistant_msg);
        let user_msg = ChatMessage::user("Continue with the next step based on the tool results above.".to_string());
        messages.push(user_msg);
    }

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