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