vtcode_core/tools/
tool_middleware.rs

1//! Composable async middleware for tool execution.
2//!
3//! Provides a chain-of-responsibility pattern for pre/post-processing
4//! tool calls with observability and error recovery.
5
6use crate::tools::{ToolCallRequest, ToolCallResponse, UnifiedToolError};
7use async_trait::async_trait;
8use std::ops::Add;
9use std::sync::Arc;
10use std::sync::atomic::{AtomicU64, Ordering};
11
12/// Request context flowing through middleware.
13pub type ToolRequest = ToolCallRequest;
14
15/// Response from tool execution.
16pub type ToolResponse = ToolCallResponse;
17
18/// Typed middleware result.
19pub type MiddlewareResult<T> = Result<T, UnifiedToolError>;
20
21/// Middleware trait for intercepting tool execution.
22#[async_trait]
23pub trait Middleware: Send + Sync {
24    /// Called before tool execution.
25    async fn before_execute(&self, _req: &ToolRequest) -> MiddlewareResult<()> {
26        Ok(())
27    }
28
29    /// Called after successful execution.
30    async fn after_execute(&self, _req: &ToolRequest, _res: &ToolResponse) -> MiddlewareResult<()> {
31        Ok(())
32    }
33
34    /// Called on execution error.
35    async fn on_error(&self, _req: &ToolRequest, _err: &UnifiedToolError) -> MiddlewareResult<()> {
36        Ok(())
37    }
38}
39
40/// Noop middleware (default).
41pub struct NoopMiddleware;
42
43#[async_trait]
44impl Middleware for NoopMiddleware {}
45
46/// Middleware composition chain.
47#[derive(Clone)]
48pub struct MiddlewareChain {
49    middlewares: Vec<Arc<dyn Middleware>>,
50}
51
52impl MiddlewareChain {
53    /// Create new chain.
54    pub fn new() -> Self {
55        Self { middlewares: Vec::new() }
56    }
57
58    /// Add middleware to chain.
59    pub fn push(mut self, mw: Arc<dyn Middleware>) -> Self {
60        self.middlewares.push(mw);
61        self
62    }
63
64    /// Execute before hooks, with optional fail-open behavior.
65    ///
66    /// When `fail_open` is `false` (the default), the first middleware error
67    /// short-circuits the chain and is returned to the caller (fail-closed).
68    /// When `fail_open` is `true`, middleware errors are logged and the chain
69    /// continues; the caller receives `Ok(())` unless every middleware fails.
70    pub async fn before_execute_opt(&self, req: &ToolRequest, fail_open: bool) -> MiddlewareResult<()> {
71        if !fail_open {
72            return self.before_execute(req).await;
73        }
74        let mut last_err = None;
75        for mw in &self.middlewares {
76            if let Err(err) = mw.before_execute(req).await {
77                tracing::warn!(
78                    tool = %req.tool_name,
79                    error = %err,
80                    "middleware before_execute failed; ignoring (fail-open)"
81                );
82                last_err = Some(err);
83            }
84        }
85        if let Some(err) = last_err { Err(err) } else { Ok(()) }
86    }
87
88    /// Execute before hooks.
89    pub async fn before_execute(&self, req: &ToolRequest) -> MiddlewareResult<()> {
90        for mw in &self.middlewares {
91            mw.before_execute(req).await?;
92        }
93        Ok(())
94    }
95
96    /// Execute after hooks.
97    pub async fn after_execute(&self, req: &ToolRequest, res: &ToolResponse) -> MiddlewareResult<()> {
98        // Run in reverse order.
99        for mw in self.middlewares.iter().rev() {
100            mw.after_execute(req, res).await?;
101        }
102        Ok(())
103    }
104
105    /// Execute error hooks.
106    pub async fn on_error(&self, req: &ToolRequest, err: &UnifiedToolError) -> MiddlewareResult<()> {
107        for mw in self.middlewares.iter().rev() {
108            if let Err(handler_err) = mw.on_error(req, err).await {
109                tracing::warn!(
110                    error = %handler_err,
111                    "Middleware error handler itself failed"
112                );
113            }
114        }
115        Ok(())
116    }
117}
118
119impl Add<Arc<dyn Middleware>> for MiddlewareChain {
120    type Output = Self;
121
122    fn add(self, rhs: Arc<dyn Middleware>) -> Self::Output {
123        self.push(rhs)
124    }
125}
126
127impl Default for MiddlewareChain {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133/// Logging middleware.
134pub struct LoggingMiddleware {
135    name: String,
136}
137
138impl LoggingMiddleware {
139    pub fn new(name: impl Into<String>) -> Arc<Self> {
140        Arc::new(Self { name: name.into() })
141    }
142}
143
144#[async_trait]
145impl Middleware for LoggingMiddleware {
146    async fn before_execute(&self, req: &ToolRequest) -> MiddlewareResult<()> {
147        tracing::info!(
148            middleware = %self.name,
149            tool = %req.tool_name,
150            "Executing tool"
151        );
152        Ok(())
153    }
154
155    async fn after_execute(&self, req: &ToolRequest, res: &ToolResponse) -> MiddlewareResult<()> {
156        let duration_ms = res.duration_ms.unwrap_or(0);
157        let cache_hit = res.cache_hit.unwrap_or(false);
158        tracing::info!(
159            middleware = %self.name,
160            tool = %req.tool_name,
161            duration_ms,
162            cache_hit,
163            "Completed tool"
164        );
165        Ok(())
166    }
167
168    async fn on_error(&self, req: &ToolRequest, err: &UnifiedToolError) -> MiddlewareResult<()> {
169        tracing::error!(
170            middleware = %self.name,
171            tool = %req.tool_name,
172            error = %err,
173            "Tool execution failed"
174        );
175        Ok(())
176    }
177}
178
179/// Metrics middleware.
180#[derive(Clone, Copy, Debug)]
181pub struct MetricsSnapshot {
182    pub total_calls: u64,
183    pub successful_calls: u64,
184    pub failed_calls: u64,
185    pub total_duration_ms: u64,
186    pub cache_hits: u64,
187}
188
189pub struct MetricsMiddleware {
190    total_calls: Arc<AtomicU64>,
191    successful_calls: Arc<AtomicU64>,
192    failed_calls: Arc<AtomicU64>,
193    total_duration_ms: Arc<AtomicU64>,
194    cache_hits: Arc<AtomicU64>,
195}
196
197impl MetricsMiddleware {
198    fn new_inner() -> Self {
199        Self {
200            total_calls: Arc::new(AtomicU64::new(0)),
201            successful_calls: Arc::new(AtomicU64::new(0)),
202            failed_calls: Arc::new(AtomicU64::new(0)),
203            total_duration_ms: Arc::new(AtomicU64::new(0)),
204            cache_hits: Arc::new(AtomicU64::new(0)),
205        }
206    }
207
208    pub fn new() -> Arc<Self> {
209        Arc::new(Self::new_inner())
210    }
211
212    pub async fn snapshot(&self) -> MetricsSnapshot {
213        MetricsSnapshot {
214            total_calls: self.total_calls.load(Ordering::Relaxed),
215            successful_calls: self.successful_calls.load(Ordering::Relaxed),
216            failed_calls: self.failed_calls.load(Ordering::Relaxed),
217            total_duration_ms: self.total_duration_ms.load(Ordering::Relaxed),
218            cache_hits: self.cache_hits.load(Ordering::Relaxed),
219        }
220    }
221}
222
223#[async_trait]
224impl Middleware for MetricsMiddleware {
225    async fn before_execute(&self, _: &ToolRequest) -> MiddlewareResult<()> {
226        self.total_calls.fetch_add(1, Ordering::Relaxed);
227        Ok(())
228    }
229
230    async fn after_execute(&self, _: &ToolRequest, res: &ToolResponse) -> MiddlewareResult<()> {
231        self.successful_calls.fetch_add(1, Ordering::Relaxed);
232        self.total_duration_ms
233            .fetch_add(res.duration_ms.unwrap_or(0), Ordering::Relaxed);
234
235        if res.cache_hit.unwrap_or(false) {
236            self.cache_hits.fetch_add(1, Ordering::Relaxed);
237        }
238
239        Ok(())
240    }
241
242    async fn on_error(&self, _: &ToolRequest, _: &UnifiedToolError) -> MiddlewareResult<()> {
243        self.failed_calls.fetch_add(1, Ordering::Relaxed);
244        Ok(())
245    }
246}
247
248impl Default for MetricsMiddleware {
249    fn default() -> Self {
250        Self::new_inner()
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use serde_json::Value;
258
259    #[tokio::test]
260    async fn test_chain_execution() {
261        let chain = MiddlewareChain::new()
262            .push(LoggingMiddleware::new("test"))
263            .push(MetricsMiddleware::new());
264
265        let req = ToolRequest {
266            id: "req-1".to_string(),
267            tool_name: "test_tool".into(),
268            args: Value::Null,
269            metadata: Some(Default::default()),
270        };
271
272        chain.before_execute(&req).await.unwrap();
273
274        let res = ToolResponse {
275            id: "req-1".to_string(),
276            success: true,
277            result: Some(Value::Null),
278            error: None,
279            duration_ms: Some(100),
280            cache_hit: Some(false),
281        };
282
283        chain.after_execute(&req, &res).await.unwrap();
284    }
285
286    #[tokio::test]
287    async fn test_metrics_tracking() {
288        let metrics = MetricsMiddleware::new();
289        let chain = MiddlewareChain::new().push(metrics.clone());
290
291        let req = ToolRequest {
292            id: "req-2".to_string(),
293            tool_name: "test".into(),
294            args: Value::Null,
295            metadata: Some(Default::default()),
296        };
297
298        for i in 0..5 {
299            chain.before_execute(&req).await.unwrap();
300            let res = ToolResponse {
301                id: format!("req-2-{i}"),
302                success: true,
303                result: Some(Value::Null),
304                error: None,
305                duration_ms: Some(10),
306                cache_hit: Some(true),
307            };
308            chain.after_execute(&req, &res).await.unwrap();
309        }
310
311        let snapshot = metrics.snapshot().await;
312        assert_eq!(snapshot.total_calls, 5);
313        assert_eq!(snapshot.successful_calls, 5);
314        assert_eq!(snapshot.cache_hits, 5);
315    }
316}
vtcode_core/tools/tool_middleware.rs

vtcode_core/tools/
tool_middleware.rs
