zesdex/apps/infrastructure/src/subagent/engine.rs

//! Subagent engine — runs an LLM-powered agent with tool execution loop.
//!
//! Flow: construct system message → call LLM → parse tool calls → execute
//! tools → continue until the model returns a final text response (no more
//! tool calls) or the iteration limit is reached.
//!
//! Progress reporting: when a `TurnEvent` queue is available via the
//! `ToolCtx`, the engine emits `AgentProgress` events so the TUI can show
//! which tool the subagent is currently executing.

use anyhow::Result;
use tracing::{debug, info, instrument};

use crate::llm::provider::LlmClient;
use crate::subagent::context::SubagentContext;
use crate::subagent::division::{tools_for, AccessTier};
use crate::tools::{tool_defs, Tool, ToolCtx};
use serde_json::Value;
use zesdex_domain::agent::progress::AgentProgress;
use zesdex_domain::core::tool_call::sanitize_tool_arguments;
use zesdex_domain::core::ChatMessage;
use zesdex_domain::subagent_directive;

/// Maximum number of tool-call iterations before the engine gives up.
const MAX_ITERATIONS: u32 = 25;

/// A single tool-result message is truncated before entering the subagent's
/// context so it cannot blow the window (matches the main turn service).
const TOOL_OUTPUT_MAX_CHARS: usize = 12_000;

/// Maximum consecutive identical tool errors before the engine injects a
/// recovery note steering the model to a different approach.
const MAX_CONSECUTIVE_TOOL_ERRORS: usize = 3;

/// Maximum number of read-only tool calls executed concurrently in a single
/// subagent batch. Read-only tools (read/grep/glob/…) block on disk I/O, so
/// running them on parallel OS threads removes the serial round-trip latency
/// for a batch of independent lookups, mirroring the main turn loop.
const MAX_PARALLEL_TOOLS: usize = 8;

/// Execute a batch of tool calls, running read-only tools concurrently when
/// the whole batch is parallel-safe.
///
/// Returns one `(tool_call_id, tool_name, result)` per call **in the original
/// call order** (OpenAI/Anthropic tool-result ordering contract). `Tool::run`
/// is synchronous, so real parallelism comes from scoped OS threads; `Tool`
/// and `ToolCtx` are `Send + Sync`, so the borrowed references can be shared
/// across the short-lived scoped threads.
///
/// If any single tool in the batch mutates state (edit/write/bash/git/…), the
/// whole batch falls back to the safe sequential path so writes never race.
fn execute_tool_batch(
    tools: &[Box<dyn Tool>],
    tool_ctx: &ToolCtx,
    tool_calls: &[zesdex_domain::core::ToolCall],
) -> Vec<(String, String, String)> {
    let parallel = tool_calls.len() > 1
        && tool_calls
            .iter()
            .all(|tc| crate::tools::tool_is_parallel_safe(&tc.function.name));

    if !parallel {
        // Sequential fallback (kept identical to the historical behavior).
        return tool_calls
            .iter()
            .map(|tc| {
                let tool_name = tc.function.name.clone();
                let args = sanitize_tool_arguments(&tc.function.arguments);
                let result = run_one_tool(tools, tool_ctx, &tool_name, &args);
                (tc.id.clone(), tool_name, result)
            })
            .collect();
    }

    // Bounded parallel path: process the batch in windows of
    // `MAX_PARALLEL_TOOLS` so concurrency stays bounded, joining each window
    // before the next so results stay in original order.
    let mut ordered = Vec::with_capacity(tool_calls.len());
    for window in tool_calls.chunks(MAX_PARALLEL_TOOLS) {
        let window_results = std::thread::scope(|s| {
            let handles: Vec<_> = window
                .iter()
                .map(|tc| {
                    let tool_name = tc.function.name.clone();
                    let args = sanitize_tool_arguments(&tc.function.arguments);
                    s.spawn(move || {
                        debug!("Subagent executing tool: {tool_name}");
                        run_one_tool(tools, tool_ctx, &tool_name, &args)
                    })
                })
                .collect();
            handles
                .into_iter()
                .map(|h| {
                    h.join()
                        .unwrap_or_else(|_| "Error: tool panicked".to_string())
                })
                .collect::<Vec<_>>()
        });
        for (tc, result) in window.iter().zip(window_results) {
            ordered.push((tc.id.clone(), tc.function.name.clone(), result));
        }
    }
    ordered
}

/// Run a single synchronous tool call and capture its result string.
fn run_one_tool(
    tools: &[Box<dyn Tool>],
    tool_ctx: &ToolCtx,
    tool_name: &str,
    args: &Value,
) -> String {
    if let Some(tool) = tools.iter().find(|t| t.name() == tool_name) {
        match tool.run(tool_ctx, args) {
            Ok(output) => output,
            Err(e) => format!("Error: {e}"),
        }
    } else {
        format!("Unknown tool: {tool_name}")
    }
}

/// Pick a `max_tokens` budget proportional to the directive's length.
fn adaptive_max_tokens(directive_len: usize) -> u32 {
    if directive_len <= 80 {
        800
    } else if directive_len <= 400 {
        1600
    } else {
        4096
    }
}

fn truncate_tool_output(output: String) -> String {
    if output.len() <= TOOL_OUTPUT_MAX_CHARS {
        return output;
    }
    let mut result: String = output.chars().take(TOOL_OUTPUT_MAX_CHARS).collect();
    result.push_str(&format!(
        "\n...[truncated {} chars]",
        output.len() - TOOL_OUTPUT_MAX_CHARS
    ));
    result
}

/// Emit an `AgentProgress` event onto the turn-event queue, if one is
/// configured in the `ToolCtx`.
fn report_progress(tool_ctx: &ToolCtx, progress: AgentProgress) {
    if let Some(ref queue) = tool_ctx.turn_events {
        if let Ok(mut q) = queue.lock() {
            q.push_back(zesdex_domain::agent::TurnEvent::AgentProgress(progress));
        }
    }
}

/// Build the system message for a subagent, including current working
/// directory and workspace root information.
fn build_system_message(directive: &str, tool_ctx: &ToolCtx) -> ChatMessage {
    let cwd = std::env::current_dir()
        .map(|p| p.to_string_lossy().to_string())
        .unwrap_or_else(|_| "unknown".to_string());
    let ws_root = tool_ctx
        .workspaces
        .first()
        .map(|p| p.to_string_lossy().to_string())
        .unwrap_or_else(|| cwd.clone());

    ChatMessage::system(subagent_directive(directive, &cwd, &ws_root))
}

/// Run an agent with a directive, access tier, and tool context.
///
/// Flow:
///   1. Resolve allowed tools for the given `access` tier.
///   2. Build a system prompt from the directive using the domain prompt module.
///   3. Loop (up to `MAX_ITERATIONS`):
///      a. Call the LLM (non-streaming) with accumulated messages + tool defs.
///      b. If the response has no tool calls → return the text content.
///      c. Otherwise execute each tool call and append the result as a
///      tool-role message.
///      d. If the response also contained text, append an assistant message.
///   4. If the loop exits naturally, return the iteration-limit message.
///
/// Progress: each tool invocation is reported via `AgentProgress` if a
/// turn-event queue is available in the `ToolCtx`.
#[instrument(skip(ctx, tool_ctx))]
pub async fn run_agent(
    ctx: SubagentContext,
    directive: &str,
    access: AccessTier,
    tool_ctx: ToolCtx,
) -> Result<String> {
    info!("Subagent starting with directive: {directive}");

    let tools = tools_for(&access);
    let defs = tool_defs(&tools);

    let sys_msg = build_system_message(directive, &tool_ctx);

    let mut messages = vec![sys_msg];

    let client = LlmClient::new(
        ctx.api_key.clone(),
        ctx.model.clone(),
        Some(ctx.base_url.clone()),
    );

    let max_tokens = adaptive_max_tokens(directive.len());

    // Track repeated tool errors so the agent can recover from a dead end.
    let mut consecutive_errors = 0usize;
    let mut last_tool = String::new();

    // Limited iteration loop so we don't run forever
    for iteration in 0..MAX_ITERATIONS {
        use zesdex_application::ports::ProviderService;
        let (response_msg, _usage) = client
            .chat(&messages, Some(defs.clone()), Some(max_tokens), Some(0.2))
            .await?;

        let content = response_msg.content.clone().unwrap_or_default();
        let tool_calls = response_msg.tool_calls.unwrap_or_default();

        // If no tool calls, we're done — return content
        if tool_calls.is_empty() {
            info!("Subagent completed after {iteration} iterations");
            report_progress(&tool_ctx, AgentProgress::completed("subagent", directive));
            return Ok(content);
        }

        // Execute tool calls — read-only batches run concurrently (bounded,
        // order preserved); any mutating tool forces the safe sequential path.
        let results = execute_tool_batch(&tools, &tool_ctx, &tool_calls);

        for (id, tool_name, result) in results {
            debug!("Subagent tool {tool_name} finished");

            report_progress(
                &tool_ctx,
                AgentProgress::running(
                    "subagent",
                    format!("{}:{tool_name}", directive),
                    Some(tool_name.clone()),
                ),
            );

            // Error-recovery: if the same tool keeps failing, inject a
            // system note steering the model to a different approach.
            if result.starts_with("Error:") {
                if last_tool.as_str() == tool_name.as_str() {
                    consecutive_errors += 1;
                } else {
                    consecutive_errors = 1;
                    last_tool = tool_name.clone();
                }
                if consecutive_errors >= MAX_CONSECUTIVE_TOOL_ERRORS {
                    messages.push(ChatMessage::system(
                        zesdex_domain::agent::prompt::error_recovery_note(&tool_name, &result),
                    ));
                    consecutive_errors = 0;
                }
            } else {
                consecutive_errors = 0;
            }

            messages.push(ChatMessage::tool(id, truncate_tool_output(result)));
        }

        // Add assistant response if there was text content
        if !content.is_empty() {
            messages.push(ChatMessage::assistant(Some(content)));
        }
    }

    info!("Subagent reached iteration limit ({MAX_ITERATIONS})");
    report_progress(
        &tool_ctx,
        AgentProgress::failed(
            "subagent",
            directive,
            format!("iteration limit ({MAX_ITERATIONS})"),
        ),
    );
    Ok(format!(
        "Subagent reached iteration limit ({MAX_ITERATIONS})"
    ))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tools::ToolCtxBuilder;
    use serde_json::json;

    /// A deterministic mock tool whose `run` returns its own name (opting into
    /// an optional sleep to make parallel-vs-sequential observable).
    struct MockTool {
        name: &'static str,
        sleep_ms: u64,
    }

    impl MockTool {
        fn new(name: &'static str, sleep_ms: u64) -> Self {
            Self { name, sleep_ms }
        }
    }

    impl Tool for MockTool {
        fn name(&self) -> &'static str {
            self.name
        }
        fn description(&self) -> &'static str {
            "mock tool for tests"
        }
        fn parameters(&self) -> Value {
            json!({"type":"object","properties":{}})
        }
        fn run(&self, _ctx: &ToolCtx, _args: &Value) -> Result<String> {
            if self.sleep_ms > 0 {
                std::thread::sleep(std::time::Duration::from_millis(self.sleep_ms));
            }
            Ok(self.name.to_string())
        }
    }

    fn tc(name: &str, id: usize) -> zesdex_domain::core::ToolCall {
        zesdex_domain::core::ToolCall {
            id: format!("call_{id}"),
            type_: "function".to_string(),
            function: zesdex_domain::core::ToolFunction {
                name: name.to_string(),
                arguments: serde_json::Value::String(String::new()),
            },
        }
    }

    fn ctx() -> ToolCtx {
        ToolCtxBuilder::default().build()
    }

    #[test]
    fn parallel_batch_preserves_original_order() {
        let tools: Vec<Box<dyn Tool>> = vec![
            Box::new(MockTool::new("read", 0)),
            Box::new(MockTool::new("grep", 0)),
        ];
        let calls = vec![tc("read", 1), tc("grep", 2), tc("read", 3)];

        let results = execute_tool_batch(&tools, &ctx(), &calls);

        // Results keep the assistant's original call order.
        let names: Vec<&str> = results.iter().map(|(_, n, _)| n.as_str()).collect();
        assert_eq!(names, vec!["read", "grep", "read"]);
        // IDs follow the same original order (ordering contract).
        let ids: Vec<&str> = results.iter().map(|(id, _, _)| id.as_str()).collect();
        assert_eq!(ids, vec!["call_1", "call_2", "call_3"]);
    }

    #[test]
    fn parallel_read_batch_is_faster_than_sequential() {
        // Both reads sleep 30ms each. Parallel should finish ~30ms (both run
        // at once), sequential would take ~60ms.
        let tools: Vec<Box<dyn Tool>> = vec![Box::new(MockTool::new("read", 30))];
        let calls = vec![tc("read", 1), tc("read", 2)];

        let started = std::time::Instant::now();
        let results = execute_tool_batch(&tools, &ctx(), &calls);
        let elapsed = started.elapsed();

        assert_eq!(results.len(), 2);
        assert!(
            elapsed < std::time::Duration::from_millis(55),
            "parallel read batch took {elapsed:?}, expected concurrent execution"
        );
        assert!(elapsed >= std::time::Duration::from_millis(25));
    }

    #[test]
    fn mutating_tool_forces_sequential_batch() {
        // A batch containing a mutating tool ("write") must NOT run in
        // parallel — the single 30ms read runs alone, then the write runs.
        let tools: Vec<Box<dyn Tool>> = vec![
            Box::new(MockTool::new("read", 30)),
            Box::new(MockTool::new("write", 0)),
        ];
        let calls = vec![tc("read", 1), tc("write", 2)];

        let results = execute_tool_batch(&tools, &ctx(), &calls);

        let names: Vec<&str> = results.iter().map(|(_, n, _)| n.as_str()).collect();
        assert_eq!(names, vec!["read", "write"]);
        let ids: Vec<&str> = results.iter().map(|(id, _, _)| id.as_str()).collect();
        assert_eq!(ids, vec!["call_1", "call_2"]);
    }
}