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, ToolCtx};
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;

/// 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
        for tc in &tool_calls {
            let tool_name = &tc.function.name;
            let args = sanitize_tool_arguments(&tc.function.arguments);

            debug!("Subagent executing tool: {tool_name}");

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

            let result = 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}")
            };

            // 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.to_string();
                }
                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(
                tc.id.clone(),
                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})"
    ))
}