1use std::collections::HashMap;
9use std::fmt;
10use std::future::Future;
11use std::pin::Pin;
12use std::sync::Arc;
13use std::sync::{LazyLock, Mutex};
14
15use regex::Regex;
16use serde::{Deserialize, Serialize};
17use serde_json::Value;
18use tokio_util::sync::CancellationToken;
19
20use crate::agent_options::{ApproveToolFn, ApproveToolFuture};
21use crate::schema::schema_for;
22use crate::transfer::TransferSignal;
23use crate::types::ContentBlock;
24
25static SCHEMA_VALIDATOR_CACHE: LazyLock<Mutex<HashMap<String, Arc<jsonschema::Validator>>>> =
26 LazyLock::new(|| Mutex::new(HashMap::new()));
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct AgentToolResult {
36 pub content: Vec<ContentBlock>,
38 pub details: Value,
40 pub is_error: bool,
42 #[serde(default, skip_serializing_if = "Option::is_none")]
44 pub transfer_signal: Option<TransferSignal>,
45}
46
47impl AgentToolResult {
48 pub fn text(text: impl Into<String>) -> Self {
50 Self {
51 content: vec![ContentBlock::Text { text: text.into() }],
52 details: Value::Null,
53 is_error: false,
54 transfer_signal: None,
55 }
56 }
57
58 pub fn error(message: impl Into<String>) -> Self {
64 Self {
65 content: vec![ContentBlock::Text {
66 text: message.into(),
67 }],
68 details: Value::Null,
69 is_error: true,
70 transfer_signal: None,
71 }
72 }
73
74 pub fn transfer(signal: TransferSignal) -> Self {
80 let text = format!("Transfer to {} initiated.", signal.target_agent());
81 Self {
82 content: vec![ContentBlock::Text { text }],
83 details: Value::Null,
84 is_error: false,
85 transfer_signal: Some(signal),
86 }
87 }
88
89 pub const fn is_transfer(&self) -> bool {
91 self.transfer_signal.is_some()
92 }
93}
94
95pub type ToolFuture<'a> = Pin<Box<dyn Future<Output = AgentToolResult> + Send + 'a>>;
97
98#[derive(Debug, Clone, Default, PartialEq, Eq)]
105pub struct ToolMetadata {
106 pub namespace: Option<String>,
108 pub version: Option<String>,
110}
111
112impl ToolMetadata {
113 #[must_use]
115 pub fn with_namespace(namespace: impl Into<String>) -> Self {
116 Self {
117 namespace: Some(namespace.into()),
118 version: None,
119 }
120 }
121
122 #[must_use]
124 pub fn with_version(mut self, version: impl Into<String>) -> Self {
125 self.version = Some(version.into());
126 self
127 }
128}
129
130pub trait AgentTool: Send + Sync {
137 fn name(&self) -> &str;
139
140 fn label(&self) -> &str;
142
143 fn description(&self) -> &str;
145
146 fn parameters_schema(&self) -> &Value;
148
149 fn requires_approval(&self) -> bool {
152 false
153 }
154
155 fn metadata(&self) -> Option<ToolMetadata> {
159 None
160 }
161
162 fn approval_context(&self, _params: &Value) -> Option<Value> {
171 None
172 }
173
174 fn auth_config(&self) -> Option<crate::credential::AuthConfig> {
180 None
181 }
182
183 fn execute(
194 &self,
195 tool_call_id: &str,
196 params: Value,
197 cancellation_token: CancellationToken,
198 on_update: Option<Box<dyn Fn(AgentToolResult) + Send + Sync>>,
199 state: Arc<std::sync::RwLock<crate::SessionState>>,
200 credential: Option<crate::credential::ResolvedCredential>,
201 ) -> ToolFuture<'_>;
202}
203
204pub trait IntoTool {
217 fn into_tool(self) -> Arc<dyn AgentTool>;
219}
220
221impl<T: AgentTool + 'static> IntoTool for T {
222 fn into_tool(self) -> Arc<dyn AgentTool> {
223 Arc::new(self)
224 }
225}
226
227pub fn validate_schema(schema: &Value) -> Result<(), String> {
240 compiled_validator(schema)?;
241 Ok(())
242}
243
244pub fn validate_tool_arguments(schema: &Value, arguments: &Value) -> Result<(), Vec<String>> {
254 let validator = compiled_validator(schema).map_err(|e| vec![e])?;
255
256 let errors: Vec<String> = validator
257 .iter_errors(arguments)
258 .map(|e| e.to_string())
259 .collect();
260
261 if errors.is_empty() {
262 Ok(())
263 } else {
264 Err(errors)
265 }
266}
267
268#[must_use]
270pub(crate) fn permissive_object_schema() -> Value {
271 serde_json::json!({
272 "type": "object",
273 "properties": {},
274 "additionalProperties": true
275 })
276}
277
278#[must_use]
280pub(crate) fn debug_validated_schema(schema: Value) -> Value {
281 debug_assert!(validate_schema(&schema).is_ok());
282 schema
283}
284
285#[must_use]
287pub(crate) fn validated_schema_for<T: schemars::JsonSchema>() -> Value {
288 debug_validated_schema(schema_for::<T>())
289}
290
291fn compiled_validator(schema: &Value) -> Result<Arc<jsonschema::Validator>, String> {
292 let cache_key = serde_json::to_string(schema).map_err(|e| e.to_string())?;
293
294 {
295 let cache = SCHEMA_VALIDATOR_CACHE
296 .lock()
297 .unwrap_or_else(std::sync::PoisonError::into_inner);
298 if let Some(validator) = cache.get(&cache_key) {
299 return Ok(Arc::clone(validator));
300 }
301 }
302
303 let compiled = Arc::new(jsonschema::validator_for(schema).map_err(|e| e.to_string())?);
304 let mut cache = SCHEMA_VALIDATOR_CACHE
305 .lock()
306 .unwrap_or_else(std::sync::PoisonError::into_inner);
307 Ok(Arc::clone(
308 cache
309 .entry(cache_key)
310 .or_insert_with(|| Arc::clone(&compiled)),
311 ))
312}
313
314#[must_use]
316pub fn unknown_tool_result(tool_name: &str) -> AgentToolResult {
317 AgentToolResult::error(format!("unknown tool: {tool_name}"))
318}
319
320#[must_use]
322pub fn validation_error_result(errors: &[String]) -> AgentToolResult {
323 let message = errors.join("\n");
324 AgentToolResult::error(message)
325}
326
327#[derive(Debug, Clone, PartialEq, Eq)]
331pub enum ToolApproval {
332 Approved,
334 Rejected,
336 ApprovedWith(serde_json::Value),
338}
339
340#[derive(Clone)]
346pub struct ToolApprovalRequest {
347 pub tool_call_id: String,
349 pub tool_name: String,
351 pub arguments: Value,
353 pub requires_approval: bool,
355 pub context: Option<Value>,
357}
358
359impl fmt::Debug for ToolApprovalRequest {
360 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
361 let redacted_context = self.context.as_ref().map(redact_sensitive_values);
362 f.debug_struct("ToolApprovalRequest")
363 .field("tool_call_id", &self.tool_call_id)
364 .field("tool_name", &self.tool_name)
365 .field("arguments", &"[REDACTED]")
366 .field("requires_approval", &self.requires_approval)
367 .field("context", &redacted_context)
368 .finish()
369 }
370}
371
372#[non_exhaustive]
374#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
375pub enum ApprovalMode {
376 Enabled,
378 #[default]
381 Smart,
382 Bypassed,
385}
386
387#[allow(clippy::type_complexity)]
392pub fn selective_approve<F>(inner: F) -> Box<ApproveToolFn>
393where
394 F: Fn(ToolApprovalRequest) -> ApproveToolFuture + Send + Sync + 'static,
395{
396 Box::new(move |req: ToolApprovalRequest| {
397 if req.requires_approval {
398 inner(req)
399 } else {
400 Box::pin(async { ToolApproval::Approved })
401 }
402 })
403}
404
405const REDACTED: &str = "[REDACTED]";
409
410const SENSITIVE_KEYS: &[&str] = &[
412 "password",
413 "secret",
414 "token",
415 "api_key",
416 "apikey",
417 "authorization",
418];
419
420#[must_use]
435pub fn redact_sensitive_values(value: &Value) -> Value {
436 redact_value(value, None)
437}
438
439fn redact_value(value: &Value, parent_key: Option<&str>) -> Value {
441 if let Some(key) = parent_key
443 && SENSITIVE_KEYS.iter().any(|&s| key.eq_ignore_ascii_case(s))
444 {
445 return Value::String(REDACTED.to_string());
446 }
447
448 match value {
449 Value::String(s) => {
450 if is_sensitive_string(s) {
451 Value::String(REDACTED.to_string())
452 } else {
453 value.clone()
454 }
455 }
456 Value::Array(arr) => Value::Array(arr.iter().map(|v| redact_value(v, None)).collect()),
457 Value::Object(map) => {
458 let redacted = map
459 .iter()
460 .map(|(k, v)| (k.clone(), redact_value(v, Some(k))))
461 .collect();
462 Value::Object(redacted)
463 }
464 _ => value.clone(),
466 }
467}
468
469fn is_sensitive_string(s: &str) -> bool {
472 if s.starts_with("sk-")
474 || s.starts_with("key-")
475 || s.starts_with("token-")
476 || s.to_ascii_lowercase().starts_with("bearer ")
477 || s.to_ascii_lowercase().starts_with("basic ")
478 {
479 return true;
480 }
481
482 thread_local! {
485 static ENV_VAR_RE: Regex =
486 Regex::new(r"^\$\{?[A-Z_][A-Z0-9_]*\}?$").expect("valid regex");
487 }
488 ENV_VAR_RE.with(|re| re.is_match(s))
489}
490
491pub trait ToolParameters {
498 fn json_schema() -> Value;
500}
501
502const _: () = {
505 const fn assert_send_sync<T: Send + Sync>() {}
506 assert_send_sync::<AgentToolResult>();
507 assert_send_sync::<ToolApproval>();
508 assert_send_sync::<ToolApprovalRequest>();
509 assert_send_sync::<ApprovalMode>();
510};
511
512#[cfg(test)]
513mod tests {
514 use serde_json::json;
515
516 use super::*;
517 use crate::FnTool;
518
519 fn stub_tool(name: &str) -> FnTool {
520 FnTool::new(name, name, "A test tool.")
521 }
522
523 #[test]
526 fn approval_request_debug_redacts_arguments_and_context() {
527 let req = ToolApprovalRequest {
528 tool_call_id: "call_1".into(),
529 tool_name: "bash".into(),
530 arguments: json!({"command": "echo secret"}),
531 requires_approval: true,
532 context: Some(json!({
533 "Authorization": "Bearer top-secret",
534 "path": "/tmp/output.txt",
535 })),
536 };
537 let debug = format!("{req:?}");
538 assert!(debug.contains("tool_call_id: \"call_1\""));
539 assert!(debug.contains("tool_name: \"bash\""));
540 assert!(debug.contains("[REDACTED]"));
541 assert!(!debug.contains("echo secret"));
542 assert!(!debug.contains("top-secret"));
543 assert!(debug.contains("/tmp/output.txt"));
544 }
545
546 #[test]
549 fn redacts_sk_prefix() {
550 let val = json!({"key": "sk-abc123"});
551 let redacted = redact_sensitive_values(&val);
552 assert_eq!(redacted["key"], json!("[REDACTED]"));
553 }
554
555 #[test]
556 fn redacts_key_prefix() {
557 let val = json!({"data": "key-live-xyz"});
558 let redacted = redact_sensitive_values(&val);
559 assert_eq!(redacted["data"], json!("[REDACTED]"));
560 }
561
562 #[test]
563 fn redacts_token_prefix() {
564 let val = json!({"tok": "token-abcdef"});
565 let redacted = redact_sensitive_values(&val);
566 assert_eq!(redacted["tok"], json!("[REDACTED]"));
567 }
568
569 #[test]
570 fn redacts_bearer_prefix_case_insensitive() {
571 let val = json!({"auth": "Bearer eyJhbGciOi..."});
572 let redacted = redact_sensitive_values(&val);
573 assert_eq!(redacted["auth"], json!("[REDACTED]"));
574
575 let val2 = json!({"auth": "bearer xyz"});
576 let redacted2 = redact_sensitive_values(&val2);
577 assert_eq!(redacted2["auth"], json!("[REDACTED]"));
578 }
579
580 #[test]
581 fn redacts_basic_prefix_case_insensitive() {
582 let val = json!({"auth": "Basic dXNlcjpwYXNz"});
583 let redacted = redact_sensitive_values(&val);
584 assert_eq!(redacted["auth"], json!("[REDACTED]"));
585
586 let val2 = json!({"auth": "basic abc"});
587 let redacted2 = redact_sensitive_values(&val2);
588 assert_eq!(redacted2["auth"], json!("[REDACTED]"));
589 }
590
591 #[test]
592 fn redacts_env_var_dollar_sign() {
593 let val = json!({"ref": "$SECRET"});
594 let redacted = redact_sensitive_values(&val);
595 assert_eq!(redacted["ref"], json!("[REDACTED]"));
596 }
597
598 #[test]
599 fn redacts_env_var_braced() {
600 let val = json!({"ref": "${API_KEY}"});
601 let redacted = redact_sensitive_values(&val);
602 assert_eq!(redacted["ref"], json!("[REDACTED]"));
603 }
604
605 #[test]
606 fn redacts_sensitive_key_password() {
607 let val = json!({"password": "hunter2"});
608 let redacted = redact_sensitive_values(&val);
609 assert_eq!(redacted["password"], json!("[REDACTED]"));
610 }
611
612 #[test]
613 fn redacts_sensitive_key_secret() {
614 let val = json!({"secret": "mysecret"});
615 let redacted = redact_sensitive_values(&val);
616 assert_eq!(redacted["secret"], json!("[REDACTED]"));
617 }
618
619 #[test]
620 fn redacts_sensitive_key_token() {
621 let val = json!({"Token": "abc"});
622 let redacted = redact_sensitive_values(&val);
623 assert_eq!(redacted["Token"], json!("[REDACTED]"));
624 }
625
626 #[test]
627 fn redacts_sensitive_key_api_key() {
628 let val = json!({"api_key": "abc"});
629 let redacted = redact_sensitive_values(&val);
630 assert_eq!(redacted["api_key"], json!("[REDACTED]"));
631 }
632
633 #[test]
634 fn redacts_sensitive_key_apikey() {
635 let val = json!({"apiKey": "abc"});
636 let redacted = redact_sensitive_values(&val);
637 assert_eq!(redacted["apiKey"], json!("[REDACTED]"));
638 }
639
640 #[test]
641 fn redacts_sensitive_key_authorization() {
642 let val = json!({"Authorization": "something"});
643 let redacted = redact_sensitive_values(&val);
644 assert_eq!(redacted["Authorization"], json!("[REDACTED]"));
645 }
646
647 #[test]
648 fn passes_through_non_sensitive_values() {
649 let val = json!({
650 "command": "echo hello",
651 "path": "/tmp/file.txt",
652 "count": 42,
653 "verbose": true,
654 "items": ["one", "two"]
655 });
656 let redacted = redact_sensitive_values(&val);
657 assert_eq!(redacted, val);
658 }
659
660 #[test]
661 fn redacts_nested_objects() {
662 let val = json!({
663 "config": {
664 "password": "secret123",
665 "host": "localhost"
666 }
667 });
668 let redacted = redact_sensitive_values(&val);
669 assert_eq!(redacted["config"]["password"], json!("[REDACTED]"));
670 assert_eq!(redacted["config"]["host"], json!("localhost"));
671 }
672
673 #[test]
674 fn redacts_values_in_arrays() {
675 let val = json!(["normal", "sk-secret", "also normal"]);
676 let redacted = redact_sensitive_values(&val);
677 assert_eq!(redacted, json!(["normal", "[REDACTED]", "also normal"]));
678 }
679
680 #[test]
681 fn handles_null_and_numbers() {
682 let val = json!({"a": null, "b": 42, "c": 2.72});
683 let redacted = redact_sensitive_values(&val);
684 assert_eq!(redacted, val);
685 }
686
687 #[test]
690 fn valid_schema_passes() {
691 let schema = json!({
692 "type": "object",
693 "properties": {
694 "name": { "type": "string" }
695 },
696 "required": ["name"]
697 });
698 assert!(validate_schema(&schema).is_ok());
699 }
700
701 #[test]
702 fn invalid_schema_returns_error() {
703 let schema = json!({
704 "type": "not_a_real_type"
705 });
706 assert!(validate_schema(&schema).is_err());
707 }
708
709 #[test]
710 fn empty_object_schema_is_valid() {
711 let schema = json!({
712 "type": "object",
713 "properties": {}
714 });
715 assert!(validate_schema(&schema).is_ok());
716 }
717
718 #[test]
721 fn approval_mode_default_is_smart() {
722 assert_eq!(ApprovalMode::default(), ApprovalMode::Smart);
723 }
724
725 #[test]
726 fn approval_mode_variants_are_distinct() {
727 assert_ne!(ApprovalMode::Enabled, ApprovalMode::Smart);
728 assert_ne!(ApprovalMode::Smart, ApprovalMode::Bypassed);
729 assert_ne!(ApprovalMode::Enabled, ApprovalMode::Bypassed);
730 }
731
732 #[test]
735 fn tool_metadata_default_is_empty() {
736 let meta = ToolMetadata::default();
737 assert_eq!(meta.namespace, None);
738 assert_eq!(meta.version, None);
739 }
740
741 #[test]
742 fn tool_metadata_builder() {
743 let meta = ToolMetadata::with_namespace("filesystem").with_version("1.2.0");
744 assert_eq!(meta.namespace.as_deref(), Some("filesystem"));
745 assert_eq!(meta.version.as_deref(), Some("1.2.0"));
746 }
747
748 #[test]
749 fn agent_tool_metadata_defaults_to_none() {
750 let tool = stub_tool("minimal");
751 assert!(tool.metadata().is_none());
752 }
753
754 #[test]
756 fn agent_tool_auth_config_defaults_to_none() {
757 let tool = stub_tool("no-auth");
758 assert!(tool.auth_config().is_none());
759 }
760
761 #[test]
764 fn approval_context_default_none() {
765 let tool = stub_tool("plain");
766 assert!(tool.approval_context(&json!({})).is_none());
767 }
768
769 #[test]
770 fn approval_context_returns_value() {
771 use crate::FnTool;
772
773 let tool = FnTool::new("ctx", "Ctx", "With context").with_approval_context(|params| {
774 Some(json!({"preview": format!("Will process: {}", params)}))
775 });
776
777 let ctx = tool.approval_context(&json!({"file": "test.txt"}));
778 assert!(ctx.is_some());
779 assert!(
780 ctx.unwrap()["preview"]
781 .as_str()
782 .unwrap()
783 .contains("test.txt")
784 );
785 }
786
787 #[test]
788 fn approval_context_panic_caught() {
789 use crate::FnTool;
790
791 let tool = FnTool::new("panicker", "Panicker", "Panics in context").with_approval_context(
792 |_params| {
793 panic!("oops");
794 },
795 );
796
797 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
798 tool.approval_context(&json!({}))
799 }));
800 assert!(result.is_err());
802 }
803
804 #[test]
805 fn approval_request_includes_context() {
806 let ctx = json!({"diff": "+new line"});
807 let req = ToolApprovalRequest {
808 tool_call_id: "call_1".into(),
809 tool_name: "write_file".into(),
810 arguments: json!({"path": "/tmp/test"}),
811 requires_approval: true,
812 context: Some(ctx.clone()),
813 };
814 assert_eq!(req.context, Some(ctx));
815 }
816
817 #[test]
820 fn transfer_constructor_sets_signal_and_text() {
821 use crate::transfer::TransferSignal;
822
823 let signal = TransferSignal::new("billing", "billing issue");
824 let result = AgentToolResult::transfer(signal);
825
826 assert!(result.is_transfer());
827 assert!(!result.is_error);
828 let text = match &result.content[0] {
829 ContentBlock::Text { text } => text.as_str(),
830 _ => panic!("expected text block"),
831 };
832 assert_eq!(text, "Transfer to billing initiated.");
833 assert!(result.transfer_signal.is_some());
834 let sig = result.transfer_signal.as_ref().unwrap();
835 assert_eq!(sig.target_agent(), "billing");
836 assert_eq!(sig.reason(), "billing issue");
837 }
838
839 #[test]
840 fn text_constructor_has_no_transfer_signal() {
841 let result = AgentToolResult::text("hello");
842 assert!(!result.is_transfer());
843 assert!(result.transfer_signal.is_none());
844 }
845
846 #[test]
847 fn error_constructor_has_no_transfer_signal() {
848 let result = AgentToolResult::error("something failed");
849 assert!(!result.is_transfer());
850 assert!(result.transfer_signal.is_none());
851 }
852
853 #[test]
854 fn deserialize_without_transfer_signal_defaults_to_none() {
855 let json = r#"{
856 "content": [{"type": "text", "text": "hello"}],
857 "details": null,
858 "is_error": false
859 }"#;
860 let result: AgentToolResult = serde_json::from_str(json).unwrap();
861 assert!(!result.is_transfer());
862 assert!(result.transfer_signal.is_none());
863 }
864
865 #[test]
866 fn transfer_signal_not_serialized_when_none() {
867 let result = AgentToolResult::text("hello");
868 let json = serde_json::to_value(&result).unwrap();
869 assert!(!json.as_object().unwrap().contains_key("transfer_signal"));
870 }
871}