embacle/

agent.rs

// ABOUTME: Configurable agent loop building on text-based tool simulation
// ABOUTME: Provides AgentExecutor with multi-turn tool calling, callbacks, and token accumulation
//
// SPDX-License-Identifier: Apache-2.0
// Copyright (c) 2026 dravr.ai

//! # Agent Loop
//!
//! [`AgentExecutor`] provides a configurable multi-turn agent loop that:
//!
//! 1. Injects a tool catalog into the conversation
//! 2. Calls `provider.complete()` and parses tool calls from the response
//! 3. Executes tools via the provided handler
//! 4. Feeds results back and repeats until no tool calls remain or
//!    `max_turns` is reached
//!
//! Builds on [`tool_simulation`](crate::tool_simulation) types and functions.
//!
//! ## Observability
//!
//! An optional [`OnTurnCallback`] is invoked after each turn with a
//! [`TurnInfo`] snapshot, enabling logging, metrics, or UI updates.

use std::sync::Arc;

use tracing::{debug, info};

use crate::tool_simulation::{
    format_tool_results_as_text, generate_tool_catalog, inject_tool_catalog,
    parse_tool_call_blocks, strip_tool_call_blocks, FunctionCall, FunctionDeclaration,
    TextToolHandler,
};
use crate::types::{ChatMessage, ChatRequest, LlmProvider, RunnerError, TokenUsage};

/// Default maximum turns for the agent loop
const DEFAULT_MAX_TURNS: u32 = 10;

/// Absolute ceiling for `max_turns` to prevent runaway loops
const MAX_TURNS_CEILING: u32 = 50;

/// Callback invoked after each agent turn for observability
pub type OnTurnCallback = Arc<dyn Fn(&TurnInfo) + Send + Sync>;

/// Information about a single agent turn, passed to [`OnTurnCallback`]
#[derive(Debug, Clone)]
pub struct TurnInfo {
    /// Turn number (1-based)
    pub turn: u32,
    /// Tool calls made during this turn
    pub tool_calls: Vec<FunctionCall>,
    /// Text content from the LLM (tool call blocks stripped)
    pub content: String,
    /// Token usage for this turn (if reported by the provider)
    pub usage: Option<TokenUsage>,
}

/// Result of an agent execution run
#[derive(Debug, Clone)]
pub struct AgentResult {
    /// Final text content from the LLM
    pub content: String,
    /// All tool calls made across all turns
    pub tool_calls: Vec<FunctionCall>,
    /// Total number of turns executed
    pub total_turns: u32,
    /// Accumulated token usage across all turns
    pub total_usage: TokenUsage,
    /// Finish reason (e.g., "stop", "`max_turns`")
    pub finish_reason: Option<String>,
}

/// Configurable agent loop with tool calling support.
///
/// # Usage
///
/// ```rust,no_run
/// # use embacle::agent::AgentExecutor;
/// # use embacle::tool_simulation::{FunctionDeclaration, FunctionResponse, TextToolHandler};
/// # use embacle::types::{ChatMessage, LlmProvider};
/// # use serde_json::json;
/// # use std::sync::Arc;
/// # async fn example(provider: &dyn LlmProvider) -> Result<(), embacle::types::RunnerError> {
/// let declarations = vec![FunctionDeclaration::new(
///     "search",
///     "Search the web",
///     Some(json!({"type": "object", "properties": {"q": {"type": "string"}}})),
/// )];
///
/// let handler: TextToolHandler = Arc::new(|name, _args| {
///     FunctionResponse { name: name.to_owned(), response: json!({"results": []}) }
/// });
///
/// let executor = AgentExecutor::new(provider, declarations, handler);
/// let messages = vec![ChatMessage::user("Search for Rust tutorials")];
/// let result = executor.run(messages).await?;
/// println!("{}", result.content);
/// # Ok(())
/// # }
/// ```
pub struct AgentExecutor<'a> {
    provider: &'a dyn LlmProvider,
    declarations: Vec<FunctionDeclaration>,
    tool_handler: TextToolHandler,
    max_turns: u32,
    on_turn: Option<OnTurnCallback>,
    request_template: Option<ChatRequest>,
}

impl<'a> AgentExecutor<'a> {
    /// Create a new agent executor with default settings (`max_turns=10`)
    pub fn new(
        provider: &'a dyn LlmProvider,
        declarations: Vec<FunctionDeclaration>,
        tool_handler: TextToolHandler,
    ) -> Self {
        Self {
            provider,
            declarations,
            tool_handler,
            max_turns: DEFAULT_MAX_TURNS,
            on_turn: None,
            request_template: None,
        }
    }

    /// Set the maximum number of turns (clamped to ceiling of 50)
    pub fn with_max_turns(mut self, max_turns: u32) -> Self {
        self.max_turns = max_turns.min(MAX_TURNS_CEILING);
        self
    }

    /// Set an observability callback invoked after each turn
    pub fn with_on_turn(mut self, callback: OnTurnCallback) -> Self {
        self.on_turn = Some(callback);
        self
    }

    /// Use a request template so each turn inherits the caller's model and
    /// sampling parameters (temperature, `max_tokens`, `top_p`, stop sequences).
    ///
    /// The template's `messages` are replaced with the running conversation on
    /// every turn; all other fields are preserved.
    pub fn with_request_template(mut self, template: ChatRequest) -> Self {
        self.request_template = Some(template);
        self
    }

    /// Build the per-turn request, applying the template if one was set.
    fn build_turn_request(&self, messages: Vec<ChatMessage>) -> ChatRequest {
        match &self.request_template {
            Some(template) => {
                let mut request = template.clone();
                request.messages = messages;
                request
            }
            None => ChatRequest::new(messages),
        }
    }

    /// Run the agent loop with the given initial messages.
    ///
    /// # Errors
    ///
    /// Returns [`RunnerError`] if any `provider.complete()` call fails.
    pub async fn run(
        &self,
        initial_messages: Vec<ChatMessage>,
    ) -> Result<AgentResult, RunnerError> {
        let mut messages = initial_messages;

        // Inject tool catalog into the conversation
        let catalog = generate_tool_catalog(&self.declarations);
        inject_tool_catalog(&mut messages, &catalog);

        debug!(
            tool_count = self.declarations.len(),
            max_turns = self.max_turns,
            "agent: starting loop"
        );

        let mut all_tool_calls: Vec<FunctionCall> = Vec::new();
        let mut total_usage = TokenUsage::default();
        let mut turn: u32 = 0;

        loop {
            turn += 1;
            if turn > self.max_turns {
                info!(max_turns = self.max_turns, "agent: max turns reached");
                return Ok(AgentResult {
                    content: String::new(),
                    tool_calls: all_tool_calls,
                    total_turns: turn - 1,
                    total_usage,
                    finish_reason: Some("max_turns".to_owned()),
                });
            }

            let request = self.build_turn_request(messages.clone());
            tracing::info!(
                turn,
                provider = self.provider.name(),
                message_count = request.messages.len(),
                "agent: dispatching turn to provider"
            );
            let response = self.provider.complete(&request).await?;

            // Accumulate token usage
            if let Some(ref usage) = response.usage {
                accumulate_usage(&mut total_usage, usage);
            }

            tracing::info!(
                turn,
                content_len = response.content.len(),
                tool_calls_inline = response.tool_calls.as_ref().map_or(0, Vec::len),
                finish_reason = response.finish_reason.as_deref().unwrap_or("none"),
                "agent: provider response received"
            );

            // Parse tool calls from the response
            let parsed_calls = parse_tool_call_blocks(&response.content);
            let content = strip_tool_call_blocks(&response.content);

            if parsed_calls.is_empty() {
                // No tool calls — final response
                let turn_info = TurnInfo {
                    turn,
                    tool_calls: vec![],
                    content: content.clone(),
                    usage: response.usage.clone(),
                };

                if let Some(ref callback) = self.on_turn {
                    callback(&turn_info);
                }

                debug!(turn, "agent: final response (no tool calls)");
                return Ok(AgentResult {
                    content,
                    tool_calls: all_tool_calls,
                    total_turns: turn,
                    total_usage,
                    finish_reason: response.finish_reason,
                });
            }

            info!(
                turn,
                call_count = parsed_calls.len(),
                "agent: executing tool calls"
            );

            // Execute tool calls
            let mut function_responses = Vec::with_capacity(parsed_calls.len());
            for call in &parsed_calls {
                let resp = (self.tool_handler)(&call.name, &call.args);
                function_responses.push(resp);
            }

            // Build turn info for callback
            let turn_info = TurnInfo {
                turn,
                tool_calls: parsed_calls.clone(),
                content: content.clone(),
                usage: response.usage,
            };

            if let Some(ref callback) = self.on_turn {
                callback(&turn_info);
            }

            all_tool_calls.extend(parsed_calls);

            // Append assistant response and tool results to conversation
            if !content.is_empty() {
                messages.push(ChatMessage::assistant(content));
            }

            let tool_results_text = format_tool_results_as_text(&function_responses);
            messages.push(ChatMessage::user(tool_results_text));
        }
    }
}

/// Add one turn's usage into the running total.
///
/// The optional categories sum as options: a category stays `None` until
/// some turn reports it, so a run that never measured cache or reasoning
/// does not claim a measured zero.
fn accumulate_usage(total: &mut TokenUsage, turn: &TokenUsage) {
    total.prompt_tokens = total.prompt_tokens.saturating_add(turn.prompt_tokens);
    total.completion_tokens = total
        .completion_tokens
        .saturating_add(turn.completion_tokens);
    total.total_tokens = total.total_tokens.saturating_add(turn.total_tokens);
    total.cached_read_tokens = add_reported(total.cached_read_tokens, turn.cached_read_tokens);
    total.cached_write_tokens = add_reported(total.cached_write_tokens, turn.cached_write_tokens);
    total.cached_write_1h_tokens =
        add_reported(total.cached_write_1h_tokens, turn.cached_write_1h_tokens);
    total.reasoning_tokens = add_reported(total.reasoning_tokens, turn.reasoning_tokens);
    total.model_requests = add_reported(total.model_requests, turn.model_requests);
}

/// Sum two optional counts, `None` only when neither was reported.
fn add_reported(total: Option<u32>, turn: Option<u32>) -> Option<u32> {
    match (total, turn) {
        (Some(a), Some(b)) => Some(a.saturating_add(b)),
        (a, b) => a.or(b),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tool_simulation::FunctionResponse;
    use crate::types::{
        ChatMessage, ChatRequest, ChatResponse, ChatStream, LlmCapabilities, LlmProvider,
        RunnerError, TokenUsage,
    };
    use async_trait::async_trait;
    use serde_json::json;
    use std::sync::atomic::{AtomicU32, Ordering};
    use std::sync::Mutex;

    struct TestProvider {
        responses: Mutex<Vec<Result<ChatResponse, RunnerError>>>,
        call_count: AtomicU32,
    }

    impl TestProvider {
        fn new(responses: Vec<Result<ChatResponse, RunnerError>>) -> Self {
            Self {
                responses: Mutex::new(responses),
                call_count: AtomicU32::new(0),
            }
        }
    }

    #[async_trait]
    impl LlmProvider for TestProvider {
        fn name(&self) -> &'static str {
            "test"
        }
        fn display_name(&self) -> &str {
            "Test Provider"
        }
        fn capabilities(&self) -> LlmCapabilities {
            LlmCapabilities::text_only()
        }
        fn default_model(&self) -> &'static str {
            "test-model"
        }
        fn available_models(&self) -> &[String] {
            &[]
        }
        async fn complete(&self, _request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
            self.call_count.fetch_add(1, Ordering::SeqCst);
            let mut responses = self.responses.lock().expect("test lock"); // Safe: test assertion
            if responses.is_empty() {
                Err(RunnerError::internal("no more test responses"))
            } else {
                responses.remove(0)
            }
        }
        async fn complete_stream(&self, _request: &ChatRequest) -> Result<ChatStream, RunnerError> {
            Err(RunnerError::internal("not supported"))
        }
        async fn health_check(&self) -> Result<bool, RunnerError> {
            Ok(true)
        }
    }

    fn make_response(content: &str, usage: Option<TokenUsage>) -> ChatResponse {
        ChatResponse {
            content: content.to_owned(),
            model: "test-model".to_owned(),
            usage,
            finish_reason: Some("stop".to_owned()),
            warnings: None,
            tool_calls: None,
        }
    }

    fn noop_handler() -> TextToolHandler {
        Arc::new(|name: &str, _args: &serde_json::Value| FunctionResponse {
            name: name.to_owned(),
            response: json!({"status": "ok"}),
        })
    }

    #[tokio::test]
    async fn single_turn_no_tool_calls() {
        let provider = TestProvider::new(vec![Ok(make_response(
            "Here is a direct answer without tool calls.",
            Some(TokenUsage::new(10, 8, 18)),
        ))]);

        let declarations = vec![FunctionDeclaration {
            name: "search".to_owned(),
            description: "Search the web".to_owned(),
            parameters: None,
            defer_loading: false,
        }];

        let executor = AgentExecutor::new(&provider, declarations, noop_handler());
        let messages = vec![ChatMessage::user("Hello")];
        let result = executor.run(messages).await.expect("should succeed"); // Safe: test assertion

        assert_eq!(
            result.content,
            "Here is a direct answer without tool calls."
        );
        assert!(result.tool_calls.is_empty());
        assert_eq!(result.total_turns, 1);
        assert_eq!(result.total_usage.prompt_tokens, 10);
        assert_eq!(result.finish_reason, Some("stop".to_owned()));
    }

    #[tokio::test]
    async fn multi_turn_with_tool_calls() {
        let provider = TestProvider::new(vec![
            // Turn 1: LLM calls a tool
            Ok(make_response(
                "Let me search for that.\n<tool_call>\n{\"name\": \"search\", \"arguments\": {\"q\": \"rust\"}}\n</tool_call>",
                Some(TokenUsage::new(10, 15, 25)),
            )),
            // Turn 2: LLM responds with the result
            Ok(make_response(
                "Based on the search results, Rust is a systems programming language.",
                Some(TokenUsage::new(30, 12, 42)),
            )),
        ]);

        let declarations = vec![FunctionDeclaration {
            name: "search".to_owned(),
            description: "Search the web".to_owned(),
            parameters: Some(json!({"type": "object", "properties": {"q": {"type": "string"}}})),
            defer_loading: false,
        }];

        let executor = AgentExecutor::new(&provider, declarations, noop_handler());
        let messages = vec![ChatMessage::user("What is Rust?")];
        let result = executor.run(messages).await.expect("should succeed"); // Safe: test assertion

        assert!(result.content.contains("systems programming"));
        assert_eq!(result.tool_calls.len(), 1);
        assert_eq!(result.tool_calls[0].name, "search");
        assert_eq!(result.total_turns, 2);
        assert_eq!(result.total_usage.prompt_tokens, 40);
        assert_eq!(result.total_usage.completion_tokens, 27);
        // Neither turn reported a breakdown, so none is claimed.
        assert_eq!(result.total_usage.cached_read_tokens, None);
        assert_eq!(result.total_usage.cached_write_tokens, None);
        assert_eq!(result.total_usage.reasoning_tokens, None);
    }

    #[tokio::test]
    async fn multi_turn_sums_cache_and_reasoning_as_reported() {
        let provider = TestProvider::new(vec![
            // Turn 1: a cold prefix — a cache write, no read, no reasoning reported
            Ok(make_response(
                "<tool_call>\n{\"name\": \"search\", \"arguments\": {}}\n</tool_call>",
                Some(TokenUsage::new(100, 10, 110).with_cache(None, Some(80))),
            )),
            // Turn 2: the warm prefix is read back, and reasoning is reported
            Ok(make_response(
                "Done.",
                Some(
                    TokenUsage::new(120, 5, 145)
                        .with_cache(Some(80), Some(0))
                        .with_reasoning(Some(20)),
                ),
            )),
        ]);

        let declarations = vec![FunctionDeclaration {
            name: "search".to_owned(),
            description: "Search the web".to_owned(),
            parameters: None,
            defer_loading: false,
        }];

        let executor = AgentExecutor::new(&provider, declarations, noop_handler());
        let result = executor
            .run(vec![ChatMessage::user("Go")])
            .await
            .expect("should succeed"); // Safe: test assertion

        assert_eq!(result.total_turns, 2);
        assert_eq!(result.total_usage.prompt_tokens, 220);
        assert_eq!(result.total_usage.completion_tokens, 15);
        assert_eq!(result.total_usage.total_tokens, 255);
        assert_eq!(result.total_usage.cached_read_tokens, Some(80));
        assert_eq!(result.total_usage.cached_write_tokens, Some(80));
        assert_eq!(result.total_usage.reasoning_tokens, Some(20));
    }

    #[test]
    fn accumulated_counts_saturate_instead_of_overflowing() {
        let mut total = TokenUsage::new(u32::MAX - 1, u32::MAX, u32::MAX).with_reasoning(Some(1));
        accumulate_usage(
            &mut total,
            &TokenUsage::new(5, 5, 10).with_reasoning(Some(u32::MAX)),
        );
        assert_eq!(total.prompt_tokens, u32::MAX);
        assert_eq!(total.completion_tokens, u32::MAX);
        assert_eq!(total.total_tokens, u32::MAX);
        assert_eq!(total.reasoning_tokens, Some(u32::MAX));
    }

    #[tokio::test]
    async fn on_turn_callback_invoked() {
        let provider = TestProvider::new(vec![
            Ok(make_response(
                "<tool_call>\n{\"name\": \"ping\", \"arguments\": {}}\n</tool_call>",
                None,
            )),
            Ok(make_response("Done!", None)),
        ]);

        let declarations = vec![FunctionDeclaration {
            name: "ping".to_owned(),
            description: "Ping".to_owned(),
            parameters: None,
            defer_loading: false,
        }];

        let turn_log: Arc<Mutex<Vec<u32>>> = Arc::new(Mutex::new(Vec::new()));
        let turn_log_clone = Arc::clone(&turn_log);

        let callback: OnTurnCallback = Arc::new(move |info: &TurnInfo| {
            turn_log_clone.lock().expect("lock").push(info.turn); // Safe: test assertion
        });

        let executor =
            AgentExecutor::new(&provider, declarations, noop_handler()).with_on_turn(callback);
        let messages = vec![ChatMessage::user("ping")];
        executor.run(messages).await.expect("should succeed"); // Safe: test assertion

        let logged = turn_log.lock().expect("lock").clone(); // Safe: test assertion
        assert_eq!(logged, vec![1, 2]);
    }

    #[tokio::test]
    async fn max_turns_exhaustion() {
        // Provider always returns tool calls — should exhaust max_turns
        let mut responses = Vec::new();
        for _ in 0..5 {
            responses.push(Ok(make_response(
                "<tool_call>\n{\"name\": \"loop\", \"arguments\": {}}\n</tool_call>",
                None,
            )));
        }
        let provider = TestProvider::new(responses);

        let declarations = vec![FunctionDeclaration {
            name: "loop".to_owned(),
            description: "Loop forever".to_owned(),
            parameters: None,
            defer_loading: false,
        }];

        let executor =
            AgentExecutor::new(&provider, declarations, noop_handler()).with_max_turns(3);
        let messages = vec![ChatMessage::user("go")];
        let result = executor.run(messages).await.expect("should not error"); // Safe: test assertion

        assert_eq!(result.finish_reason, Some("max_turns".to_owned()));
        assert_eq!(result.total_turns, 3);
        assert_eq!(result.tool_calls.len(), 3);
    }

    #[tokio::test]
    async fn request_template_preserves_params() {
        // Provider captures the request it receives so we can assert on it.
        struct CapturingProvider {
            seen_model: Mutex<Option<String>>,
            seen_temp: Mutex<Option<f32>>,
        }

        #[async_trait]
        impl LlmProvider for CapturingProvider {
            fn name(&self) -> &'static str {
                "capture"
            }
            fn display_name(&self) -> &str {
                "Capture"
            }
            fn capabilities(&self) -> LlmCapabilities {
                LlmCapabilities::text_only()
            }
            fn default_model(&self) -> &'static str {
                "default-model"
            }
            fn available_models(&self) -> &[String] {
                &[]
            }
            async fn complete(&self, request: &ChatRequest) -> Result<ChatResponse, RunnerError> {
                *self.seen_model.lock().expect("lock") = request.model.clone(); // Safe: test assertion
                *self.seen_temp.lock().expect("lock") = request.temperature; // Safe: test assertion
                Ok(make_response("done", None))
            }
            async fn complete_stream(
                &self,
                _request: &ChatRequest,
            ) -> Result<ChatStream, RunnerError> {
                Err(RunnerError::internal("not supported"))
            }
            async fn health_check(&self) -> Result<bool, RunnerError> {
                Ok(true)
            }
        }

        let provider = CapturingProvider {
            seen_model: Mutex::new(None),
            seen_temp: Mutex::new(None),
        };

        let mut template = ChatRequest::new(vec![]);
        template.model = Some("gpt-5.4".to_owned());
        template.temperature = Some(0.3);

        let executor =
            AgentExecutor::new(&provider, vec![], noop_handler()).with_request_template(template);
        executor
            .run(vec![ChatMessage::user("hi")])
            .await
            .expect("should succeed"); // Safe: test assertion

        assert_eq!(
            provider.seen_model.lock().expect("lock").as_deref(), // Safe: test assertion
            Some("gpt-5.4")
        );
        assert_eq!(*provider.seen_temp.lock().expect("lock"), Some(0.3)); // Safe: test assertion
    }

    #[tokio::test]
    async fn token_accumulation() {
        let provider = TestProvider::new(vec![
            Ok(make_response(
                "<tool_call>\n{\"name\": \"a\", \"arguments\": {}}\n</tool_call>",
                Some(TokenUsage::new(10, 5, 15)),
            )),
            Ok(make_response("final", Some(TokenUsage::new(20, 3, 23)))),
        ]);

        let declarations = vec![FunctionDeclaration {
            name: "a".to_owned(),
            description: "Tool A".to_owned(),
            parameters: None,
            defer_loading: false,
        }];

        let executor = AgentExecutor::new(&provider, declarations, noop_handler());
        let messages = vec![ChatMessage::user("go")];
        let result = executor.run(messages).await.expect("should succeed"); // Safe: test assertion

        assert_eq!(result.total_usage.prompt_tokens, 30);
        assert_eq!(result.total_usage.completion_tokens, 8);
        assert_eq!(result.total_usage.total_tokens, 38);
    }
}