1mod callable;
7pub mod context;
8mod definition;
9pub mod envelope;
10mod from_fn;
11pub mod hitl;
12pub mod io;
13pub mod registry;
14pub mod result;
15
16pub use callable::Callable;
17pub use context::Context;
18pub use definition::{Annotations, Definition, ExecutionHint};
19pub use envelope::{
20 DataClassification, Envelope, FilesystemMode, FilesystemRule, NetworkPolicy, NetworkRule,
21 Safety, SensitiveMatcher, SensitivePredicate, SubprocessRule,
22};
23pub use from_fn::{from_fn, from_fn_simple};
24pub use hitl::{
25 Decision, DenyHumanApproval, DenyOnSensitive, HumanApproval, SensitivityEvaluator, ToolCall,
26 denied_error,
27};
28pub use io::{ToolInput, ToolMetadata, ToolOutput, ToolSchema};
29pub use registry::{BatchOptions, Registry};
30pub use result::{ToolResult, error_result, json_result, text_result};
31
32#[cfg(test)]
33mod tests {
34 use super::*;
35 use rskit_errors::{AppResult, ErrorCode};
36 use schemars::JsonSchema;
37 use serde::{Deserialize, Serialize};
38
39 #[derive(Deserialize, JsonSchema)]
40 struct AddInput {
41 a: i32,
42 b: i32,
43 }
44
45 #[derive(Serialize)]
46 struct AddOutput {
47 sum: i32,
48 }
49
50 #[tokio::test]
51 async fn test_from_fn_basic() {
52 let tool = from_fn(
53 "add",
54 "Add two numbers",
55 |_ctx: Context, input: AddInput| async move {
56 Ok(text_result(&format!("{}", input.a + input.b)))
57 },
58 )
59 .unwrap();
60
61 assert_eq!(tool.definition().name, "add");
62 assert_eq!(tool.definition().description, "Add two numbers");
63
64 let ctx = Context::new();
65 let result = tool
66 .call(
67 &ctx,
68 ToolInput::new(serde_json::json!({"a": 1, "b": 2})).unwrap(),
69 )
70 .await
71 .unwrap();
72 assert_eq!(result.text(), "3");
73 }
74
75 #[tokio::test]
76 async fn test_from_fn_simple_basic() {
77 let tool = from_fn_simple("add", "Add", |input: AddInput| async move {
78 Ok(AddOutput {
79 sum: input.a + input.b,
80 })
81 })
82 .unwrap();
83
84 let ctx = Context::new();
85 let result = tool
86 .call(
87 &ctx,
88 ToolInput::new(serde_json::json!({"a": 1, "b": 2})).unwrap(),
89 )
90 .await
91 .unwrap();
92 assert!(result.output.is_some());
93 assert_eq!(result.output.unwrap()["sum"], 3);
94 }
95
96 #[tokio::test]
97 async fn from_fn_simple_rejects_schema_invalid_input_before_handler_runs() {
98 use std::sync::Arc;
99 use std::sync::atomic::{AtomicUsize, Ordering};
100
101 let calls = Arc::new(AtomicUsize::new(0));
102 let handler_calls = Arc::clone(&calls);
103 let tool = from_fn_simple("add", "Add", move |input: AddInput| {
104 let handler_calls = Arc::clone(&handler_calls);
105 async move {
106 handler_calls.fetch_add(1, Ordering::SeqCst);
107 Ok(AddOutput {
108 sum: input.a + input.b,
109 })
110 }
111 })
112 .unwrap();
113
114 let ctx = Context::new();
115 let err = tool
116 .call(
117 &ctx,
118 ToolInput::new(serde_json::json!({"a": "not-an-int", "b": 2})).unwrap(),
119 )
120 .await
121 .expect_err("untrusted input that violates the schema must fail closed");
122 assert_eq!(err.code(), ErrorCode::InvalidInput);
123 assert!(err.message().contains("invalid tool input"));
124 assert_eq!(
125 calls.load(Ordering::SeqCst),
126 0,
127 "handler must not run on schema-invalid input",
128 );
129 }
130
131 #[tokio::test]
132 async fn from_fn_simple_validate_reports_schema_conformance() {
133 let tool = from_fn_simple("add", "Add", |input: AddInput| async move {
134 Ok(AddOutput {
135 sum: input.a + input.b,
136 })
137 })
138 .unwrap();
139
140 assert!(
141 tool.validate(&ToolInput::new(serde_json::json!({"a": 1, "b": 2})).unwrap())
142 .valid
143 );
144 assert!(
145 !tool
146 .validate(&ToolInput::new(serde_json::json!({"a": "bad", "b": 2})).unwrap())
147 .valid
148 );
149 }
150
151 #[tokio::test]
152 async fn test_from_fn_schema_generated() {
153 let tool = from_fn("add", "Add", |_ctx: Context, _input: AddInput| async move {
154 Ok(text_result("0"))
155 })
156 .unwrap();
157
158 let schema = &tool.definition().input_schema;
159 assert!(schema.is_object());
160 let obj = schema.as_object().unwrap();
161 let props = obj.get("properties").unwrap().as_object().unwrap();
162 assert!(props.contains_key("a"));
163 assert!(props.contains_key("b"));
164 }
165
166 #[tokio::test]
167 async fn test_from_fn_invalid_input() {
168 let tool = from_fn("add", "Add", |_ctx: Context, _input: AddInput| async move {
169 Ok(text_result("0"))
170 })
171 .unwrap();
172
173 let ctx = Context::new();
174 let result = tool
175 .call(&ctx, ToolInput::new(serde_json::json!({"x": 1})).unwrap())
176 .await;
177 assert!(result.is_err());
178 }
179
180 #[tokio::test]
181 async fn test_validate() {
182 let tool = from_fn("add", "Add", |_ctx: Context, _input: AddInput| async move {
183 Ok(text_result("0"))
184 })
185 .unwrap();
186
187 let valid = tool.validate(&ToolInput::new(serde_json::json!({"a": 1, "b": 2})).unwrap());
188 assert!(valid.valid);
189
190 let invalid =
191 tool.validate(&ToolInput::new(serde_json::json!({"a": "bad", "b": 2})).unwrap());
192 assert!(!invalid.valid);
193 }
194
195 #[tokio::test]
196 async fn test_registry_operations() {
197 let registry = Registry::new();
198 assert!(registry.is_empty());
199
200 let tool = from_fn(
201 "add",
202 "Add two numbers",
203 |_ctx: Context, input: AddInput| async move {
204 Ok(text_result(&format!("{}", input.a + input.b)))
205 },
206 )
207 .unwrap();
208
209 registry.register(tool).unwrap();
210 assert_eq!(registry.len(), 1);
211 assert!(registry.contains("add"));
212 assert!(!registry.contains("missing"));
213
214 let ctx = Context::new();
215 let result = registry
216 .call(
217 "add",
218 &ctx,
219 ToolInput::new(serde_json::json!({"a": 3, "b": 4})).unwrap(),
220 )
221 .await
222 .unwrap();
223 assert_eq!(result.text(), "7");
224 }
225
226 #[tokio::test]
227 async fn test_registry_rejects_schema_invalid_input() {
228 let registry = Registry::new();
229 let tool = from_fn(
230 "add",
231 "Add two numbers",
232 |_ctx: Context, input: AddInput| async move {
233 Ok(text_result(&format!("{}", input.a + input.b)))
234 },
235 )
236 .unwrap();
237 registry.register(tool).unwrap();
238
239 let err = registry
240 .call(
241 "add",
242 &Context::new(),
243 ToolInput::new(serde_json::json!({"a": "bad", "b": 2})).unwrap(),
244 )
245 .await
246 .unwrap_err();
247
248 assert_eq!(err.code(), ErrorCode::InvalidInput);
249 assert!(err.message().contains("invalid tool input"));
250 }
251
252 #[tokio::test]
253 async fn test_registry_duplicate() {
254 let registry = Registry::new();
255
256 let t1 = from_fn("dup", "First", |_ctx: Context, _: AddInput| async move {
257 Ok(text_result("1"))
258 })
259 .unwrap();
260 let t2 = from_fn("dup", "Second", |_ctx: Context, _: AddInput| async move {
261 Ok(text_result("2"))
262 })
263 .unwrap();
264
265 registry.register(t1).unwrap();
266 assert!(registry.register(t2).is_err());
267 }
268
269 #[tokio::test]
270 async fn test_registry_not_found() {
271 let registry = Registry::new();
272 let ctx = Context::new();
273 let result = registry
274 .call(
275 "missing",
276 &ctx,
277 ToolInput::new(serde_json::json!({})).unwrap(),
278 )
279 .await;
280 assert!(result.is_err());
281 }
282
283 #[tokio::test]
284 async fn test_registry_list() {
285 let registry = Registry::new();
286
287 let t1 = from_fn("alpha", "A tool", |_ctx: Context, _: AddInput| async move {
288 Ok(text_result("a"))
289 })
290 .unwrap();
291 let t2 = from_fn("beta", "B tool", |_ctx: Context, _: AddInput| async move {
292 Ok(text_result("b"))
293 })
294 .unwrap();
295
296 registry.register(t1).unwrap();
297 registry.register(t2).unwrap();
298
299 let defs = registry.list();
300 assert_eq!(defs.len(), 2);
301
302 let mut names: Vec<_> = registry.names();
303 names.sort();
304 assert_eq!(names, vec!["alpha", "beta"]);
305 }
306
307 #[tokio::test]
308 async fn test_registry_search() {
309 let registry = Registry::new();
310 registry
311 .register(
312 from_fn(
313 "file_read",
314 "Read a file",
315 |_ctx: Context, _: AddInput| async move { Ok(text_result("")) },
316 )
317 .unwrap(),
318 )
319 .unwrap();
320 registry
321 .register(
322 from_fn(
323 "web_search",
324 "Search the web",
325 |_ctx: Context, _: AddInput| async move { Ok(text_result("")) },
326 )
327 .unwrap(),
328 )
329 .unwrap();
330
331 let results = registry.search("file");
332 assert_eq!(results.len(), 1);
333 assert_eq!(results[0].name, "file_read");
334
335 let results = registry.search("search");
336 assert_eq!(results.len(), 1);
337 assert_eq!(results[0].name, "web_search");
338 }
339
340 #[test]
341 fn test_definition_serialization() {
342 let def = Definition {
343 name: "test".to_string(),
344 description: "Test tool".to_string(),
345 input_schema: ToolSchema::new(serde_json::json!({"type": "object"})).unwrap(),
346 output_schema: None,
347 annotations: Annotations {
348 title: "Test".to_string(),
349 ..Default::default()
350 },
351 envelope: Envelope::default(),
352 };
353
354 let json = serde_json::to_value(&def).unwrap();
355 assert_eq!(json["name"], "test");
356 assert_eq!(json["annotations"]["title"], "Test");
357 assert!(json.get("output_schema").is_none());
358 }
359
360 #[test]
361 fn test_annotations_execution_hint() {
362 let ann = Annotations {
363 execution_hint: ExecutionHint::Ui,
364 ..Default::default()
365 };
366 assert_eq!(ann.execution_hint, ExecutionHint::Ui);
367
368 let default_ann = Annotations::default();
369 assert_eq!(default_ann.execution_hint, ExecutionHint::Backend);
370 }
371
372 #[test]
373 fn test_execution_hint_serialization() {
374 let ann = Annotations {
375 title: "My Tool".to_string(),
376 execution_hint: ExecutionHint::Hybrid,
377 ..Default::default()
378 };
379 let json = serde_json::to_value(&ann).unwrap();
380 assert_eq!(json["execution_hint"], "hybrid");
381 assert_eq!(json["title"], "My Tool");
382
383 let ann_default = Annotations {
384 title: "Other".to_string(),
385 ..Default::default()
386 };
387 let json_default = serde_json::to_value(&ann_default).unwrap();
388 assert_eq!(json_default["execution_hint"], "backend");
389 }
390
391 #[test]
392 fn test_execution_hint_deserialization() {
393 let json = serde_json::json!({
394 "title": "T",
395 "execution_hint": "backend"
396 });
397 let ann: Annotations = serde_json::from_value(json).unwrap();
398 assert_eq!(ann.execution_hint, ExecutionHint::Backend);
399
400 let json_missing = serde_json::json!({"title": "T"});
401 let ann2: Annotations = serde_json::from_value(json_missing).unwrap();
402 assert_eq!(ann2.execution_hint, ExecutionHint::Backend);
403 }
404
405 struct StubTool(Definition);
407
408 #[async_trait::async_trait]
409 impl Callable for StubTool {
410 fn definition(&self) -> &Definition {
411 &self.0
412 }
413 fn validate(&self, _input: &ToolInput) -> rskit_schema::ValidationResult {
414 rskit_schema::ValidationResult {
415 valid: true,
416 errors: vec![],
417 }
418 }
419 async fn call(&self, _ctx: &Context, _input: ToolInput) -> AppResult<ToolResult> {
420 Ok(text_result("stub"))
421 }
422 }
423
424 fn stub_def(name: &str, annotations: Annotations) -> Box<dyn Callable> {
425 Box::new(StubTool(Definition {
426 name: name.to_string(),
427 description: name.to_string(),
428 input_schema: ToolSchema::new(serde_json::json!({"type": "object"})).unwrap(),
429 output_schema: None,
430 annotations,
431 envelope: Envelope::default(),
432 }))
433 }
434
435 #[tokio::test]
436 async fn test_registry_filter_by_execution_hint() {
437 let registry = Registry::new();
438
439 registry
440 .register(stub_def(
441 "validate_form",
442 Annotations {
443 execution_hint: ExecutionHint::Ui,
444 ..Default::default()
445 },
446 ))
447 .unwrap();
448
449 registry
450 .register(stub_def(
451 "run_query",
452 Annotations {
453 execution_hint: ExecutionHint::Backend,
454 ..Default::default()
455 },
456 ))
457 .unwrap();
458
459 registry
460 .register(stub_def("noop", Annotations::default()))
461 .unwrap();
462
463 let ui = registry.filter_by_execution_hint(ExecutionHint::Ui);
464 assert_eq!(ui.len(), 1);
465 assert_eq!(ui[0].name, "validate_form");
466
467 let backend = registry.filter_by_execution_hint(ExecutionHint::Backend);
468 let mut backend_names = backend
469 .iter()
470 .map(|def| def.name.as_str())
471 .collect::<Vec<_>>();
472 backend_names.sort_unstable();
473 assert_eq!(backend_names, vec!["noop", "run_query"]);
474
475 let hybrid = registry.filter_by_execution_hint(ExecutionHint::Hybrid);
476 assert!(hybrid.is_empty());
477 }
478
479 #[test]
480 fn test_context_metadata() {
481 let mut ctx = Context::new();
482 ctx.set("key", serde_json::json!("value").into());
483 assert_eq!(ctx.get("key").unwrap(), &serde_json::json!("value"));
484 assert!(ctx.get("missing").is_none());
485 }
486
487 #[test]
488 fn test_context_cancellation() {
489 let token = tokio_util::sync::CancellationToken::new();
490 let ctx = Context::with_cancellation(token.clone());
491 assert!(!ctx.is_cancelled());
492 assert!(!ctx.cancel_token().is_cancelled());
493 token.cancel();
494 assert!(ctx.is_cancelled());
495 }
496
497 #[test]
498 fn test_context_default_and_debug_are_metadata_safe() {
499 let mut ctx = Context::default();
500 ctx.request_id = "req-1".to_owned();
501 ctx.tool_use_id = "tool-1".to_owned();
502 ctx.max_result_size = 1024;
503 ctx.set("secret", serde_json::json!("redacted-by-key-only").into());
504
505 let debug = format!("{ctx:?}");
506 assert!(debug.contains("req-1"));
507 assert!(debug.contains("tool-1"));
508 assert!(debug.contains("metadata_keys"));
509 assert!(debug.contains("secret"));
510 assert!(!debug.contains("redacted-by-key-only"));
511 }
512
513 #[test]
514 fn test_tool_result_text() {
515 let r = text_result("hello");
516 assert_eq!(r.text(), "hello");
517 assert!(!r.is_error);
518 }
519
520 #[test]
521 fn test_tool_result_error() {
522 let r = error_result("something broke");
523 assert_eq!(r.text(), "something broke");
524 assert!(r.is_error);
525 }
526
527 #[test]
528 fn test_tool_result_json() {
529 let r = json_result(&serde_json::json!({"x": 1})).unwrap();
530 assert!(!r.is_error);
531 assert!(r.output.is_some());
532 assert_eq!(r.output.unwrap()["x"], 1);
533 }
534
535 #[test]
536 fn test_tool_result_block_conversion_preserves_error_state() {
537 let result = error_result("denied");
538 let block = result.to_block("tool-use-1");
539 assert_eq!(block.id, "tool-use-1");
540 assert_eq!(block.content, "denied");
541 assert!(block.is_error);
542 }
543
544 #[test]
545 fn test_tool_result_metadata() {
546 let mut r = text_result("hi");
547 r.set_meta("timing_ms", serde_json::json!(42).into());
548 assert_eq!(r.metadata.get("timing_ms").unwrap(), &serde_json::json!(42));
549 }
550}