behest_runtime/
tool_runtime.rs

1//! Tool runtime layer with validation, timeout, and bounded parallelism.
2//!
3//! [`ToolRuntime`] wraps a [`ToolRegistry`] and enforces [`RuntimePolicy`]
4//! constraints during tool execution: JSON schema validation, per-tool
5//! timeouts, bounded concurrency, and execution recording to an
6//! [`ExecutionStore`].
7
8use std::sync::Arc;
9use std::time::Instant;
10
11use behest_tool::Tool;
12use futures_util::future::join_all;
13use serde_json::Value;
14use tokio::sync::Semaphore;
15use tokio::time::timeout;
16use tracing::{debug, warn};
17use uuid::Uuid;
18
19use super::error::{RuntimeError, RuntimeResult};
20use super::policy::RuntimePolicy;
21use super::tool_output::{self, ToolOutputConfig};
22use super::tool_scope::ScopedToolRegistry;
23use behest_core::error::ToolError;
24use behest_provider::{Message, ToolCall};
25use behest_store::{ExecutionStore, ToolExecution};
26use behest_tool::{ToolOutput, ToolRegistry};
27
28/// Outcome of running a single tool call within the runtime.
29///
30/// Carries the original [`ToolCall`], the execution result (success or error),
31/// and the [`Message`] produced for the conversation history.
32#[derive(Debug)]
33pub struct ToolExecutionOutcome {
34    /// The original tool call.
35    pub call: ToolCall,
36    /// Output if successful, or the tool error if failed.
37    pub output: Result<ToolOutput, ToolError>,
38    /// The resulting message for the conversation.
39    pub message: Message,
40}
41
42/// Runtime layer that orchestrates tool execution with policy enforcement.
43pub struct ToolRuntime {
44    registry: Arc<ScopedToolRegistry>,
45    policy: RuntimePolicy,
46    semaphore: Arc<Semaphore>,
47}
48
49impl ToolRuntime {
50    /// Creates a new tool runtime wrapping the given registry.
51    #[must_use]
52    pub fn new(registry: ToolRegistry, policy: RuntimePolicy) -> Self {
53        let mut policy = policy;
54        policy.max_tool_concurrency = policy.max_tool_concurrency.max(1);
55        let semaphore = Arc::new(Semaphore::new(policy.max_tool_concurrency));
56        Self {
57            registry: Arc::new(ScopedToolRegistry::new(registry)),
58            policy,
59            semaphore,
60        }
61    }
62
63    /// Returns a reference to the underlying scoped tool registry.
64    #[must_use]
65    pub fn registry(&self) -> &Arc<ScopedToolRegistry> {
66        &self.registry
67    }
68
69    /// Returns a mutable reference to the underlying scoped tool registry.
70    pub fn registry_mut(&mut self) -> &mut Arc<ScopedToolRegistry> {
71        &mut self.registry
72    }
73
74    /// Returns the runtime policy.
75    #[must_use]
76    pub fn policy(&self) -> &RuntimePolicy {
77        &self.policy
78    }
79
80    /// Registers a tool at the base level of the scoped registry.
81    ///
82    /// Returns the previously registered tool with the same name, if any.
83    pub fn register_tool(&self, tool: Arc<dyn Tool>) -> Option<Arc<dyn Tool>> {
84        self.registry.base().register_arc(tool)
85    }
86
87    /// Unregisters a tool from the base level of the scoped registry.
88    ///
89    /// Returns the removed tool if it existed, or [`None`] if no tool with
90    /// `name` was registered.
91    #[must_use]
92    pub fn unregister_tool(&self, name: &str) -> Option<Arc<dyn Tool>> {
93        self.registry.unregister_from_base(name)
94    }
95
96    /// Executes a batch of tool calls with bounded parallelism, timeout,
97    /// validation, and optional execution recording.
98    ///
99    /// Tool calls are partitioned by [`Tool::is_concurrency_safe`]:
100    /// concurrent-safe tools execute in parallel (bounded by semaphore),
101    /// while exclusive tools execute sequentially, one at a time.
102    /// Results are merged in the original call order.
103    ///
104    /// Each tool call is executed independently. If `execution_store` is
105    /// provided, each execution is recorded.
106    ///
107    /// When `continue_on_tool_failure` is true (from policy), failed tool
108    /// calls produce error messages but do not abort the batch.
109    ///
110    /// # Errors
111    ///
112    /// Returns [`RuntimeError::ToolTimeout`] when the semaphore is exhausted,
113    /// or propagates errors from individual tool executions when the policy
114    /// does not allow continuation on failure.
115    pub async fn execute_batch(
116        &self,
117        calls: Vec<ToolCall>,
118        session_id: Uuid,
119        message_id: Uuid,
120        execution_store: Option<&dyn ExecutionStore>,
121    ) -> RuntimeResult<Vec<ToolExecutionOutcome>> {
122        if calls.is_empty() {
123            return Ok(Vec::new());
124        }
125
126        let call_count = calls.len();
127        let indexed_calls: Vec<(usize, ToolCall)> = calls.into_iter().enumerate().collect();
128
129        let (concurrent_group, exclusive_group): (Vec<_>, Vec<_>) =
130            indexed_calls.into_iter().partition(|(_, call)| {
131                self.registry
132                    .get(&call.name)
133                    .is_some_and(|tool| tool.is_concurrency_safe())
134            });
135
136        let mut results: Vec<Option<ToolExecutionOutcome>> =
137            (0..call_count).map(|_| None).collect();
138
139        if !concurrent_group.is_empty() {
140            let concurrent_results = {
141                let futures: Vec<_> = concurrent_group
142                    .into_iter()
143                    .map(|(idx, call)| {
144                        let sem = Arc::clone(&self.semaphore);
145                        let registry = self.registry.clone();
146                        let tool_timeout = self.policy.tool_timeout;
147                        let continue_on_failure = self.policy.continue_on_tool_failure;
148                        let truncation_config = self.policy.tool_output.clone();
149
150                        async move {
151                            let _permit =
152                                sem.acquire().await.map_err(|_| RuntimeError::ToolTimeout {
153                                    tool: call.name.clone(),
154                                })?;
155
156                            let outcome = Self::execute_single(
157                                &registry,
158                                &call,
159                                tool_timeout,
160                                continue_on_failure,
161                                &truncation_config,
162                            )
163                            .await?;
164                            Ok::<(usize, ToolExecutionOutcome), RuntimeError>((idx, outcome))
165                        }
166                    })
167                    .collect();
168
169                join_all(futures).await
170            };
171
172            for res in concurrent_results {
173                let (idx, outcome) = res?;
174                if let Some(store) = execution_store {
175                    Self::record_execution(store, session_id, message_id, &outcome.call, &outcome)
176                        .await;
177                }
178                results[idx] = Some(outcome);
179            }
180        }
181
182        for (idx, call) in exclusive_group {
183            let _permit =
184                self.semaphore
185                    .acquire()
186                    .await
187                    .map_err(|_| RuntimeError::ToolTimeout {
188                        tool: call.name.clone(),
189                    })?;
190
191            let outcome = Self::execute_single(
192                &self.registry,
193                &call,
194                self.policy.tool_timeout,
195                self.policy.continue_on_tool_failure,
196                &self.policy.tool_output,
197            )
198            .await?;
199
200            if let Some(store) = execution_store {
201                Self::record_execution(store, session_id, message_id, &call, &outcome).await;
202            }
203
204            results[idx] = Some(outcome);
205        }
206
207        Ok(results
208            .into_iter()
209            .map(|r| r.unwrap_or_else(|| unreachable!("all indices populated")))
210            .collect())
211    }
212
213    async fn execute_single(
214        registry: &ScopedToolRegistry,
215        call: &ToolCall,
216        tool_timeout: std::time::Duration,
217        continue_on_failure: bool,
218        truncation_config: &ToolOutputConfig,
219    ) -> RuntimeResult<ToolExecutionOutcome> {
220        let Some(tool) = registry.get(&call.name) else {
221            let error_msg = format!("tool not found: {}", call.name);
222            warn!(tool = %call.name, "tool not found");
223            let error = ToolError::NotFound {
224                name: call.name.clone(),
225            };
226            if !continue_on_failure {
227                return Err(error.into());
228            }
229
230            return Ok(ToolExecutionOutcome {
231                call: call.clone(),
232                output: Err(error),
233                message: Message::tool_text(
234                    call.id.clone(),
235                    call.name.clone(),
236                    format!("{{\"error\":\"{error_msg}\"}}"),
237                ),
238            });
239        };
240
241        if let Err(validation_error) =
242            Self::validate_arguments(&tool.parameters_schema(), &call.arguments)
243        {
244            debug!(tool = %call.name, error = %validation_error, "schema validation failed");
245            let err = ToolError::InvalidArguments {
246                name: call.name.clone(),
247                message: validation_error.clone(),
248            };
249            if !continue_on_failure {
250                return Err(err.into());
251            }
252            return Ok(ToolExecutionOutcome {
253                call: call.clone(),
254                output: Err(err),
255                message: Message::tool_text(
256                    call.id.clone(),
257                    call.name.clone(),
258                    format!("{{\"error\":\"{validation_error}\"}}"),
259                ),
260            });
261        }
262
263        let start = Instant::now();
264        match timeout(tool_timeout, tool.execute(call.arguments.clone())).await {
265            Ok(Ok(output)) => {
266                let duration = start.elapsed();
267                debug!(tool = %call.name, ?duration, "tool executed successfully");
268                let raw_text = output.value.to_string();
269                let truncated =
270                    tool_output::truncate_output(&raw_text, truncation_config, Some(&call.name));
271                let msg = Message::tool_text(call.id.clone(), call.name.clone(), truncated.text);
272                Ok(ToolExecutionOutcome {
273                    call: call.clone(),
274                    output: Ok(output),
275                    message: msg,
276                })
277            }
278            Ok(Err(tool_error)) => {
279                let duration = start.elapsed();
280                let error_msg = tool_error.to_string();
281                warn!(tool = %call.name, ?duration, error = %error_msg, "tool execution failed");
282                if !continue_on_failure {
283                    return Err(tool_error.into());
284                }
285
286                Ok(ToolExecutionOutcome {
287                    call: call.clone(),
288                    output: Err(tool_error),
289                    message: Message::tool_text(
290                        call.id.clone(),
291                        call.name.clone(),
292                        format!("{{\"error\":\"{error_msg}\"}}"),
293                    ),
294                })
295            }
296            Err(_) => {
297                let error_msg = format!("tool execution timeout after {tool_timeout:?}");
298                warn!(tool = %call.name, "tool execution timed out");
299                if !continue_on_failure {
300                    return Err(RuntimeError::ToolTimeout {
301                        tool: call.name.clone(),
302                    });
303                }
304
305                Ok(ToolExecutionOutcome {
306                    call: call.clone(),
307                    output: Err(ToolError::Execution {
308                        name: call.name.clone(),
309                        message: error_msg.clone(),
310                    }),
311                    message: Message::tool_text(
312                        call.id.clone(),
313                        call.name.clone(),
314                        format!("{{\"error\":\"{error_msg}\"}}"),
315                    ),
316                })
317            }
318        }
319    }
320
321    fn validate_arguments(schema: &Value, arguments: &Value) -> Result<(), String> {
322        if schema.is_null() || schema.as_object().is_none() {
323            return Ok(());
324        }
325
326        let evaluation = jsonschema::evaluate(schema, arguments);
327        let flag = evaluation.flag();
328
329        if flag.valid {
330            Ok(())
331        } else {
332            let errors: Vec<String> = evaluation
333                .iter_errors()
334                .map(|e| e.error.to_string())
335                .collect();
336            Err(format!("validation errors: {}", errors.join("; ")))
337        }
338    }
339
340    async fn record_execution(
341        store: &dyn ExecutionStore,
342        session_id: Uuid,
343        message_id: Uuid,
344        call: &ToolCall,
345        outcome: &ToolExecutionOutcome,
346    ) {
347        let mut execution = ToolExecution::new(
348            session_id,
349            message_id,
350            &call.id,
351            &call.name,
352            call.arguments.clone(),
353        );
354
355        match &outcome.output {
356            Ok(output) => {
357                execution = execution.with_success(output.value.clone(), std::time::Duration::ZERO);
358            }
359            Err(error) => {
360                execution = execution.with_failure(error.to_string(), std::time::Duration::ZERO);
361            }
362        }
363
364        if let Err(e) = store.record_execution(execution).await {
365            warn!(error = %e, "failed to record tool execution");
366        }
367    }
368}
369
370#[cfg(test)]
371#[allow(clippy::unwrap_used, clippy::type_complexity, clippy::expect_used)]
372mod tests {
373    use super::*;
374    use behest_store::memory::MemoryExecutionStore;
375    use behest_tool::FunctionTool;
376    use serde_json::json;
377
378    fn echo_tool() -> FunctionTool {
379        FunctionTool::new(
380            "echo",
381            "Echoes input",
382            json!({
383                "type": "object",
384                "properties": {
385                    "message": { "type": "string" }
386                },
387                "required": ["message"]
388            }),
389            |args: Value| -> std::pin::Pin<
390                Box<dyn std::future::Future<Output = behest_tool::ToolResult<Value>> + Send>,
391            > {
392                Box::pin(async move { Ok(args.get("message").cloned().unwrap_or(Value::Null)) })
393            },
394        )
395    }
396
397    fn failing_tool() -> FunctionTool {
398        FunctionTool::new(
399            "fail",
400            "Always fails",
401            json!({"type": "object"}),
402            |_args: Value| -> std::pin::Pin<
403                Box<dyn std::future::Future<Output = behest_tool::ToolResult<Value>> + Send>,
404            > {
405                Box::pin(async move {
406                    Err(ToolError::Execution {
407                        name: "fail".to_owned(),
408                        message: "intentional failure".to_owned(),
409                    })
410                })
411            },
412        )
413    }
414
415    #[tokio::test]
416    async fn execute_batch_should_run_tools_in_parallel() {
417        let registry = ToolRegistry::new();
418        registry.register(echo_tool());
419        let policy = RuntimePolicy::new().with_max_tool_concurrency(2);
420        let runtime = ToolRuntime::new(registry, policy);
421
422        let calls = vec![
423            ToolCall::new("call_1", "echo", json!({"message": "hello"})),
424            ToolCall::new("call_2", "echo", json!({"message": "world"})),
425        ];
426
427        let session_id = Uuid::now_v7();
428        let message_id = Uuid::now_v7();
429        let outcomes = runtime
430            .execute_batch(calls, session_id, message_id, None)
431            .await
432            .unwrap();
433
434        assert_eq!(outcomes.len(), 2);
435        assert!(outcomes[0].output.is_ok());
436        assert!(outcomes[1].output.is_ok());
437    }
438
439    #[tokio::test]
440    async fn execute_batch_should_handle_unknown_tool() {
441        let registry = ToolRegistry::new();
442        let policy = RuntimePolicy::new();
443        let runtime = ToolRuntime::new(registry, policy);
444
445        let calls = vec![ToolCall::new("call_1", "unknown", json!({}))];
446        let session_id = Uuid::now_v7();
447        let message_id = Uuid::now_v7();
448        let outcomes = runtime
449            .execute_batch(calls, session_id, message_id, None)
450            .await
451            .unwrap();
452
453        assert_eq!(outcomes.len(), 1);
454        assert!(outcomes[0].output.is_err());
455    }
456
457    #[tokio::test]
458    async fn execute_batch_should_validate_schema() {
459        let registry = ToolRegistry::new();
460        registry.register(echo_tool());
461        let policy = RuntimePolicy::new();
462        let runtime = ToolRuntime::new(registry, policy);
463
464        let calls = vec![ToolCall::new("call_1", "echo", json!({"message": 123}))];
465        let session_id = Uuid::now_v7();
466        let message_id = Uuid::now_v7();
467        let outcomes = runtime
468            .execute_batch(calls, session_id, message_id, None)
469            .await
470            .unwrap();
471
472        assert_eq!(outcomes.len(), 1);
473        assert!(outcomes[0].output.is_err());
474        let err = outcomes[0].output.as_ref().unwrap_err();
475        assert!(err.to_string().contains("validation"));
476    }
477
478    #[tokio::test]
479    async fn execute_batch_should_record_to_execution_store() {
480        let registry = ToolRegistry::new();
481        registry.register(echo_tool());
482        let policy = RuntimePolicy::new();
483        let runtime = ToolRuntime::new(registry, policy);
484
485        let store = MemoryExecutionStore::new();
486        let calls = vec![ToolCall::new("call_1", "echo", json!({"message": "test"}))];
487        let session_id = Uuid::now_v7();
488        let message_id = Uuid::now_v7();
489
490        runtime
491            .execute_batch(calls, session_id, message_id, Some(&store))
492            .await
493            .unwrap();
494
495        let executions = store.list_executions(&session_id).await.unwrap();
496        assert_eq!(executions.len(), 1);
497        assert_eq!(executions[0].tool_name, "echo");
498    }
499
500    #[tokio::test]
501    async fn execute_batch_should_handle_tool_failure() {
502        let registry = ToolRegistry::new();
503        registry.register(failing_tool());
504        let policy = RuntimePolicy::new().with_continue_on_tool_failure(true);
505        let runtime = ToolRuntime::new(registry, policy);
506
507        let calls = vec![ToolCall::new("call_1", "fail", json!({}))];
508        let session_id = Uuid::now_v7();
509        let message_id = Uuid::now_v7();
510        let outcomes = runtime
511            .execute_batch(calls, session_id, message_id, None)
512            .await
513            .unwrap();
514
515        assert_eq!(outcomes.len(), 1);
516        assert!(outcomes[0].output.is_err());
517    }
518
519    #[tokio::test]
520    async fn execute_batch_empty_returns_empty() {
521        let registry = ToolRegistry::new();
522        let policy = RuntimePolicy::new();
523        let runtime = ToolRuntime::new(registry, policy);
524
525        let session_id = Uuid::now_v7();
526        let message_id = Uuid::now_v7();
527        let outcomes = runtime
528            .execute_batch(Vec::new(), session_id, message_id, None)
529            .await
530            .unwrap();
531
532        assert!(outcomes.is_empty());
533    }
534
535    #[tokio::test]
536    async fn execute_batch_partitions_concurrent_and_exclusive() {
537        use std::sync::Arc;
538        use std::sync::atomic::{AtomicBool, Ordering};
539
540        let concurrent_flag = Arc::new(AtomicBool::new(false));
541        let exclusive_flag = Arc::new(AtomicBool::new(false));
542
543        let con_flag = Arc::clone(&concurrent_flag);
544        let concurrent_tool = FunctionTool::new(
545            "concurrent",
546            "Safe for concurrency",
547            json!({"type": "object"}),
548            move |_args: Value| -> std::pin::Pin<
549                Box<dyn std::future::Future<Output = behest_tool::ToolResult<Value>> + Send>,
550            > {
551                let flag = Arc::clone(&con_flag);
552                Box::pin(async move {
553                    flag.store(true, Ordering::SeqCst);
554                    Ok(Value::Null)
555                })
556            },
557        )
558        .concurrency_safe();
559
560        let exc_flag = Arc::clone(&exclusive_flag);
561        let exclusive_tool = FunctionTool::new(
562            "exclusive",
563            "Not safe for concurrency",
564            json!({"type": "object"}),
565            move |_args: Value| -> std::pin::Pin<
566                Box<dyn std::future::Future<Output = behest_tool::ToolResult<Value>> + Send>,
567            > {
568                let flag = Arc::clone(&exc_flag);
569                Box::pin(async move {
570                    flag.store(true, Ordering::SeqCst);
571                    Ok(Value::Null)
572                })
573            },
574        );
575        // exclusive_tool is NOT .concurrency_safe() — defaults to false
576
577        let registry = ToolRegistry::new();
578        registry.register(concurrent_tool);
579        registry.register(exclusive_tool);
580
581        let policy = RuntimePolicy::new().with_max_tool_concurrency(4);
582        let runtime = ToolRuntime::new(registry, policy);
583
584        let calls = vec![
585            ToolCall::new("call_1", "concurrent", json!({})),
586            ToolCall::new("call_2", "exclusive", json!({})),
587        ];
588
589        let session_id = Uuid::now_v7();
590        let message_id = Uuid::now_v7();
591        let outcomes = runtime
592            .execute_batch(calls, session_id, message_id, None)
593            .await
594            .unwrap();
595
596        assert_eq!(outcomes.len(), 2);
597        assert!(outcomes[0].output.is_ok());
598        assert!(outcomes[1].output.is_ok());
599        assert!(concurrent_flag.load(Ordering::SeqCst));
600        assert!(exclusive_flag.load(Ordering::SeqCst));
601    }
602
603    #[tokio::test]
604    async fn execute_batch_preserves_result_order() {
605        let tool_a = FunctionTool::new(
606            "a",
607            "Tool A",
608            json!({"type": "object"}),
609            |_args: Value| -> std::pin::Pin<
610                Box<dyn std::future::Future<Output = behest_tool::ToolResult<Value>> + Send>,
611            > { Box::pin(async move { Ok(Value::String("A".into())) }) },
612        )
613        .concurrency_safe();
614
615        let tool_b = FunctionTool::new(
616            "b",
617            "Tool B",
618            json!({"type": "object"}),
619            |_args: Value| -> std::pin::Pin<
620                Box<dyn std::future::Future<Output = behest_tool::ToolResult<Value>> + Send>,
621            > { Box::pin(async move { Ok(Value::String("B".into())) }) },
622        );
623        // tool_b is NOT concurrency_safe
624
625        let registry = ToolRegistry::new();
626        registry.register(tool_a);
627        registry.register(tool_b);
628
629        let policy = RuntimePolicy::new().with_max_tool_concurrency(4);
630        let runtime = ToolRuntime::new(registry, policy);
631
632        let calls = vec![
633            ToolCall::new("call_1", "a", json!({})),
634            ToolCall::new("call_2", "b", json!({})),
635        ];
636
637        let session_id = Uuid::now_v7();
638        let message_id = Uuid::now_v7();
639        let outcomes = runtime
640            .execute_batch(calls, session_id, message_id, None)
641            .await
642            .unwrap();
643
644        assert_eq!(outcomes.len(), 2);
645        // Results must be in original call order: a first, then b
646        assert_eq!(outcomes[0].call.name, "a");
647        assert_eq!(outcomes[1].call.name, "b");
648    }
649
650    #[tokio::test]
651    async fn execute_batch_should_propagate_tool_failure_when_continue_disabled() {
652        let registry = ToolRegistry::new();
653        registry.register(failing_tool());
654        let policy = RuntimePolicy::new().with_continue_on_tool_failure(false);
655        let runtime = ToolRuntime::new(registry, policy);
656
657        let calls = vec![ToolCall::new("call_1", "fail", json!({}))];
658        let session_id = Uuid::now_v7();
659        let message_id = Uuid::now_v7();
660        let result = runtime
661            .execute_batch(calls, session_id, message_id, None)
662            .await;
663
664        assert!(matches!(result, Err(RuntimeError::Tool(_))));
665    }
666
667    #[tokio::test]
668    async fn execute_batch_should_propagate_invalid_arguments_when_continue_disabled() {
669        let registry = ToolRegistry::new();
670        registry.register(echo_tool());
671        let policy = RuntimePolicy::new().with_continue_on_tool_failure(false);
672        let runtime = ToolRuntime::new(registry, policy);
673
674        let calls = vec![ToolCall::new("call_1", "echo", json!({"message": 123}))];
675        let session_id = Uuid::now_v7();
676        let message_id = Uuid::now_v7();
677        let result = runtime
678            .execute_batch(calls, session_id, message_id, None)
679            .await;
680
681        assert!(matches!(result, Err(RuntimeError::Tool(_))));
682    }
683
684    #[tokio::test]
685    async fn execute_batch_with_zero_concurrency_should_not_hang() {
686        let registry = ToolRegistry::new();
687        registry.register(echo_tool());
688        let policy = RuntimePolicy::new().with_max_tool_concurrency(0);
689        let runtime = ToolRuntime::new(registry, policy);
690
691        let calls = vec![ToolCall::new("call_1", "echo", json!({"message": "hello"}))];
692        let session_id = Uuid::now_v7();
693        let message_id = Uuid::now_v7();
694        let result = tokio::time::timeout(
695            std::time::Duration::from_millis(100),
696            runtime.execute_batch(calls, session_id, message_id, None),
697        )
698        .await;
699
700        let outcomes = result.expect("tool runtime should not hang").unwrap();
701        assert_eq!(outcomes.len(), 1);
702        assert!(outcomes[0].output.is_ok());
703    }
704}
behest_runtime/tool_runtime.rs

behest_runtime/
tool_runtime.rs
