1use std::any::Any;
22use std::collections::{BTreeMap, BTreeSet};
23use std::fmt;
24use std::path::{Path, PathBuf};
25use std::sync::atomic::{AtomicU64, Ordering};
26use std::sync::{Arc, Mutex, OnceLock};
27use std::time::Duration;
28
29use agentkit_capabilities::{
30 CapabilityContext, CapabilityError, CapabilityName, CapabilityProvider, Invocable,
31 InvocableOutput, InvocableRequest, InvocableResult, InvocableSpec, PromptProvider,
32 ResourceProvider,
33};
34use agentkit_core::{
35 ApprovalId, Item, ItemKind, MetadataMap, Part, SessionId, TaskId, ToolCallId, ToolOutput,
36 ToolResultPart, TurnCancellation, TurnId,
37};
38use async_trait::async_trait;
39use serde::{Deserialize, Serialize};
40use serde_json::{Value, json};
41use thiserror::Error;
42
43#[doc(hidden)]
47pub mod __private_async_trait {
48 pub use async_trait::async_trait;
49}
50
51#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
69pub struct ToolName(pub String);
70
71impl ToolName {
72 pub fn new(value: impl Into<String>) -> Self {
74 Self(value.into())
75 }
76}
77
78impl fmt::Display for ToolName {
79 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80 self.0.fmt(f)
81 }
82}
83
84impl From<&str> for ToolName {
85 fn from(value: &str) -> Self {
86 Self::new(value)
87 }
88}
89
90#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
97pub struct ToolAnnotations {
98 pub read_only_hint: bool,
100 pub destructive_hint: bool,
102 pub idempotent_hint: bool,
104 pub needs_approval_hint: bool,
106 pub supports_streaming_hint: bool,
108}
109
110impl ToolAnnotations {
111 pub fn new() -> Self {
113 Self::default()
114 }
115
116 pub fn read_only() -> Self {
118 Self::default().with_read_only(true)
119 }
120
121 pub fn destructive() -> Self {
123 Self::default().with_destructive(true)
124 }
125
126 pub fn needs_approval() -> Self {
128 Self::default().with_needs_approval(true)
129 }
130
131 pub fn streaming() -> Self {
133 Self::default().with_supports_streaming(true)
134 }
135
136 pub fn with_read_only(mut self, read_only_hint: bool) -> Self {
137 self.read_only_hint = read_only_hint;
138 self
139 }
140
141 pub fn with_destructive(mut self, destructive_hint: bool) -> Self {
142 self.destructive_hint = destructive_hint;
143 self
144 }
145
146 pub fn with_idempotent(mut self, idempotent_hint: bool) -> Self {
147 self.idempotent_hint = idempotent_hint;
148 self
149 }
150
151 pub fn with_needs_approval(mut self, needs_approval_hint: bool) -> Self {
152 self.needs_approval_hint = needs_approval_hint;
153 self
154 }
155
156 pub fn with_supports_streaming(mut self, supports_streaming_hint: bool) -> Self {
157 self.supports_streaming_hint = supports_streaming_hint;
158 self
159 }
160}
161
162pub const TOOL_OUTPUT_LIMIT_METADATA_KEY: &str = "agentkit.tool_output_limit";
167
168#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
171#[serde(rename_all = "snake_case")]
172pub enum ToolOutputOverflowAction {
173 Fail,
177 InlineClip,
179 StoreForReadback,
183}
184
185#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
187pub struct ToolOutputLimit {
188 pub max_bytes: usize,
190 pub action: ToolOutputOverflowAction,
192}
193
194impl ToolOutputLimit {
195 pub fn fail(max_bytes: usize) -> Self {
197 Self {
198 max_bytes,
199 action: ToolOutputOverflowAction::Fail,
200 }
201 }
202
203 pub fn inline_clip(max_bytes: usize) -> Self {
205 Self {
206 max_bytes,
207 action: ToolOutputOverflowAction::InlineClip,
208 }
209 }
210
211 pub fn store_for_readback(max_bytes: usize) -> Self {
214 Self {
215 max_bytes,
216 action: ToolOutputOverflowAction::StoreForReadback,
217 }
218 }
219
220 fn to_metadata_value(&self) -> Value {
221 serde_json::to_value(self).expect("ToolOutputLimit serializes")
222 }
223
224 fn from_metadata(metadata: &MetadataMap) -> Option<Self> {
225 metadata
226 .get(TOOL_OUTPUT_LIMIT_METADATA_KEY)
227 .and_then(|value| serde_json::from_value(value.clone()).ok())
228 }
229}
230
231#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
257pub struct ToolSpec {
258 pub name: ToolName,
260 pub description: String,
262 pub input_schema: Value,
264 #[serde(skip_serializing_if = "Option::is_none", default)]
274 pub output_schema: Option<Value>,
275 pub annotations: ToolAnnotations,
277 pub metadata: MetadataMap,
279}
280
281#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
287pub struct ToolCatalogEvent {
288 pub source: String,
290 pub added: Vec<String>,
292 pub removed: Vec<String>,
294 pub changed: Vec<String>,
296}
297
298impl ToolCatalogEvent {
299 pub fn new(source: impl Into<String>) -> Self {
301 Self {
302 source: source.into(),
303 added: Vec::new(),
304 removed: Vec::new(),
305 changed: Vec::new(),
306 }
307 }
308
309 pub fn for_each_name_mut(&mut self, mut f: impl FnMut(&mut String)) {
311 for vec in [&mut self.added, &mut self.removed, &mut self.changed] {
312 for name in vec.iter_mut() {
313 f(name);
314 }
315 }
316 }
317
318 pub fn retain_names(&mut self, mut predicate: impl FnMut(&str) -> bool) {
321 self.added.retain(|n| predicate(n));
322 self.removed.retain(|n| predicate(n));
323 self.changed.retain(|n| predicate(n));
324 }
325}
326
327impl ToolSpec {
328 pub fn new(
330 name: impl Into<ToolName>,
331 description: impl Into<String>,
332 input_schema: Value,
333 ) -> Self {
334 Self {
335 name: name.into(),
336 description: description.into(),
337 input_schema,
338 output_schema: None,
339 annotations: ToolAnnotations::default(),
340 metadata: MetadataMap::new(),
341 }
342 }
343
344 pub fn with_output_schema(mut self, schema: Value) -> Self {
347 self.output_schema = Some(schema);
348 self
349 }
350
351 pub fn with_annotations(mut self, annotations: ToolAnnotations) -> Self {
353 self.annotations = annotations;
354 self
355 }
356
357 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
359 self.metadata = metadata;
360 self
361 }
362
363 pub fn with_output_limit(mut self, limit: ToolOutputLimit) -> Self {
369 self.metadata.insert(
370 TOOL_OUTPUT_LIMIT_METADATA_KEY.to_string(),
371 limit.to_metadata_value(),
372 );
373 self
374 }
375}
376
377#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
383pub struct ToolRequest {
384 pub call_id: ToolCallId,
386 pub tool_name: ToolName,
388 pub input: Value,
390 pub session_id: SessionId,
392 pub turn_id: TurnId,
394 pub metadata: MetadataMap,
396}
397
398impl ToolRequest {
399 pub fn new(
401 call_id: impl Into<ToolCallId>,
402 tool_name: impl Into<ToolName>,
403 input: Value,
404 session_id: impl Into<SessionId>,
405 turn_id: impl Into<TurnId>,
406 ) -> Self {
407 Self {
408 call_id: call_id.into(),
409 tool_name: tool_name.into(),
410 input,
411 session_id: session_id.into(),
412 turn_id: turn_id.into(),
413 metadata: MetadataMap::new(),
414 }
415 }
416
417 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
419 self.metadata = metadata;
420 self
421 }
422}
423
424#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
429pub struct ToolResult {
430 pub result: ToolResultPart,
432 pub duration: Option<Duration>,
434 pub metadata: MetadataMap,
436}
437
438impl ToolResult {
439 pub fn new(result: ToolResultPart) -> Self {
441 Self {
442 result,
443 duration: None,
444 metadata: MetadataMap::new(),
445 }
446 }
447
448 pub fn with_duration(mut self, duration: Duration) -> Self {
450 self.duration = Some(duration);
451 self
452 }
453
454 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
456 self.metadata = metadata;
457 self
458 }
459}
460
461pub trait ToolResources: Send + Sync {
487 fn as_any(&self) -> &dyn Any;
490}
491
492impl ToolResources for () {
493 fn as_any(&self) -> &dyn Any {
494 self
495 }
496}
497
498pub struct ToolContext<'a> {
504 pub capability: CapabilityContext<'a>,
506 pub permissions: &'a dyn PermissionChecker,
508 pub resources: &'a dyn ToolResources,
510 pub cancellation: Option<TurnCancellation>,
512 pub execution_scope: Option<ToolExecutionScope>,
515 pub approved_request: Option<ApprovalRequest>,
518}
519
520#[derive(Clone)]
526pub struct ToolExecutionScope {
527 pub executor: Arc<dyn ToolExecutor>,
528 pub session_id: SessionId,
529 pub turn_id: TurnId,
530 pub permissions: Arc<dyn PermissionChecker>,
531 pub resources: Arc<dyn ToolResources>,
532 pub cancellation: Option<TurnCancellation>,
533}
534
535impl ToolExecutionScope {
536 pub fn nested_context(&self, metadata: MetadataMap) -> OwnedToolContext {
538 OwnedToolContext {
539 session_id: self.session_id.clone(),
540 turn_id: self.turn_id.clone(),
541 metadata,
542 permissions: self.permissions.clone(),
543 resources: self.resources.clone(),
544 cancellation: self.cancellation.clone(),
545 execution_scope: Some(self.clone()),
546 approved_request: None,
547 }
548 }
549
550 pub async fn execute_child(&self, request: ToolRequest) -> ToolExecutionOutcome {
552 let ctx = self.nested_context(request.metadata.clone());
553 self.executor.execute_owned(request, ctx).await
554 }
555
556 pub async fn execute_approved_child(
559 &self,
560 request: ToolRequest,
561 approval: &ApprovalRequest,
562 ) -> ToolExecutionOutcome {
563 let ctx = self.nested_context(request.metadata.clone());
564 self.executor
565 .execute_approved_owned(request, approval, ctx)
566 .await
567 }
568}
569
570#[derive(Clone)]
576pub struct OwnedToolContext {
577 pub session_id: SessionId,
579 pub turn_id: TurnId,
581 pub metadata: MetadataMap,
583 pub permissions: Arc<dyn PermissionChecker>,
585 pub resources: Arc<dyn ToolResources>,
587 pub cancellation: Option<TurnCancellation>,
589 pub execution_scope: Option<ToolExecutionScope>,
591 pub approved_request: Option<ApprovalRequest>,
593}
594
595impl OwnedToolContext {
596 pub fn borrowed(&self) -> ToolContext<'_> {
598 ToolContext {
599 capability: CapabilityContext {
600 session_id: Some(&self.session_id),
601 turn_id: Some(&self.turn_id),
602 metadata: &self.metadata,
603 },
604 permissions: self.permissions.as_ref(),
605 resources: self.resources.as_ref(),
606 cancellation: self.cancellation.clone(),
607 execution_scope: self.execution_scope.clone(),
608 approved_request: self.approved_request.clone(),
609 }
610 }
611}
612
613#[derive(Clone, Debug)]
616pub struct ToolOutputTruncationContext {
617 pub tool_name: ToolName,
618 pub call_id: ToolCallId,
619 pub session_id: SessionId,
620 pub turn_id: TurnId,
621 pub tool_spec: ToolSpec,
622}
623
624impl From<(&ToolRequest, ToolSpec)> for ToolOutputTruncationContext {
625 fn from((request, tool_spec): (&ToolRequest, ToolSpec)) -> Self {
626 Self {
627 tool_name: request.tool_name.clone(),
628 call_id: request.call_id.clone(),
629 session_id: request.session_id.clone(),
630 turn_id: request.turn_id.clone(),
631 tool_spec,
632 }
633 }
634}
635
636#[async_trait]
641pub trait ToolOutputTruncationStrategy: Send + Sync {
642 async fn apply(
643 &self,
644 ctx: ToolOutputTruncationContext,
645 output: ToolOutput,
646 ) -> Result<ToolOutput, ToolError>;
647}
648
649#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
651pub struct ToolOutputArtifactId(pub String);
652
653impl fmt::Display for ToolOutputArtifactId {
654 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
655 self.0.fmt(f)
656 }
657}
658
659#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
661pub struct ToolOutputArtifact {
662 pub id: ToolOutputArtifactId,
663 pub tool_name: ToolName,
664 pub call_id: ToolCallId,
665 pub session_id: SessionId,
666 pub turn_id: TurnId,
667 pub original_bytes: usize,
668 pub body: String,
669}
670
671#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
673pub struct ToolOutputArtifactSlice {
674 pub id: ToolOutputArtifactId,
675 pub offset: usize,
676 pub next_offset: usize,
677 pub original_bytes: usize,
678 pub eof: bool,
679 pub content: String,
680}
681
682#[async_trait]
683pub trait ToolOutputArtifactStore: Send + Sync {
684 async fn put(
685 &self,
686 ctx: &ToolOutputTruncationContext,
687 body: String,
688 original_bytes: usize,
689 ) -> Result<ToolOutputArtifact, ToolError>;
690
691 async fn read(
692 &self,
693 id: &ToolOutputArtifactId,
694 offset: usize,
695 max_bytes: usize,
696 ) -> Result<ToolOutputArtifactSlice, ToolError>;
697}
698
699#[derive(Debug, Default)]
701pub struct InMemoryToolOutputArtifactStore {
702 next_id: AtomicU64,
703 artifacts: Mutex<BTreeMap<ToolOutputArtifactId, ToolOutputArtifact>>,
704}
705
706impl InMemoryToolOutputArtifactStore {
707 pub fn new() -> Self {
708 Self::default()
709 }
710}
711
712#[async_trait]
713impl ToolOutputArtifactStore for InMemoryToolOutputArtifactStore {
714 async fn put(
715 &self,
716 ctx: &ToolOutputTruncationContext,
717 body: String,
718 original_bytes: usize,
719 ) -> Result<ToolOutputArtifact, ToolError> {
720 let n = self.next_id.fetch_add(1, Ordering::Relaxed);
721 let id = ToolOutputArtifactId(format!(
722 "{}:{}:{}",
723 sanitize_artifact_id_component(ctx.session_id.0.as_str()),
724 sanitize_artifact_id_component(ctx.call_id.0.as_str()),
725 n
726 ));
727 let artifact = ToolOutputArtifact {
728 id: id.clone(),
729 tool_name: ctx.tool_name.clone(),
730 call_id: ctx.call_id.clone(),
731 session_id: ctx.session_id.clone(),
732 turn_id: ctx.turn_id.clone(),
733 original_bytes,
734 body,
735 };
736 self.artifacts
737 .lock()
738 .unwrap_or_else(|e| e.into_inner())
739 .insert(id, artifact.clone());
740 Ok(artifact)
741 }
742
743 async fn read(
744 &self,
745 id: &ToolOutputArtifactId,
746 offset: usize,
747 max_bytes: usize,
748 ) -> Result<ToolOutputArtifactSlice, ToolError> {
749 let artifact = self
750 .artifacts
751 .lock()
752 .unwrap_or_else(|e| e.into_inner())
753 .get(id)
754 .cloned()
755 .ok_or_else(|| {
756 ToolError::InvalidInput(format!("unknown tool result artifact: {id}"))
757 })?;
758 let body = artifact.body;
759 if offset > body.len() || !body.is_char_boundary(offset) {
760 return Err(ToolError::InvalidInput(format!(
761 "offset {offset} is not a UTF-8 boundary in tool result artifact {id}"
762 )));
763 }
764 let requested_end = offset.saturating_add(max_bytes).min(body.len());
765 let end = body.floor_char_boundary(requested_end);
766 Ok(ToolOutputArtifactSlice {
767 id: id.clone(),
768 offset,
769 next_offset: end,
770 original_bytes: artifact.original_bytes,
771 eof: end == body.len(),
772 content: body[offset..end].to_string(),
773 })
774 }
775}
776
777fn sanitize_artifact_id_component(s: &str) -> String {
778 let cleaned: String = s
779 .chars()
780 .map(|c| {
781 if c.is_ascii_alphanumeric() || c == '-' || c == '_' {
782 c
783 } else {
784 '_'
785 }
786 })
787 .take(64)
788 .collect();
789 if cleaned.is_empty() {
790 "_".to_string()
791 } else {
792 cleaned
793 }
794}
795
796pub struct ConfigurableToolOutputTruncationStrategy {
799 default_limit: Option<ToolOutputLimit>,
800 per_tool_limits: BTreeMap<ToolName, ToolOutputLimit>,
801 use_tool_metadata: bool,
802 store: Arc<dyn ToolOutputArtifactStore>,
803}
804
805impl ConfigurableToolOutputTruncationStrategy {
806 pub fn new(store: Arc<dyn ToolOutputArtifactStore>) -> Self {
807 Self {
808 default_limit: None,
809 per_tool_limits: BTreeMap::new(),
810 use_tool_metadata: true,
811 store,
812 }
813 }
814
815 pub fn with_default_limit(mut self, limit: ToolOutputLimit) -> Self {
816 self.default_limit = Some(limit);
817 self
818 }
819
820 pub fn with_tool_limit(
821 mut self,
822 tool_name: impl Into<ToolName>,
823 limit: ToolOutputLimit,
824 ) -> Self {
825 self.per_tool_limits.insert(tool_name.into(), limit);
826 self
827 }
828
829 pub fn use_tool_metadata(mut self, value: bool) -> Self {
830 self.use_tool_metadata = value;
831 self
832 }
833
834 fn limit_for(&self, ctx: &ToolOutputTruncationContext) -> Option<ToolOutputLimit> {
835 self.per_tool_limits
836 .get(&ctx.tool_name)
837 .cloned()
838 .or_else(|| {
839 self.use_tool_metadata
840 .then(|| ToolOutputLimit::from_metadata(&ctx.tool_spec.metadata))
841 .flatten()
842 })
843 .or_else(|| self.default_limit.clone())
844 }
845}
846
847#[async_trait]
848impl ToolOutputTruncationStrategy for ConfigurableToolOutputTruncationStrategy {
849 async fn apply(
850 &self,
851 ctx: ToolOutputTruncationContext,
852 output: ToolOutput,
853 ) -> Result<ToolOutput, ToolError> {
854 let Some(limit) = self.limit_for(&ctx) else {
855 return Ok(output);
856 };
857 let model_bytes = tool_output_model_bytes(&output);
858 if model_bytes <= limit.max_bytes {
859 return Ok(output);
860 }
861
862 match limit.action {
863 ToolOutputOverflowAction::Fail => Err(ToolError::ExecutionFailed(format!(
864 "tool {} produced {model_bytes} bytes, exceeding configured limit of {} bytes",
865 ctx.tool_name, limit.max_bytes
866 ))),
867 ToolOutputOverflowAction::InlineClip => Ok(clip_tool_output_inline(
868 output,
869 limit.max_bytes,
870 model_bytes,
871 )),
872 ToolOutputOverflowAction::StoreForReadback => {
873 let body = tool_output_readback_body(&output);
874 let artifact = self.store.put(&ctx, body, model_bytes).await?;
875 Ok(fit_structured_tool_output(
876 json!({
877 "truncated": true,
878 "tool_result_id": artifact.id.0,
879 "read_tool": TOOL_RESULT_READ_TOOL_NAME,
880 "read_args": {
881 "id": artifact.id.0,
882 "offset": 0,
883 "limit": limit.max_bytes
884 },
885 "original_bytes": artifact.original_bytes,
886 }),
887 limit.max_bytes,
888 ))
889 }
890 }
891 }
892}
893
894fn tool_output_model_bytes(output: &ToolOutput) -> usize {
895 match output {
896 ToolOutput::Text(s) => s.len(),
897 other => serde_json::to_string(other)
898 .map(|s| s.len())
899 .unwrap_or(usize::MAX),
900 }
901}
902
903fn tool_output_readback_body(output: &ToolOutput) -> String {
904 match output {
905 ToolOutput::Text(s) => s.clone(),
906 ToolOutput::Structured(value) => {
907 serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string())
908 }
909 ToolOutput::Parts(parts) => serde_json::to_string_pretty(parts).unwrap_or_default(),
910 ToolOutput::Files(files) => serde_json::to_string_pretty(files).unwrap_or_default(),
911 }
912}
913
914fn clip_tool_output_inline(
915 output: ToolOutput,
916 max_bytes: usize,
917 original_bytes: usize,
918) -> ToolOutput {
919 match output {
920 ToolOutput::Text(s) => {
921 ToolOutput::Text(clip_string_with_marker(&s, max_bytes, original_bytes))
922 }
923 other => {
924 let body = tool_output_readback_body(&other);
925 fit_structured_tool_output(
926 json!({
927 "truncated": true,
928 "original_bytes": original_bytes,
929 "content": body,
930 }),
931 max_bytes,
932 )
933 }
934 }
935}
936
937fn clip_string_with_marker(s: &str, max_bytes: usize, original_bytes: usize) -> String {
938 let marker = format!("\n[tool output truncated: original_bytes={original_bytes}]");
939 if marker.len() >= max_bytes {
940 let cut = marker.floor_char_boundary(max_bytes.min(marker.len()));
941 return marker[..cut].to_string();
942 }
943 let keep_bytes = max_bytes.saturating_sub(marker.len());
944 let cut = s.floor_char_boundary(keep_bytes.min(s.len()));
945 format!("{}{}", &s[..cut], marker)
946}
947
948fn fit_structured_tool_output(mut value: Value, max_bytes: usize) -> ToolOutput {
949 loop {
950 let output = ToolOutput::Structured(value.clone());
951 if tool_output_model_bytes(&output) <= max_bytes {
952 return output;
953 }
954
955 let Some(Value::String(content)) = value.get_mut("content") else {
956 return ToolOutput::Structured(json!({
957 "truncated": true,
958 "error": "tool output metadata exceeded configured max_bytes"
959 }));
960 };
961 if content.is_empty() {
962 return ToolOutput::Structured(json!({
963 "truncated": true,
964 "error": "tool output metadata exceeded configured max_bytes"
965 }));
966 }
967
968 let current_len = content.len();
969 let shrink_by = tool_output_model_bytes(&output)
970 .saturating_sub(max_bytes)
971 .saturating_add(32)
972 .min(current_len);
973 let new_len = content.floor_char_boundary(current_len - shrink_by);
974 content.truncate(new_len);
975 }
976}
977
978pub const TOOL_RESULT_READ_TOOL_NAME: &str = "tool_result_read";
979const TOOL_RESULT_READ_OUTPUT_ENVELOPE_BYTES: usize = 4096;
980const TOOL_RESULT_READ_JSON_ESCAPE_BYTES_PER_INPUT_BYTE: usize = 6;
981
982#[derive(Clone)]
985pub struct ToolResultReadTool {
986 spec: ToolSpec,
987 store: Arc<dyn ToolOutputArtifactStore>,
988 max_read_bytes: usize,
989}
990
991impl ToolResultReadTool {
992 pub fn new(store: Arc<dyn ToolOutputArtifactStore>, max_read_bytes: usize) -> Self {
993 Self {
994 spec: ToolSpec::new(
995 TOOL_RESULT_READ_TOOL_NAME,
996 "Read a bounded UTF-8 byte slice from a stored oversized tool result.",
997 json!({
998 "type": "object",
999 "properties": {
1000 "id": { "type": "string" },
1001 "offset": { "type": "integer", "minimum": 0 },
1002 "limit": { "type": "integer", "minimum": 1 }
1003 },
1004 "required": ["id", "offset", "limit"],
1005 "additionalProperties": false
1006 }),
1007 )
1008 .with_annotations(ToolAnnotations {
1009 read_only_hint: true,
1010 idempotent_hint: true,
1011 ..ToolAnnotations::default()
1012 })
1013 .with_output_limit(ToolOutputLimit::fail(
1014 max_read_bytes
1015 .saturating_mul(TOOL_RESULT_READ_JSON_ESCAPE_BYTES_PER_INPUT_BYTE)
1016 .saturating_add(TOOL_RESULT_READ_OUTPUT_ENVELOPE_BYTES),
1017 )),
1018 store,
1019 max_read_bytes,
1020 }
1021 }
1022}
1023
1024#[derive(Deserialize)]
1025struct ToolResultReadInput {
1026 id: String,
1027 offset: usize,
1028 limit: usize,
1029}
1030
1031#[async_trait]
1032impl Tool for ToolResultReadTool {
1033 fn spec(&self) -> &ToolSpec {
1034 &self.spec
1035 }
1036
1037 async fn invoke(
1038 &self,
1039 request: ToolRequest,
1040 _ctx: &mut ToolContext<'_>,
1041 ) -> Result<ToolResult, ToolError> {
1042 let input: ToolResultReadInput = serde_json::from_value(request.input.clone())
1043 .map_err(|error| ToolError::InvalidInput(format!("invalid tool input: {error}")))?;
1044 if input.limit == 0 {
1045 return Err(ToolError::InvalidInput(
1046 "limit must be greater than 0".to_string(),
1047 ));
1048 }
1049 if input.limit > self.max_read_bytes {
1050 return Err(ToolError::InvalidInput(format!(
1051 "limit {} exceeds maximum read size of {} bytes",
1052 input.limit, self.max_read_bytes
1053 )));
1054 }
1055 let slice = self
1056 .store
1057 .read(&ToolOutputArtifactId(input.id), input.offset, input.limit)
1058 .await?;
1059 Ok(ToolResult::new(ToolResultPart::success(
1060 request.call_id,
1061 ToolOutput::Structured(json!({
1062 "id": slice.id.0,
1063 "offset": slice.offset,
1064 "next_offset": slice.next_offset,
1065 "original_bytes": slice.original_bytes,
1066 "eof": slice.eof,
1067 "content": slice.content,
1068 })),
1069 )))
1070 }
1071}
1072
1073pub fn tool_result_readback_registry(
1075 store: Arc<dyn ToolOutputArtifactStore>,
1076 max_read_bytes: usize,
1077) -> ToolRegistry {
1078 ToolRegistry::new().with(ToolResultReadTool::new(store, max_read_bytes))
1079}
1080
1081pub trait PermissionRequest: Send + Sync {
1110 fn kind(&self) -> &'static str;
1112 fn summary(&self) -> String;
1114 fn metadata(&self) -> &MetadataMap;
1116 fn as_any(&self) -> &dyn Any;
1118}
1119
1120#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1125pub enum PermissionCode {
1126 PathNotAllowed,
1128 CommandNotAllowed,
1130 NetworkNotAllowed,
1132 ServerNotTrusted,
1134 AuthScopeNotAllowed,
1136 CustomPolicyDenied,
1138 UnknownRequest,
1140}
1141
1142#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1147pub struct PermissionDenial {
1148 pub code: PermissionCode,
1150 pub message: String,
1152 pub metadata: MetadataMap,
1154}
1155
1156#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1161pub enum ApprovalReason {
1162 PolicyRequiresConfirmation,
1164 EscalatedRisk,
1166 UnknownTarget,
1168 SensitivePath,
1170 SensitiveCommand,
1172 SensitiveServer,
1174 SensitiveAuthScope,
1176}
1177
1178#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1183pub struct ApprovalRequest {
1184 pub task_id: Option<TaskId>,
1186 pub call_id: Option<ToolCallId>,
1189 pub id: ApprovalId,
1191 pub request_kind: String,
1193 pub reason: ApprovalReason,
1195 pub summary: String,
1197 pub metadata: MetadataMap,
1199}
1200
1201impl ApprovalRequest {
1202 pub fn new(
1204 id: impl Into<ApprovalId>,
1205 request_kind: impl Into<String>,
1206 reason: ApprovalReason,
1207 summary: impl Into<String>,
1208 ) -> Self {
1209 Self {
1210 task_id: None,
1211 call_id: None,
1212 id: id.into(),
1213 request_kind: request_kind.into(),
1214 reason,
1215 summary: summary.into(),
1216 metadata: MetadataMap::new(),
1217 }
1218 }
1219
1220 pub fn with_task_id(mut self, task_id: impl Into<TaskId>) -> Self {
1222 self.task_id = Some(task_id.into());
1223 self
1224 }
1225
1226 pub fn with_call_id(mut self, call_id: impl Into<ToolCallId>) -> Self {
1228 self.call_id = Some(call_id.into());
1229 self
1230 }
1231
1232 pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
1234 self.metadata = metadata;
1235 self
1236 }
1237}
1238
1239#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1241pub enum ApprovalDecision {
1242 Approve,
1244 Deny {
1246 reason: Option<String>,
1248 },
1249}
1250
1251#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1258pub enum ToolInterruption {
1259 ApprovalRequired(ApprovalRequest),
1261}
1262
1263#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1265pub enum PermissionDecision {
1266 Allow,
1268 Deny(PermissionDenial),
1270 RequireApproval(ApprovalRequest),
1272}
1273
1274pub trait PermissionChecker: Send + Sync {
1298 fn evaluate(&self, request: &dyn PermissionRequest) -> PermissionDecision;
1300}
1301
1302#[derive(Copy, Clone, Debug, Default)]
1307pub struct AllowAllPermissions;
1308
1309impl PermissionChecker for AllowAllPermissions {
1310 fn evaluate(&self, _request: &dyn PermissionRequest) -> PermissionDecision {
1311 PermissionDecision::Allow
1312 }
1313}
1314
1315#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1321pub enum PolicyMatch {
1322 NoOpinion,
1324 Allow,
1326 Deny(PermissionDenial),
1328 RequireApproval(ApprovalRequest),
1330}
1331
1332pub trait PermissionPolicy: Send + Sync {
1341 fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch;
1343}
1344
1345pub struct CompositePermissionChecker {
1365 policies: Vec<Box<dyn PermissionPolicy>>,
1366 fallback: PermissionDecision,
1367}
1368
1369impl CompositePermissionChecker {
1370 pub fn new(fallback: PermissionDecision) -> Self {
1378 Self {
1379 policies: Vec::new(),
1380 fallback,
1381 }
1382 }
1383
1384 pub fn with_policy(mut self, policy: impl PermissionPolicy + 'static) -> Self {
1386 self.policies.push(Box::new(policy));
1387 self
1388 }
1389}
1390
1391impl PermissionChecker for CompositePermissionChecker {
1392 fn evaluate(&self, request: &dyn PermissionRequest) -> PermissionDecision {
1393 let mut saw_allow = false;
1394 let mut approval = None;
1395
1396 for policy in &self.policies {
1397 match policy.evaluate(request) {
1398 PolicyMatch::NoOpinion => {}
1399 PolicyMatch::Allow => saw_allow = true,
1400 PolicyMatch::Deny(denial) => return PermissionDecision::Deny(denial),
1401 PolicyMatch::RequireApproval(req) => approval = Some(req),
1402 }
1403 }
1404
1405 if let Some(req) = approval {
1406 PermissionDecision::RequireApproval(req)
1407 } else if saw_allow {
1408 PermissionDecision::Allow
1409 } else {
1410 self.fallback.clone()
1411 }
1412 }
1413}
1414
1415#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1420pub struct ShellPermissionRequest {
1421 pub executable: String,
1423 pub argv: Vec<String>,
1425 pub cwd: Option<PathBuf>,
1427 pub env_keys: Vec<String>,
1429 pub metadata: MetadataMap,
1431}
1432
1433impl PermissionRequest for ShellPermissionRequest {
1434 fn kind(&self) -> &'static str {
1435 "shell.command"
1436 }
1437
1438 fn summary(&self) -> String {
1439 if self.argv.is_empty() {
1440 self.executable.clone()
1441 } else {
1442 format!("{} {}", self.executable, self.argv.join(" "))
1443 }
1444 }
1445
1446 fn metadata(&self) -> &MetadataMap {
1447 &self.metadata
1448 }
1449
1450 fn as_any(&self) -> &dyn Any {
1451 self
1452 }
1453}
1454
1455#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1460pub enum FileSystemPermissionRequest {
1461 Read {
1463 path: PathBuf,
1464 metadata: MetadataMap,
1465 },
1466 Write {
1468 path: PathBuf,
1469 metadata: MetadataMap,
1470 },
1471 Edit {
1473 path: PathBuf,
1474 metadata: MetadataMap,
1475 },
1476 Delete {
1478 path: PathBuf,
1479 metadata: MetadataMap,
1480 },
1481 Move {
1483 from: PathBuf,
1484 to: PathBuf,
1485 metadata: MetadataMap,
1486 },
1487 List {
1489 path: PathBuf,
1490 metadata: MetadataMap,
1491 },
1492 CreateDir {
1494 path: PathBuf,
1495 metadata: MetadataMap,
1496 },
1497}
1498
1499impl FileSystemPermissionRequest {
1500 fn metadata_map(&self) -> &MetadataMap {
1501 match self {
1502 Self::Read { metadata, .. }
1503 | Self::Write { metadata, .. }
1504 | Self::Edit { metadata, .. }
1505 | Self::Delete { metadata, .. }
1506 | Self::Move { metadata, .. }
1507 | Self::List { metadata, .. }
1508 | Self::CreateDir { metadata, .. } => metadata,
1509 }
1510 }
1511}
1512
1513impl PermissionRequest for FileSystemPermissionRequest {
1514 fn kind(&self) -> &'static str {
1515 match self {
1516 Self::Read { .. } => "filesystem.read",
1517 Self::Write { .. } => "filesystem.write",
1518 Self::Edit { .. } => "filesystem.edit",
1519 Self::Delete { .. } => "filesystem.delete",
1520 Self::Move { .. } => "filesystem.move",
1521 Self::List { .. } => "filesystem.list",
1522 Self::CreateDir { .. } => "filesystem.mkdir",
1523 }
1524 }
1525
1526 fn summary(&self) -> String {
1527 match self {
1528 Self::Read { path, .. } => format!("Read {}", path.display()),
1529 Self::Write { path, .. } => format!("Write {}", path.display()),
1530 Self::Edit { path, .. } => format!("Edit {}", path.display()),
1531 Self::Delete { path, .. } => format!("Delete {}", path.display()),
1532 Self::Move { from, to, .. } => {
1533 format!("Move {} to {}", from.display(), to.display())
1534 }
1535 Self::List { path, .. } => format!("List {}", path.display()),
1536 Self::CreateDir { path, .. } => format!("Create directory {}", path.display()),
1537 }
1538 }
1539
1540 fn metadata(&self) -> &MetadataMap {
1541 self.metadata_map()
1542 }
1543
1544 fn as_any(&self) -> &dyn Any {
1545 self
1546 }
1547}
1548
1549#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1554pub enum McpPermissionRequest {
1555 Connect {
1557 server_id: String,
1558 metadata: MetadataMap,
1559 },
1560 InvokeTool {
1562 server_id: String,
1563 tool_name: String,
1564 metadata: MetadataMap,
1565 },
1566 ReadResource {
1568 server_id: String,
1569 resource_id: String,
1570 metadata: MetadataMap,
1571 },
1572 FetchPrompt {
1574 server_id: String,
1575 prompt_id: String,
1576 metadata: MetadataMap,
1577 },
1578 UseAuthScope {
1580 server_id: String,
1581 scope: String,
1582 metadata: MetadataMap,
1583 },
1584}
1585
1586impl McpPermissionRequest {
1587 fn metadata_map(&self) -> &MetadataMap {
1588 match self {
1589 Self::Connect { metadata, .. }
1590 | Self::InvokeTool { metadata, .. }
1591 | Self::ReadResource { metadata, .. }
1592 | Self::FetchPrompt { metadata, .. }
1593 | Self::UseAuthScope { metadata, .. } => metadata,
1594 }
1595 }
1596}
1597
1598impl PermissionRequest for McpPermissionRequest {
1599 fn kind(&self) -> &'static str {
1600 match self {
1601 Self::Connect { .. } => "mcp.connect",
1602 Self::InvokeTool { .. } => "mcp.invoke_tool",
1603 Self::ReadResource { .. } => "mcp.read_resource",
1604 Self::FetchPrompt { .. } => "mcp.fetch_prompt",
1605 Self::UseAuthScope { .. } => "mcp.use_auth_scope",
1606 }
1607 }
1608
1609 fn summary(&self) -> String {
1610 match self {
1611 Self::Connect { server_id, .. } => format!("Connect MCP server {server_id}"),
1612 Self::InvokeTool {
1613 server_id,
1614 tool_name,
1615 ..
1616 } => format!("Invoke MCP tool {server_id}.{tool_name}"),
1617 Self::ReadResource {
1618 server_id,
1619 resource_id,
1620 ..
1621 } => format!("Read MCP resource {server_id}:{resource_id}"),
1622 Self::FetchPrompt {
1623 server_id,
1624 prompt_id,
1625 ..
1626 } => format!("Fetch MCP prompt {server_id}:{prompt_id}"),
1627 Self::UseAuthScope {
1628 server_id, scope, ..
1629 } => format!("Use MCP auth scope {server_id}:{scope}"),
1630 }
1631 }
1632
1633 fn metadata(&self) -> &MetadataMap {
1634 self.metadata_map()
1635 }
1636
1637 fn as_any(&self) -> &dyn Any {
1638 self
1639 }
1640}
1641
1642pub struct CustomKindPolicy {
1658 allowed_kinds: BTreeSet<String>,
1659 denied_kinds: BTreeSet<String>,
1660 require_approval_by_default: bool,
1661}
1662
1663impl CustomKindPolicy {
1664 pub fn new(require_approval_by_default: bool) -> Self {
1671 Self {
1672 allowed_kinds: BTreeSet::new(),
1673 denied_kinds: BTreeSet::new(),
1674 require_approval_by_default,
1675 }
1676 }
1677
1678 pub fn allow_kind(mut self, kind: impl Into<String>) -> Self {
1680 self.allowed_kinds.insert(kind.into());
1681 self
1682 }
1683
1684 pub fn deny_kind(mut self, kind: impl Into<String>) -> Self {
1686 self.denied_kinds.insert(kind.into());
1687 self
1688 }
1689}
1690
1691impl PermissionPolicy for CustomKindPolicy {
1692 fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
1693 let kind = request.kind();
1694 if !kind.starts_with("custom.") {
1695 return PolicyMatch::NoOpinion;
1696 }
1697 if self.denied_kinds.contains(kind) {
1698 return PolicyMatch::Deny(PermissionDenial {
1699 code: PermissionCode::CustomPolicyDenied,
1700 message: format!("custom permission kind {kind} is denied"),
1701 metadata: request.metadata().clone(),
1702 });
1703 }
1704 if self.allowed_kinds.contains(kind) {
1705 return PolicyMatch::Allow;
1706 }
1707 if self.require_approval_by_default {
1708 PolicyMatch::RequireApproval(ApprovalRequest {
1709 task_id: None,
1710 call_id: None,
1711 id: ApprovalId::new(format!("approval:{kind}")),
1712 request_kind: kind.to_string(),
1713 reason: ApprovalReason::PolicyRequiresConfirmation,
1714 summary: request.summary(),
1715 metadata: request.metadata().clone(),
1716 })
1717 } else {
1718 PolicyMatch::NoOpinion
1719 }
1720 }
1721}
1722
1723pub struct PathPolicy {
1743 allowed_roots: Vec<CanonicalRoot>,
1744 read_only_roots: Vec<CanonicalRoot>,
1745 protected_roots: Vec<CanonicalRoot>,
1746 require_approval_outside_allowed: bool,
1747}
1748
1749impl PathPolicy {
1750 pub fn new() -> Self {
1753 Self {
1754 allowed_roots: Vec::new(),
1755 read_only_roots: Vec::new(),
1756 protected_roots: Vec::new(),
1757 require_approval_outside_allowed: true,
1758 }
1759 }
1760
1761 pub fn allow_root(mut self, root: impl Into<PathBuf>) -> Self {
1763 self.allowed_roots.push(CanonicalRoot::new(root.into()));
1764 self
1765 }
1766
1767 pub fn read_only_root(mut self, root: impl Into<PathBuf>) -> Self {
1769 self.read_only_roots.push(CanonicalRoot::new(root.into()));
1770 self
1771 }
1772
1773 pub fn protect_root(mut self, root: impl Into<PathBuf>) -> Self {
1775 self.protected_roots.push(CanonicalRoot::new(root.into()));
1776 self
1777 }
1778
1779 pub fn require_approval_outside_allowed(mut self, value: bool) -> Self {
1782 self.require_approval_outside_allowed = value;
1783 self
1784 }
1785}
1786
1787impl Default for PathPolicy {
1788 fn default() -> Self {
1789 Self::new()
1790 }
1791}
1792
1793fn resolve_canonical(path: &Path) -> PathBuf {
1797 let abs = std::path::absolute(path).unwrap_or_else(|_| path.to_path_buf());
1798 canonicalize_with_partial_fallback(&abs).unwrap_or(abs)
1799}
1800
1801fn canonicalize_with_partial_fallback(abs: &Path) -> Option<PathBuf> {
1802 if let Ok(canonical) = std::fs::canonicalize(abs) {
1803 return Some(canonical);
1804 }
1805 let mut tail: Vec<std::ffi::OsString> = Vec::new();
1806 let mut current = abs.to_path_buf();
1807 loop {
1808 let name = current.file_name().map(|n| n.to_os_string())?;
1809 tail.push(name);
1810 if !current.pop() {
1811 return None;
1812 }
1813 if let Ok(canonical) = std::fs::canonicalize(¤t) {
1814 let mut out = canonical;
1815 for seg in tail.iter().rev() {
1816 out.push(seg);
1817 }
1818 return Some(out);
1819 }
1820 }
1821}
1822
1823struct CanonicalRoot {
1829 lexical: PathBuf,
1830 canonical: OnceLock<PathBuf>,
1831}
1832
1833impl CanonicalRoot {
1834 fn new(lexical: PathBuf) -> Self {
1835 Self {
1836 lexical,
1837 canonical: OnceLock::new(),
1838 }
1839 }
1840
1841 fn resolve(&self) -> std::borrow::Cow<'_, Path> {
1842 if let Some(canonical) = self.canonical.get() {
1843 return std::borrow::Cow::Borrowed(canonical);
1844 }
1845 let abs = std::path::absolute(&self.lexical).unwrap_or_else(|_| self.lexical.clone());
1846 if let Ok(canonical) = std::fs::canonicalize(&abs) {
1847 let _ = self.canonical.set(canonical);
1848 return std::borrow::Cow::Borrowed(self.canonical.get().unwrap());
1849 }
1850 std::borrow::Cow::Owned(canonicalize_with_partial_fallback(&abs).unwrap_or(abs))
1851 }
1852}
1853
1854impl PermissionPolicy for PathPolicy {
1855 fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
1856 let Some(fs) = request
1857 .as_any()
1858 .downcast_ref::<FileSystemPermissionRequest>()
1859 else {
1860 return PolicyMatch::NoOpinion;
1861 };
1862
1863 let raw_paths: Vec<&Path> = match fs {
1864 FileSystemPermissionRequest::Move { from, to, .. } => {
1865 vec![from.as_path(), to.as_path()]
1866 }
1867 FileSystemPermissionRequest::Read { path, .. }
1868 | FileSystemPermissionRequest::Write { path, .. }
1869 | FileSystemPermissionRequest::Edit { path, .. }
1870 | FileSystemPermissionRequest::Delete { path, .. }
1871 | FileSystemPermissionRequest::List { path, .. }
1872 | FileSystemPermissionRequest::CreateDir { path, .. } => vec![path.as_path()],
1873 };
1874
1875 let candidate_paths: Vec<PathBuf> =
1876 raw_paths.iter().map(|p| resolve_canonical(p)).collect();
1877
1878 let mutates = matches!(
1879 fs,
1880 FileSystemPermissionRequest::Write { .. }
1881 | FileSystemPermissionRequest::Edit { .. }
1882 | FileSystemPermissionRequest::Delete { .. }
1883 | FileSystemPermissionRequest::Move { .. }
1884 | FileSystemPermissionRequest::CreateDir { .. }
1885 );
1886
1887 if candidate_paths.iter().any(|path| {
1888 self.protected_roots
1889 .iter()
1890 .any(|root| path.starts_with(root.resolve().as_ref()))
1891 }) {
1892 return PolicyMatch::Deny(PermissionDenial {
1893 code: PermissionCode::PathNotAllowed,
1894 message: format!("path access denied for {}", fs.summary()),
1895 metadata: fs.metadata().clone(),
1896 });
1897 }
1898
1899 if mutates
1900 && candidate_paths.iter().any(|path| {
1901 self.read_only_roots
1902 .iter()
1903 .any(|root| path.starts_with(root.resolve().as_ref()))
1904 })
1905 {
1906 return PolicyMatch::Deny(PermissionDenial {
1907 code: PermissionCode::PathNotAllowed,
1908 message: format!("path is read-only for {}", fs.summary()),
1909 metadata: fs.metadata().clone(),
1910 });
1911 }
1912
1913 if self.allowed_roots.is_empty() {
1914 return PolicyMatch::NoOpinion;
1915 }
1916
1917 let all_allowed = candidate_paths.iter().all(|path| {
1918 self.allowed_roots
1919 .iter()
1920 .any(|root| path.starts_with(root.resolve().as_ref()))
1921 });
1922
1923 if all_allowed {
1924 PolicyMatch::Allow
1925 } else if self.require_approval_outside_allowed {
1926 PolicyMatch::RequireApproval(ApprovalRequest {
1927 task_id: None,
1928 call_id: None,
1929 id: ApprovalId::new(format!("approval:{}", fs.kind())),
1930 request_kind: fs.kind().to_string(),
1931 reason: ApprovalReason::SensitivePath,
1932 summary: fs.summary(),
1933 metadata: fs.metadata().clone(),
1934 })
1935 } else {
1936 PolicyMatch::Deny(PermissionDenial {
1937 code: PermissionCode::PathNotAllowed,
1938 message: format!("path outside allowed roots for {}", fs.summary()),
1939 metadata: fs.metadata().clone(),
1940 })
1941 }
1942 }
1943}
1944
1945pub struct CommandPolicy {
1966 allowed_executables: BTreeSet<String>,
1967 denied_executables: BTreeSet<String>,
1968 allowed_cwds: Vec<PathBuf>,
1969 denied_env_keys: BTreeSet<String>,
1970 require_approval_for_unknown: bool,
1971}
1972
1973impl CommandPolicy {
1974 pub fn new() -> Self {
1977 Self {
1978 allowed_executables: BTreeSet::new(),
1979 denied_executables: BTreeSet::new(),
1980 allowed_cwds: Vec::new(),
1981 denied_env_keys: BTreeSet::new(),
1982 require_approval_for_unknown: true,
1983 }
1984 }
1985
1986 pub fn allow_executable(mut self, executable: impl Into<String>) -> Self {
1988 self.allowed_executables.insert(executable.into());
1989 self
1990 }
1991
1992 pub fn deny_executable(mut self, executable: impl Into<String>) -> Self {
1994 self.denied_executables.insert(executable.into());
1995 self
1996 }
1997
1998 pub fn allow_cwd(mut self, cwd: impl Into<PathBuf>) -> Self {
2000 self.allowed_cwds.push(cwd.into());
2001 self
2002 }
2003
2004 pub fn deny_env_key(mut self, key: impl Into<String>) -> Self {
2006 self.denied_env_keys.insert(key.into());
2007 self
2008 }
2009
2010 pub fn require_approval_for_unknown(mut self, value: bool) -> Self {
2013 self.require_approval_for_unknown = value;
2014 self
2015 }
2016}
2017
2018impl Default for CommandPolicy {
2019 fn default() -> Self {
2020 Self::new()
2021 }
2022}
2023
2024impl PermissionPolicy for CommandPolicy {
2025 fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
2026 let Some(shell) = request.as_any().downcast_ref::<ShellPermissionRequest>() else {
2027 return PolicyMatch::NoOpinion;
2028 };
2029
2030 if self.denied_executables.contains(&shell.executable)
2031 || shell
2032 .env_keys
2033 .iter()
2034 .any(|key| self.denied_env_keys.contains(key))
2035 {
2036 return PolicyMatch::Deny(PermissionDenial {
2037 code: PermissionCode::CommandNotAllowed,
2038 message: format!("command denied for {}", shell.summary()),
2039 metadata: shell.metadata().clone(),
2040 });
2041 }
2042
2043 if let Some(cwd) = &shell.cwd
2044 && !self.allowed_cwds.is_empty()
2045 && !self.allowed_cwds.iter().any(|root| cwd.starts_with(root))
2046 {
2047 return PolicyMatch::RequireApproval(ApprovalRequest {
2048 task_id: None,
2049 call_id: None,
2050 id: ApprovalId::new("approval:shell.cwd"),
2051 request_kind: shell.kind().to_string(),
2052 reason: ApprovalReason::SensitiveCommand,
2053 summary: shell.summary(),
2054 metadata: shell.metadata().clone(),
2055 });
2056 }
2057
2058 if self.allowed_executables.is_empty()
2059 || self.allowed_executables.contains(&shell.executable)
2060 {
2061 PolicyMatch::Allow
2062 } else if self.require_approval_for_unknown {
2063 PolicyMatch::RequireApproval(ApprovalRequest {
2064 task_id: None,
2065 call_id: None,
2066 id: ApprovalId::new("approval:shell.command"),
2067 request_kind: shell.kind().to_string(),
2068 reason: ApprovalReason::SensitiveCommand,
2069 summary: shell.summary(),
2070 metadata: shell.metadata().clone(),
2071 })
2072 } else {
2073 PolicyMatch::Deny(PermissionDenial {
2074 code: PermissionCode::CommandNotAllowed,
2075 message: format!("executable {} is not allowed", shell.executable),
2076 metadata: shell.metadata().clone(),
2077 })
2078 }
2079 }
2080}
2081
2082pub struct McpServerPolicy {
2096 trusted_servers: BTreeSet<String>,
2097 allowed_auth_scopes: BTreeSet<String>,
2098 require_approval_for_untrusted: bool,
2099}
2100
2101impl McpServerPolicy {
2102 pub fn new() -> Self {
2105 Self {
2106 trusted_servers: BTreeSet::new(),
2107 allowed_auth_scopes: BTreeSet::new(),
2108 require_approval_for_untrusted: true,
2109 }
2110 }
2111
2112 pub fn trust_server(mut self, server_id: impl Into<String>) -> Self {
2114 self.trusted_servers.insert(server_id.into());
2115 self
2116 }
2117
2118 pub fn allow_auth_scope(mut self, scope: impl Into<String>) -> Self {
2120 self.allowed_auth_scopes.insert(scope.into());
2121 self
2122 }
2123}
2124
2125impl Default for McpServerPolicy {
2126 fn default() -> Self {
2127 Self::new()
2128 }
2129}
2130
2131impl PermissionPolicy for McpServerPolicy {
2132 fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
2133 let Some(mcp) = request.as_any().downcast_ref::<McpPermissionRequest>() else {
2134 return PolicyMatch::NoOpinion;
2135 };
2136
2137 let server_id = match mcp {
2138 McpPermissionRequest::Connect { server_id, .. }
2139 | McpPermissionRequest::InvokeTool { server_id, .. }
2140 | McpPermissionRequest::ReadResource { server_id, .. }
2141 | McpPermissionRequest::FetchPrompt { server_id, .. }
2142 | McpPermissionRequest::UseAuthScope { server_id, .. } => server_id,
2143 };
2144
2145 if !self.trusted_servers.is_empty() && !self.trusted_servers.contains(server_id) {
2146 return if self.require_approval_for_untrusted {
2147 PolicyMatch::RequireApproval(ApprovalRequest {
2148 task_id: None,
2149 call_id: None,
2150 id: ApprovalId::new(format!("approval:mcp:{server_id}")),
2151 request_kind: mcp.kind().to_string(),
2152 reason: ApprovalReason::SensitiveServer,
2153 summary: mcp.summary(),
2154 metadata: mcp.metadata().clone(),
2155 })
2156 } else {
2157 PolicyMatch::Deny(PermissionDenial {
2158 code: PermissionCode::ServerNotTrusted,
2159 message: format!("MCP server {server_id} is not trusted"),
2160 metadata: mcp.metadata().clone(),
2161 })
2162 };
2163 }
2164
2165 if let McpPermissionRequest::UseAuthScope { scope, .. } = mcp
2166 && !self.allowed_auth_scopes.is_empty()
2167 && !self.allowed_auth_scopes.contains(scope)
2168 {
2169 return PolicyMatch::Deny(PermissionDenial {
2170 code: PermissionCode::AuthScopeNotAllowed,
2171 message: format!("MCP auth scope {scope} is not allowed"),
2172 metadata: mcp.metadata().clone(),
2173 });
2174 }
2175
2176 PolicyMatch::Allow
2177 }
2178}
2179
2180#[async_trait]
2232pub trait Tool: Send + Sync {
2233 fn spec(&self) -> &ToolSpec;
2235
2236 fn current_spec(&self) -> Option<ToolSpec> {
2244 Some(self.spec().clone())
2245 }
2246
2247 fn proposed_requests(
2259 &self,
2260 _request: &ToolRequest,
2261 ) -> Result<Vec<Box<dyn PermissionRequest>>, ToolError> {
2262 Ok(Vec::new())
2263 }
2264
2265 async fn invoke(
2274 &self,
2275 request: ToolRequest,
2276 ctx: &mut ToolContext<'_>,
2277 ) -> Result<ToolResult, ToolError>;
2278
2279 async fn invoke_outcome(
2285 &self,
2286 request: ToolRequest,
2287 ctx: &mut ToolContext<'_>,
2288 ) -> ToolExecutionOutcome {
2289 match self.invoke(request, ctx).await {
2290 Ok(result) => ToolExecutionOutcome::Completed(result),
2291 Err(error) => ToolExecutionOutcome::Failed(error),
2292 }
2293 }
2294}
2295
2296#[derive(Clone, Default)]
2327pub struct ToolRegistry {
2328 tools: BTreeMap<ToolName, Arc<dyn Tool>>,
2329}
2330
2331impl ToolRegistry {
2332 pub fn new() -> Self {
2334 Self::default()
2335 }
2336
2337 pub fn register<T>(&mut self, tool: T) -> &mut Self
2339 where
2340 T: Tool + 'static,
2341 {
2342 self.tools.insert(tool.spec().name.clone(), Arc::new(tool));
2343 self
2344 }
2345
2346 pub fn with<T>(mut self, tool: T) -> Self
2348 where
2349 T: Tool + 'static,
2350 {
2351 self.register(tool);
2352 self
2353 }
2354
2355 pub fn register_arc(&mut self, tool: Arc<dyn Tool>) -> &mut Self {
2357 self.tools.insert(tool.spec().name.clone(), tool);
2358 self
2359 }
2360
2361 pub fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2363 self.tools.get(name).cloned()
2364 }
2365
2366 pub fn tools(&self) -> Vec<Arc<dyn Tool>> {
2368 self.tools.values().cloned().collect()
2369 }
2370
2371 pub fn merge(mut self, other: Self) -> Self {
2380 self.tools.extend(other.tools);
2381 self
2382 }
2383
2384 pub fn specs(&self) -> Vec<ToolSpec> {
2386 self.tools
2387 .values()
2388 .filter_map(|tool| tool.current_spec())
2389 .collect()
2390 }
2391}
2392
2393pub trait ToolSource: Send + Sync {
2401 fn specs(&self) -> Vec<ToolSpec>;
2403
2404 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>>;
2406
2407 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2411 Vec::new()
2412 }
2413
2414 fn prefixed(self, prefix: impl Into<String>) -> Prefixed<Self>
2424 where
2425 Self: Sized,
2426 {
2427 Prefixed::new(self, prefix)
2428 }
2429
2430 fn filtered<F>(self, predicate: F) -> Filtered<Self, F>
2436 where
2437 Self: Sized,
2438 F: Fn(&ToolName) -> bool + Send + Sync + 'static,
2439 {
2440 Filtered::new(self, predicate)
2441 }
2442
2443 fn renamed<I>(self, mapping: I) -> Renamed<Self>
2450 where
2451 Self: Sized,
2452 I: IntoIterator<Item = (ToolName, ToolName)>,
2453 {
2454 Renamed::new(self, mapping)
2455 }
2456
2457 fn unadvertised(self) -> Unadvertised<Self>
2468 where
2469 Self: Sized,
2470 {
2471 Unadvertised::new(self)
2472 }
2473}
2474
2475impl ToolSource for ToolRegistry {
2476 fn specs(&self) -> Vec<ToolSpec> {
2477 ToolRegistry::specs(self)
2478 }
2479
2480 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2481 ToolRegistry::get(self, name)
2482 }
2483}
2484
2485impl<S> ToolSource for Arc<S>
2486where
2487 S: ToolSource + ?Sized,
2488{
2489 fn specs(&self) -> Vec<ToolSpec> {
2490 (**self).specs()
2491 }
2492
2493 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2494 (**self).get(name)
2495 }
2496
2497 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2498 (**self).drain_catalog_events()
2499 }
2500}
2501
2502pub struct Prefixed<S> {
2505 inner: S,
2506 prefix: String,
2507}
2508
2509impl<S> Prefixed<S> {
2510 pub fn new(inner: S, prefix: impl Into<String>) -> Self {
2512 Self {
2513 inner,
2514 prefix: prefix.into(),
2515 }
2516 }
2517
2518 fn rewrite(&self, name: &str) -> String {
2519 format!("{}_{}", self.prefix, name)
2520 }
2521
2522 fn strip<'a>(&self, name: &'a str) -> Option<&'a str> {
2523 name.strip_prefix(self.prefix.as_str())
2524 .and_then(|rest| rest.strip_prefix('_'))
2525 }
2526}
2527
2528impl<S> ToolSource for Prefixed<S>
2529where
2530 S: ToolSource,
2531{
2532 fn specs(&self) -> Vec<ToolSpec> {
2533 self.inner
2534 .specs()
2535 .into_iter()
2536 .map(|mut spec| {
2537 spec.name = ToolName::new(self.rewrite(spec.name.0.as_str()));
2538 spec
2539 })
2540 .collect()
2541 }
2542
2543 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2544 let original = self.strip(name.0.as_str())?;
2545 let inner_name = ToolName::new(original);
2546 let inner_tool = self.inner.get(&inner_name)?;
2547 let mut public_spec = inner_tool.spec().clone();
2548 public_spec.name = name.clone();
2549 Some(Arc::new(RewrittenTool {
2550 inner: inner_tool,
2551 inner_name,
2552 public_spec,
2553 }))
2554 }
2555
2556 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2557 self.inner
2558 .drain_catalog_events()
2559 .into_iter()
2560 .map(|mut event| {
2561 event.for_each_name_mut(|name| *name = self.rewrite(name.as_str()));
2562 event
2563 })
2564 .collect()
2565 }
2566}
2567
2568pub struct Filtered<S, F> {
2571 inner: S,
2572 predicate: F,
2573}
2574
2575impl<S, F> Filtered<S, F> {
2576 pub fn new(inner: S, predicate: F) -> Self {
2578 Self { inner, predicate }
2579 }
2580}
2581
2582impl<S, F> ToolSource for Filtered<S, F>
2583where
2584 S: ToolSource,
2585 F: Fn(&ToolName) -> bool + Send + Sync + 'static,
2586{
2587 fn specs(&self) -> Vec<ToolSpec> {
2588 self.inner
2589 .specs()
2590 .into_iter()
2591 .filter(|spec| (self.predicate)(&spec.name))
2592 .collect()
2593 }
2594
2595 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2596 if !(self.predicate)(name) {
2597 return None;
2598 }
2599 self.inner.get(name)
2600 }
2601
2602 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2603 self.inner
2604 .drain_catalog_events()
2605 .into_iter()
2606 .map(|mut event| {
2607 event.retain_names(|n| (self.predicate)(&ToolName::new(n)));
2608 event
2609 })
2610 .collect()
2611 }
2612}
2613
2614pub struct Unadvertised<S> {
2621 inner: S,
2622}
2623
2624impl<S> Unadvertised<S> {
2625 pub fn new(inner: S) -> Self {
2627 Self { inner }
2628 }
2629}
2630
2631impl<S> ToolSource for Unadvertised<S>
2632where
2633 S: ToolSource,
2634{
2635 fn specs(&self) -> Vec<ToolSpec> {
2636 Vec::new()
2637 }
2638
2639 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2640 self.inner.get(name)
2641 }
2642
2643 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2644 let _ = self.inner.drain_catalog_events();
2647 Vec::new()
2648 }
2649}
2650
2651pub struct Renamed<S> {
2658 inner: S,
2659 forward: BTreeMap<ToolName, ToolName>,
2660 backward: BTreeMap<ToolName, ToolName>,
2661}
2662
2663impl<S> Renamed<S> {
2664 pub fn new<I>(inner: S, mapping: I) -> Self
2666 where
2667 I: IntoIterator<Item = (ToolName, ToolName)>,
2668 {
2669 let forward: BTreeMap<ToolName, ToolName> = mapping.into_iter().collect();
2670 let backward = forward
2671 .iter()
2672 .map(|(k, v)| (v.clone(), k.clone()))
2673 .collect();
2674 Self {
2675 inner,
2676 forward,
2677 backward,
2678 }
2679 }
2680}
2681
2682impl<S> ToolSource for Renamed<S>
2683where
2684 S: ToolSource,
2685{
2686 fn specs(&self) -> Vec<ToolSpec> {
2687 self.inner
2688 .specs()
2689 .into_iter()
2690 .map(|mut spec| {
2691 if let Some(new_name) = self.forward.get(&spec.name) {
2692 spec.name = new_name.clone();
2693 }
2694 spec
2695 })
2696 .collect()
2697 }
2698
2699 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2700 if let Some(original) = self.backward.get(name) {
2701 let inner_tool = self.inner.get(original)?;
2702 let mut public_spec = inner_tool.spec().clone();
2703 public_spec.name = name.clone();
2704 Some(Arc::new(RewrittenTool {
2705 inner: inner_tool,
2706 inner_name: original.clone(),
2707 public_spec,
2708 }))
2709 } else if self.forward.contains_key(name) {
2710 None
2712 } else {
2713 self.inner.get(name)
2714 }
2715 }
2716
2717 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2718 self.inner
2719 .drain_catalog_events()
2720 .into_iter()
2721 .map(|mut event| {
2722 event.for_each_name_mut(|name| {
2723 if let Some(new) = self.forward.get(&ToolName::new(name.as_str())) {
2724 *name = new.0.clone();
2725 }
2726 });
2727 event
2728 })
2729 .collect()
2730 }
2731}
2732
2733#[cfg(feature = "schemars")]
2759pub fn schema_for<T: schemars::JsonSchema>() -> Value {
2760 let schema = schemars::schema_for!(T);
2761 serde_json::to_value(schema)
2762 .expect("schemars produces valid JSON; this conversion is infallible")
2763}
2764
2765#[cfg(feature = "schemars")]
2783pub fn tool_spec_for<T: schemars::JsonSchema>(
2784 name: impl Into<ToolName>,
2785 description: impl Into<String>,
2786) -> ToolSpec {
2787 ToolSpec::new(name, description, schema_for::<T>())
2788}
2789
2790struct RewrittenTool {
2796 inner: Arc<dyn Tool>,
2797 inner_name: ToolName,
2798 public_spec: ToolSpec,
2799}
2800
2801#[async_trait]
2802impl Tool for RewrittenTool {
2803 fn spec(&self) -> &ToolSpec {
2804 &self.public_spec
2805 }
2806
2807 fn current_spec(&self) -> Option<ToolSpec> {
2808 let inner_current = self.inner.current_spec()?;
2809 Some(ToolSpec {
2810 name: self.public_spec.name.clone(),
2811 description: inner_current.description,
2812 input_schema: inner_current.input_schema,
2813 output_schema: inner_current.output_schema,
2814 annotations: inner_current.annotations,
2815 metadata: inner_current.metadata,
2816 })
2817 }
2818
2819 fn proposed_requests(
2820 &self,
2821 request: &ToolRequest,
2822 ) -> Result<Vec<Box<dyn PermissionRequest>>, ToolError> {
2823 let mut inner_request = request.clone();
2824 inner_request.tool_name = self.inner_name.clone();
2825 self.inner.proposed_requests(&inner_request)
2826 }
2827
2828 async fn invoke(
2829 &self,
2830 mut request: ToolRequest,
2831 ctx: &mut ToolContext<'_>,
2832 ) -> Result<ToolResult, ToolError> {
2833 request.tool_name = self.inner_name.clone();
2834 self.inner.invoke(request, ctx).await
2835 }
2836
2837 async fn invoke_outcome(
2838 &self,
2839 mut request: ToolRequest,
2840 ctx: &mut ToolContext<'_>,
2841 ) -> ToolExecutionOutcome {
2842 request.tool_name = self.inner_name.clone();
2843 self.inner.invoke_outcome(request, ctx).await
2844 }
2845}
2846
2847struct ToolMap {
2865 inner: std::sync::RwLock<BTreeMap<ToolName, Arc<dyn Tool>>>,
2866}
2867
2868impl ToolMap {
2869 fn new() -> Self {
2870 Self {
2871 inner: std::sync::RwLock::new(BTreeMap::new()),
2872 }
2873 }
2874
2875 fn read(&self) -> std::sync::RwLockReadGuard<'_, BTreeMap<ToolName, Arc<dyn Tool>>> {
2876 self.inner.read().unwrap_or_else(|e| e.into_inner())
2877 }
2878
2879 fn write(&self) -> std::sync::RwLockWriteGuard<'_, BTreeMap<ToolName, Arc<dyn Tool>>> {
2880 self.inner.write().unwrap_or_else(|e| e.into_inner())
2881 }
2882}
2883
2884struct DynamicCatalogInner {
2887 source_id: String,
2888 tools: ToolMap,
2889 events_tx: tokio::sync::broadcast::Sender<ToolCatalogEvent>,
2890}
2891
2892pub fn dynamic_catalog(source_id: impl Into<String>) -> (CatalogWriter, CatalogReader) {
2910 let (events_tx, events_rx) = tokio::sync::broadcast::channel(128);
2911 let inner = Arc::new(DynamicCatalogInner {
2912 source_id: source_id.into(),
2913 tools: ToolMap::new(),
2914 events_tx,
2915 });
2916 (
2917 CatalogWriter {
2918 inner: Arc::clone(&inner),
2919 },
2920 CatalogReader {
2921 inner,
2922 events_rx: std::sync::Mutex::new(events_rx),
2923 },
2924 )
2925}
2926
2927pub struct CatalogWriter {
2934 inner: Arc<DynamicCatalogInner>,
2935}
2936
2937impl CatalogWriter {
2938 pub fn source_id(&self) -> &str {
2940 &self.inner.source_id
2941 }
2942
2943 pub fn reader(&self) -> CatalogReader {
2947 CatalogReader {
2948 inner: Arc::clone(&self.inner),
2949 events_rx: std::sync::Mutex::new(self.inner.events_tx.subscribe()),
2950 }
2951 }
2952
2953 pub fn upsert(&self, tool: Arc<dyn Tool>) {
2956 let name = tool.spec().name.clone();
2957 let mut guard = self.inner.tools.write();
2958 let existed = guard.insert(name.clone(), tool).is_some();
2959 drop(guard);
2960 let mut event = ToolCatalogEvent::new(self.inner.source_id.clone());
2961 if existed {
2962 event.changed.push(name.0);
2963 } else {
2964 event.added.push(name.0);
2965 }
2966 let _ = self.inner.events_tx.send(event);
2967 }
2968
2969 pub fn remove(&self, name: &ToolName) -> bool {
2972 let mut guard = self.inner.tools.write();
2973 let removed = guard.remove(name).is_some();
2974 drop(guard);
2975 if removed {
2976 let mut event = ToolCatalogEvent::new(self.inner.source_id.clone());
2977 event.removed.push(name.0.clone());
2978 let _ = self.inner.events_tx.send(event);
2979 }
2980 removed
2981 }
2982
2983 pub fn replace_all(&self, tools: impl IntoIterator<Item = Arc<dyn Tool>>) {
2986 let new_map: BTreeMap<ToolName, Arc<dyn Tool>> = tools
2987 .into_iter()
2988 .map(|tool| (tool.spec().name.clone(), tool))
2989 .collect();
2990
2991 let mut guard = self.inner.tools.write();
2992 let mut event = ToolCatalogEvent::new(self.inner.source_id.clone());
2993
2994 for (name, new_tool) in new_map.iter() {
2995 match guard.get(name) {
2996 None => event.added.push(name.0.clone()),
2997 Some(existing)
2998 if !Arc::ptr_eq(existing, new_tool)
2999 && existing.current_spec() != new_tool.current_spec() =>
3000 {
3001 event.changed.push(name.0.clone());
3002 }
3003 Some(_) => {}
3004 }
3005 }
3006 for name in guard.keys() {
3007 if !new_map.contains_key(name) {
3008 event.removed.push(name.0.clone());
3009 }
3010 }
3011
3012 *guard = new_map;
3013 drop(guard);
3014
3015 if !event.added.is_empty() || !event.removed.is_empty() || !event.changed.is_empty() {
3016 let _ = self.inner.events_tx.send(event);
3017 }
3018 }
3019
3020 pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<ToolCatalogEvent> {
3024 self.inner.events_tx.subscribe()
3025 }
3026}
3027
3028pub struct CatalogReader {
3033 inner: Arc<DynamicCatalogInner>,
3034 events_rx: std::sync::Mutex<tokio::sync::broadcast::Receiver<ToolCatalogEvent>>,
3035}
3036
3037impl CatalogReader {
3038 pub fn source_id(&self) -> &str {
3040 &self.inner.source_id
3041 }
3042
3043 pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<ToolCatalogEvent> {
3046 self.inner.events_tx.subscribe()
3047 }
3048}
3049
3050impl Clone for CatalogReader {
3051 fn clone(&self) -> Self {
3052 Self {
3053 inner: Arc::clone(&self.inner),
3054 events_rx: std::sync::Mutex::new(self.inner.events_tx.subscribe()),
3055 }
3056 }
3057}
3058
3059impl ToolSource for CatalogReader {
3060 fn specs(&self) -> Vec<ToolSpec> {
3061 self.inner
3062 .tools
3063 .read()
3064 .values()
3065 .filter_map(|tool| tool.current_spec())
3066 .collect()
3067 }
3068
3069 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
3070 self.inner.tools.read().get(name).cloned()
3071 }
3072
3073 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3074 let mut rx = self.events_rx.lock().unwrap_or_else(|e| e.into_inner());
3079 let mut out = Vec::new();
3080 loop {
3081 match rx.try_recv() {
3082 Ok(event) => out.push(event),
3083 Err(tokio::sync::broadcast::error::TryRecvError::Empty) => break,
3084 Err(tokio::sync::broadcast::error::TryRecvError::Closed) => break,
3085 Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
3086 }
3087 }
3088 out
3089 }
3090}
3091
3092impl ToolSpec {
3093 pub fn as_invocable_spec(&self) -> InvocableSpec {
3096 InvocableSpec::new(
3097 CapabilityName::new(self.name.0.clone()),
3098 self.description.clone(),
3099 self.input_schema.clone(),
3100 )
3101 .with_metadata(self.metadata.clone())
3102 }
3103}
3104
3105pub struct ToolInvocableAdapter {
3111 spec: InvocableSpec,
3112 tool: Arc<dyn Tool>,
3113 permissions: Arc<dyn PermissionChecker>,
3114 resources: Arc<dyn ToolResources>,
3115 next_call_id: AtomicU64,
3116}
3117
3118impl ToolInvocableAdapter {
3119 pub fn new(
3122 tool: Arc<dyn Tool>,
3123 permissions: Arc<dyn PermissionChecker>,
3124 resources: Arc<dyn ToolResources>,
3125 ) -> Option<Self> {
3126 let spec = tool.current_spec()?.as_invocable_spec();
3127 Some(Self {
3128 spec,
3129 tool,
3130 permissions,
3131 resources,
3132 next_call_id: AtomicU64::new(1),
3133 })
3134 }
3135}
3136
3137#[async_trait]
3138impl Invocable for ToolInvocableAdapter {
3139 fn spec(&self) -> &InvocableSpec {
3140 &self.spec
3141 }
3142
3143 async fn invoke(
3144 &self,
3145 request: InvocableRequest,
3146 ctx: &mut CapabilityContext<'_>,
3147 ) -> Result<InvocableResult, CapabilityError> {
3148 let tool_request = ToolRequest {
3149 call_id: ToolCallId::new(format!(
3150 "tool-call-{}",
3151 self.next_call_id.fetch_add(1, Ordering::Relaxed)
3152 )),
3153 tool_name: self.tool.spec().name.clone(),
3154 input: request.input,
3155 session_id: ctx
3156 .session_id
3157 .cloned()
3158 .unwrap_or_else(|| SessionId::new("capability-session")),
3159 turn_id: ctx
3160 .turn_id
3161 .cloned()
3162 .unwrap_or_else(|| TurnId::new("capability-turn")),
3163 metadata: request.metadata,
3164 };
3165
3166 for permission_request in self
3167 .tool
3168 .proposed_requests(&tool_request)
3169 .map_err(|error| CapabilityError::InvalidInput(error.to_string()))?
3170 {
3171 match self.permissions.evaluate(permission_request.as_ref()) {
3172 PermissionDecision::Allow => {}
3173 PermissionDecision::Deny(denial) => {
3174 return Err(CapabilityError::ExecutionFailed(format!(
3175 "tool permission denied: {denial:?}"
3176 )));
3177 }
3178 PermissionDecision::RequireApproval(req) => {
3179 return Err(CapabilityError::Unavailable(format!(
3180 "tool invocation requires approval: {}",
3181 req.summary
3182 )));
3183 }
3184 }
3185 }
3186
3187 let mut tool_ctx = ToolContext {
3188 capability: CapabilityContext {
3189 session_id: ctx.session_id,
3190 turn_id: ctx.turn_id,
3191 metadata: ctx.metadata,
3192 },
3193 permissions: self.permissions.as_ref(),
3194 resources: self.resources.as_ref(),
3195 cancellation: None,
3196 execution_scope: None,
3197 approved_request: None,
3198 };
3199
3200 let result = self
3201 .tool
3202 .invoke(tool_request, &mut tool_ctx)
3203 .await
3204 .map_err(|error| CapabilityError::ExecutionFailed(error.to_string()))?;
3205
3206 Ok(InvocableResult {
3207 output: match result.result.output {
3208 ToolOutput::Text(text) => InvocableOutput::Text(text),
3209 ToolOutput::Structured(value) => InvocableOutput::Structured(value),
3210 ToolOutput::Parts(parts) => InvocableOutput::Items(vec![Item {
3211 id: None,
3212 kind: ItemKind::Tool,
3213 parts,
3214 metadata: MetadataMap::new(),
3215 usage: None,
3216 finish_reason: None,
3217 created_at: None,
3218 }]),
3219 ToolOutput::Files(files) => {
3220 let parts = files.into_iter().map(Part::File).collect();
3221 InvocableOutput::Items(vec![Item {
3222 id: None,
3223 kind: ItemKind::Tool,
3224 parts,
3225 metadata: MetadataMap::new(),
3226 usage: None,
3227 finish_reason: None,
3228 created_at: None,
3229 }])
3230 }
3231 },
3232 metadata: result.metadata,
3233 })
3234 }
3235}
3236
3237pub struct ToolCapabilityProvider {
3243 invocables: Vec<Arc<dyn Invocable>>,
3244}
3245
3246impl ToolCapabilityProvider {
3247 pub fn from_registry(
3250 registry: &ToolRegistry,
3251 permissions: Arc<dyn PermissionChecker>,
3252 resources: Arc<dyn ToolResources>,
3253 ) -> Self {
3254 let invocables = registry
3255 .tools()
3256 .into_iter()
3257 .filter_map(|tool| {
3258 ToolInvocableAdapter::new(tool, permissions.clone(), resources.clone())
3259 .map(|adapter| Arc::new(adapter) as Arc<dyn Invocable>)
3260 })
3261 .collect();
3262
3263 Self { invocables }
3264 }
3265}
3266
3267impl CapabilityProvider for ToolCapabilityProvider {
3268 fn invocables(&self) -> Vec<Arc<dyn Invocable>> {
3269 self.invocables.clone()
3270 }
3271
3272 fn resources(&self) -> Vec<Arc<dyn ResourceProvider>> {
3273 Vec::new()
3274 }
3275
3276 fn prompts(&self) -> Vec<Arc<dyn PromptProvider>> {
3277 Vec::new()
3278 }
3279}
3280
3281#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3287pub enum ToolExecutionOutcome {
3288 Completed(ToolResult),
3290 Interrupted(ToolInterruption),
3292 FailedBeforeInvocation(ToolError),
3294 Failed(ToolError),
3296}
3297
3298#[async_trait]
3306pub trait ToolExecutor: Send + Sync {
3307 fn specs(&self) -> Vec<ToolSpec>;
3309
3310 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3315 Vec::new()
3316 }
3317
3318 async fn execute(
3320 &self,
3321 request: ToolRequest,
3322 ctx: &mut ToolContext<'_>,
3323 ) -> ToolExecutionOutcome;
3324
3325 async fn execute_owned(
3328 &self,
3329 request: ToolRequest,
3330 ctx: OwnedToolContext,
3331 ) -> ToolExecutionOutcome {
3332 let mut borrowed = ctx.borrowed();
3333 self.execute(request, &mut borrowed).await
3334 }
3335
3336 async fn execute_approved(
3342 &self,
3343 request: ToolRequest,
3344 approved_request: &ApprovalRequest,
3345 ctx: &mut ToolContext<'_>,
3346 ) -> ToolExecutionOutcome {
3347 let _ = approved_request;
3348 self.execute(request, ctx).await
3349 }
3350
3351 async fn execute_approved_owned(
3354 &self,
3355 request: ToolRequest,
3356 approved_request: &ApprovalRequest,
3357 mut ctx: OwnedToolContext,
3358 ) -> ToolExecutionOutcome {
3359 ctx.approved_request = Some(approved_request.clone());
3360 let mut borrowed = ctx.borrowed();
3361 self.execute_approved(request, approved_request, &mut borrowed)
3362 .await
3363 }
3364}
3365
3366#[async_trait]
3367impl<T> ToolExecutor for Arc<T>
3368where
3369 T: ToolExecutor + ?Sized,
3370{
3371 fn specs(&self) -> Vec<ToolSpec> {
3372 (**self).specs()
3373 }
3374
3375 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3376 (**self).drain_catalog_events()
3377 }
3378
3379 async fn execute(
3380 &self,
3381 request: ToolRequest,
3382 ctx: &mut ToolContext<'_>,
3383 ) -> ToolExecutionOutcome {
3384 (**self).execute(request, ctx).await
3385 }
3386
3387 async fn execute_approved(
3388 &self,
3389 request: ToolRequest,
3390 approved_request: &ApprovalRequest,
3391 ctx: &mut ToolContext<'_>,
3392 ) -> ToolExecutionOutcome {
3393 (**self)
3394 .execute_approved(request, approved_request, ctx)
3395 .await
3396 }
3397}
3398
3399#[derive(Clone, Debug, Default, PartialEq, Eq)]
3402pub enum CollisionPolicy {
3403 #[default]
3406 FirstWins,
3407 LastWins,
3409}
3410
3411pub struct BasicToolExecutor {
3430 sources: Vec<Arc<dyn ToolSource>>,
3431 collision: CollisionPolicy,
3432 output_truncation: Option<Arc<dyn ToolOutputTruncationStrategy>>,
3433}
3434
3435impl BasicToolExecutor {
3436 pub fn new(sources: impl IntoIterator<Item = Arc<dyn ToolSource>>) -> Self {
3438 Self {
3439 sources: sources.into_iter().collect(),
3440 collision: CollisionPolicy::default(),
3441 output_truncation: None,
3442 }
3443 }
3444
3445 pub fn from_registry(registry: ToolRegistry) -> Self {
3447 Self::new([Arc::new(registry) as Arc<dyn ToolSource>])
3448 }
3449
3450 pub fn with_collision_policy(mut self, policy: CollisionPolicy) -> Self {
3453 self.collision = policy;
3454 self
3455 }
3456
3457 pub fn with_output_truncation_strategy(
3461 mut self,
3462 strategy: impl ToolOutputTruncationStrategy + 'static,
3463 ) -> Self {
3464 self.output_truncation = Some(Arc::new(strategy));
3465 self
3466 }
3467
3468 pub fn with_output_truncation_strategy_arc(
3470 mut self,
3471 strategy: Arc<dyn ToolOutputTruncationStrategy>,
3472 ) -> Self {
3473 self.output_truncation = Some(strategy);
3474 self
3475 }
3476
3477 pub fn specs(&self) -> Vec<ToolSpec> {
3480 let mut seen = BTreeSet::new();
3481 let mut out = Vec::new();
3482 let iter: Box<dyn Iterator<Item = &Arc<dyn ToolSource>>> = match self.collision {
3483 CollisionPolicy::FirstWins => Box::new(self.sources.iter()),
3484 CollisionPolicy::LastWins => Box::new(self.sources.iter().rev()),
3485 };
3486 for source in iter {
3487 for spec in source.specs() {
3488 if seen.insert(spec.name.clone()) {
3489 out.push(spec);
3490 }
3491 }
3492 }
3493 out
3494 }
3495
3496 fn lookup(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
3497 match self.collision {
3498 CollisionPolicy::FirstWins => self.sources.iter().find_map(|s| s.get(name)),
3499 CollisionPolicy::LastWins => self.sources.iter().rev().find_map(|s| s.get(name)),
3500 }
3501 }
3502
3503 async fn execute_inner(
3504 &self,
3505 request: ToolRequest,
3506 approved_request_id: Option<&ApprovalId>,
3507 ctx: &mut ToolContext<'_>,
3508 ) -> ToolExecutionOutcome {
3509 let Some(tool) = self.lookup(&request.tool_name) else {
3510 return ToolExecutionOutcome::FailedBeforeInvocation(ToolError::NotFound(
3511 request.tool_name,
3512 ));
3513 };
3514
3515 match tool.proposed_requests(&request) {
3516 Ok(requests) => {
3517 for permission_request in requests {
3518 match ctx.permissions.evaluate(permission_request.as_ref()) {
3519 PermissionDecision::Allow => {}
3520 PermissionDecision::Deny(denial) => {
3521 return ToolExecutionOutcome::FailedBeforeInvocation(
3522 ToolError::PermissionDenied(denial),
3523 );
3524 }
3525 PermissionDecision::RequireApproval(mut req) => {
3526 req.call_id = Some(request.call_id.clone());
3527 if approved_request_id != Some(&req.id) {
3528 return ToolExecutionOutcome::Interrupted(
3529 ToolInterruption::ApprovalRequired(req),
3530 );
3531 }
3532 }
3533 }
3534 }
3535 }
3536 Err(error) => return ToolExecutionOutcome::FailedBeforeInvocation(error),
3537 }
3538
3539 let truncation_ctx = ToolOutputTruncationContext::from((&request, tool.spec().clone()));
3540 match tool.invoke_outcome(request, ctx).await {
3541 ToolExecutionOutcome::Completed(mut result) => {
3542 if let Some(strategy) = &self.output_truncation {
3543 match strategy.apply(truncation_ctx, result.result.output).await {
3544 Ok(output) => {
3545 result.result.output = output;
3546 }
3547 Err(error) => return ToolExecutionOutcome::Failed(error),
3548 }
3549 }
3550 ToolExecutionOutcome::Completed(result)
3551 }
3552 other => other,
3553 }
3554 }
3555}
3556
3557#[async_trait]
3558impl ToolExecutor for BasicToolExecutor {
3559 fn specs(&self) -> Vec<ToolSpec> {
3560 BasicToolExecutor::specs(self)
3561 }
3562
3563 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3564 self.sources
3565 .iter()
3566 .flat_map(|s| s.drain_catalog_events())
3567 .collect()
3568 }
3569
3570 async fn execute(
3571 &self,
3572 request: ToolRequest,
3573 ctx: &mut ToolContext<'_>,
3574 ) -> ToolExecutionOutcome {
3575 self.execute_inner(request, None, ctx).await
3576 }
3577
3578 async fn execute_approved(
3579 &self,
3580 request: ToolRequest,
3581 approved_request: &ApprovalRequest,
3582 ctx: &mut ToolContext<'_>,
3583 ) -> ToolExecutionOutcome {
3584 let previous = ctx.approved_request.replace(approved_request.clone());
3585 let outcome = self
3586 .execute_inner(request, Some(&approved_request.id), ctx)
3587 .await;
3588 ctx.approved_request = previous;
3589 outcome
3590 }
3591}
3592
3593#[derive(Debug, Error, Clone, PartialEq, Serialize, Deserialize)]
3598pub enum ToolError {
3599 #[error("tool not found: {0}")]
3601 NotFound(ToolName),
3602 #[error("invalid tool input: {0}")]
3604 InvalidInput(String),
3605 #[error("tool permission denied: {0:?}")]
3607 PermissionDenied(PermissionDenial),
3608 #[error("tool execution failed: {0}")]
3610 ExecutionFailed(String),
3611 #[error("tool unavailable: {0}")]
3613 Unavailable(String),
3614 #[error("tool execution cancelled")]
3616 Cancelled,
3617 #[error("internal tool error: {0}")]
3619 Internal(String),
3620}
3621
3622impl ToolError {
3623 pub fn permission_denied(denial: PermissionDenial) -> Self {
3625 Self::PermissionDenied(denial)
3626 }
3627}
3628
3629impl From<PermissionDenial> for ToolError {
3630 fn from(value: PermissionDenial) -> Self {
3631 Self::permission_denied(value)
3632 }
3633}
3634
3635#[cfg(test)]
3636mod tests {
3637 use super::*;
3638 use async_trait::async_trait;
3639 use serde_json::json;
3640
3641 #[test]
3642 fn command_policy_can_deny_unknown_executables_without_approval() {
3643 let policy = CommandPolicy::new()
3644 .allow_executable("pwd")
3645 .require_approval_for_unknown(false);
3646 let request = ShellPermissionRequest {
3647 executable: "rm".into(),
3648 argv: vec!["-rf".into(), "/tmp/demo".into()],
3649 cwd: None,
3650 env_keys: Vec::new(),
3651 metadata: MetadataMap::new(),
3652 };
3653
3654 match policy.evaluate(&request) {
3655 PolicyMatch::Deny(denial) => {
3656 assert_eq!(denial.code, PermissionCode::CommandNotAllowed);
3657 }
3658 other => panic!("unexpected policy match: {other:?}"),
3659 }
3660 }
3661
3662 #[test]
3663 fn path_policy_allows_reads_under_read_only_roots() {
3664 let policy = PathPolicy::new().read_only_root("/workspace/vendor");
3665 let request = FileSystemPermissionRequest::Read {
3666 path: PathBuf::from("/workspace/vendor/lib.rs"),
3667 metadata: MetadataMap::new(),
3668 };
3669
3670 match policy.evaluate(&request) {
3671 PolicyMatch::NoOpinion | PolicyMatch::Allow => {}
3672 other => panic!("unexpected policy match: {other:?}"),
3673 }
3674 }
3675
3676 #[test]
3677 fn path_policy_denies_mutations_under_read_only_roots() {
3678 let policy = PathPolicy::new().read_only_root("/workspace/vendor");
3679 let request = FileSystemPermissionRequest::Edit {
3680 path: PathBuf::from("/workspace/vendor/lib.rs"),
3681 metadata: MetadataMap::new(),
3682 };
3683
3684 match policy.evaluate(&request) {
3685 PolicyMatch::Deny(denial) => {
3686 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3687 assert!(denial.message.contains("read-only"));
3688 }
3689 other => panic!("unexpected policy match: {other:?}"),
3690 }
3691 }
3692
3693 #[test]
3694 fn path_policy_denies_moves_into_read_only_roots() {
3695 let policy = PathPolicy::new().read_only_root("/workspace/vendor");
3696 let request = FileSystemPermissionRequest::Move {
3697 from: PathBuf::from("/workspace/src/lib.rs"),
3698 to: PathBuf::from("/workspace/vendor/lib.rs"),
3699 metadata: MetadataMap::new(),
3700 };
3701
3702 match policy.evaluate(&request) {
3703 PolicyMatch::Deny(denial) => {
3704 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3705 assert!(denial.message.contains("read-only"));
3706 }
3707 other => panic!("unexpected policy match: {other:?}"),
3708 }
3709 }
3710
3711 #[cfg(unix)]
3712 struct SymlinkTmpDir(PathBuf);
3713
3714 #[cfg(unix)]
3715 impl SymlinkTmpDir {
3716 fn new(label: &str) -> Self {
3717 use std::time::{SystemTime, UNIX_EPOCH};
3718 let nanos = SystemTime::now()
3719 .duration_since(UNIX_EPOCH)
3720 .unwrap()
3721 .as_nanos();
3722 let dir = std::env::temp_dir().join(format!(
3723 "agentkit-pathpolicy-{}-{}-{}",
3724 label,
3725 std::process::id(),
3726 nanos
3727 ));
3728 std::fs::create_dir_all(&dir).unwrap();
3729 Self(std::fs::canonicalize(&dir).unwrap())
3732 }
3733
3734 fn path(&self) -> &Path {
3735 &self.0
3736 }
3737 }
3738
3739 #[cfg(unix)]
3740 impl Drop for SymlinkTmpDir {
3741 fn drop(&mut self) {
3742 let _ = std::fs::remove_dir_all(&self.0);
3743 }
3744 }
3745
3746 #[cfg(unix)]
3747 fn assert_path_denied(
3748 policy: &PathPolicy,
3749 request: FileSystemPermissionRequest,
3750 ) -> PermissionDenial {
3751 match policy.evaluate(&request) {
3752 PolicyMatch::Deny(denial) => denial,
3753 other => panic!("expected deny, got: {other:?}"),
3754 }
3755 }
3756
3757 #[cfg(unix)]
3758 #[test]
3759 fn path_policy_blocks_symlink_escape_from_allowed_root() {
3760 let tmp = SymlinkTmpDir::new("allow-escape");
3761 let allowed = tmp.path().join("workspace");
3762 let outside = tmp.path().join("outside");
3763 std::fs::create_dir_all(&allowed).unwrap();
3764 std::fs::create_dir_all(&outside).unwrap();
3765 let secret = outside.join("secret.txt");
3766 std::fs::write(&secret, b"top-secret").unwrap();
3767 let escape = allowed.join("leak");
3768 std::os::unix::fs::symlink(&secret, &escape).unwrap();
3769
3770 let policy = PathPolicy::new()
3771 .allow_root(&allowed)
3772 .require_approval_outside_allowed(false);
3773 let denial = assert_path_denied(
3774 &policy,
3775 FileSystemPermissionRequest::Read {
3776 path: escape,
3777 metadata: MetadataMap::new(),
3778 },
3779 );
3780 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3781 }
3782
3783 #[cfg(unix)]
3784 #[test]
3785 fn path_policy_blocks_symlink_into_protected_root() {
3786 let tmp = SymlinkTmpDir::new("protect-bypass");
3787 let workspace = tmp.path().join("workspace");
3788 std::fs::create_dir_all(&workspace).unwrap();
3789 let secret = workspace.join(".env");
3790 std::fs::write(&secret, b"API_KEY=xxx").unwrap();
3791 let alias = workspace.join("config");
3792 std::os::unix::fs::symlink(&secret, &alias).unwrap();
3793
3794 let policy = PathPolicy::new()
3795 .allow_root(&workspace)
3796 .protect_root(&secret);
3797 let denial = assert_path_denied(
3798 &policy,
3799 FileSystemPermissionRequest::Read {
3800 path: alias,
3801 metadata: MetadataMap::new(),
3802 },
3803 );
3804 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3805 assert!(denial.message.contains("denied"));
3806 }
3807
3808 #[cfg(unix)]
3809 #[test]
3810 fn path_policy_blocks_symlink_write_into_read_only_root() {
3811 let tmp = SymlinkTmpDir::new("readonly-bypass");
3812 let workspace = tmp.path().join("workspace");
3813 let vendor = workspace.join("vendor");
3814 std::fs::create_dir_all(&vendor).unwrap();
3815 let target = vendor.join("lib.rs");
3816 std::fs::write(&target, b"// vendored").unwrap();
3817 let writable_alias = workspace.join("writable");
3818 std::os::unix::fs::symlink(&target, &writable_alias).unwrap();
3819
3820 let policy = PathPolicy::new()
3821 .allow_root(&workspace)
3822 .read_only_root(&vendor);
3823 let denial = assert_path_denied(
3824 &policy,
3825 FileSystemPermissionRequest::Edit {
3826 path: writable_alias,
3827 metadata: MetadataMap::new(),
3828 },
3829 );
3830 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3831 assert!(denial.message.contains("read-only"));
3832 }
3833
3834 #[cfg(unix)]
3835 #[test]
3836 fn path_policy_resolves_symlink_parent_for_nonexistent_leaf() {
3837 let tmp = SymlinkTmpDir::new("create-escape");
3838 let allowed = tmp.path().join("workspace");
3839 let outside = tmp.path().join("outside");
3840 std::fs::create_dir_all(&allowed).unwrap();
3841 std::fs::create_dir_all(&outside).unwrap();
3842 let escape_dir = allowed.join("escape");
3843 std::os::unix::fs::symlink(&outside, &escape_dir).unwrap();
3844 let new_file = escape_dir.join("new.txt");
3845
3846 let policy = PathPolicy::new()
3847 .allow_root(&allowed)
3848 .require_approval_outside_allowed(false);
3849 let denial = assert_path_denied(
3850 &policy,
3851 FileSystemPermissionRequest::Write {
3852 path: new_file,
3853 metadata: MetadataMap::new(),
3854 },
3855 );
3856 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3857 }
3858
3859 #[derive(Clone)]
3860 struct HiddenTool {
3861 spec: ToolSpec,
3862 }
3863
3864 impl HiddenTool {
3865 fn new() -> Self {
3866 Self {
3867 spec: ToolSpec {
3868 name: ToolName::new("hidden"),
3869 description: "hidden".into(),
3870 input_schema: json!({"type": "object"}),
3871 output_schema: None,
3872 annotations: ToolAnnotations::default(),
3873 metadata: MetadataMap::new(),
3874 },
3875 }
3876 }
3877 }
3878
3879 #[async_trait]
3880 impl Tool for HiddenTool {
3881 fn spec(&self) -> &ToolSpec {
3882 &self.spec
3883 }
3884
3885 fn current_spec(&self) -> Option<ToolSpec> {
3886 None
3887 }
3888
3889 async fn invoke(
3890 &self,
3891 request: ToolRequest,
3892 _ctx: &mut ToolContext<'_>,
3893 ) -> Result<ToolResult, ToolError> {
3894 Ok(ToolResult {
3895 result: ToolResultPart {
3896 call_id: request.call_id,
3897 output: ToolOutput::Text("hidden".into()),
3898 is_error: false,
3899 metadata: MetadataMap::new(),
3900 },
3901 duration: None,
3902 metadata: MetadataMap::new(),
3903 })
3904 }
3905 }
3906
3907 #[test]
3908 fn hidden_tools_are_omitted_from_specs_and_capabilities() {
3909 let registry = ToolRegistry::new().with(HiddenTool::new());
3910
3911 assert!(registry.specs().is_empty());
3912
3913 let provider = ToolCapabilityProvider::from_registry(
3914 ®istry,
3915 Arc::new(AllowAllPermissionChecker),
3916 Arc::new(()),
3917 );
3918 assert!(provider.invocables().is_empty());
3919 }
3920
3921 struct AllowAllPermissionChecker;
3922
3923 impl PermissionChecker for AllowAllPermissionChecker {
3924 fn evaluate(&self, _request: &dyn PermissionRequest) -> PermissionDecision {
3925 PermissionDecision::Allow
3926 }
3927 }
3928
3929 #[derive(Clone)]
3932 struct PanickingSpecTool {
3933 spec: ToolSpec,
3934 }
3935
3936 impl PanickingSpecTool {
3937 fn new(name: &str) -> Self {
3938 Self {
3939 spec: ToolSpec {
3940 name: ToolName::new(name),
3941 description: "panics on current_spec".into(),
3942 input_schema: json!({"type": "object"}),
3943 output_schema: None,
3944 annotations: ToolAnnotations::default(),
3945 metadata: MetadataMap::new(),
3946 },
3947 }
3948 }
3949 }
3950
3951 #[async_trait]
3952 impl Tool for PanickingSpecTool {
3953 fn spec(&self) -> &ToolSpec {
3954 &self.spec
3955 }
3956
3957 fn current_spec(&self) -> Option<ToolSpec> {
3958 panic!("PanickingSpecTool::current_spec");
3959 }
3960
3961 async fn invoke(
3962 &self,
3963 request: ToolRequest,
3964 _ctx: &mut ToolContext<'_>,
3965 ) -> Result<ToolResult, ToolError> {
3966 Ok(ToolResult {
3967 result: ToolResultPart {
3968 call_id: request.call_id,
3969 output: ToolOutput::Text("never".into()),
3970 is_error: false,
3971 metadata: MetadataMap::new(),
3972 },
3973 duration: None,
3974 metadata: MetadataMap::new(),
3975 })
3976 }
3977 }
3978
3979 #[test]
3990 fn catalog_recovers_from_panicked_writer() {
3991 let (writer, reader) = dynamic_catalog("test");
3992
3993 writer.upsert(Arc::new(PanickingSpecTool::new("boom")));
3997 let _ = reader.drain_catalog_events();
3998
3999 let panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4004 writer.replace_all(vec![
4005 Arc::new(PanickingSpecTool::new("boom")) as Arc<dyn Tool>
4006 ]);
4007 }));
4008 assert!(
4009 panic_result.is_err(),
4010 "PanickingSpecTool::current_spec must propagate"
4011 );
4012
4013 assert!(
4017 reader.get(&ToolName::new("boom")).is_some(),
4018 "catalog still readable after poisoning panic"
4019 );
4020
4021 assert!(writer.remove(&ToolName::new("boom")));
4024
4025 writer.upsert(Arc::new(HiddenTool::new()));
4029 assert!(
4030 reader.get(&ToolName::new("hidden")).is_some(),
4031 "catalog usable for further writes + reads"
4032 );
4033 }
4034
4035 #[derive(Clone)]
4036 struct EchoTool {
4037 spec: ToolSpec,
4038 }
4039
4040 impl EchoTool {
4041 fn new(name: &str) -> Self {
4042 Self {
4043 spec: ToolSpec {
4044 name: ToolName::new(name),
4045 description: format!("echo {name}"),
4046 input_schema: json!({"type": "object"}),
4047 output_schema: None,
4048 annotations: ToolAnnotations::default(),
4049 metadata: MetadataMap::new(),
4050 },
4051 }
4052 }
4053 }
4054
4055 #[async_trait]
4056 impl Tool for EchoTool {
4057 fn spec(&self) -> &ToolSpec {
4058 &self.spec
4059 }
4060
4061 async fn invoke(
4062 &self,
4063 request: ToolRequest,
4064 _ctx: &mut ToolContext<'_>,
4065 ) -> Result<ToolResult, ToolError> {
4066 Ok(ToolResult::new(ToolResultPart::success(
4067 request.call_id,
4068 ToolOutput::text(request.tool_name.0.clone()),
4069 )))
4070 }
4071 }
4072
4073 fn registry_with(names: &[&str]) -> ToolRegistry {
4074 names.iter().fold(ToolRegistry::new(), |reg, name| {
4075 reg.with(EchoTool::new(name))
4076 })
4077 }
4078
4079 #[test]
4080 fn prefixed_rewrites_specs_and_resolves_lookups() {
4081 let source = registry_with(&["get_temp", "get_humidity"]).prefixed("weather");
4082 let names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4083 assert_eq!(names, vec!["weather_get_humidity", "weather_get_temp"]);
4084
4085 assert!(source.get(&ToolName::new("weather_get_temp")).is_some());
4086 assert!(
4087 source.get(&ToolName::new("get_temp")).is_none(),
4088 "original name must not resolve when prefixed"
4089 );
4090 assert!(source.get(&ToolName::new("unknown")).is_none());
4091 }
4092
4093 #[tokio::test]
4094 async fn prefixed_invoke_sees_inner_name_on_request() {
4095 let source = registry_with(&["get_temp"]).prefixed("weather");
4096 let tool = source.get(&ToolName::new("weather_get_temp")).unwrap();
4097
4098 assert_eq!(tool.spec().name.0, "weather_get_temp");
4100
4101 let owned = OwnedToolContext {
4103 session_id: SessionId::new("s"),
4104 turn_id: TurnId::new("t"),
4105 metadata: MetadataMap::new(),
4106 permissions: Arc::new(AllowAllPermissions),
4107 resources: Arc::new(()),
4108 cancellation: None,
4109 execution_scope: None,
4110 approved_request: None,
4111 };
4112 let mut ctx = owned.borrowed();
4113 let request = ToolRequest {
4114 call_id: ToolCallId::new("c"),
4115 tool_name: ToolName::new("weather_get_temp"),
4116 input: json!({}),
4117 session_id: SessionId::new("s"),
4118 turn_id: TurnId::new("t"),
4119 metadata: MetadataMap::new(),
4120 };
4121 let result = tool.invoke(request, &mut ctx).await.unwrap();
4122 match result.result.output {
4123 ToolOutput::Text(text) => assert_eq!(text, "get_temp"),
4124 other => panic!("unexpected output: {other:?}"),
4125 }
4126 }
4127
4128 #[derive(Clone)]
4129 struct StaticOutputTool {
4130 spec: ToolSpec,
4131 output: ToolOutput,
4132 }
4133
4134 impl StaticOutputTool {
4135 fn new(name: &str, output: ToolOutput) -> Self {
4136 Self {
4137 spec: ToolSpec::new(name, format!("static {name}"), json!({"type": "object"})),
4138 output,
4139 }
4140 }
4141
4142 fn with_output_limit(mut self, limit: ToolOutputLimit) -> Self {
4143 self.spec = self.spec.with_output_limit(limit);
4144 self
4145 }
4146 }
4147
4148 #[async_trait]
4149 impl Tool for StaticOutputTool {
4150 fn spec(&self) -> &ToolSpec {
4151 &self.spec
4152 }
4153
4154 async fn invoke(
4155 &self,
4156 request: ToolRequest,
4157 _ctx: &mut ToolContext<'_>,
4158 ) -> Result<ToolResult, ToolError> {
4159 Ok(ToolResult::new(ToolResultPart::success(
4160 request.call_id,
4161 self.output.clone(),
4162 )))
4163 }
4164 }
4165
4166 struct ApprovedContextTool {
4167 spec: ToolSpec,
4168 }
4169
4170 impl ApprovedContextTool {
4171 fn new() -> Self {
4172 Self {
4173 spec: ToolSpec::new(
4174 "approved_context",
4175 "approved context",
4176 json!({"type": "object"}),
4177 ),
4178 }
4179 }
4180 }
4181
4182 #[async_trait]
4183 impl Tool for ApprovedContextTool {
4184 fn spec(&self) -> &ToolSpec {
4185 &self.spec
4186 }
4187
4188 async fn invoke(
4189 &self,
4190 request: ToolRequest,
4191 ctx: &mut ToolContext<'_>,
4192 ) -> Result<ToolResult, ToolError> {
4193 Ok(ToolResult::new(ToolResultPart::success(
4194 request.call_id,
4195 ToolOutput::structured(json!({
4196 "approved": ctx.approved_request.is_some()
4197 })),
4198 )))
4199 }
4200 }
4201
4202 struct ScopeChildTool {
4203 spec: ToolSpec,
4204 }
4205
4206 impl ScopeChildTool {
4207 fn new() -> Self {
4208 Self {
4209 spec: ToolSpec::new("scope_child", "scope child", json!({"type": "object"})),
4210 }
4211 }
4212 }
4213
4214 #[async_trait]
4215 impl Tool for ScopeChildTool {
4216 fn spec(&self) -> &ToolSpec {
4217 &self.spec
4218 }
4219
4220 async fn invoke(
4221 &self,
4222 request: ToolRequest,
4223 _ctx: &mut ToolContext<'_>,
4224 ) -> Result<ToolResult, ToolError> {
4225 Ok(ToolResult::new(ToolResultPart::success(
4226 request.call_id,
4227 ToolOutput::structured(json!({ "child": request.input })),
4228 )))
4229 }
4230 }
4231
4232 struct ScopeParentTool {
4233 spec: ToolSpec,
4234 }
4235
4236 impl ScopeParentTool {
4237 fn new() -> Self {
4238 Self {
4239 spec: ToolSpec::new("scope_parent", "scope parent", json!({"type": "object"})),
4240 }
4241 }
4242 }
4243
4244 #[async_trait]
4245 impl Tool for ScopeParentTool {
4246 fn spec(&self) -> &ToolSpec {
4247 &self.spec
4248 }
4249
4250 async fn invoke(
4251 &self,
4252 request: ToolRequest,
4253 _ctx: &mut ToolContext<'_>,
4254 ) -> Result<ToolResult, ToolError> {
4255 Ok(ToolResult::new(ToolResultPart::success(
4256 request.call_id,
4257 ToolOutput::text("unused"),
4258 )))
4259 }
4260
4261 async fn invoke_outcome(
4262 &self,
4263 request: ToolRequest,
4264 ctx: &mut ToolContext<'_>,
4265 ) -> ToolExecutionOutcome {
4266 let Some(scope) = ctx.execution_scope.clone() else {
4267 return ToolExecutionOutcome::Failed(ToolError::Internal(
4268 "missing execution scope".into(),
4269 ));
4270 };
4271 let child = ToolRequest::new(
4272 "child-call",
4273 "scope_child",
4274 request.input.clone(),
4275 request.session_id.clone(),
4276 request.turn_id.clone(),
4277 );
4278 match scope.execute_child(child).await {
4279 ToolExecutionOutcome::Completed(child_result) => {
4280 ToolExecutionOutcome::Completed(ToolResult::new(ToolResultPart::success(
4281 request.call_id,
4282 child_result.result.output,
4283 )))
4284 }
4285 other => other,
4286 }
4287 }
4288 }
4289
4290 fn test_context() -> OwnedToolContext {
4291 OwnedToolContext {
4292 session_id: SessionId::new("s"),
4293 turn_id: TurnId::new("t"),
4294 metadata: MetadataMap::new(),
4295 permissions: Arc::new(AllowAllPermissions),
4296 resources: Arc::new(()),
4297 cancellation: None,
4298 execution_scope: None,
4299 approved_request: None,
4300 }
4301 }
4302
4303 fn test_context_with_scope(executor: Arc<dyn ToolExecutor>) -> OwnedToolContext {
4304 let session_id = SessionId::new("s");
4305 let turn_id = TurnId::new("t");
4306 let metadata = MetadataMap::new();
4307 let permissions: Arc<dyn PermissionChecker> = Arc::new(AllowAllPermissions);
4308 let resources: Arc<dyn ToolResources> = Arc::new(());
4309 let scope = ToolExecutionScope {
4310 executor,
4311 session_id: session_id.clone(),
4312 turn_id: turn_id.clone(),
4313 permissions: permissions.clone(),
4314 resources: resources.clone(),
4315 cancellation: None,
4316 };
4317 OwnedToolContext {
4318 session_id,
4319 turn_id,
4320 metadata,
4321 permissions,
4322 resources,
4323 cancellation: None,
4324 execution_scope: Some(scope),
4325 approved_request: None,
4326 }
4327 }
4328
4329 #[tokio::test]
4330 async fn default_invoke_outcome_wraps_invoke_success() {
4331 let executor = BasicToolExecutor::from_registry(ToolRegistry::new().with(
4332 StaticOutputTool::new("plain", ToolOutput::structured(json!({"ok": true}))),
4333 ));
4334 let outcome = executor
4335 .execute_owned(
4336 ToolRequest::new("call", "plain", json!({}), "s", "t"),
4337 test_context(),
4338 )
4339 .await;
4340
4341 let ToolExecutionOutcome::Completed(result) = outcome else {
4342 panic!("expected completed outcome, got {outcome:?}");
4343 };
4344 assert_eq!(
4345 result.result.output,
4346 ToolOutput::structured(json!({"ok": true}))
4347 );
4348 }
4349
4350 #[tokio::test]
4351 async fn execute_approved_passes_approval_context_to_tool() {
4352 let executor =
4353 BasicToolExecutor::from_registry(ToolRegistry::new().with(ApprovedContextTool::new()));
4354 let approval = ApprovalRequest {
4355 task_id: None,
4356 call_id: Some(ToolCallId::new("call")),
4357 id: ApprovalId::new("approval"),
4358 request_kind: "test.approval".into(),
4359 reason: ApprovalReason::PolicyRequiresConfirmation,
4360 summary: "approve".into(),
4361 metadata: MetadataMap::new(),
4362 };
4363 let outcome = executor
4364 .execute_approved_owned(
4365 ToolRequest::new("call", "approved_context", json!({}), "s", "t"),
4366 &approval,
4367 test_context(),
4368 )
4369 .await;
4370
4371 let ToolExecutionOutcome::Completed(result) = outcome else {
4372 panic!("expected completed outcome, got {outcome:?}");
4373 };
4374 assert_eq!(
4375 result.result.output,
4376 ToolOutput::structured(json!({"approved": true}))
4377 );
4378 }
4379
4380 #[tokio::test]
4381 async fn execution_scope_invokes_child_through_executor() {
4382 let executor: Arc<dyn ToolExecutor> = Arc::new(BasicToolExecutor::from_registry(
4383 ToolRegistry::new()
4384 .with(ScopeParentTool::new())
4385 .with(ScopeChildTool::new()),
4386 ));
4387 let outcome = executor
4388 .execute_owned(
4389 ToolRequest::new("parent-call", "scope_parent", json!({"value": 3}), "s", "t"),
4390 test_context_with_scope(executor.clone()),
4391 )
4392 .await;
4393
4394 let ToolExecutionOutcome::Completed(result) = outcome else {
4395 panic!("expected completed outcome, got {outcome:?}");
4396 };
4397 assert_eq!(
4398 result.result.output,
4399 ToolOutput::structured(json!({ "child": { "value": 3 } }))
4400 );
4401 }
4402
4403 #[tokio::test]
4404 async fn executor_stores_oversized_output_using_tool_metadata_limit() {
4405 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4406 let strategy = ConfigurableToolOutputTruncationStrategy::new(store.clone());
4407 let tool = StaticOutputTool::new("big", ToolOutput::text("x".repeat(500)))
4408 .with_output_limit(ToolOutputLimit::store_for_readback(300));
4409 let executor = BasicToolExecutor::from_registry(ToolRegistry::new().with(tool))
4410 .with_output_truncation_strategy(strategy);
4411
4412 let outcome = executor
4413 .execute_owned(
4414 ToolRequest::new(
4415 "call",
4416 "big",
4417 json!({}),
4418 SessionId::new("s"),
4419 TurnId::new("t"),
4420 ),
4421 test_context(),
4422 )
4423 .await;
4424
4425 let ToolExecutionOutcome::Completed(result) = outcome else {
4426 panic!("expected completed outcome, got {outcome:?}");
4427 };
4428 let ToolOutput::Structured(envelope) = result.result.output else {
4429 panic!("expected truncation envelope");
4430 };
4431 assert_eq!(envelope["truncated"], true);
4432 assert_eq!(envelope["read_tool"], TOOL_RESULT_READ_TOOL_NAME);
4433 let id = envelope["tool_result_id"].as_str().expect("tool_result_id");
4434
4435 let slice = store
4436 .read(&ToolOutputArtifactId(id.to_string()), 0, 50)
4437 .await
4438 .expect("read artifact");
4439 assert_eq!(slice.content, "x".repeat(50));
4440 assert_eq!(slice.next_offset, 50);
4441 assert!(!slice.eof);
4442 }
4443
4444 #[tokio::test]
4445 async fn tool_result_read_enforces_explicit_max_read_size() {
4446 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4447 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4448 let request = ToolRequest::new(
4449 "call",
4450 "big",
4451 json!({}),
4452 SessionId::new("s"),
4453 TurnId::new("t"),
4454 );
4455 let ctx = ToolOutputTruncationContext::from((&request, spec));
4456 let artifact = store
4457 .put(&ctx, "abcdef".to_string(), 6)
4458 .await
4459 .expect("store artifact");
4460 let tool = ToolResultReadTool::new(store, 4);
4461 let owned_ctx = test_context();
4462 let mut tool_ctx = owned_ctx.borrowed();
4463
4464 let err = tool
4465 .invoke(
4466 ToolRequest::new(
4467 "read-call",
4468 TOOL_RESULT_READ_TOOL_NAME,
4469 json!({"id": artifact.id.0, "offset": 0, "limit": 5}),
4470 SessionId::new("s"),
4471 TurnId::new("t"),
4472 ),
4473 &mut tool_ctx,
4474 )
4475 .await
4476 .expect_err("read past max must fail");
4477 match err {
4478 ToolError::InvalidInput(message) => assert!(message.contains("exceeds maximum")),
4479 other => panic!("expected InvalidInput, got {other:?}"),
4480 }
4481 }
4482
4483 #[tokio::test]
4484 async fn tool_result_read_rejects_zero_limit() {
4485 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4486 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4487 let request = ToolRequest::new(
4488 "call",
4489 "big",
4490 json!({}),
4491 SessionId::new("s"),
4492 TurnId::new("t"),
4493 );
4494 let ctx = ToolOutputTruncationContext::from((&request, spec));
4495 let artifact = store
4496 .put(&ctx, "abcdef".to_string(), 6)
4497 .await
4498 .expect("store artifact");
4499 let tool = ToolResultReadTool::new(store, 4);
4500 let owned_ctx = test_context();
4501 let mut tool_ctx = owned_ctx.borrowed();
4502
4503 let err = tool
4504 .invoke(
4505 ToolRequest::new(
4506 "read-call",
4507 TOOL_RESULT_READ_TOOL_NAME,
4508 json!({"id": artifact.id.0, "offset": 0, "limit": 0}),
4509 SessionId::new("s"),
4510 TurnId::new("t"),
4511 ),
4512 &mut tool_ctx,
4513 )
4514 .await
4515 .expect_err("zero limit must fail");
4516 match err {
4517 ToolError::InvalidInput(message) => assert!(message.contains("greater than 0")),
4518 other => panic!("expected InvalidInput, got {other:?}"),
4519 }
4520 }
4521
4522 #[tokio::test]
4523 async fn tool_result_read_executor_allows_full_content_limit_with_envelope() {
4524 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4525 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4526 let request = ToolRequest::new(
4527 "call",
4528 "big",
4529 json!({}),
4530 SessionId::new("s"),
4531 TurnId::new("t"),
4532 );
4533 let ctx = ToolOutputTruncationContext::from((&request, spec));
4534 let artifact = store
4535 .put(&ctx, "abcd".to_string(), 4)
4536 .await
4537 .expect("store artifact");
4538 let executor = BasicToolExecutor::from_registry(
4539 ToolRegistry::new().with(ToolResultReadTool::new(store.clone(), 4)),
4540 )
4541 .with_output_truncation_strategy(ConfigurableToolOutputTruncationStrategy::new(store));
4542
4543 let outcome = executor
4544 .execute_owned(
4545 ToolRequest::new(
4546 "read-call",
4547 TOOL_RESULT_READ_TOOL_NAME,
4548 json!({"id": artifact.id.0, "offset": 0, "limit": 4}),
4549 SessionId::new("s"),
4550 TurnId::new("t"),
4551 ),
4552 test_context(),
4553 )
4554 .await;
4555
4556 let ToolExecutionOutcome::Completed(result) = outcome else {
4557 panic!("expected completed outcome, got {outcome:?}");
4558 };
4559 let ToolOutput::Structured(output) = result.result.output else {
4560 panic!("expected structured readback output");
4561 };
4562 assert_eq!(output["content"], "abcd");
4563 assert_eq!(output["eof"], true);
4564 }
4565
4566 #[tokio::test]
4567 async fn tool_result_read_executor_allows_json_escaped_full_content_limit() {
4568 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4569 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4570 let request = ToolRequest::new(
4571 "call",
4572 "big",
4573 json!({}),
4574 SessionId::new("s"),
4575 TurnId::new("t"),
4576 );
4577 let ctx = ToolOutputTruncationContext::from((&request, spec));
4578 let content = "\0".repeat(4);
4579 let artifact = store
4580 .put(&ctx, content.clone(), content.len())
4581 .await
4582 .expect("store artifact");
4583 let executor = BasicToolExecutor::from_registry(
4584 ToolRegistry::new().with(ToolResultReadTool::new(store.clone(), 4)),
4585 )
4586 .with_output_truncation_strategy(ConfigurableToolOutputTruncationStrategy::new(store));
4587
4588 let outcome = executor
4589 .execute_owned(
4590 ToolRequest::new(
4591 "read-call",
4592 TOOL_RESULT_READ_TOOL_NAME,
4593 json!({"id": artifact.id.0, "offset": 0, "limit": 4}),
4594 SessionId::new("s"),
4595 TurnId::new("t"),
4596 ),
4597 test_context(),
4598 )
4599 .await;
4600
4601 let ToolExecutionOutcome::Completed(result) = outcome else {
4602 panic!("expected completed outcome, got {outcome:?}");
4603 };
4604 let ToolOutput::Structured(output) = result.result.output else {
4605 panic!("expected structured readback output");
4606 };
4607 assert_eq!(output["content"], content);
4608 assert_eq!(output["eof"], true);
4609 }
4610
4611 #[test]
4612 fn inline_clip_respects_limit_when_marker_exceeds_budget() {
4613 let clipped = clip_string_with_marker("abcdef", 8, 1000);
4614
4615 assert!(clipped.len() <= 8);
4616 assert!(clipped.is_char_boundary(clipped.len()));
4617 }
4618
4619 #[test]
4620 fn filtered_hides_tools_rejected_by_predicate() {
4621 let source = registry_with(&["safe", "danger_drop", "danger_delete"])
4622 .filtered(|name| !name.0.starts_with("danger_"));
4623 let names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4624 assert_eq!(names, vec!["safe"]);
4625
4626 assert!(source.get(&ToolName::new("safe")).is_some());
4627 assert!(source.get(&ToolName::new("danger_drop")).is_none());
4628 }
4629
4630 #[test]
4631 fn unadvertised_dispatches_without_advertising() {
4632 let source = registry_with(&["hidden_a", "hidden_b"]).unadvertised();
4633 assert!(source.specs().is_empty());
4634 assert!(source.get(&ToolName::new("hidden_a")).is_some());
4635 assert!(source.get(&ToolName::new("hidden_b")).is_some());
4636 assert!(source.get(&ToolName::new("missing")).is_none());
4637 assert!(source.drain_catalog_events().is_empty());
4638 }
4639
4640 #[test]
4641 fn renamed_remaps_specs_and_lookups() {
4642 let source = registry_with(&["legacy_name", "passthrough"])
4643 .renamed([(ToolName::new("legacy_name"), ToolName::new("modern_name"))]);
4644 let mut names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4645 names.sort();
4646 assert_eq!(names, vec!["modern_name", "passthrough"]);
4647
4648 assert!(source.get(&ToolName::new("modern_name")).is_some());
4649 assert!(
4650 source.get(&ToolName::new("legacy_name")).is_none(),
4651 "original name is hidden after renaming"
4652 );
4653 assert!(source.get(&ToolName::new("passthrough")).is_some());
4654 }
4655
4656 #[cfg(feature = "schemars")]
4657 mod schemars_helpers {
4658 use super::*;
4659 use schemars::JsonSchema;
4660 use serde::Deserialize;
4661
4662 #[derive(JsonSchema, Deserialize)]
4663 #[allow(dead_code)]
4664 struct WeatherInput {
4665 location: String,
4667 #[serde(default)]
4669 celsius: bool,
4670 }
4671
4672 #[test]
4673 fn schema_for_emits_object_schema_with_typed_fields() {
4674 let schema = schema_for::<WeatherInput>();
4675 let obj = schema.as_object().expect("schema is a JSON object");
4676 assert_eq!(
4677 obj.get("type").and_then(|v| v.as_str()),
4678 Some("object"),
4679 "root type should be object"
4680 );
4681 let properties = obj
4682 .get("properties")
4683 .and_then(|v| v.as_object())
4684 .expect("properties block");
4685 assert!(properties.contains_key("location"));
4686 assert!(properties.contains_key("celsius"));
4687 }
4688
4689 #[test]
4690 fn tool_spec_for_carries_schema_name_and_description() {
4691 let spec = tool_spec_for::<WeatherInput>("get_weather", "Fetch current weather");
4692 assert_eq!(spec.name.0, "get_weather");
4693 assert_eq!(spec.description, "Fetch current weather");
4694 assert!(spec.input_schema.is_object());
4695 }
4696 }
4697
4698 #[test]
4699 fn transforms_compose_via_chained_methods() {
4700 let source = registry_with(&["read_file", "write_file", "delete_file"])
4701 .filtered(|name| name.0 != "delete_file")
4702 .prefixed("fs");
4703 let mut names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4704 names.sort();
4705 assert_eq!(names, vec!["fs_read_file", "fs_write_file"]);
4706 }
4707}