Skip to main content

vv_agent/tools/
executor.rs

1use std::any::Any;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5use std::time::Duration;
6
7use crate::context::RunContext;
8use crate::tools::{ToolContext, ToolMetadata, ToolSpec};
9use crate::types::{Metadata, ToolArguments, ToolCall, ToolExecutionResult};
10
11pub type ToolFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, ToolError>> + Send + 'a>>;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum ToolExposure {
15    /// Model-visible and executable through the runtime.
16    Direct,
17    /// Not model-visible, but still available for explicit runtime invocation.
18    Hidden,
19}
20
21#[derive(Clone, Default)]
22pub struct ToolEnablementContext {
23    pub run: RunContext,
24    pub app_state: Option<Arc<dyn Any + Send + Sync>>,
25}
26
27impl ToolEnablementContext {
28    pub fn new(run: RunContext) -> Self {
29        Self {
30            run,
31            app_state: None,
32        }
33    }
34
35    pub fn with_app_state(mut self, app_state: Option<Arc<dyn Any + Send + Sync>>) -> Self {
36        self.app_state = app_state;
37        self
38    }
39
40    pub fn app_state<T: Send + Sync + 'static>(&self) -> Option<&T> {
41        self.app_state.as_ref()?.downcast_ref::<T>()
42    }
43}
44
45pub type ToolEnablementPredicate =
46    Arc<dyn Fn(&ToolEnablementContext) -> bool + Send + Sync + 'static>;
47
48#[derive(Clone)]
49pub enum ToolEnablementRule {
50    Static(bool),
51    Predicate(ToolEnablementPredicate),
52}
53
54impl ToolEnablementRule {
55    pub fn predicate<F>(predicate: F) -> Self
56    where
57        F: Fn(&ToolEnablementContext) -> bool + Send + Sync + 'static,
58    {
59        Self::Predicate(Arc::new(predicate))
60    }
61
62    pub fn is_enabled(&self, context: &ToolEnablementContext) -> bool {
63        match self {
64            Self::Static(enabled) => *enabled,
65            Self::Predicate(predicate) => predicate(context),
66        }
67    }
68}
69
70impl Default for ToolEnablementRule {
71    fn default() -> Self {
72        Self::Static(true)
73    }
74}
75
76impl From<bool> for ToolEnablementRule {
77    fn from(enabled: bool) -> Self {
78        Self::Static(enabled)
79    }
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub enum ApprovalRequirement {
84    NotRequired,
85    Required,
86    Provider,
87}
88
89pub type ApprovalPredicate =
90    Arc<dyn Fn(&ToolContext, &ToolArguments) -> bool + Send + Sync + 'static>;
91
92#[derive(Clone)]
93pub enum ToolApprovalRule {
94    Static(ApprovalRequirement),
95    Predicate(ApprovalPredicate),
96}
97
98impl ToolApprovalRule {
99    pub fn predicate<F>(predicate: F) -> Self
100    where
101        F: Fn(&ToolContext, &ToolArguments) -> bool + Send + Sync + 'static,
102    {
103        Self::Predicate(Arc::new(predicate))
104    }
105
106    pub fn requirement(
107        &self,
108        context: &ToolContext,
109        arguments: &ToolArguments,
110    ) -> ApprovalRequirement {
111        match self {
112            Self::Static(requirement) => *requirement,
113            Self::Predicate(predicate) => {
114                if predicate(context, arguments) {
115                    ApprovalRequirement::Required
116                } else {
117                    ApprovalRequirement::NotRequired
118                }
119            }
120        }
121    }
122}
123
124impl Default for ToolApprovalRule {
125    fn default() -> Self {
126        Self::Static(ApprovalRequirement::NotRequired)
127    }
128}
129
130impl From<bool> for ToolApprovalRule {
131    fn from(needs_approval: bool) -> Self {
132        Self::Static(if needs_approval {
133            ApprovalRequirement::Required
134        } else {
135            ApprovalRequirement::NotRequired
136        })
137    }
138}
139
140impl From<ApprovalRequirement> for ToolApprovalRule {
141    fn from(requirement: ApprovalRequirement) -> Self {
142        Self::Static(requirement)
143    }
144}
145
146#[derive(Debug, Clone, Default)]
147pub struct ToolSpecContext;
148
149pub struct ToolRunContext<'a> {
150    pub context: &'a mut ToolContext,
151}
152
153impl<'a> ToolRunContext<'a> {
154    pub fn new(context: &'a mut ToolContext) -> Self {
155        Self { context }
156    }
157}
158
159#[derive(Debug, Clone, PartialEq, Eq)]
160pub struct ToolError {
161    message: String,
162}
163
164impl ToolError {
165    pub fn new(message: impl Into<String>) -> Self {
166        Self {
167            message: message.into(),
168        }
169    }
170}
171
172impl std::fmt::Display for ToolError {
173    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        formatter.write_str(&self.message)
175    }
176}
177
178impl std::error::Error for ToolError {}
179
180pub trait ToolExecutor: Send + Sync {
181    fn name(&self) -> &str;
182    fn description(&self) -> &str;
183    fn metadata(&self) -> &Metadata {
184        static EMPTY_METADATA: std::sync::OnceLock<Metadata> = std::sync::OnceLock::new();
185        EMPTY_METADATA.get_or_init(Metadata::new)
186    }
187    fn exposure(&self) -> ToolExposure {
188        ToolExposure::Direct
189    }
190    fn timeout(&self) -> Option<Duration> {
191        None
192    }
193    fn tool_metadata(&self) -> Option<&ToolMetadata> {
194        None
195    }
196    fn spec(&self, _ctx: &ToolSpecContext) -> Result<ToolSpec, ToolError>;
197    fn approval_requirement(
198        &self,
199        _call: &ToolCall,
200        _ctx: &ToolRunContext<'_>,
201    ) -> ApprovalRequirement {
202        ApprovalRequirement::NotRequired
203    }
204    fn validate_arguments(
205        &self,
206        call: &ToolCall,
207    ) -> Result<Option<ToolExecutionResult>, ToolError> {
208        let spec = self.spec(&ToolSpecContext)?;
209        let schema = crate::tools::argument_validation::close_object_schemas(&spec.schema);
210        let validator = crate::tools::argument_validation::validator_for_tool_schema(&schema)
211            .map_err(ToolError::new)?;
212        Ok(crate::tools::argument_validation::invalid_tool_arguments_result(&validator, call))
213    }
214    fn run<'a>(
215        &'a self,
216        call: ToolCall,
217        ctx: ToolRunContext<'a>,
218    ) -> ToolFuture<'a, ToolExecutionResult>;
219}
220
221#[derive(Clone)]
222pub struct ToolSpecExecutor {
223    spec: ToolSpec,
224    exposure: ToolExposure,
225    argument_validator: Result<jsonschema::Validator, String>,
226}
227
228impl ToolSpecExecutor {
229    pub fn new(mut spec: ToolSpec) -> Self {
230        spec.schema = crate::tools::argument_validation::close_object_schemas(&spec.schema);
231        let exposure = spec.exposure;
232        let argument_validator =
233            crate::tools::argument_validation::validator_for_tool_schema(&spec.schema);
234        Self {
235            spec,
236            exposure,
237            argument_validator,
238        }
239    }
240
241    pub fn with_exposure(mut self, exposure: ToolExposure) -> Self {
242        self.exposure = exposure;
243        self
244    }
245
246    pub fn into_arc(self) -> Arc<dyn ToolExecutor> {
247        Arc::new(self)
248    }
249}
250
251impl ToolExecutor for ToolSpecExecutor {
252    fn name(&self) -> &str {
253        &self.spec.name
254    }
255
256    fn description(&self) -> &str {
257        &self.spec.description
258    }
259
260    fn metadata(&self) -> &Metadata {
261        &self.spec.metadata
262    }
263
264    fn exposure(&self) -> ToolExposure {
265        self.exposure
266    }
267
268    fn timeout(&self) -> Option<Duration> {
269        self.spec.timeout
270    }
271
272    fn tool_metadata(&self) -> Option<&ToolMetadata> {
273        self.spec.tool_metadata.as_ref()
274    }
275
276    fn spec(&self, _ctx: &ToolSpecContext) -> Result<ToolSpec, ToolError> {
277        Ok(self.spec.clone())
278    }
279
280    fn approval_requirement(
281        &self,
282        call: &ToolCall,
283        ctx: &ToolRunContext<'_>,
284    ) -> ApprovalRequirement {
285        self.spec.approval.requirement(ctx.context, &call.arguments)
286    }
287
288    fn validate_arguments(
289        &self,
290        call: &ToolCall,
291    ) -> Result<Option<ToolExecutionResult>, ToolError> {
292        let validator = self
293            .argument_validator
294            .as_ref()
295            .map_err(|error| ToolError::new(error.clone()))?;
296        Ok(crate::tools::argument_validation::invalid_tool_arguments_result(validator, call))
297    }
298
299    fn run<'a>(
300        &'a self,
301        call: ToolCall,
302        ctx: ToolRunContext<'a>,
303    ) -> ToolFuture<'a, ToolExecutionResult> {
304        let handler = self.spec.handler.clone();
305        Box::pin(async move {
306            if let Some(result) = self.validate_arguments(&call)? {
307                return Ok(result);
308            }
309            ctx.context.begin_tool_call(&call);
310            if self.spec.timeout.is_none() {
311                let mut result = handler(ctx.context, &call.arguments);
312                if result.tool_call_id.trim().is_empty() || result.tool_call_id == "pending" {
313                    result.tool_call_id = call.id;
314                }
315                return Ok(result);
316            }
317
318            let arguments = call.arguments.clone();
319            let mut isolated_context = ctx.context.clone();
320            let (mut result, updated_context) = tokio::task::spawn_blocking(move || {
321                let result = handler(&mut isolated_context, &arguments);
322                (result, isolated_context)
323            })
324            .await
325            .map_err(|error| ToolError::new(format!("tool task failed: {error}")))?;
326            *ctx.context = updated_context;
327            if result.tool_call_id.trim().is_empty() || result.tool_call_id == "pending" {
328                result.tool_call_id = call.id;
329            }
330            Ok(result)
331        })
332    }
333}