agent_base/engine/
pipeline.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use async_trait::async_trait;
5use serde_json::Value;
6
7use crate::engine::EventBus;
8use crate::tool::{Tool, ToolContext, ToolControlFlow, ToolOutput, ToolPolicy, TruncationInfo};
9use crate::types::{RuntimeEvent, AgentResult};
10
11/// Pure execution pipeline — cares about *how* to safely execute a tool.
12///
13/// Responsibilities: policy hooks, timeout, output truncation.
14/// Does NOT do: tool lookup, event emission, user-event forwarding.
15#[async_trait]
16pub trait ToolExecutionPipeline: Send + Sync {
17    async fn execute(
18        &self,
19        tool: &dyn Tool,
20        args: &Value,
21        ctx: &ToolContext,
22    ) -> AgentResult<ToolOutput>;
23}
24
25/// Default pipeline: ToolPolicy hooks + timeout + output truncation.
26#[derive(Clone)]
27pub struct DefaultPipeline {
28    tool_policy: Option<Arc<dyn ToolPolicy>>,
29    tool_timeout_ms: Option<u64>,
30    max_output_chars: Option<usize>,
31}
32
33impl DefaultPipeline {
34    pub fn new(
35        tool_policy: Option<Arc<dyn ToolPolicy>>,
36        tool_timeout_ms: Option<u64>,
37        max_output_chars: Option<usize>,
38    ) -> Self {
39        Self {
40            tool_policy,
41            tool_timeout_ms,
42            max_output_chars,
43        }
44    }
45
46    pub fn policy(&self) -> Option<Arc<dyn ToolPolicy>> {
47        self.tool_policy.clone()
48    }
49}
50
51#[async_trait]
52impl ToolExecutionPipeline for DefaultPipeline {
53    async fn execute(
54        &self,
55        tool: &dyn Tool,
56        args: &Value,
57        ctx: &ToolContext,
58    ) -> AgentResult<ToolOutput> {
59        // 1. before_call hook
60        if let Some(policy) = &self.tool_policy {
61            policy.before_call(tool.name(), args, ctx)?;
62        }
63
64        // 2. Execute with optional timeout
65        let result = if let Some(timeout_ms) = self.tool_timeout_ms {
66            match tokio::time::timeout(
67                Duration::from_millis(timeout_ms),
68                tool.call(args, ctx),
69            )
70            .await
71            {
72                Ok(result) => result,
73                Err(_) => {
74                    tracing::warn!(
75                        tool = tool.name(),
76                        timeout_ms = timeout_ms,
77                        "tool execution timed out"
78                    );
79                    return Ok(ToolOutput {
80                        summary: "[Tool Timeout]".to_string(),
81                        control_flow: ToolControlFlow::Continue,
82                        ..Default::default()
83                    });
84                }
85            }
86        } else {
87            tool.call(args, ctx).await
88        };
89
90        let mut output = result?;
91
92        // 3. Output truncation
93        if let Some(max_chars) = self.max_output_chars {
94            if output.summary.len() > max_chars {
95                let original_summary_len = output.summary.len();
96                let original_raw_len = output.raw.as_ref().map(|v| v.to_string().len());
97                let suffix = "...(truncated)";
98                let keep = max_chars.saturating_sub(suffix.len());
99                if keep > 0 {
100                    // Use floor_char_boundary to avoid panicking on multi-byte
101                    // UTF-8 characters (e.g. CJK, emoji) where `keep` falls
102                    // in the middle of a character.
103                    let truncate_at = output.summary.floor_char_boundary(keep);
104                    output.summary.truncate(truncate_at);
105                    output.summary.push_str(suffix);
106                } else {
107                    output.summary = suffix[..max_chars].to_string();
108                }
109                output.truncation = Some(TruncationInfo {
110                    original_summary_len,
111                    original_raw_len,
112                    max_allowed_chars: max_chars,
113                });
114                tracing::debug!(
115                    tool = tool.name(),
116                    original_summary_len = original_summary_len,
117                    original_raw_len = original_raw_len,
118                    max_allowed_chars = max_chars,
119                    "tool output truncated"
120                );
121            }
122        }
123
124        // 4. after_call hook
125        if let Some(policy) = &self.tool_policy {
126            policy.after_call(tool.name(), args, &output, ctx)?;
127        }
128
129        Ok(output)
130    }
131}
132
133/// Event-emitting decorator — wraps any pipeline to emit ToolCallStarted/Finished
134/// events on the internal [`EventBus`] and forward [`UserEvent`]s from tools.
135pub(crate) struct EventEmittingPipeline<P: ToolExecutionPipeline> {
136    inner: P,
137    event_bus: EventBus,
138}
139
140impl<P: ToolExecutionPipeline> EventEmittingPipeline<P> {
141    pub fn new(inner: P, event_bus: EventBus) -> Self {
142        Self { inner, event_bus }
143    }
144}
145
146#[async_trait]
147impl<P: ToolExecutionPipeline + Send + Sync> ToolExecutionPipeline for EventEmittingPipeline<P> {
148    async fn execute(
149        &self,
150        tool: &dyn Tool,
151        args: &Value,
152        ctx: &ToolContext,
153    ) -> AgentResult<ToolOutput> {
154        self.event_bus.emit(RuntimeEvent::ToolCallStarted {
155            session_id: ctx.session_id.clone(),
156            tool_name: tool.name().to_string(),
157            args_json: args.to_string(),
158        });
159
160        let result = self.inner.execute(tool, args, ctx).await;
161
162        let summary = match &result {
163            Ok(output) => output.summary.clone(),
164            Err(e) => e.to_string(),
165        };
166        self.event_bus.emit(RuntimeEvent::ToolCallFinished {
167            session_id: ctx.session_id.clone(),
168            tool_name: tool.name().to_string(),
169            summary,
170        });
171
172        result
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use serde_json::json;
180    use std::sync::atomic::{AtomicU32, Ordering};
181
182    use crate::tool::ToolRegistry;
183    use crate::types::{AgentError, Language, SessionId};
184
185    // ── Test helpers ──
186
187    use tokio::sync::mpsc;
188
189    fn test_ctx() -> ToolContext {
190        let (tx, _rx) = mpsc::unbounded_channel();
191        ToolContext {
192            session_id: SessionId::new(1),
193            user_event_tx: tx,
194            llm_client: None,
195            session_store: None,
196            language: Language::En,
197            cancel_token: tokio_util::sync::CancellationToken::new(),
198        }
199    }
200
201    struct EchoTool;
202    #[async_trait]
203    impl Tool for EchoTool {
204        fn name(&self) -> &'static str { "echo" }
205        fn definition(&self) -> Value { json!({}) }
206        async fn call(&self, args: &Value, _ctx: &ToolContext) -> AgentResult<ToolOutput> {
207            Ok(ToolOutput {
208                summary: args.get("msg").and_then(|v| v.as_str()).unwrap_or("ok").to_string(),
209                ..Default::default()
210            })
211        }
212    }
213
214    struct SlowTool;
215    #[async_trait]
216    impl Tool for SlowTool {
217        fn name(&self) -> &'static str { "slow" }
218        fn definition(&self) -> Value { json!({}) }
219        async fn call(&self, _args: &Value, _ctx: &ToolContext) -> AgentResult<ToolOutput> {
220            tokio::time::sleep(Duration::from_secs(10)).await;
221            Ok(ToolOutput { summary: "done".to_string(), ..Default::default() })
222        }
223    }
224
225    struct FailingTool;
226    #[async_trait]
227    impl Tool for FailingTool {
228        fn name(&self) -> &'static str { "fail" }
229        fn definition(&self) -> Value { json!({}) }
230        async fn call(&self, _args: &Value, _ctx: &ToolContext) -> AgentResult<ToolOutput> {
231            Err(AgentError::tool_not_found("intentional"))
232        }
233    }
234
235    struct TrackingPolicy {
236        before_count: AtomicU32,
237        after_count: AtomicU32,
238        fail_before: bool,
239    }
240    impl TrackingPolicy {
241        fn new() -> Self {
242            Self { before_count: AtomicU32::new(0), after_count: AtomicU32::new(0), fail_before: false }
243        }
244        fn fail_before_call() -> Self {
245            Self { before_count: AtomicU32::new(0), after_count: AtomicU32::new(0), fail_before: true }
246        }
247    }
248    #[async_trait]
249    impl ToolPolicy for TrackingPolicy {
250        async fn evaluate_approval(&self, _: &str, _: &Value) -> Option<crate::types::ApprovalRequest> { None }
251        fn before_call(&self, _name: &str, _args: &Value, _ctx: &ToolContext) -> AgentResult<()> {
252            self.before_count.fetch_add(1, Ordering::SeqCst);
253            if self.fail_before {
254                return Err(AgentError::internal("before_call denied"));
255            }
256            Ok(())
257        }
258        fn after_call(&self, _name: &str, _args: &Value, _output: &ToolOutput, _ctx: &ToolContext) -> AgentResult<()> {
259            self.after_count.fetch_add(1, Ordering::SeqCst);
260            Ok(())
261        }
262    }
263
264    // ── DefaultPipeline tests ──
265
266    #[tokio::test]
267    async fn basic_execution() {
268        let pipeline = DefaultPipeline::new(None, None, None);
269        let output = pipeline.execute(&EchoTool, &json!({"msg": "hello"}), &test_ctx()).await.unwrap();
270        assert_eq!(output.summary, "hello");
271    }
272
273    #[tokio::test]
274    async fn policy_before_and_after_called() {
275        let policy = Arc::new(TrackingPolicy::new());
276        let pipeline = DefaultPipeline::new(Some(policy.clone()), None, None);
277
278        pipeline.execute(&EchoTool, &json!({}), &test_ctx()).await.unwrap();
279
280        assert_eq!(policy.before_count.load(Ordering::SeqCst), 1);
281        assert_eq!(policy.after_count.load(Ordering::SeqCst), 1);
282    }
283
284    #[tokio::test]
285    async fn policy_before_call_aborts() {
286        let policy = Arc::new(TrackingPolicy::fail_before_call());
287        let pipeline = DefaultPipeline::new(Some(policy.clone()), None, None);
288
289        let result = pipeline.execute(&EchoTool, &json!({}), &test_ctx()).await;
290        assert!(result.is_err());
291        // after_call should NOT be called
292        assert_eq!(policy.after_count.load(Ordering::SeqCst), 0);
293    }
294
295    #[tokio::test]
296    async fn timeout_fires() {
297        let pipeline = DefaultPipeline::new(None, Some(50), None); // 50ms timeout
298        let output = pipeline.execute(&SlowTool, &json!({}), &test_ctx()).await.unwrap();
299        assert_eq!(output.summary, "[Tool Timeout]");
300    }
301
302    #[tokio::test]
303    async fn no_timeout_when_tool_fast() {
304        let pipeline = DefaultPipeline::new(None, Some(5000), None);
305        let output = pipeline.execute(&EchoTool, &json!({"msg": "fast"}), &test_ctx()).await.unwrap();
306        assert_eq!(output.summary, "fast");
307    }
308
309    #[tokio::test]
310    async fn truncation_applies() {
311        let pipeline = DefaultPipeline::new(None, None, Some(10)); // max 10 chars
312        let output = pipeline.execute(&EchoTool, &json!({"msg": "this is a very long message"}), &test_ctx()).await.unwrap();
313        assert!(output.summary.len() <= 10);
314        assert!(output.truncation.is_some());
315        let t = output.truncation.unwrap();
316        assert_eq!(t.original_summary_len, 27);
317        assert_eq!(t.max_allowed_chars, 10);
318    }
319
320    #[tokio::test]
321    async fn no_truncation_when_short() {
322        let pipeline = DefaultPipeline::new(None, None, Some(100));
323        let output = pipeline.execute(&EchoTool, &json!({"msg": "short"}), &test_ctx()).await.unwrap();
324        assert_eq!(output.summary, "short");
325        assert!(output.truncation.is_none());
326    }
327
328    #[tokio::test]
329    async fn truncation_cjk_no_panic() {
330        // CJK chars are 3 bytes each. With max_chars=20, keep=6 bytes,
331        // which falls inside a 3-byte CJK char. floor_char_boundary
332        // should round down to the nearest char boundary.
333        let pipeline = DefaultPipeline::new(None, None, Some(20));
334        let output = pipeline.execute(&EchoTool, &json!({"msg": "这是一个很长的中文消息，用于测试多字节字符的截断处理"}), &test_ctx()).await.unwrap();
335        assert!(output.summary.len() <= 20);
336        assert!(output.summary.ends_with("...(truncated)") || output.summary.len() <= 20);
337        assert!(output.truncation.is_some());
338        // Must not panic — that's the main assertion
339    }
340
341    #[tokio::test]
342    async fn timeout_plus_truncation() {
343        let pipeline = DefaultPipeline::new(None, Some(50), Some(100));
344        let output = pipeline.execute(&SlowTool, &json!({}), &test_ctx()).await.unwrap();
345        assert_eq!(output.summary, "[Tool Timeout]");
346        assert!(output.truncation.is_none()); // timeout output is short
347    }
348
349    #[tokio::test]
350    async fn tool_error_propagates() {
351        let pipeline = DefaultPipeline::new(None, None, None);
352        let result = pipeline.execute(&FailingTool, &json!({}), &test_ctx()).await;
353        assert!(result.is_err());
354    }
355
356    // ── EventEmittingPipeline tests ──
357
358    #[tokio::test]
359    async fn emits_start_and_finish_events() {
360        let inner = DefaultPipeline::new(None, None, None);
361        let event_bus = EventBus::new(64);
362        let mut rx = event_bus.subscribe();
363        let pipeline = EventEmittingPipeline::new(inner, event_bus);
364
365        let _ = pipeline.execute(&EchoTool, &json!({"msg": "test"}), &test_ctx()).await;
366
367        let mut events = Vec::new();
368        while let Ok(event) = rx.try_recv() {
369            events.push(event);
370        }
371        assert_eq!(events.len(), 2);
372
373        match &events[0] {
374            RuntimeEvent::ToolCallStarted { tool_name, .. } => assert_eq!(tool_name, "echo"),
375            _ => panic!("expected ToolCallStarted"),
376        }
377        match &events[1] {
378            RuntimeEvent::ToolCallFinished { tool_name, summary, .. } => {
379                assert_eq!(tool_name, "echo");
380                assert_eq!(summary, "test");
381            }
382            _ => panic!("expected ToolCallFinished"),
383        }
384    }
385
386    #[tokio::test]
387    async fn emits_finish_with_error_on_failure() {
388        let inner = DefaultPipeline::new(None, None, None);
389        let event_bus = EventBus::new(64);
390        let mut rx = event_bus.subscribe();
391        let pipeline = EventEmittingPipeline::new(inner, event_bus);
392
393        let _ = pipeline.execute(&FailingTool, &json!({}), &test_ctx()).await;
394
395        let mut events = Vec::new();
396        while let Ok(event) = rx.try_recv() {
397            events.push(event);
398        }
399        assert_eq!(events.len(), 2);
400
401        match &events[1] {
402            RuntimeEvent::ToolCallFinished { summary, .. } => {
403                assert!(summary.contains("intentional"));
404            }
405            _ => panic!("expected ToolCallFinished"),
406        }
407    }
408
409    #[tokio::test]
410    async fn event_emitting_delegates_to_inner() {
411        let policy = Arc::new(TrackingPolicy::new());
412        let inner = DefaultPipeline::new(Some(policy.clone()), None, None);
413        let event_bus = EventBus::new(64);
414        let pipeline = EventEmittingPipeline::new(inner, event_bus);
415
416        let output = pipeline.execute(&EchoTool, &json!({"msg": "delegated"}), &test_ctx()).await.unwrap();
417        assert_eq!(output.summary, "delegated");
418        assert_eq!(policy.before_count.load(Ordering::SeqCst), 1);
419        assert_eq!(policy.after_count.load(Ordering::SeqCst), 1);
420    }
421}
agent_base/engine/pipeline.rs

agent_base/engine/
pipeline.rs
