tool_call_batcher/
lib.rs

1/*!
2tool-call-batcher: queue pending tool calls and flush them as a batch.
3
4When an agent produces multiple tool calls in one turn, this crate lets
5you collect them, inspect them, and flush them in a controlled way.
6
7```rust
8use tool_call_batcher::CallBatcher;
9use serde_json::json;
10
11let mut b = CallBatcher::new();
12b.enqueue("search", json!({"q": "rust"}));
13b.enqueue("fetch",  json!({"url": "http://example.com"}));
14let batch = b.flush();
15assert_eq!(batch.len(), 2);
16assert!(b.is_empty());
17```
18*/
19
20use serde_json::Value;
21
22#[derive(Debug, Clone)]
23pub struct PendingCall {
24    pub id: usize,
25    pub tool: String,
26    pub args: Value,
27    pub priority: i32,
28}
29
30#[derive(Debug, Default)]
31pub struct CallBatcher {
32    queue: Vec<PendingCall>,
33    next_id: usize,
34}
35
36impl CallBatcher {
37    pub fn new() -> Self { Self::default() }
38
39    /// Add a call to the queue. Returns its id.
40    pub fn enqueue(&mut self, tool: impl Into<String>, args: Value) -> usize {
41        self.enqueue_with_priority(tool, args, 0)
42    }
43
44    pub fn enqueue_with_priority(&mut self, tool: impl Into<String>, args: Value, priority: i32) -> usize {
45        let id = self.next_id;
46        self.next_id += 1;
47        self.queue.push(PendingCall { id, tool: tool.into(), args, priority });
48        id
49    }
50
51    /// Remove and return all queued calls (in enqueue order).
52    pub fn flush(&mut self) -> Vec<PendingCall> {
53        std::mem::take(&mut self.queue)
54    }
55
56    /// Remove and return calls sorted by priority descending.
57    pub fn flush_by_priority(&mut self) -> Vec<PendingCall> {
58        let mut calls = std::mem::take(&mut self.queue);
59        calls.sort_by(|a, b| b.priority.cmp(&a.priority));
60        calls
61    }
62
63    /// Peek at calls without removing them.
64    pub fn peek(&self) -> &[PendingCall] { &self.queue }
65
66    /// Remove a specific call by id.
67    pub fn cancel(&mut self, id: usize) -> bool {
68        let before = self.queue.len();
69        self.queue.retain(|c| c.id != id);
70        self.queue.len() < before
71    }
72
73    /// Calls for a specific tool (peeked, not removed).
74    pub fn calls_for(&self, tool: &str) -> Vec<&PendingCall> {
75        self.queue.iter().filter(|c| c.tool == tool).collect()
76    }
77
78    pub fn len(&self) -> usize { self.queue.len() }
79    pub fn is_empty(&self) -> bool { self.queue.is_empty() }
80    pub fn clear(&mut self) { self.queue.clear(); }
81
82    /// Build Anthropic-style tool_use content blocks.
83    pub fn to_tool_use_blocks(&self) -> Vec<Value> {
84        self.queue.iter().map(|c| {
85            serde_json::json!({
86                "type": "tool_use",
87                "id": format!("call_{}", c.id),
88                "name": c.tool,
89                "input": c.args,
90            })
91        }).collect()
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use serde_json::json;
99
100    #[test]
101    fn enqueue_and_flush() {
102        let mut b = CallBatcher::new();
103        b.enqueue("search", json!({}));
104        b.enqueue("fetch", json!({}));
105        let calls = b.flush();
106        assert_eq!(calls.len(), 2);
107        assert!(b.is_empty());
108    }
109
110    #[test]
111    fn flush_clears_queue() {
112        let mut b = CallBatcher::new();
113        b.enqueue("t", json!(null));
114        b.flush();
115        assert!(b.is_empty());
116    }
117
118    #[test]
119    fn flush_by_priority() {
120        let mut b = CallBatcher::new();
121        b.enqueue_with_priority("low", json!({}), 1);
122        b.enqueue_with_priority("high", json!({}), 10);
123        let calls = b.flush_by_priority();
124        assert_eq!(calls[0].tool, "high");
125    }
126
127    #[test]
128    fn cancel_removes_call() {
129        let mut b = CallBatcher::new();
130        let id = b.enqueue("x", json!({}));
131        assert!(b.cancel(id));
132        assert!(b.is_empty());
133    }
134
135    #[test]
136    fn cancel_nonexistent_returns_false() {
137        let mut b = CallBatcher::new();
138        assert!(!b.cancel(999));
139    }
140
141    #[test]
142    fn calls_for_filter() {
143        let mut b = CallBatcher::new();
144        b.enqueue("search", json!({}));
145        b.enqueue("fetch", json!({}));
146        b.enqueue("search", json!({}));
147        assert_eq!(b.calls_for("search").len(), 2);
148    }
149
150    #[test]
151    fn ids_are_unique() {
152        let mut b = CallBatcher::new();
153        let a = b.enqueue("t", json!({}));
154        let c = b.enqueue("t", json!({}));
155        assert_ne!(a, c);
156    }
157
158    #[test]
159    fn peek_does_not_remove() {
160        let mut b = CallBatcher::new();
161        b.enqueue("t", json!({}));
162        let _ = b.peek();
163        assert_eq!(b.len(), 1);
164    }
165
166    #[test]
167    fn to_tool_use_blocks() {
168        let mut b = CallBatcher::new();
169        b.enqueue("search", json!({"q": "rust"}));
170        let blocks = b.to_tool_use_blocks();
171        assert_eq!(blocks.len(), 1);
172        assert_eq!(blocks[0]["type"], "tool_use");
173        assert_eq!(blocks[0]["name"], "search");
174    }
175
176    #[test]
177    fn clear_empties_queue() {
178        let mut b = CallBatcher::new();
179        b.enqueue("t", json!({}));
180        b.clear();
181        assert!(b.is_empty());
182    }
183
184    #[test]
185    fn flush_preserves_order() {
186        let mut b = CallBatcher::new();
187        b.enqueue("a", json!({}));
188        b.enqueue("b", json!({}));
189        b.enqueue("c", json!({}));
190        let calls = b.flush();
191        let tools: Vec<&str> = calls.iter().map(|c| c.tool.as_str()).collect();
192        assert_eq!(tools, vec!["a", "b", "c"]);
193    }
194}
tool_call_batcher/lib.rs

tool_call_batcher/
lib.rs
