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.

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::core::tool_call::sanitize_tool_arguments;
use zesdex_domain::core::ChatMessage;

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

/// 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.
///   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.
#[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 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());

    let mut messages = vec![ChatMessage::system(format!(
        "You are a focused subagent.\n\n\
         Current directory (PWD): {cwd}\n\
         Workspace root: {ws_root}\n\n\
         Your directive:\n{directive}\n\n\
         Complete the directive autonomously using the tools available to you. \
         Return your final answer when done."
    ))];

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

    // Limited iteration loop so we don't run forever
    for iteration in 0..MAX_ITERATIONS {
        let (response_msg, _usage) = client.chat_with_tools_non_streaming(
            &messages,
            Some(defs.clone()),
            Some(4096),
            None,
            None,
        )?;

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

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

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

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

    Ok("Subagent reached iteration limit".to_string())
}