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