1use std::future::Future;
2use std::marker::PhantomData;
3use std::panic::{catch_unwind, AssertUnwindSafe};
4use std::sync::Arc;
5use std::sync::Mutex;
6use std::time::Duration;
7
8use serde::de::DeserializeOwned;
9use serde_json::Value;
10
11use crate::context::{RunContext, ToolCallContext};
12use crate::tools::{
13 Tool, ToolApprovalRule, ToolContext, ToolEnablementContext, ToolEnablementRule, ToolExposure,
14 ToolMetadata, ToolOutput, ToolSpec, ToolSpecExecutor,
15};
16use crate::types::{Metadata, ToolArguments};
17
18type DynamicFunctionHandler =
19 Arc<dyn Fn(ToolCallContext, Value) -> Result<ToolOutput, String> + Send + Sync>;
20pub type ToolErrorMapper = Arc<dyn Fn(&str) -> String + Send + Sync + 'static>;
21
22pub struct FunctionTool<Args = Value> {
23 name: String,
24 description: String,
25 parameters_schema: Value,
26 handler: DynamicFunctionHandler,
27 strict_schema: bool,
28 exposure: ToolExposure,
29 timeout: Option<Duration>,
30 approval: ToolApprovalRule,
31 tool_metadata: Option<ToolMetadata>,
32 enablement: ToolEnablementRule,
33 error_mapper: Option<ToolErrorMapper>,
34 metadata: Metadata,
35 _args: PhantomData<fn() -> Args>,
36}
37
38impl FunctionTool<Value> {
39 pub fn builder(name: impl Into<String>) -> FunctionToolBuilder<Value> {
40 FunctionToolBuilder {
41 name: name.into(),
42 description: String::new(),
43 parameters_schema: serde_json::json!({
44 "type": "object",
45 "properties": {},
46 "required": []
47 }),
48 handler: None,
49 strict_schema: true,
50 exposure: ToolExposure::Direct,
51 timeout: None,
52 approval: ToolApprovalRule::default(),
53 tool_metadata: None,
54 enablement: ToolEnablementRule::default(),
55 error_mapper: None,
56 metadata: Metadata::new(),
57 _args: PhantomData,
58 }
59 }
60}
61
62impl<Args> Clone for FunctionTool<Args> {
63 fn clone(&self) -> Self {
64 Self {
65 name: self.name.clone(),
66 description: self.description.clone(),
67 parameters_schema: self.parameters_schema.clone(),
68 handler: self.handler.clone(),
69 strict_schema: self.strict_schema,
70 exposure: self.exposure,
71 timeout: self.timeout,
72 approval: self.approval.clone(),
73 tool_metadata: self.tool_metadata.clone(),
74 enablement: self.enablement.clone(),
75 error_mapper: self.error_mapper.clone(),
76 metadata: self.metadata.clone(),
77 _args: PhantomData,
78 }
79 }
80}
81
82impl<Args> FunctionTool<Args> {
83 pub fn metadata(&self) -> &Metadata {
84 &self.metadata
85 }
86}
87
88impl<Args> FunctionTool<Args>
89where
90 Args: Send + Sync + 'static,
91{
92 pub fn to_executor(&self) -> std::sync::Arc<dyn crate::tools::ToolExecutor> {
93 ToolSpecExecutor::new(self.as_tool_spec()).into_arc()
94 }
95}
96
97impl<Args> Tool for FunctionTool<Args>
98where
99 Args: Send + Sync + 'static,
100{
101 fn name(&self) -> &str {
102 &self.name
103 }
104
105 fn description(&self) -> &str {
106 &self.description
107 }
108
109 fn parameters_schema(&self) -> &Value {
110 &self.parameters_schema
111 }
112
113 fn strict_schema(&self) -> bool {
114 self.strict_schema
115 }
116
117 fn exposure(&self) -> ToolExposure {
118 self.exposure
119 }
120
121 fn timeout(&self) -> Option<Duration> {
122 self.timeout
123 }
124
125 fn approval_rule(&self) -> ToolApprovalRule {
126 self.approval.clone()
127 }
128
129 fn tool_metadata(&self) -> Option<&ToolMetadata> {
130 self.tool_metadata.as_ref()
131 }
132
133 fn is_enabled(&self, context: &ToolEnablementContext) -> bool {
134 self.enablement.is_enabled(context)
135 }
136
137 fn as_tool_spec(&self) -> ToolSpec {
138 let name = self.name.clone();
139 let description = self.description.clone();
140 let parameters_schema = self.parameters_schema.clone();
141 let handler = self.handler.clone();
142 let error_mapper = self.error_mapper.clone();
143 let mut spec = ToolSpec::new(
144 name.clone(),
145 description.clone(),
146 Arc::new(
147 move |context: &mut ToolContext, arguments: &ToolArguments| {
148 let raw_arguments = Value::Object(arguments.clone().into_iter().collect());
149 let shared_state = Arc::new(Mutex::new(context.shared_state.clone()));
150 let run_context = context.run_context.clone().unwrap_or_else(|| RunContext {
151 run_id: context.task_id.clone(),
152 agent_name: context
153 .metadata
154 .get("agent_name")
155 .and_then(Value::as_str)
156 .unwrap_or_default()
157 .to_string(),
158 model: None,
159 workspace: Some(context.workspace.clone()),
160 metadata: context.metadata.clone(),
161 app_state: context.app_state.clone(),
162 });
163 let call_context = ToolCallContext {
164 run: run_context,
165 tool_call_id: context.tool_call_id.clone(),
166 tool_name: context.tool_name.clone(),
167 raw_arguments: raw_arguments.clone(),
168 idempotency_key: context.idempotency_key.clone(),
169 metadata: context.metadata.clone(),
170 app_state: context.app_state.clone(),
171 shared_state: shared_state.clone(),
172 };
173 let outcome =
174 catch_unwind(AssertUnwindSafe(|| handler(call_context, raw_arguments)));
175 context.shared_state = shared_state
176 .lock()
177 .unwrap_or_else(std::sync::PoisonError::into_inner)
178 .clone();
179 match outcome {
180 Ok(Ok(output)) => output.to_result(&context.tool_call_id),
181 Ok(Err(error)) => function_error_result(
182 &name,
183 &context.tool_call_id,
184 &error,
185 error_mapper.as_ref(),
186 ),
187 Err(payload) => function_error_result(
188 &name,
189 &context.tool_call_id,
190 &panic_message(payload),
191 error_mapper.as_ref(),
192 ),
193 }
194 },
195 ),
196 );
197 spec.schema = serde_json::json!({
198 "type": "function",
199 "function": {
200 "name": self.name,
201 "description": self.description,
202 "parameters": parameters_schema,
203 "strict": self.strict_schema,
204 }
205 });
206 spec.strict_schema = self.strict_schema;
207 spec.exposure = self.exposure;
208 spec.timeout = self.timeout;
209 spec.approval = self.approval.clone();
210 spec.tool_metadata = self.tool_metadata.clone();
211 spec.metadata = self.metadata.clone();
212 spec
213 }
214}
215
216pub struct FunctionToolBuilder<Args = Value> {
217 name: String,
218 description: String,
219 parameters_schema: Value,
220 handler: Option<DynamicFunctionHandler>,
221 strict_schema: bool,
222 exposure: ToolExposure,
223 timeout: Option<Duration>,
224 approval: ToolApprovalRule,
225 tool_metadata: Option<ToolMetadata>,
226 enablement: ToolEnablementRule,
227 error_mapper: Option<ToolErrorMapper>,
228 metadata: Metadata,
229 _args: PhantomData<fn() -> Args>,
230}
231
232impl<Args> FunctionToolBuilder<Args> {
233 pub fn description(mut self, description: impl Into<String>) -> Self {
234 self.description = description.into();
235 self
236 }
237
238 pub fn json_schema(mut self, schema: Value) -> Self {
239 self.parameters_schema = schema;
240 self
241 }
242
243 pub fn strict_schema(mut self, strict_schema: bool) -> Self {
244 self.strict_schema = strict_schema;
245 self
246 }
247
248 pub fn exposure(mut self, exposure: ToolExposure) -> Self {
249 self.exposure = exposure;
250 self
251 }
252
253 pub fn timeout(mut self, timeout: Duration) -> Self {
254 self.timeout = Some(timeout);
255 self
256 }
257
258 pub fn needs_approval(mut self, approval: impl Into<ToolApprovalRule>) -> Self {
259 self.approval = approval.into();
260 self
261 }
262
263 pub fn needs_approval_if<F>(mut self, predicate: F) -> Self
264 where
265 F: Fn(&ToolContext, &ToolArguments) -> bool + Send + Sync + 'static,
266 {
267 self.approval = ToolApprovalRule::predicate(predicate);
268 self
269 }
270
271 pub fn tool_metadata(mut self, tool_metadata: ToolMetadata) -> Self {
272 self.tool_metadata = Some(tool_metadata);
273 self
274 }
275
276 pub fn enabled(mut self, enabled: bool) -> Self {
277 self.enablement = ToolEnablementRule::Static(enabled);
278 self
279 }
280
281 pub fn enabled_if<F>(mut self, predicate: F) -> Self
282 where
283 F: Fn(&ToolEnablementContext) -> bool + Send + Sync + 'static,
284 {
285 self.enablement = ToolEnablementRule::predicate(predicate);
286 self
287 }
288
289 pub fn failure_error_function<F>(mut self, mapper: F) -> Self
290 where
291 F: Fn(&str) -> String + Send + Sync + 'static,
292 {
293 self.error_mapper = Some(Arc::new(mapper));
294 self
295 }
296
297 pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
298 self.metadata.insert(key.into(), value);
299 self
300 }
301
302 pub fn handler<NextArgs, F, Fut>(self, handler: F) -> FunctionToolBuilder<NextArgs>
303 where
304 NextArgs: DeserializeOwned + Send + Sync + 'static,
305 F: Fn(ToolCallContext, NextArgs) -> Fut + Send + Sync + 'static,
306 Fut: Future<Output = Result<ToolOutput, String>> + Send + 'static,
307 {
308 let handler = Arc::new(handler);
309 FunctionToolBuilder {
310 name: self.name,
311 description: self.description,
312 parameters_schema: self.parameters_schema,
313 handler: Some(Arc::new(move |context, raw_arguments| {
314 let args = serde_json::from_value::<NextArgs>(raw_arguments)
315 .map_err(|error| format!("invalid tool arguments: {error}"))?;
316 block_on_tool_future(handler(context, args))
317 })),
318 strict_schema: self.strict_schema,
319 exposure: self.exposure,
320 timeout: self.timeout,
321 approval: self.approval,
322 tool_metadata: self.tool_metadata,
323 enablement: self.enablement,
324 error_mapper: self.error_mapper,
325 metadata: self.metadata,
326 _args: PhantomData,
327 }
328 }
329
330 pub fn build(self) -> Result<FunctionTool<Args>, String> {
331 if self.name.trim().is_empty() {
332 return Err("tool name cannot be empty".to_string());
333 }
334 if self.timeout.is_some_and(|timeout| timeout.is_zero()) {
335 return Err("tool timeout must be greater than zero".to_string());
336 }
337 let Some(handler) = self.handler else {
338 return Err("tool handler is required".to_string());
339 };
340 let tool_metadata = self
341 .tool_metadata
342 .as_ref()
343 .map(ToolMetadata::normalized)
344 .transpose()
345 .map_err(|error| error.to_string())?;
346 let parameters_schema =
347 crate::tools::argument_validation::close_object_schemas(&self.parameters_schema);
348 crate::tools::argument_validation::validator_for_parameters(¶meters_schema)?;
349 Ok(FunctionTool {
350 name: self.name,
351 description: self.description,
352 parameters_schema,
353 handler,
354 strict_schema: self.strict_schema,
355 exposure: self.exposure,
356 timeout: self.timeout,
357 approval: self.approval,
358 tool_metadata,
359 enablement: self.enablement,
360 error_mapper: self.error_mapper,
361 metadata: self.metadata,
362 _args: PhantomData,
363 })
364 }
365}
366
367fn function_error_result(
368 tool_name: &str,
369 tool_call_id: &str,
370 error: &str,
371 mapper: Option<&ToolErrorMapper>,
372) -> crate::types::ToolExecutionResult {
373 let message = mapper
374 .map(|mapper| mapper(error))
375 .unwrap_or_else(|| format!("Tool execution failed ({tool_name}): {error}"));
376 ToolOutput::error(message)
377 .with_code("tool_execution_failed")
378 .to_result(tool_call_id)
379}
380
381fn panic_message(payload: Box<dyn std::any::Any + Send>) -> String {
382 if let Some(message) = payload.downcast_ref::<&str>() {
383 return (*message).to_string();
384 }
385 if let Some(message) = payload.downcast_ref::<String>() {
386 return message.clone();
387 }
388 "tool handler panicked".to_string()
389}
390
391fn block_on_tool_future<F>(future: F) -> Result<ToolOutput, String>
392where
393 F: Future<Output = Result<ToolOutput, String>> + Send + 'static,
394{
395 if let Ok(handle) = tokio::runtime::Handle::try_current() {
396 if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
397 tokio::task::block_in_place(|| handle.block_on(future))
398 } else {
399 std::thread::spawn(move || {
400 tokio::runtime::Builder::new_current_thread()
401 .enable_all()
402 .build()
403 .map_err(|error| error.to_string())?
404 .block_on(future)
405 })
406 .join()
407 .map_err(|_| "tool handler thread panicked".to_string())?
408 }
409 } else {
410 tokio::runtime::Builder::new_current_thread()
411 .enable_all()
412 .build()
413 .map_err(|error| error.to_string())?
414 .block_on(future)
415 }
416}