llm_tool/registry/batch.rs
1//! Concurrent and sequential batch tool execution for [`super::ToolRegistry`].
2
3use alloc::{boxed::Box, string::String, vec::Vec};
4use core::future;
5
6use serde::{Deserialize, Serialize};
7
8use super::ToolRegistry;
9use crate::{
10 rust_tool::BoxToolFuture,
11 types::{ToolContext, ToolError, ToolOutput},
12};
13
14/// Execution strategy for a batch of tool invocations.
15#[derive(
16 Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, schemars::JsonSchema,
17)]
18#[serde(rename_all = "snake_case")]
19pub enum BatchExecutionMode {
20 /// Execute all tool calls concurrently and collect outcomes in original call order.
21 ///
22 /// If any tool in the batch has [`crate::types::ToolEffect::Destructive`], its execution
23 /// will fail with a [`ToolError`] unless [`AllowDestructiveBatch`] is present in context
24 /// or [`BatchExecutionMode::ParallelAllowDestructive`] is used.
25 #[default]
26 Parallel,
27 /// Execute all tool calls concurrently, explicitly confirming and allowing destructive tools.
28 ParallelAllowDestructive,
29 /// Execute tool calls sequentially in call order, halting immediately on the first error.
30 SequentialStopOnError,
31}
32
33/// Marker context extension to explicitly allow destructive tools in parallel batch execution.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub struct AllowDestructiveBatch(pub bool);
36
37/// Owned tool invocation specification within a batch request.
38#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
39pub struct BatchToolCall {
40 /// Caller-assigned tool call identifier.
41 #[serde(default)]
42 pub id: String,
43 /// Target registered tool name.
44 pub name: String,
45 /// JSON arguments object passed to the tool.
46 #[serde(default)]
47 pub arguments: serde_json::Value,
48}
49
50/// Borrowed string-based tool invocation specification for zero-DOM-allocation batch dispatch.
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub struct BatchToolCallStr<'a> {
53 /// Caller-assigned tool call identifier.
54 pub id: &'a str,
55 /// Target registered tool name.
56 pub name: &'a str,
57 /// Raw JSON arguments string.
58 pub arguments_json: &'a str,
59}
60
61/// Outcome of an individual tool invocation within a batch.
62#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
63pub struct BatchToolCallOutcome {
64 /// Caller-assigned tool call identifier.
65 pub id: String,
66 /// Target tool name.
67 pub name: String,
68 /// Result produced by the tool handler.
69 pub result: Result<ToolOutput, ToolError>,
70}
71
72impl ToolRegistry {
73 /// Dispatches a batch of [`BatchToolCall`] items according to `mode`, preserving call order.
74 ///
75 /// - [`BatchExecutionMode::Parallel`]: polls all tool futures concurrently using
76 /// [`futures_util::future::join_all`], returning every outcome in exact input slice order.
77 /// Destructive tools require explicit confirmation via [`BatchExecutionMode::ParallelAllowDestructive`]
78 /// or [`AllowDestructiveBatch`] on `ctx`.
79 /// - [`BatchExecutionMode::SequentialStopOnError`]: runs calls one by one and stops after
80 /// the first `Err(ToolError)`.
81 pub async fn dispatch_batch(
82 &self,
83 calls: &[BatchToolCall],
84 mode: BatchExecutionMode,
85 ctx: &ToolContext,
86 ) -> Vec<BatchToolCallOutcome> {
87 match mode {
88 BatchExecutionMode::Parallel | BatchExecutionMode::ParallelAllowDestructive => {
89 let allow_destructive =
90 matches!(mode, BatchExecutionMode::ParallelAllowDestructive)
91 || ctx
92 .get_ext::<AllowDestructiveBatch>()
93 .is_some_and(|ext| ext.0);
94
95 let futures: Vec<BoxToolFuture<'_>> = calls
96 .iter()
97 .map(|call| -> BoxToolFuture<'_> {
98 if !allow_destructive && self.is_destructive(&call.name) {
99 tracing::warn!(
100 tool = %call.name,
101 "Destructive tool rejected in parallel batch execution without explicit confirmation"
102 );
103 let fut: BoxToolFuture<'_> = Box::pin(future::ready(Err(
104 ToolError::destructive_in_parallel_batch(&call.name),
105 )));
106 return fut;
107 }
108 match self.dispatch_boxed(&call.name, call.arguments.clone(), ctx) {
109 Ok(fut) => fut,
110 Err(err) => {
111 let fut: BoxToolFuture<'_> = Box::pin(future::ready(Err(err)));
112 fut
113 }
114 }
115 })
116 .collect();
117
118 let results = futures_util::future::join_all(futures).await;
119 calls
120 .iter()
121 .zip(results)
122 .map(|(call, result)| BatchToolCallOutcome {
123 id: call.id.clone(),
124 name: call.name.clone(),
125 result,
126 })
127 .collect()
128 }
129 BatchExecutionMode::SequentialStopOnError => {
130 let mut outcomes = Vec::with_capacity(calls.len());
131 for call in calls {
132 let result = self.dispatch(&call.name, call.arguments.clone(), ctx).await;
133 let should_stop = match &result {
134 Ok(_) => false,
135 Err(err) => {
136 tracing::debug!(
137 error = %err,
138 tool = %call.name,
139 "Sequential batch execution stopped due to error"
140 );
141 true
142 }
143 };
144 outcomes.push(BatchToolCallOutcome {
145 id: call.id.clone(),
146 name: call.name.clone(),
147 result,
148 });
149 if should_stop {
150 break;
151 }
152 }
153 outcomes
154 }
155 }
156 }
157
158 /// Dispatches a batch of borrowed [`BatchToolCallStr`] items according to `mode`,
159 /// deserializing directly from raw JSON strings without intermediate DOM allocation.
160 pub async fn dispatch_batch_str(
161 &self,
162 calls: &[BatchToolCallStr<'_>],
163 mode: BatchExecutionMode,
164 ctx: &ToolContext,
165 ) -> Vec<BatchToolCallOutcome> {
166 match mode {
167 BatchExecutionMode::Parallel | BatchExecutionMode::ParallelAllowDestructive => {
168 let allow_destructive =
169 matches!(mode, BatchExecutionMode::ParallelAllowDestructive)
170 || ctx
171 .get_ext::<AllowDestructiveBatch>()
172 .is_some_and(|ext| ext.0);
173
174 let futures: Vec<BoxToolFuture<'_>> = calls
175 .iter()
176 .map(|call| -> BoxToolFuture<'_> {
177 if !allow_destructive && self.is_destructive(call.name) {
178 tracing::warn!(
179 tool = %call.name,
180 "Destructive tool rejected in parallel batch execution without explicit confirmation"
181 );
182 let fut: BoxToolFuture<'_> = Box::pin(future::ready(Err(
183 ToolError::destructive_in_parallel_batch(call.name),
184 )));
185 return fut;
186 }
187 match self.dispatch_boxed_str(call.name, call.arguments_json, ctx) {
188 Ok(fut) => fut,
189 Err(err) => Box::pin(future::ready(Err(err))),
190 }
191 })
192 .collect();
193
194 let results = futures_util::future::join_all(futures).await;
195 calls
196 .iter()
197 .zip(results)
198 .map(|(call, result)| BatchToolCallOutcome {
199 id: String::from(call.id),
200 name: String::from(call.name),
201 result,
202 })
203 .collect()
204 }
205 BatchExecutionMode::SequentialStopOnError => {
206 let mut outcomes = Vec::with_capacity(calls.len());
207 for call in calls {
208 let result = self.dispatch_str(call.name, call.arguments_json, ctx).await;
209 let should_stop = match &result {
210 Ok(_) => false,
211 Err(err) => {
212 tracing::debug!(
213 error = %err,
214 tool = %call.name,
215 "Sequential batch execution stopped due to error"
216 );
217 true
218 }
219 };
220 outcomes.push(BatchToolCallOutcome {
221 id: String::from(call.id),
222 name: String::from(call.name),
223 result,
224 });
225 if should_stop {
226 break;
227 }
228 }
229 outcomes
230 }
231 }
232 }
233}