lc-agents 0.22.4

// lc-agents/src/executor/tools.rs
//! Tool execution helpers shared by the streaming and non-streaming paths.

use super::AgentError;
use crate::types::{AgentAction, ToolInput};
use lc_core::tools::{BaseTool, ToolError};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Semaphore;

/// A11: build an O(1) name → tool index from a tool list.
///
/// The prior lookups were `tools.iter().find(|t| t.name() == name)` — O(n) on every
/// tool call. `entry().or_insert_with` preserves the original "first match wins"
/// semantics exactly: the linear scan returned the *first* tool whose name matched,
/// and `or_insert_with` keeps the first-inserted entry when names collide.
pub(crate) fn index_tools(tools: &[Arc<dyn BaseTool>]) -> HashMap<String, Arc<dyn BaseTool>> {
    let mut map: HashMap<String, Arc<dyn BaseTool>> = HashMap::with_capacity(tools.len());
    for tool in tools {
        map.entry(tool.name().to_string())
            .or_insert_with(|| tool.clone());
    }
    map
}

/// A1 Spotlighting — open marker wrapping untrusted tool output.
///
/// Mirrors `lc_guardrails`' `DEFAULT_OPEN_MARKER`/`DEFAULT_CLOSE_MARKER` and its escaping
/// (a backslash-prefixed close tag) so output wrapped here parse back cleanly through
/// [`lc_guardrails::spotlighting::unwrap`]. `lc-agents` cannot import `lc-guardrails`
/// (the dependency points the other way), so the constants are mirrored rather than reused.
pub(crate) const UNTRUSTED_OPEN: &str = "<untrusted_data>";
pub(crate) const UNTRUSTED_CLOSE: &str = "</untrusted_data>";

/// A1: wraps untrusted tool output in the data delimiters, escaping any embedded close
/// marker so a hostile tool result cannot forge a premature `</untrusted_data>`.
pub(crate) fn wrap_tool_output(output: &str) -> String {
    let escaped = output.replace(UNTRUSTED_CLOSE, "\\</untrusted_data>");
    format!("{UNTRUSTED_OPEN}{escaped}{UNTRUSTED_CLOSE}")
}

/// Tool-**execution** error → observation text, fed back to the loop so the agent can
/// recover on its own. 0.20.0 S3.1 unified all four execution paths (invoke/stream ×
/// single/parallel) to this soft-fail semantics — the sequential `invoke` single-tool
/// path previously hard-failed upward and no longer does.
///
/// Only `AgentError::ToolExecutionError` (the tool ran and failed) is routed here.
/// Framework guardrails that reject a call *before* execution — `ToolNotFound`, tool
/// permission policy, hook `Reject`, and `ToolError::ControlAbort` (e.g. the handoff
/// cycle / depth guard) — are **not** soft-failed: the agent cannot recover from them
/// by re-planning, so they propagate hard.
pub(crate) fn tool_error_observation(err: &AgentError) -> String {
    format!("[Tool execution error: {err}]")
}

/// Executes a tool with an optional timeout.
///
/// With `Some(d)`, the tool call is cancelled (and errors) if it exceeds `d`.
/// Shared by both the non-streaming and streaming execution paths.
pub(crate) async fn run_tool_with_timeout(
    tool: &Arc<dyn BaseTool>,
    input: String,
    timeout: Option<Duration>,
) -> Result<String, ToolError> {
    let fut = tool.run(input);
    match timeout {
        Some(d) => match tokio::time::timeout(d, fut).await {
            Ok(result) => result,
            Err(_) => Err(ToolError::Timeout(d.as_secs())),
        },
        None => fut.await,
    }
}

/// Helper: execute a single tool for streaming (no RunTree dependency).
///
/// A11: takes the prebuilt name → tool index (an O(1) lookup) instead of a raw slice.
/// `spotlight`/`rule_of_two` mirror the executor's A1/A2 toggles so the streaming path
/// (which cannot read `AgentExecutor` fields) applies the same guards as invoke.
pub(crate) async fn execute_tool_for_stream(
    tools: &HashMap<String, Arc<dyn BaseTool>>,
    action: &AgentAction,
    timeout: Option<Duration>,
    spotlight: bool,
    rule_of_two: bool,
) -> Result<String, AgentError> {
    let tool = tools
        .get(&action.tool)
        .ok_or_else(|| AgentError::ToolNotFound(action.tool.clone()))?;

    // A2: block a tool that arms all three risk properties (v0.22.1 §S8).
    if rule_of_two && tool.risk().count_armed() >= 3 {
        return Ok(
            "[BLOCKED by Rule of Two: tool declares untrusted-input + sensitive-access + state-changing]"
                .to_string(),
        );
    }

    let input_str = match &action.tool_input {
        ToolInput::String { value: s } => s.clone(),
        ToolInput::Object { value: v } => serde_json::to_string(v)
            .map_err(|e| AgentError::Other(format!("Failed to serialize tool input: {}", e)))?,
    };

    let output = run_tool_with_timeout(tool, input_str, timeout)
        .await
        .map_err(|e| match e {
            // 0.20.0 A-H1: keep `ControlAbort` (handoff cycle / depth guard) distinct
            // from a plain execution failure, mirroring `execute_tool_inner`. The
            // streaming caller must be able to tell "the agent cannot recover, stop"
            // apart from "the tool ran and failed, re-plan".
            ToolError::ControlAbort(msg) => AgentError::Other(format!("Tool call aborted: {msg}")),
            other => AgentError::ToolExecutionError(other.to_string()),
        })?;

    // A1: wrap untrusted tool output when spotlighting is on (v0.22.1 §S8).
    Ok(if spotlight {
        wrap_tool_output(&output)
    } else {
        output
    })
}

/// Helper: execute multiple tools in parallel for streaming.
///
/// Concurrency is capped at `max_concurrency` via a local semaphore.
///
/// 0.20.0 A-H1: mirrors the non-streaming parallel path — only a tool-**execution**
/// error becomes an observation; a framework guardrail (`ToolNotFound` /
/// `ControlAbort` / input serialization) in any one tool propagates hard as `Err` so
/// the caller ends the stream instead of feeding a re-plan loop that cannot recover.
pub(crate) async fn execute_tools_parallel_for_stream(
    tools: &HashMap<String, Arc<dyn BaseTool>>,
    actions: &[AgentAction],
    timeout: Option<Duration>,
    max_concurrency: usize,
    spotlight: bool,
    rule_of_two: bool,
) -> Result<Vec<String>, AgentError> {
    use futures_util::future::join_all;

    let sem = Arc::new(Semaphore::new(max_concurrency));
    let futures = actions.iter().map(|action| {
        let sem = sem.clone();
        async move {
            let _permit = sem
                .acquire_owned()
                .await
                .map_err(|e| AgentError::Other(format!("concurrency semaphore closed: {e}")))?;
            execute_tool_for_stream(tools, action, timeout, spotlight, rule_of_two).await
        }
    });

    let results = join_all(futures).await;

    let mut observations = Vec::with_capacity(results.len());
    for result in results {
        match result {
            Ok(output) => observations.push(output),
            Err(e @ AgentError::ToolExecutionError(_)) => {
                observations.push(tool_error_observation(&e))
            }
            Err(e) => return Err(e),
        }
    }
    Ok(observations)
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;

    /// Minimal named tool for the A11 index tests.
    struct NamedTool(&'static str);

    #[async_trait]
    impl BaseTool for NamedTool {
        fn name(&self) -> &str {
            self.0
        }
        fn description(&self) -> &str {
            "A11 index test tool"
        }
        async fn run(&self, input: String) -> Result<String, ToolError> {
            Ok(input)
        }
    }

    fn named(name: &'static str) -> Arc<dyn BaseTool> {
        Arc::new(NamedTool(name))
    }

    #[test]
    fn index_resolves_each_tool_by_name() {
        let tools: Vec<Arc<dyn BaseTool>> = vec![named("alpha"), named("beta"), named("gamma")];
        let idx = index_tools(&tools);
        assert!(idx.contains_key("alpha"));
        assert!(idx.contains_key("beta"));
        assert!(idx.contains_key("gamma"));
        assert!(!idx.contains_key("delta"));
        assert_eq!(idx.len(), 3);
    }

    #[test]
    fn index_preserves_first_match_on_name_collision() {
        // The old linear `find` returned the *first* tool whose name matched; the
        // index uses `or_insert_with`, so the first-inserted entry must win too —
        // otherwise the O(1) optimization would silently change which tool runs.
        let first = named("dup");
        let second = named("dup");
        let tools: Vec<Arc<dyn BaseTool>> = vec![first.clone(), second.clone()];
        let idx = index_tools(&tools);
        let resolved = idx.get("dup").unwrap();
        assert!(
            Arc::ptr_eq(resolved, &first),
            "first-inserted tool must win on name collision"
        );
    }
}