synaptic_middleware/
tool_call_limit.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3use async_trait::async_trait;
4use serde_json::Value;
5use synaptic_core::SynapticError;
6
7use crate::{AgentMiddleware, ToolCallRequest, ToolCaller};
8
9/// Limits the number of tool invocations during a single agent run.
10///
11/// When the limit is exceeded, `wrap_tool_call` returns a
12/// `SynapticError::MaxStepsExceeded` error.
13pub struct ToolCallLimitMiddleware {
14    max_calls: usize,
15    count: AtomicUsize,
16}
17
18impl ToolCallLimitMiddleware {
19    pub fn new(max_calls: usize) -> Self {
20        Self {
21            max_calls,
22            count: AtomicUsize::new(0),
23        }
24    }
25
26    pub fn call_count(&self) -> usize {
27        self.count.load(Ordering::SeqCst)
28    }
29
30    pub fn reset(&self) {
31        self.count.store(0, Ordering::SeqCst);
32    }
33}
34
35#[async_trait]
36impl AgentMiddleware for ToolCallLimitMiddleware {
37    async fn wrap_tool_call(
38        &self,
39        request: ToolCallRequest,
40        next: &dyn ToolCaller,
41    ) -> Result<Value, SynapticError> {
42        let current = self.count.fetch_add(1, Ordering::SeqCst);
43        if current >= self.max_calls {
44            return Err(SynapticError::MaxStepsExceeded {
45                max_steps: self.max_calls,
46            });
47        }
48        next.call(request).await
49    }
50}
synaptic_middleware/tool_call_limit.rs

synaptic_middleware/
tool_call_limit.rs
