1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use serde_json::{Value, json};
6use tokio::sync::mpsc;
7
8use crate::engine::SessionStore;
9use crate::types::{AgentError, AgentResult, SessionId, UserEvent};
10
11pub mod auto_continue;
12pub mod policy;
13pub mod update_plan;
14
15pub use auto_continue::AutoContinueTool;
16pub use update_plan::UpdatePlanTool;
17
18pub use policy::{DenyAllToolPolicy, ToolDecision, ToolPolicy};
19
20pub use agent_types::{
22 ActivationContext, Content, ToolExposure, ToolMetadata, content_details, content_text,
23};
24
25#[derive(Clone)]
26pub struct ToolContext {
27 pub session_id: SessionId,
28 pub user_event_tx: mpsc::UnboundedSender<UserEvent>,
31 pub llm_client: Option<Arc<dyn llm_trait::LlmProvider>>,
32 pub session_store: Option<Arc<dyn SessionStore>>,
33 pub language: crate::types::Language,
36 pub cancel_token: tokio_util::sync::CancellationToken,
38 pub max_output_chars: Option<usize>,
43 pub(crate) event_bus: crate::engine::EventBus,
46}
47
48impl ToolContext {
49 pub fn emit_user_event(&self, event: UserEvent) {
51 let _ = self.user_event_tx.send(event);
52 }
53
54 pub fn emit_progress(&self, text: impl Into<String>) {
56 self.emit_user_event(UserEvent::Progress { text: text.into() });
57 }
58
59 pub fn emit_partial_result(
62 &self,
63 tool_call_id: &str,
64 content: impl Into<String>,
65 is_partial: bool,
66 ) {
67 self.emit_user_event(UserEvent::ToolPartialResult {
68 tool_call_id: tool_call_id.to_string(),
69 content: content.into(),
70 is_partial,
71 });
72 }
73
74 pub fn for_test() -> Self {
78 let (tx, _rx) = mpsc::unbounded_channel();
79 ToolContext {
80 session_id: SessionId::new(0),
81 user_event_tx: tx,
82 llm_client: None,
83 session_store: None,
84 language: crate::types::Language::En,
85 cancel_token: tokio_util::sync::CancellationToken::new(),
86 max_output_chars: None,
87 event_bus: crate::engine::EventBus::new(1),
88 }
89 }
90}
91
92#[async_trait]
93pub trait Tool: Send + Sync {
94 fn name(&self) -> &'static str;
95 fn description(&self) -> &'static str;
97 fn schema(&self) -> Value;
100 async fn call(&self, args: &Value, ctx: &ToolContext) -> AgentResult<Vec<Content>>;
101
102 fn timeout_ms(&self) -> Option<u64> {
108 None }
110
111 fn metadata(&self) -> ToolMetadata {
118 ToolMetadata {
119 name: self.name().to_string(),
120 description: self.description().to_string(),
121 origin: "custom".to_string(),
122 version: "unknown".to_string(),
123 requirements: vec![],
124 }
125 }
126
127 fn exposure(&self) -> ToolExposure {
133 ToolExposure::Direct
134 }
135
136 fn should_activate(&self, _ctx: &ActivationContext) -> bool {
142 true
143 }
144}
145
146#[async_trait]
147pub trait TypedTool: Send + Sync {
148 type Args: serde::de::DeserializeOwned + schemars::JsonSchema;
149 type Output: serde::Serialize;
150
151 fn name(&self) -> &'static str;
152 fn description(&self) -> &'static str;
153 async fn call_typed(&self, args: Self::Args, ctx: &ToolContext) -> AgentResult<Self::Output>;
154
155 fn format_output(&self, output: Self::Output) -> Content {
156 match serde_json::to_value(&output) {
161 Ok(serde_json::Value::String(s)) => Content::text(s),
162 Ok(other) => Content::text(other.to_string()),
163 Err(_) => Content::text(String::new()),
164 }
165 }
166
167 fn origin(&self) -> &'static str {
169 "custom"
170 }
171
172 fn version(&self) -> &'static str {
174 "unknown"
175 }
176
177 fn exposure(&self) -> ToolExposure {
179 ToolExposure::Direct
180 }
181
182 fn should_activate(&self, _ctx: &ActivationContext) -> bool {
184 true
185 }
186}
187
188#[async_trait]
189impl<T: TypedTool + Send + Sync + 'static> Tool for T {
190 fn name(&self) -> &'static str {
191 TypedTool::name(self)
192 }
193
194 fn description(&self) -> &'static str {
195 TypedTool::description(self)
196 }
197
198 fn schema(&self) -> Value {
199 let settings = schemars::generate::SchemaSettings::draft07().with(|s| {
205 s.inline_subschemas = true;
206 s.meta_schema = None;
207 });
208 let generator = schemars::SchemaGenerator::new(settings);
209 let schema = generator.into_root_schema_for::<T::Args>();
210 serde_json::to_value(schema).unwrap_or(Value::Null)
211 }
212
213 fn metadata(&self) -> ToolMetadata {
214 ToolMetadata {
215 name: self.name().to_string(),
216 description: self.description().to_string(),
217 origin: self.origin().to_string(),
218 version: self.version().to_string(),
219 requirements: vec![],
220 }
221 }
222
223 fn exposure(&self) -> ToolExposure {
224 TypedTool::exposure(self)
225 }
226
227 fn should_activate(&self, ctx: &ActivationContext) -> bool {
228 TypedTool::should_activate(self, ctx)
229 }
230
231 async fn call(&self, args: &Value, ctx: &ToolContext) -> AgentResult<Vec<Content>> {
232 let typed_args: T::Args =
233 serde_json::from_value(args.clone()).map_err(|_| AgentError::ToolArgsInvalid {
234 name: self.name().to_string(),
235 raw: args.to_string(),
236 })?;
237 let output = self.call_typed(typed_args, ctx).await?;
238 Ok(vec![self.format_output(output)])
239 }
240}
241
242pub fn render_tool_definition(tool: &dyn Tool) -> Value {
247 json!({
248 "type": "function",
249 "function": {
250 "name": tool.name(),
251 "description": tool.description(),
252 "parameters": tool.schema(),
253 }
254 })
255}
256
257pub(crate) type ToolRef = Arc<dyn Tool>;
258
259#[derive(Clone, Default)]
260pub struct ToolRegistry {
261 tools: HashMap<String, ToolRef>,
262}
263
264impl ToolRegistry {
265 pub fn register(&mut self, tool: impl Tool + 'static) {
266 self.tools.insert(tool.name().to_string(), Arc::new(tool));
267 }
268
269 pub fn register_arc(&mut self, tool: Arc<dyn Tool>) {
270 self.tools.insert(tool.name().to_string(), tool);
271 }
272
273 pub fn remove(&mut self, name: &str) {
275 self.tools.remove(name);
276 }
277
278 pub fn get(&self, name: &str) -> Option<ToolRef> {
279 self.tools.get(name).cloned()
280 }
281
282 pub fn definitions(&self) -> Vec<Value> {
283 let mut tools: Vec<_> = self.tools.values().collect();
284 tools.sort_by_key(|t| t.name());
285 tools
286 .into_iter()
287 .map(|t| render_tool_definition(t.as_ref()))
288 .collect()
289 }
290
291 pub fn definitions_filtered(&self, ctx: &ActivationContext) -> Vec<Value> {
297 let mut tools: Vec<_> = self.tools.values().collect();
298 tools.sort_by_key(|t| t.name());
299
300 let direct_names: Vec<String> = tools
306 .iter()
307 .filter(|t| t.exposure() == ToolExposure::Direct)
308 .map(|t| t.name().to_string())
309 .collect();
310
311 let mut activated_names = direct_names.clone();
312 for t in &tools {
313 if t.exposure() == ToolExposure::Deferred {
314 let mut ctx_with_tools = ctx.clone();
315 ctx_with_tools.current_tools = activated_names.clone();
316 if t.should_activate(&ctx_with_tools) {
317 activated_names.push(t.name().to_string());
318 }
319 }
320 }
321
322 tools
323 .into_iter()
324 .filter(|t| match t.exposure() {
325 ToolExposure::Direct => true,
326 ToolExposure::Deferred => activated_names.contains(&t.name().to_string()),
327 ToolExposure::Hidden => false,
328 })
329 .map(|t| render_tool_definition(t.as_ref()))
330 .collect()
331 }
332
333 pub fn len(&self) -> usize {
334 self.tools.len()
335 }
336
337 pub fn is_empty(&self) -> bool {
338 self.tools.is_empty()
339 }
340
341 pub fn metadatas(&self) -> Vec<ToolMetadata> {
347 let mut list: Vec<_> = self.tools.values().map(|tool| tool.metadata()).collect();
348 list.sort_by(|a, b| a.name.cmp(&b.name));
349 list
350 }
351}
352
353#[cfg(test)]
354mod tests {
355 use super::*;
356
357 #[test]
358 fn content_text_ctor_and_into_vec() {
359 let c = Content::text("hello");
360 let v: Vec<Content> = c.clone().into();
361 assert_eq!(v.len(), 1);
362 assert!(matches!(v[0], Content::Text { .. }));
363 assert!(matches!(&c, Content::Text { text } if text == "hello"));
364 }
365
366 #[test]
367 fn content_serializes_with_type_tag() {
368 let c = Content::text("hi");
369 let j = serde_json::to_value(&c).unwrap();
370 assert_eq!(j["type"], "text");
371 assert_eq!(j["text"], "hi");
372 }
373
374 #[test]
375 fn tool_context_for_test_constructs() {
376 let ctx = ToolContext::for_test();
377 assert!(ctx.llm_client.is_none());
378 assert!(ctx.session_store.is_none());
379 assert!(!ctx.cancel_token.is_cancelled());
380 ctx.emit_progress("hello");
381 }
382
383 #[test]
384 fn typed_tool_schema_is_derived_from_args() {
385 #[derive(schemars::JsonSchema, serde::Deserialize, serde::Serialize)]
386 struct GreetArgs {
387 name: String,
388 #[serde(default)]
389 times: u32,
390 }
391
392 let schema = schemars::schema_for!(GreetArgs);
393 let j = serde_json::to_value(&schema).unwrap();
394 assert!(j["properties"]["name"].is_object());
396 assert!(j["properties"]["times"].is_object());
397 }
398
399 #[test]
400 fn typed_tool_schema_is_provider_safe_for_nested_enum() {
401 #[derive(schemars::JsonSchema, serde::Deserialize, serde::Serialize)]
402 enum Status {
403 Active,
404 Paused,
405 }
406
407 #[derive(schemars::JsonSchema, serde::Deserialize, serde::Serialize)]
408 struct Args {
409 name: String,
410 status: Status,
411 }
412
413 #[derive(Default)]
414 struct NestedTool;
415 #[async_trait]
416 impl TypedTool for NestedTool {
417 type Args = Args;
418 type Output = String;
419 fn name(&self) -> &'static str {
420 "nested"
421 }
422 fn description(&self) -> &'static str {
423 ""
424 }
425 async fn call_typed(
426 &self,
427 _args: Args,
428 _ctx: &ToolContext,
429 ) -> crate::types::AgentResult<String> {
430 Ok(String::new())
431 }
432 }
433
434 let schema = Tool::schema(&NestedTool);
435 let raw = schema.to_string();
436 assert!(!raw.contains("$ref"), "schema contains $ref: {raw}");
439 assert!(!raw.contains("$defs"), "schema contains $defs: {raw}");
440 assert!(
441 !raw.contains("definitions"),
442 "schema has definitions: {raw}"
443 );
444 assert!(schema.get("$schema").is_none(), "schema has $schema key");
445
446 let variants: Vec<&str> = schema["properties"]["status"]["enum"]
448 .as_array()
449 .unwrap()
450 .iter()
451 .map(|v| v.as_str().unwrap())
452 .collect();
453 assert!(variants.contains(&"Active"), "missing Active: {variants:?}");
454 assert!(variants.contains(&"Paused"), "missing Paused: {variants:?}");
455 }
456
457 #[test]
458 fn definitions_are_sorted_by_name() {
459 struct NamedTool(&'static str);
460 #[async_trait::async_trait]
461 impl Tool for NamedTool {
462 fn name(&self) -> &'static str {
463 self.0
464 }
465 fn description(&self) -> &'static str {
466 ""
467 }
468 fn schema(&self) -> serde_json::Value {
469 serde_json::Value::Null
470 }
471 async fn call(
472 &self,
473 _args: &serde_json::Value,
474 _ctx: &ToolContext,
475 ) -> crate::types::AgentResult<Vec<Content>> {
476 Ok(vec![])
477 }
478 }
479
480 let mut registry = ToolRegistry::default();
481 registry.register(NamedTool("zeta"));
482 registry.register(NamedTool("alpha"));
483 registry.register(NamedTool("mike"));
484
485 let defs = registry.definitions();
486 let names: Vec<&str> = defs
487 .iter()
488 .map(|d| d["function"]["name"].as_str().unwrap())
489 .collect();
490 assert_eq!(names, vec!["alpha", "mike", "zeta"]);
491 }
492
493 #[test]
496 fn content_image_and_content_text_skips_images() {
497 let img = Content::image("base64data", "image/png");
498 assert!(
499 matches!(&img, Content::Image { data, mime_type } if data == "base64data" && mime_type == "image/png")
500 );
501
502 let text = content_text(&[
503 Content::text("a"),
504 Content::image("b", "image/png"),
505 Content::text("c"),
506 ]);
507 assert_eq!(text, "a\nc");
508 }
509
510 #[test]
511 fn emit_partial_result_sends_event() {
512 let (tx, mut rx) = mpsc::unbounded_channel();
513 let ctx = ToolContext {
514 session_id: SessionId::new(0),
515 user_event_tx: tx,
516 llm_client: None,
517 session_store: None,
518 language: crate::types::Language::En,
519 cancel_token: tokio_util::sync::CancellationToken::new(),
520 max_output_chars: None,
521 event_bus: crate::engine::EventBus::new(1),
522 };
523 ctx.emit_partial_result("tc1", "partial", true);
524 match rx.try_recv().unwrap() {
525 UserEvent::ToolPartialResult {
526 tool_call_id,
527 content,
528 is_partial,
529 } => {
530 assert_eq!(tool_call_id, "tc1");
531 assert_eq!(content, "partial");
532 assert!(is_partial);
533 }
534 other => panic!("unexpected event: {other:?}"),
535 }
536 }
537
538 #[derive(schemars::JsonSchema, serde::Deserialize, serde::Serialize)]
540 struct GreetArgs {
541 name: String,
542 }
543
544 struct GreetTool;
545 #[async_trait]
546 impl TypedTool for GreetTool {
547 type Args = GreetArgs;
548 type Output = String;
549 fn name(&self) -> &'static str {
550 "greet"
551 }
552 fn description(&self) -> &'static str {
553 "greets a name"
554 }
555 fn origin(&self) -> &'static str {
556 "test-crate"
557 }
558 fn version(&self) -> &'static str {
559 "1.0.0"
560 }
561 async fn call_typed(&self, args: GreetArgs, _ctx: &ToolContext) -> AgentResult<String> {
562 Ok(format!("Hello, {}!", args.name))
563 }
564 }
565
566 #[test]
567 fn typed_tool_blanket_delegates_name_description() {
568 let t = GreetTool;
569 assert_eq!(Tool::name(&t), "greet");
570 assert_eq!(Tool::description(&t), "greets a name");
571 }
572
573 #[test]
574 fn typed_tool_metadata_uses_origin_and_version() {
575 let m = Tool::metadata(&GreetTool);
576 assert_eq!(m.name, "greet");
577 assert_eq!(m.description, "greets a name");
578 assert_eq!(m.origin, "test-crate");
579 assert_eq!(m.version, "1.0.0");
580 assert!(m.requirements.is_empty());
581 }
582
583 #[tokio::test]
584 async fn typed_tool_call_deserializes_and_formats() {
585 let ctx = ToolContext::for_test();
586 let out = Tool::call(&GreetTool, &json!({"name": "world"}), &ctx)
587 .await
588 .unwrap();
589 assert_eq!(content_text(&out), "Hello, world!");
591 }
592
593 #[derive(serde::Serialize)]
595 struct GreetResult {
596 message: String,
597 }
598
599 struct GreetStructTool;
600 #[async_trait]
601 impl TypedTool for GreetStructTool {
602 type Args = GreetArgs;
603 type Output = GreetResult;
604 fn name(&self) -> &'static str {
605 "greet_struct"
606 }
607 fn description(&self) -> &'static str {
608 "greets as json"
609 }
610 async fn call_typed(
611 &self,
612 args: GreetArgs,
613 _ctx: &ToolContext,
614 ) -> AgentResult<GreetResult> {
615 Ok(GreetResult {
616 message: format!("Hello, {}!", args.name),
617 })
618 }
619 }
620
621 #[test]
622 fn format_output_json_serializes_struct() {
623 let out = GreetStructTool.format_output(GreetResult {
624 message: "hi".into(),
625 });
626 assert_eq!(content_text(&[out]), r#"{"message":"hi"}"#);
627 }
628
629 #[tokio::test]
630 async fn typed_tool_call_invalid_args_is_tool_args_invalid() {
631 let ctx = ToolContext::for_test();
632 let err = Tool::call(&GreetTool, &json!({"nope": 1}), &ctx)
633 .await
634 .unwrap_err();
635 assert!(matches!(err, AgentError::ToolArgsInvalid { .. }));
636 }
637
638 struct NamedTool(&'static str);
639 #[async_trait]
640 impl Tool for NamedTool {
641 fn name(&self) -> &'static str {
642 self.0
643 }
644 fn description(&self) -> &'static str {
645 ""
646 }
647 fn schema(&self) -> serde_json::Value {
648 serde_json::Value::Null
649 }
650 async fn call(
651 &self,
652 _args: &serde_json::Value,
653 _ctx: &ToolContext,
654 ) -> AgentResult<Vec<Content>> {
655 Ok(vec![])
656 }
657 }
658
659 #[test]
660 fn registry_register_arc_get_remove_len_is_empty() {
661 let mut r = ToolRegistry::default();
662 assert!(r.is_empty());
663 assert_eq!(r.len(), 0);
664
665 let t: Arc<dyn Tool> = Arc::new(NamedTool("x"));
666 r.register_arc(t);
667 assert!(!r.is_empty());
668 assert_eq!(r.len(), 1);
669 assert!(r.get("x").is_some());
670 assert!(r.get("missing").is_none());
671
672 r.remove("x");
673 assert!(r.is_empty());
674 }
675
676 #[test]
677 fn metadatas_are_sorted_by_name() {
678 let mut r = ToolRegistry::default();
679 r.register(NamedTool("zeta"));
680 r.register(NamedTool("alpha"));
681
682 let metas = r.metadatas();
683 let names: Vec<&str> = metas.iter().map(|m| m.name.as_str()).collect();
684 assert_eq!(names, vec!["alpha", "zeta"]);
685 assert_eq!(metas[0].origin, "custom");
686 assert_eq!(metas[0].version, "unknown");
687 }
688
689 struct ExposureTool(&'static str, ToolExposure);
692 #[async_trait]
693 impl Tool for ExposureTool {
694 fn name(&self) -> &'static str {
695 self.0
696 }
697 fn description(&self) -> &'static str {
698 ""
699 }
700 fn schema(&self) -> serde_json::Value {
701 serde_json::Value::Null
702 }
703 async fn call(
704 &self,
705 _args: &serde_json::Value,
706 _ctx: &ToolContext,
707 ) -> AgentResult<Vec<Content>> {
708 Ok(vec![])
709 }
710 fn exposure(&self) -> ToolExposure {
711 self.1.clone()
712 }
713 }
714
715 struct ConditionalTool(&'static str, bool);
716 #[async_trait]
717 impl Tool for ConditionalTool {
718 fn name(&self) -> &'static str {
719 self.0
720 }
721 fn description(&self) -> &'static str {
722 ""
723 }
724 fn schema(&self) -> serde_json::Value {
725 serde_json::Value::Null
726 }
727 async fn call(
728 &self,
729 _args: &serde_json::Value,
730 _ctx: &ToolContext,
731 ) -> AgentResult<Vec<Content>> {
732 Ok(vec![])
733 }
734 fn exposure(&self) -> ToolExposure {
735 ToolExposure::Deferred
736 }
737 fn should_activate(&self, _ctx: &ActivationContext) -> bool {
738 self.1
739 }
740 }
741
742 fn default_ctx() -> ActivationContext {
743 ActivationContext {
744 session_id: crate::types::SessionId::new(0),
745 current_tools: vec![],
746 workspace: std::path::PathBuf::from("/tmp"),
747 }
748 }
749
750 #[test]
751 fn tool_exposure_default_is_direct() {
752 let t = NamedTool("x");
754 assert_eq!(t.exposure(), ToolExposure::Direct);
755 }
756
757 #[test]
758 fn tool_should_activate_default_is_true() {
759 let t = NamedTool("x");
760 assert!(t.should_activate(&default_ctx()));
761 }
762
763 #[test]
764 fn definitions_filtered_includes_direct_excludes_hidden() {
765 let mut r = ToolRegistry::default();
766 r.register(ExposureTool("direct_a", ToolExposure::Direct));
767 r.register(ExposureTool("hidden_a", ToolExposure::Hidden));
768 r.register(ExposureTool("direct_b", ToolExposure::Direct));
769
770 let defs = r.definitions_filtered(&default_ctx());
771 let names: Vec<&str> = defs
772 .iter()
773 .map(|d| d["function"]["name"].as_str().unwrap())
774 .collect();
775 assert_eq!(names, vec!["direct_a", "direct_b"]);
776 }
777
778 #[test]
779 fn definitions_filtered_includes_deferred_when_activated() {
780 let mut r = ToolRegistry::default();
781 r.register(ExposureTool("always", ToolExposure::Direct));
782 r.register(ConditionalTool("maybe", true)); r.register(ConditionalTool("never", false)); let defs = r.definitions_filtered(&default_ctx());
786 let names: Vec<&str> = defs
787 .iter()
788 .map(|d| d["function"]["name"].as_str().unwrap())
789 .collect();
790 assert_eq!(names, vec!["always", "maybe"]);
791 }
792
793 #[test]
794 fn definitions_filtered_all_hidden_returns_empty() {
795 let mut r = ToolRegistry::default();
796 r.register(ExposureTool("h1", ToolExposure::Hidden));
797 r.register(ExposureTool("h2", ToolExposure::Hidden));
798
799 let defs = r.definitions_filtered(&default_ctx());
800 assert!(defs.is_empty());
801 }
802
803 #[test]
804 fn definitions_filtered_empty_registry() {
805 let r = ToolRegistry::default();
806 let defs = r.definitions_filtered(&default_ctx());
807 assert!(defs.is_empty());
808 }
809
810 #[test]
811 fn definitions_unfiltered_includes_hidden() {
812 let mut r = ToolRegistry::default();
814 r.register(ExposureTool("visible", ToolExposure::Direct));
815 r.register(ExposureTool("secret", ToolExposure::Hidden));
816
817 let defs = r.definitions();
818 assert_eq!(defs.len(), 2);
819 }
820}