agent_base/engine/

pipeline.rs

use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use serde_json::Value;

use crate::tool::{Content, Tool, ToolContext, ToolPolicy, content_text};
use crate::types::{AgentError, AgentResult};

/// Pure execution pipeline — cares about *how* to safely execute a tool.
///
/// Responsibilities: policy hooks, timeout, output size limit.
/// Does NOT do: tool lookup, event emission, user-event forwarding.
#[async_trait]
pub trait ToolExecutionPipeline: Send + Sync {
    async fn execute(
        &self,
        tool: &dyn Tool,
        args: &Value,
        ctx: &ToolContext,
    ) -> AgentResult<Vec<Content>>;
}

/// Default pipeline: ToolPolicy hooks + timeout + output size limit.
#[derive(Clone)]
pub struct DefaultPipeline {
    tool_policy: Option<Arc<dyn ToolPolicy>>,
    tool_timeout_ms: Option<u64>,
    max_output_chars: Option<usize>,
}

impl DefaultPipeline {
    pub fn new(
        tool_policy: Option<Arc<dyn ToolPolicy>>,
        tool_timeout_ms: Option<u64>,
        max_output_chars: Option<usize>,
    ) -> Self {
        Self {
            tool_policy,
            tool_timeout_ms,
            max_output_chars,
        }
    }

    pub fn policy(&self) -> Option<Arc<dyn ToolPolicy>> {
        self.tool_policy.clone()
    }
}

#[async_trait]
impl ToolExecutionPipeline for DefaultPipeline {
    async fn execute(
        &self,
        tool: &dyn Tool,
        args: &Value,
        ctx: &ToolContext,
    ) -> AgentResult<Vec<Content>> {
        // 1. before_call hook
        if let Some(policy) = &self.tool_policy {
            policy.before_call(tool.name(), args, ctx)?;
        }

        // 2. Execute with optional timeout
        let output = if let Some(timeout_ms) = self.tool_timeout_ms {
            match tokio::time::timeout(Duration::from_millis(timeout_ms), tool.call(args, ctx))
                .await
            {
                Ok(result) => result?,
                Err(_) => {
                    tracing::warn!(
                        tool = tool.name(),
                        timeout_ms = timeout_ms,
                        "tool execution timed out"
                    );
                    return Ok(vec![Content::text("[Tool Timeout]")]);
                }
            }
        } else {
            tool.call(args, ctx).await?
        };

        // 3. Output size limit — reject by default rather than silently
        // truncating (design §6.5). Tools that want to return a bounded
        // subset should do their own explicit truncation before returning.
        if let Some(max_chars) = self.max_output_chars {
            let text_len = content_text(&output).chars().count();
            if text_len > max_chars {
                return Err(AgentError::ToolOutputTooLarge {
                    name: tool.name().to_string(),
                    max_chars,
                });
            }
        }

        // 4. after_call hook
        if let Some(policy) = &self.tool_policy {
            policy.after_call(tool.name(), args, &output, ctx)?;
        }

        Ok(output)
    }
}

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

    use crate::types::{ApprovalRequest, Language, SessionId};

    // ── Test helpers ──

    use tokio::sync::mpsc;

    fn test_ctx() -> ToolContext {
        let (tx, _rx) = mpsc::unbounded_channel();
        ToolContext {
            session_id: SessionId::new(1),
            user_event_tx: tx,
            llm_client: None,
            session_store: None,
            language: Language::En,
            cancel_token: tokio_util::sync::CancellationToken::new(),
            max_output_chars: None,
            event_bus: crate::engine::EventBus::new(1),
        }
    }

    struct EchoTool;
    #[async_trait]
    impl Tool for EchoTool {
        fn name(&self) -> &'static str {
            "echo"
        }
        fn description(&self) -> &'static str {
            "echo"
        }
        fn schema(&self) -> Value {
            json!({})
        }
        async fn call(&self, args: &Value, _ctx: &ToolContext) -> AgentResult<Vec<Content>> {
            Ok(vec![Content::text(
                args.get("msg").and_then(|v| v.as_str()).unwrap_or("ok"),
            )])
        }
    }

    struct SlowTool;
    #[async_trait]
    impl Tool for SlowTool {
        fn name(&self) -> &'static str {
            "slow"
        }
        fn description(&self) -> &'static str {
            "slow"
        }
        fn schema(&self) -> Value {
            json!({})
        }
        async fn call(&self, _args: &Value, _ctx: &ToolContext) -> AgentResult<Vec<Content>> {
            tokio::time::sleep(Duration::from_secs(10)).await;
            Ok(vec![Content::text("done")])
        }
    }

    struct FailingTool;
    #[async_trait]
    impl Tool for FailingTool {
        fn name(&self) -> &'static str {
            "fail"
        }
        fn description(&self) -> &'static str {
            "fail"
        }
        fn schema(&self) -> Value {
            json!({})
        }
        async fn call(&self, _args: &Value, _ctx: &ToolContext) -> AgentResult<Vec<Content>> {
            Err(AgentError::tool_not_found("intentional"))
        }
    }

    struct TrackingPolicy {
        before_count: std::sync::atomic::AtomicU32,
        after_count: std::sync::atomic::AtomicU32,
        fail_before: bool,
    }
    impl TrackingPolicy {
        fn new() -> Self {
            Self {
                before_count: std::sync::atomic::AtomicU32::new(0),
                after_count: std::sync::atomic::AtomicU32::new(0),
                fail_before: false,
            }
        }
        fn fail_before_call() -> Self {
            Self {
                before_count: std::sync::atomic::AtomicU32::new(0),
                after_count: std::sync::atomic::AtomicU32::new(0),
                fail_before: true,
            }
        }
    }
    #[async_trait]
    impl ToolPolicy for TrackingPolicy {
        async fn evaluate_approval(&self, _: &str, _: &Value) -> Option<ApprovalRequest> {
            None
        }
        fn before_call(&self, _name: &str, _args: &Value, _ctx: &ToolContext) -> AgentResult<()> {
            self.before_count
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            if self.fail_before {
                return Err(AgentError::internal("before_call denied"));
            }
            Ok(())
        }
        fn after_call(
            &self,
            _name: &str,
            _args: &Value,
            _output: &[Content],
            _ctx: &ToolContext,
        ) -> AgentResult<()> {
            self.after_count
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            Ok(())
        }
    }

    // ── DefaultPipeline tests ──

    #[tokio::test]
    async fn basic_execution() {
        let pipeline = DefaultPipeline::new(None, None, None);
        let output = pipeline
            .execute(&EchoTool, &json!({"msg": "hello"}), &test_ctx())
            .await
            .unwrap();
        assert_eq!(content_text(&output), "hello");
    }

    #[tokio::test]
    async fn policy_before_and_after_called() {
        let policy = Arc::new(TrackingPolicy::new());
        let pipeline = DefaultPipeline::new(Some(policy.clone()), None, None);

        pipeline
            .execute(&EchoTool, &json!({}), &test_ctx())
            .await
            .unwrap();

        assert_eq!(
            policy
                .before_count
                .load(std::sync::atomic::Ordering::SeqCst),
            1
        );
        assert_eq!(
            policy.after_count.load(std::sync::atomic::Ordering::SeqCst),
            1
        );
    }

    #[tokio::test]
    async fn policy_before_call_aborts() {
        let policy = Arc::new(TrackingPolicy::fail_before_call());
        let pipeline = DefaultPipeline::new(Some(policy.clone()), None, None);

        let result = pipeline.execute(&EchoTool, &json!({}), &test_ctx()).await;
        assert!(result.is_err());
        // after_call should NOT be called
        assert_eq!(
            policy.after_count.load(std::sync::atomic::Ordering::SeqCst),
            0
        );
    }

    #[tokio::test]
    async fn timeout_fires() {
        let pipeline = DefaultPipeline::new(None, Some(50), None); // 50ms timeout
        let output = pipeline
            .execute(&SlowTool, &json!({}), &test_ctx())
            .await
            .unwrap();
        assert_eq!(content_text(&output), "[Tool Timeout]");
    }

    #[tokio::test]
    async fn no_timeout_when_tool_fast() {
        let pipeline = DefaultPipeline::new(None, Some(5000), None);
        let output = pipeline
            .execute(&EchoTool, &json!({"msg": "fast"}), &test_ctx())
            .await
            .unwrap();
        assert_eq!(content_text(&output), "fast");
    }

    #[tokio::test]
    async fn output_over_limit_is_rejected() {
        let pipeline = DefaultPipeline::new(None, None, Some(10)); // max 10 chars
        let result = pipeline
            .execute(
                &EchoTool,
                &json!({"msg": "this is a very long message"}),
                &test_ctx(),
            )
            .await;
        assert!(matches!(
            result,
            Err(AgentError::ToolOutputTooLarge { max_chars: 10, .. })
        ));
    }

    #[tokio::test]
    async fn no_rejection_when_short() {
        let pipeline = DefaultPipeline::new(None, None, Some(100));
        let output = pipeline
            .execute(&EchoTool, &json!({"msg": "short"}), &test_ctx())
            .await
            .unwrap();
        assert_eq!(content_text(&output), "short");
    }

    #[tokio::test]
    async fn output_over_limit_cjk_rejected() {
        // CJK chars are 3 bytes each. The size check counts chars, not bytes,
        // so a long Chinese string must still be rejected without panicking.
        let pipeline = DefaultPipeline::new(None, None, Some(20));
        let result = pipeline
            .execute(
                &EchoTool,
                &json!({"msg": "这是一个很长的中文消息，用于测试多字节字符的超限处理"}),
                &test_ctx(),
            )
            .await;
        assert!(matches!(
            result,
            Err(AgentError::ToolOutputTooLarge { max_chars: 20, .. })
        ));
    }

    #[tokio::test]
    async fn cjk_within_char_limit_is_not_rejected() {
        // 17 CJK chars = 51 bytes. With a 20-char limit, byte-counting would
        // wrongly reject (51 > 20), but char-counting must accept (17 <= 20).
        let pipeline = DefaultPipeline::new(None, None, Some(20));
        let output = pipeline
            .execute(
                &EchoTool,
                &json!({"msg": "这是一个用于验证字符计数的中文消息"}),
                &test_ctx(),
            )
            .await
            .unwrap();
        assert_eq!(content_text(&output).chars().count(), 17);
    }

    #[tokio::test]
    async fn timeout_plus_limit() {
        let pipeline = DefaultPipeline::new(None, Some(50), Some(100));
        let output = pipeline
            .execute(&SlowTool, &json!({}), &test_ctx())
            .await
            .unwrap();
        // timeout output is short, so it does not trip the size limit
        assert_eq!(content_text(&output), "[Tool Timeout]");
    }

    #[tokio::test]
    async fn tool_error_propagates() {
        let pipeline = DefaultPipeline::new(None, None, None);
        let result = pipeline
            .execute(&FailingTool, &json!({}), &test_ctx())
            .await;
        assert!(result.is_err());
    }
}