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
2458impl ToolSource for ToolRegistry {
2459 fn specs(&self) -> Vec<ToolSpec> {
2460 ToolRegistry::specs(self)
2461 }
2462
2463 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2464 ToolRegistry::get(self, name)
2465 }
2466}
2467
2468impl<S> ToolSource for Arc<S>
2469where
2470 S: ToolSource + ?Sized,
2471{
2472 fn specs(&self) -> Vec<ToolSpec> {
2473 (**self).specs()
2474 }
2475
2476 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2477 (**self).get(name)
2478 }
2479
2480 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2481 (**self).drain_catalog_events()
2482 }
2483}
2484
2485pub struct Prefixed<S> {
2488 inner: S,
2489 prefix: String,
2490}
2491
2492impl<S> Prefixed<S> {
2493 pub fn new(inner: S, prefix: impl Into<String>) -> Self {
2495 Self {
2496 inner,
2497 prefix: prefix.into(),
2498 }
2499 }
2500
2501 fn rewrite(&self, name: &str) -> String {
2502 format!("{}_{}", self.prefix, name)
2503 }
2504
2505 fn strip<'a>(&self, name: &'a str) -> Option<&'a str> {
2506 name.strip_prefix(self.prefix.as_str())
2507 .and_then(|rest| rest.strip_prefix('_'))
2508 }
2509}
2510
2511impl<S> ToolSource for Prefixed<S>
2512where
2513 S: ToolSource,
2514{
2515 fn specs(&self) -> Vec<ToolSpec> {
2516 self.inner
2517 .specs()
2518 .into_iter()
2519 .map(|mut spec| {
2520 spec.name = ToolName::new(self.rewrite(spec.name.0.as_str()));
2521 spec
2522 })
2523 .collect()
2524 }
2525
2526 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2527 let original = self.strip(name.0.as_str())?;
2528 let inner_name = ToolName::new(original);
2529 let inner_tool = self.inner.get(&inner_name)?;
2530 let mut public_spec = inner_tool.spec().clone();
2531 public_spec.name = name.clone();
2532 Some(Arc::new(RewrittenTool {
2533 inner: inner_tool,
2534 inner_name,
2535 public_spec,
2536 }))
2537 }
2538
2539 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2540 self.inner
2541 .drain_catalog_events()
2542 .into_iter()
2543 .map(|mut event| {
2544 event.for_each_name_mut(|name| *name = self.rewrite(name.as_str()));
2545 event
2546 })
2547 .collect()
2548 }
2549}
2550
2551pub struct Filtered<S, F> {
2554 inner: S,
2555 predicate: F,
2556}
2557
2558impl<S, F> Filtered<S, F> {
2559 pub fn new(inner: S, predicate: F) -> Self {
2561 Self { inner, predicate }
2562 }
2563}
2564
2565impl<S, F> ToolSource for Filtered<S, F>
2566where
2567 S: ToolSource,
2568 F: Fn(&ToolName) -> bool + Send + Sync + 'static,
2569{
2570 fn specs(&self) -> Vec<ToolSpec> {
2571 self.inner
2572 .specs()
2573 .into_iter()
2574 .filter(|spec| (self.predicate)(&spec.name))
2575 .collect()
2576 }
2577
2578 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2579 if !(self.predicate)(name) {
2580 return None;
2581 }
2582 self.inner.get(name)
2583 }
2584
2585 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2586 self.inner
2587 .drain_catalog_events()
2588 .into_iter()
2589 .map(|mut event| {
2590 event.retain_names(|n| (self.predicate)(&ToolName::new(n)));
2591 event
2592 })
2593 .collect()
2594 }
2595}
2596
2597pub struct Renamed<S> {
2604 inner: S,
2605 forward: BTreeMap<ToolName, ToolName>,
2606 backward: BTreeMap<ToolName, ToolName>,
2607}
2608
2609impl<S> Renamed<S> {
2610 pub fn new<I>(inner: S, mapping: I) -> Self
2612 where
2613 I: IntoIterator<Item = (ToolName, ToolName)>,
2614 {
2615 let forward: BTreeMap<ToolName, ToolName> = mapping.into_iter().collect();
2616 let backward = forward
2617 .iter()
2618 .map(|(k, v)| (v.clone(), k.clone()))
2619 .collect();
2620 Self {
2621 inner,
2622 forward,
2623 backward,
2624 }
2625 }
2626}
2627
2628impl<S> ToolSource for Renamed<S>
2629where
2630 S: ToolSource,
2631{
2632 fn specs(&self) -> Vec<ToolSpec> {
2633 self.inner
2634 .specs()
2635 .into_iter()
2636 .map(|mut spec| {
2637 if let Some(new_name) = self.forward.get(&spec.name) {
2638 spec.name = new_name.clone();
2639 }
2640 spec
2641 })
2642 .collect()
2643 }
2644
2645 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
2646 if let Some(original) = self.backward.get(name) {
2647 let inner_tool = self.inner.get(original)?;
2648 let mut public_spec = inner_tool.spec().clone();
2649 public_spec.name = name.clone();
2650 Some(Arc::new(RewrittenTool {
2651 inner: inner_tool,
2652 inner_name: original.clone(),
2653 public_spec,
2654 }))
2655 } else if self.forward.contains_key(name) {
2656 None
2658 } else {
2659 self.inner.get(name)
2660 }
2661 }
2662
2663 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
2664 self.inner
2665 .drain_catalog_events()
2666 .into_iter()
2667 .map(|mut event| {
2668 event.for_each_name_mut(|name| {
2669 if let Some(new) = self.forward.get(&ToolName::new(name.as_str())) {
2670 *name = new.0.clone();
2671 }
2672 });
2673 event
2674 })
2675 .collect()
2676 }
2677}
2678
2679#[cfg(feature = "schemars")]
2705pub fn schema_for<T: schemars::JsonSchema>() -> Value {
2706 let schema = schemars::schema_for!(T);
2707 serde_json::to_value(schema)
2708 .expect("schemars produces valid JSON; this conversion is infallible")
2709}
2710
2711#[cfg(feature = "schemars")]
2729pub fn tool_spec_for<T: schemars::JsonSchema>(
2730 name: impl Into<ToolName>,
2731 description: impl Into<String>,
2732) -> ToolSpec {
2733 ToolSpec::new(name, description, schema_for::<T>())
2734}
2735
2736struct RewrittenTool {
2742 inner: Arc<dyn Tool>,
2743 inner_name: ToolName,
2744 public_spec: ToolSpec,
2745}
2746
2747#[async_trait]
2748impl Tool for RewrittenTool {
2749 fn spec(&self) -> &ToolSpec {
2750 &self.public_spec
2751 }
2752
2753 fn current_spec(&self) -> Option<ToolSpec> {
2754 let inner_current = self.inner.current_spec()?;
2755 Some(ToolSpec {
2756 name: self.public_spec.name.clone(),
2757 description: inner_current.description,
2758 input_schema: inner_current.input_schema,
2759 output_schema: inner_current.output_schema,
2760 annotations: inner_current.annotations,
2761 metadata: inner_current.metadata,
2762 })
2763 }
2764
2765 fn proposed_requests(
2766 &self,
2767 request: &ToolRequest,
2768 ) -> Result<Vec<Box<dyn PermissionRequest>>, ToolError> {
2769 let mut inner_request = request.clone();
2770 inner_request.tool_name = self.inner_name.clone();
2771 self.inner.proposed_requests(&inner_request)
2772 }
2773
2774 async fn invoke(
2775 &self,
2776 mut request: ToolRequest,
2777 ctx: &mut ToolContext<'_>,
2778 ) -> Result<ToolResult, ToolError> {
2779 request.tool_name = self.inner_name.clone();
2780 self.inner.invoke(request, ctx).await
2781 }
2782
2783 async fn invoke_outcome(
2784 &self,
2785 mut request: ToolRequest,
2786 ctx: &mut ToolContext<'_>,
2787 ) -> ToolExecutionOutcome {
2788 request.tool_name = self.inner_name.clone();
2789 self.inner.invoke_outcome(request, ctx).await
2790 }
2791}
2792
2793struct ToolMap {
2811 inner: std::sync::RwLock<BTreeMap<ToolName, Arc<dyn Tool>>>,
2812}
2813
2814impl ToolMap {
2815 fn new() -> Self {
2816 Self {
2817 inner: std::sync::RwLock::new(BTreeMap::new()),
2818 }
2819 }
2820
2821 fn read(&self) -> std::sync::RwLockReadGuard<'_, BTreeMap<ToolName, Arc<dyn Tool>>> {
2822 self.inner.read().unwrap_or_else(|e| e.into_inner())
2823 }
2824
2825 fn write(&self) -> std::sync::RwLockWriteGuard<'_, BTreeMap<ToolName, Arc<dyn Tool>>> {
2826 self.inner.write().unwrap_or_else(|e| e.into_inner())
2827 }
2828}
2829
2830struct DynamicCatalogInner {
2833 source_id: String,
2834 tools: ToolMap,
2835 events_tx: tokio::sync::broadcast::Sender<ToolCatalogEvent>,
2836}
2837
2838pub fn dynamic_catalog(source_id: impl Into<String>) -> (CatalogWriter, CatalogReader) {
2856 let (events_tx, events_rx) = tokio::sync::broadcast::channel(128);
2857 let inner = Arc::new(DynamicCatalogInner {
2858 source_id: source_id.into(),
2859 tools: ToolMap::new(),
2860 events_tx,
2861 });
2862 (
2863 CatalogWriter {
2864 inner: Arc::clone(&inner),
2865 },
2866 CatalogReader {
2867 inner,
2868 events_rx: std::sync::Mutex::new(events_rx),
2869 },
2870 )
2871}
2872
2873pub struct CatalogWriter {
2880 inner: Arc<DynamicCatalogInner>,
2881}
2882
2883impl CatalogWriter {
2884 pub fn source_id(&self) -> &str {
2886 &self.inner.source_id
2887 }
2888
2889 pub fn reader(&self) -> CatalogReader {
2893 CatalogReader {
2894 inner: Arc::clone(&self.inner),
2895 events_rx: std::sync::Mutex::new(self.inner.events_tx.subscribe()),
2896 }
2897 }
2898
2899 pub fn upsert(&self, tool: Arc<dyn Tool>) {
2902 let name = tool.spec().name.clone();
2903 let mut guard = self.inner.tools.write();
2904 let existed = guard.insert(name.clone(), tool).is_some();
2905 drop(guard);
2906 let mut event = ToolCatalogEvent::new(self.inner.source_id.clone());
2907 if existed {
2908 event.changed.push(name.0);
2909 } else {
2910 event.added.push(name.0);
2911 }
2912 let _ = self.inner.events_tx.send(event);
2913 }
2914
2915 pub fn remove(&self, name: &ToolName) -> bool {
2918 let mut guard = self.inner.tools.write();
2919 let removed = guard.remove(name).is_some();
2920 drop(guard);
2921 if removed {
2922 let mut event = ToolCatalogEvent::new(self.inner.source_id.clone());
2923 event.removed.push(name.0.clone());
2924 let _ = self.inner.events_tx.send(event);
2925 }
2926 removed
2927 }
2928
2929 pub fn replace_all(&self, tools: impl IntoIterator<Item = Arc<dyn Tool>>) {
2932 let new_map: BTreeMap<ToolName, Arc<dyn Tool>> = tools
2933 .into_iter()
2934 .map(|tool| (tool.spec().name.clone(), tool))
2935 .collect();
2936
2937 let mut guard = self.inner.tools.write();
2938 let mut event = ToolCatalogEvent::new(self.inner.source_id.clone());
2939
2940 for (name, new_tool) in new_map.iter() {
2941 match guard.get(name) {
2942 None => event.added.push(name.0.clone()),
2943 Some(existing)
2944 if !Arc::ptr_eq(existing, new_tool)
2945 && existing.current_spec() != new_tool.current_spec() =>
2946 {
2947 event.changed.push(name.0.clone());
2948 }
2949 Some(_) => {}
2950 }
2951 }
2952 for name in guard.keys() {
2953 if !new_map.contains_key(name) {
2954 event.removed.push(name.0.clone());
2955 }
2956 }
2957
2958 *guard = new_map;
2959 drop(guard);
2960
2961 if !event.added.is_empty() || !event.removed.is_empty() || !event.changed.is_empty() {
2962 let _ = self.inner.events_tx.send(event);
2963 }
2964 }
2965
2966 pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<ToolCatalogEvent> {
2970 self.inner.events_tx.subscribe()
2971 }
2972}
2973
2974pub struct CatalogReader {
2979 inner: Arc<DynamicCatalogInner>,
2980 events_rx: std::sync::Mutex<tokio::sync::broadcast::Receiver<ToolCatalogEvent>>,
2981}
2982
2983impl CatalogReader {
2984 pub fn source_id(&self) -> &str {
2986 &self.inner.source_id
2987 }
2988
2989 pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<ToolCatalogEvent> {
2992 self.inner.events_tx.subscribe()
2993 }
2994}
2995
2996impl Clone for CatalogReader {
2997 fn clone(&self) -> Self {
2998 Self {
2999 inner: Arc::clone(&self.inner),
3000 events_rx: std::sync::Mutex::new(self.inner.events_tx.subscribe()),
3001 }
3002 }
3003}
3004
3005impl ToolSource for CatalogReader {
3006 fn specs(&self) -> Vec<ToolSpec> {
3007 self.inner
3008 .tools
3009 .read()
3010 .values()
3011 .filter_map(|tool| tool.current_spec())
3012 .collect()
3013 }
3014
3015 fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
3016 self.inner.tools.read().get(name).cloned()
3017 }
3018
3019 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3020 let mut rx = self.events_rx.lock().unwrap_or_else(|e| e.into_inner());
3025 let mut out = Vec::new();
3026 loop {
3027 match rx.try_recv() {
3028 Ok(event) => out.push(event),
3029 Err(tokio::sync::broadcast::error::TryRecvError::Empty) => break,
3030 Err(tokio::sync::broadcast::error::TryRecvError::Closed) => break,
3031 Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => continue,
3032 }
3033 }
3034 out
3035 }
3036}
3037
3038impl ToolSpec {
3039 pub fn as_invocable_spec(&self) -> InvocableSpec {
3042 InvocableSpec::new(
3043 CapabilityName::new(self.name.0.clone()),
3044 self.description.clone(),
3045 self.input_schema.clone(),
3046 )
3047 .with_metadata(self.metadata.clone())
3048 }
3049}
3050
3051pub struct ToolInvocableAdapter {
3057 spec: InvocableSpec,
3058 tool: Arc<dyn Tool>,
3059 permissions: Arc<dyn PermissionChecker>,
3060 resources: Arc<dyn ToolResources>,
3061 next_call_id: AtomicU64,
3062}
3063
3064impl ToolInvocableAdapter {
3065 pub fn new(
3068 tool: Arc<dyn Tool>,
3069 permissions: Arc<dyn PermissionChecker>,
3070 resources: Arc<dyn ToolResources>,
3071 ) -> Option<Self> {
3072 let spec = tool.current_spec()?.as_invocable_spec();
3073 Some(Self {
3074 spec,
3075 tool,
3076 permissions,
3077 resources,
3078 next_call_id: AtomicU64::new(1),
3079 })
3080 }
3081}
3082
3083#[async_trait]
3084impl Invocable for ToolInvocableAdapter {
3085 fn spec(&self) -> &InvocableSpec {
3086 &self.spec
3087 }
3088
3089 async fn invoke(
3090 &self,
3091 request: InvocableRequest,
3092 ctx: &mut CapabilityContext<'_>,
3093 ) -> Result<InvocableResult, CapabilityError> {
3094 let tool_request = ToolRequest {
3095 call_id: ToolCallId::new(format!(
3096 "tool-call-{}",
3097 self.next_call_id.fetch_add(1, Ordering::Relaxed)
3098 )),
3099 tool_name: self.tool.spec().name.clone(),
3100 input: request.input,
3101 session_id: ctx
3102 .session_id
3103 .cloned()
3104 .unwrap_or_else(|| SessionId::new("capability-session")),
3105 turn_id: ctx
3106 .turn_id
3107 .cloned()
3108 .unwrap_or_else(|| TurnId::new("capability-turn")),
3109 metadata: request.metadata,
3110 };
3111
3112 for permission_request in self
3113 .tool
3114 .proposed_requests(&tool_request)
3115 .map_err(|error| CapabilityError::InvalidInput(error.to_string()))?
3116 {
3117 match self.permissions.evaluate(permission_request.as_ref()) {
3118 PermissionDecision::Allow => {}
3119 PermissionDecision::Deny(denial) => {
3120 return Err(CapabilityError::ExecutionFailed(format!(
3121 "tool permission denied: {denial:?}"
3122 )));
3123 }
3124 PermissionDecision::RequireApproval(req) => {
3125 return Err(CapabilityError::Unavailable(format!(
3126 "tool invocation requires approval: {}",
3127 req.summary
3128 )));
3129 }
3130 }
3131 }
3132
3133 let mut tool_ctx = ToolContext {
3134 capability: CapabilityContext {
3135 session_id: ctx.session_id,
3136 turn_id: ctx.turn_id,
3137 metadata: ctx.metadata,
3138 },
3139 permissions: self.permissions.as_ref(),
3140 resources: self.resources.as_ref(),
3141 cancellation: None,
3142 execution_scope: None,
3143 approved_request: None,
3144 };
3145
3146 let result = self
3147 .tool
3148 .invoke(tool_request, &mut tool_ctx)
3149 .await
3150 .map_err(|error| CapabilityError::ExecutionFailed(error.to_string()))?;
3151
3152 Ok(InvocableResult {
3153 output: match result.result.output {
3154 ToolOutput::Text(text) => InvocableOutput::Text(text),
3155 ToolOutput::Structured(value) => InvocableOutput::Structured(value),
3156 ToolOutput::Parts(parts) => InvocableOutput::Items(vec![Item {
3157 id: None,
3158 kind: ItemKind::Tool,
3159 parts,
3160 metadata: MetadataMap::new(),
3161 usage: None,
3162 finish_reason: None,
3163 created_at: None,
3164 }]),
3165 ToolOutput::Files(files) => {
3166 let parts = files.into_iter().map(Part::File).collect();
3167 InvocableOutput::Items(vec![Item {
3168 id: None,
3169 kind: ItemKind::Tool,
3170 parts,
3171 metadata: MetadataMap::new(),
3172 usage: None,
3173 finish_reason: None,
3174 created_at: None,
3175 }])
3176 }
3177 },
3178 metadata: result.metadata,
3179 })
3180 }
3181}
3182
3183pub struct ToolCapabilityProvider {
3189 invocables: Vec<Arc<dyn Invocable>>,
3190}
3191
3192impl ToolCapabilityProvider {
3193 pub fn from_registry(
3196 registry: &ToolRegistry,
3197 permissions: Arc<dyn PermissionChecker>,
3198 resources: Arc<dyn ToolResources>,
3199 ) -> Self {
3200 let invocables = registry
3201 .tools()
3202 .into_iter()
3203 .filter_map(|tool| {
3204 ToolInvocableAdapter::new(tool, permissions.clone(), resources.clone())
3205 .map(|adapter| Arc::new(adapter) as Arc<dyn Invocable>)
3206 })
3207 .collect();
3208
3209 Self { invocables }
3210 }
3211}
3212
3213impl CapabilityProvider for ToolCapabilityProvider {
3214 fn invocables(&self) -> Vec<Arc<dyn Invocable>> {
3215 self.invocables.clone()
3216 }
3217
3218 fn resources(&self) -> Vec<Arc<dyn ResourceProvider>> {
3219 Vec::new()
3220 }
3221
3222 fn prompts(&self) -> Vec<Arc<dyn PromptProvider>> {
3223 Vec::new()
3224 }
3225}
3226
3227#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
3233pub enum ToolExecutionOutcome {
3234 Completed(ToolResult),
3236 Interrupted(ToolInterruption),
3238 Failed(ToolError),
3240}
3241
3242#[async_trait]
3250pub trait ToolExecutor: Send + Sync {
3251 fn specs(&self) -> Vec<ToolSpec>;
3253
3254 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3259 Vec::new()
3260 }
3261
3262 async fn execute(
3264 &self,
3265 request: ToolRequest,
3266 ctx: &mut ToolContext<'_>,
3267 ) -> ToolExecutionOutcome;
3268
3269 async fn execute_owned(
3272 &self,
3273 request: ToolRequest,
3274 ctx: OwnedToolContext,
3275 ) -> ToolExecutionOutcome {
3276 let mut borrowed = ctx.borrowed();
3277 self.execute(request, &mut borrowed).await
3278 }
3279
3280 async fn execute_approved(
3286 &self,
3287 request: ToolRequest,
3288 approved_request: &ApprovalRequest,
3289 ctx: &mut ToolContext<'_>,
3290 ) -> ToolExecutionOutcome {
3291 let _ = approved_request;
3292 self.execute(request, ctx).await
3293 }
3294
3295 async fn execute_approved_owned(
3298 &self,
3299 request: ToolRequest,
3300 approved_request: &ApprovalRequest,
3301 mut ctx: OwnedToolContext,
3302 ) -> ToolExecutionOutcome {
3303 ctx.approved_request = Some(approved_request.clone());
3304 let mut borrowed = ctx.borrowed();
3305 self.execute_approved(request, approved_request, &mut borrowed)
3306 .await
3307 }
3308}
3309
3310#[async_trait]
3311impl<T> ToolExecutor for Arc<T>
3312where
3313 T: ToolExecutor + ?Sized,
3314{
3315 fn specs(&self) -> Vec<ToolSpec> {
3316 (**self).specs()
3317 }
3318
3319 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3320 (**self).drain_catalog_events()
3321 }
3322
3323 async fn execute(
3324 &self,
3325 request: ToolRequest,
3326 ctx: &mut ToolContext<'_>,
3327 ) -> ToolExecutionOutcome {
3328 (**self).execute(request, ctx).await
3329 }
3330
3331 async fn execute_approved(
3332 &self,
3333 request: ToolRequest,
3334 approved_request: &ApprovalRequest,
3335 ctx: &mut ToolContext<'_>,
3336 ) -> ToolExecutionOutcome {
3337 (**self)
3338 .execute_approved(request, approved_request, ctx)
3339 .await
3340 }
3341}
3342
3343#[derive(Clone, Debug, Default, PartialEq, Eq)]
3346pub enum CollisionPolicy {
3347 #[default]
3350 FirstWins,
3351 LastWins,
3353}
3354
3355pub struct BasicToolExecutor {
3374 sources: Vec<Arc<dyn ToolSource>>,
3375 collision: CollisionPolicy,
3376 output_truncation: Option<Arc<dyn ToolOutputTruncationStrategy>>,
3377}
3378
3379impl BasicToolExecutor {
3380 pub fn new(sources: impl IntoIterator<Item = Arc<dyn ToolSource>>) -> Self {
3382 Self {
3383 sources: sources.into_iter().collect(),
3384 collision: CollisionPolicy::default(),
3385 output_truncation: None,
3386 }
3387 }
3388
3389 pub fn from_registry(registry: ToolRegistry) -> Self {
3391 Self::new([Arc::new(registry) as Arc<dyn ToolSource>])
3392 }
3393
3394 pub fn with_collision_policy(mut self, policy: CollisionPolicy) -> Self {
3397 self.collision = policy;
3398 self
3399 }
3400
3401 pub fn with_output_truncation_strategy(
3405 mut self,
3406 strategy: impl ToolOutputTruncationStrategy + 'static,
3407 ) -> Self {
3408 self.output_truncation = Some(Arc::new(strategy));
3409 self
3410 }
3411
3412 pub fn with_output_truncation_strategy_arc(
3414 mut self,
3415 strategy: Arc<dyn ToolOutputTruncationStrategy>,
3416 ) -> Self {
3417 self.output_truncation = Some(strategy);
3418 self
3419 }
3420
3421 pub fn specs(&self) -> Vec<ToolSpec> {
3424 let mut seen = BTreeSet::new();
3425 let mut out = Vec::new();
3426 let iter: Box<dyn Iterator<Item = &Arc<dyn ToolSource>>> = match self.collision {
3427 CollisionPolicy::FirstWins => Box::new(self.sources.iter()),
3428 CollisionPolicy::LastWins => Box::new(self.sources.iter().rev()),
3429 };
3430 for source in iter {
3431 for spec in source.specs() {
3432 if seen.insert(spec.name.clone()) {
3433 out.push(spec);
3434 }
3435 }
3436 }
3437 out
3438 }
3439
3440 fn lookup(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
3441 match self.collision {
3442 CollisionPolicy::FirstWins => self.sources.iter().find_map(|s| s.get(name)),
3443 CollisionPolicy::LastWins => self.sources.iter().rev().find_map(|s| s.get(name)),
3444 }
3445 }
3446
3447 async fn execute_inner(
3448 &self,
3449 request: ToolRequest,
3450 approved_request_id: Option<&ApprovalId>,
3451 ctx: &mut ToolContext<'_>,
3452 ) -> ToolExecutionOutcome {
3453 let Some(tool) = self.lookup(&request.tool_name) else {
3454 return ToolExecutionOutcome::Failed(ToolError::NotFound(request.tool_name));
3455 };
3456
3457 match tool.proposed_requests(&request) {
3458 Ok(requests) => {
3459 for permission_request in requests {
3460 match ctx.permissions.evaluate(permission_request.as_ref()) {
3461 PermissionDecision::Allow => {}
3462 PermissionDecision::Deny(denial) => {
3463 return ToolExecutionOutcome::Failed(ToolError::PermissionDenied(
3464 denial,
3465 ));
3466 }
3467 PermissionDecision::RequireApproval(mut req) => {
3468 req.call_id = Some(request.call_id.clone());
3469 if approved_request_id != Some(&req.id) {
3470 return ToolExecutionOutcome::Interrupted(
3471 ToolInterruption::ApprovalRequired(req),
3472 );
3473 }
3474 }
3475 }
3476 }
3477 }
3478 Err(error) => return ToolExecutionOutcome::Failed(error),
3479 }
3480
3481 let truncation_ctx = ToolOutputTruncationContext::from((&request, tool.spec().clone()));
3482 match tool.invoke_outcome(request, ctx).await {
3483 ToolExecutionOutcome::Completed(mut result) => {
3484 if let Some(strategy) = &self.output_truncation {
3485 match strategy.apply(truncation_ctx, result.result.output).await {
3486 Ok(output) => {
3487 result.result.output = output;
3488 }
3489 Err(error) => return ToolExecutionOutcome::Failed(error),
3490 }
3491 }
3492 ToolExecutionOutcome::Completed(result)
3493 }
3494 other => other,
3495 }
3496 }
3497}
3498
3499#[async_trait]
3500impl ToolExecutor for BasicToolExecutor {
3501 fn specs(&self) -> Vec<ToolSpec> {
3502 BasicToolExecutor::specs(self)
3503 }
3504
3505 fn drain_catalog_events(&self) -> Vec<ToolCatalogEvent> {
3506 self.sources
3507 .iter()
3508 .flat_map(|s| s.drain_catalog_events())
3509 .collect()
3510 }
3511
3512 async fn execute(
3513 &self,
3514 request: ToolRequest,
3515 ctx: &mut ToolContext<'_>,
3516 ) -> ToolExecutionOutcome {
3517 self.execute_inner(request, None, ctx).await
3518 }
3519
3520 async fn execute_approved(
3521 &self,
3522 request: ToolRequest,
3523 approved_request: &ApprovalRequest,
3524 ctx: &mut ToolContext<'_>,
3525 ) -> ToolExecutionOutcome {
3526 let previous = ctx.approved_request.replace(approved_request.clone());
3527 let outcome = self
3528 .execute_inner(request, Some(&approved_request.id), ctx)
3529 .await;
3530 ctx.approved_request = previous;
3531 outcome
3532 }
3533}
3534
3535#[derive(Debug, Error, Clone, PartialEq, Serialize, Deserialize)]
3540pub enum ToolError {
3541 #[error("tool not found: {0}")]
3543 NotFound(ToolName),
3544 #[error("invalid tool input: {0}")]
3546 InvalidInput(String),
3547 #[error("tool permission denied: {0:?}")]
3549 PermissionDenied(PermissionDenial),
3550 #[error("tool execution failed: {0}")]
3552 ExecutionFailed(String),
3553 #[error("tool unavailable: {0}")]
3555 Unavailable(String),
3556 #[error("tool execution cancelled")]
3558 Cancelled,
3559 #[error("internal tool error: {0}")]
3561 Internal(String),
3562}
3563
3564impl ToolError {
3565 pub fn permission_denied(denial: PermissionDenial) -> Self {
3567 Self::PermissionDenied(denial)
3568 }
3569}
3570
3571impl From<PermissionDenial> for ToolError {
3572 fn from(value: PermissionDenial) -> Self {
3573 Self::permission_denied(value)
3574 }
3575}
3576
3577#[cfg(test)]
3578mod tests {
3579 use super::*;
3580 use async_trait::async_trait;
3581 use serde_json::json;
3582
3583 #[test]
3584 fn command_policy_can_deny_unknown_executables_without_approval() {
3585 let policy = CommandPolicy::new()
3586 .allow_executable("pwd")
3587 .require_approval_for_unknown(false);
3588 let request = ShellPermissionRequest {
3589 executable: "rm".into(),
3590 argv: vec!["-rf".into(), "/tmp/demo".into()],
3591 cwd: None,
3592 env_keys: Vec::new(),
3593 metadata: MetadataMap::new(),
3594 };
3595
3596 match policy.evaluate(&request) {
3597 PolicyMatch::Deny(denial) => {
3598 assert_eq!(denial.code, PermissionCode::CommandNotAllowed);
3599 }
3600 other => panic!("unexpected policy match: {other:?}"),
3601 }
3602 }
3603
3604 #[test]
3605 fn path_policy_allows_reads_under_read_only_roots() {
3606 let policy = PathPolicy::new().read_only_root("/workspace/vendor");
3607 let request = FileSystemPermissionRequest::Read {
3608 path: PathBuf::from("/workspace/vendor/lib.rs"),
3609 metadata: MetadataMap::new(),
3610 };
3611
3612 match policy.evaluate(&request) {
3613 PolicyMatch::NoOpinion | PolicyMatch::Allow => {}
3614 other => panic!("unexpected policy match: {other:?}"),
3615 }
3616 }
3617
3618 #[test]
3619 fn path_policy_denies_mutations_under_read_only_roots() {
3620 let policy = PathPolicy::new().read_only_root("/workspace/vendor");
3621 let request = FileSystemPermissionRequest::Edit {
3622 path: PathBuf::from("/workspace/vendor/lib.rs"),
3623 metadata: MetadataMap::new(),
3624 };
3625
3626 match policy.evaluate(&request) {
3627 PolicyMatch::Deny(denial) => {
3628 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3629 assert!(denial.message.contains("read-only"));
3630 }
3631 other => panic!("unexpected policy match: {other:?}"),
3632 }
3633 }
3634
3635 #[test]
3636 fn path_policy_denies_moves_into_read_only_roots() {
3637 let policy = PathPolicy::new().read_only_root("/workspace/vendor");
3638 let request = FileSystemPermissionRequest::Move {
3639 from: PathBuf::from("/workspace/src/lib.rs"),
3640 to: PathBuf::from("/workspace/vendor/lib.rs"),
3641 metadata: MetadataMap::new(),
3642 };
3643
3644 match policy.evaluate(&request) {
3645 PolicyMatch::Deny(denial) => {
3646 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3647 assert!(denial.message.contains("read-only"));
3648 }
3649 other => panic!("unexpected policy match: {other:?}"),
3650 }
3651 }
3652
3653 #[cfg(unix)]
3654 struct SymlinkTmpDir(PathBuf);
3655
3656 #[cfg(unix)]
3657 impl SymlinkTmpDir {
3658 fn new(label: &str) -> Self {
3659 use std::time::{SystemTime, UNIX_EPOCH};
3660 let nanos = SystemTime::now()
3661 .duration_since(UNIX_EPOCH)
3662 .unwrap()
3663 .as_nanos();
3664 let dir = std::env::temp_dir().join(format!(
3665 "agentkit-pathpolicy-{}-{}-{}",
3666 label,
3667 std::process::id(),
3668 nanos
3669 ));
3670 std::fs::create_dir_all(&dir).unwrap();
3671 Self(std::fs::canonicalize(&dir).unwrap())
3674 }
3675
3676 fn path(&self) -> &Path {
3677 &self.0
3678 }
3679 }
3680
3681 #[cfg(unix)]
3682 impl Drop for SymlinkTmpDir {
3683 fn drop(&mut self) {
3684 let _ = std::fs::remove_dir_all(&self.0);
3685 }
3686 }
3687
3688 #[cfg(unix)]
3689 fn assert_path_denied(
3690 policy: &PathPolicy,
3691 request: FileSystemPermissionRequest,
3692 ) -> PermissionDenial {
3693 match policy.evaluate(&request) {
3694 PolicyMatch::Deny(denial) => denial,
3695 other => panic!("expected deny, got: {other:?}"),
3696 }
3697 }
3698
3699 #[cfg(unix)]
3700 #[test]
3701 fn path_policy_blocks_symlink_escape_from_allowed_root() {
3702 let tmp = SymlinkTmpDir::new("allow-escape");
3703 let allowed = tmp.path().join("workspace");
3704 let outside = tmp.path().join("outside");
3705 std::fs::create_dir_all(&allowed).unwrap();
3706 std::fs::create_dir_all(&outside).unwrap();
3707 let secret = outside.join("secret.txt");
3708 std::fs::write(&secret, b"top-secret").unwrap();
3709 let escape = allowed.join("leak");
3710 std::os::unix::fs::symlink(&secret, &escape).unwrap();
3711
3712 let policy = PathPolicy::new()
3713 .allow_root(&allowed)
3714 .require_approval_outside_allowed(false);
3715 let denial = assert_path_denied(
3716 &policy,
3717 FileSystemPermissionRequest::Read {
3718 path: escape,
3719 metadata: MetadataMap::new(),
3720 },
3721 );
3722 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3723 }
3724
3725 #[cfg(unix)]
3726 #[test]
3727 fn path_policy_blocks_symlink_into_protected_root() {
3728 let tmp = SymlinkTmpDir::new("protect-bypass");
3729 let workspace = tmp.path().join("workspace");
3730 std::fs::create_dir_all(&workspace).unwrap();
3731 let secret = workspace.join(".env");
3732 std::fs::write(&secret, b"API_KEY=xxx").unwrap();
3733 let alias = workspace.join("config");
3734 std::os::unix::fs::symlink(&secret, &alias).unwrap();
3735
3736 let policy = PathPolicy::new()
3737 .allow_root(&workspace)
3738 .protect_root(&secret);
3739 let denial = assert_path_denied(
3740 &policy,
3741 FileSystemPermissionRequest::Read {
3742 path: alias,
3743 metadata: MetadataMap::new(),
3744 },
3745 );
3746 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3747 assert!(denial.message.contains("denied"));
3748 }
3749
3750 #[cfg(unix)]
3751 #[test]
3752 fn path_policy_blocks_symlink_write_into_read_only_root() {
3753 let tmp = SymlinkTmpDir::new("readonly-bypass");
3754 let workspace = tmp.path().join("workspace");
3755 let vendor = workspace.join("vendor");
3756 std::fs::create_dir_all(&vendor).unwrap();
3757 let target = vendor.join("lib.rs");
3758 std::fs::write(&target, b"// vendored").unwrap();
3759 let writable_alias = workspace.join("writable");
3760 std::os::unix::fs::symlink(&target, &writable_alias).unwrap();
3761
3762 let policy = PathPolicy::new()
3763 .allow_root(&workspace)
3764 .read_only_root(&vendor);
3765 let denial = assert_path_denied(
3766 &policy,
3767 FileSystemPermissionRequest::Edit {
3768 path: writable_alias,
3769 metadata: MetadataMap::new(),
3770 },
3771 );
3772 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3773 assert!(denial.message.contains("read-only"));
3774 }
3775
3776 #[cfg(unix)]
3777 #[test]
3778 fn path_policy_resolves_symlink_parent_for_nonexistent_leaf() {
3779 let tmp = SymlinkTmpDir::new("create-escape");
3780 let allowed = tmp.path().join("workspace");
3781 let outside = tmp.path().join("outside");
3782 std::fs::create_dir_all(&allowed).unwrap();
3783 std::fs::create_dir_all(&outside).unwrap();
3784 let escape_dir = allowed.join("escape");
3785 std::os::unix::fs::symlink(&outside, &escape_dir).unwrap();
3786 let new_file = escape_dir.join("new.txt");
3787
3788 let policy = PathPolicy::new()
3789 .allow_root(&allowed)
3790 .require_approval_outside_allowed(false);
3791 let denial = assert_path_denied(
3792 &policy,
3793 FileSystemPermissionRequest::Write {
3794 path: new_file,
3795 metadata: MetadataMap::new(),
3796 },
3797 );
3798 assert_eq!(denial.code, PermissionCode::PathNotAllowed);
3799 }
3800
3801 #[derive(Clone)]
3802 struct HiddenTool {
3803 spec: ToolSpec,
3804 }
3805
3806 impl HiddenTool {
3807 fn new() -> Self {
3808 Self {
3809 spec: ToolSpec {
3810 name: ToolName::new("hidden"),
3811 description: "hidden".into(),
3812 input_schema: json!({"type": "object"}),
3813 output_schema: None,
3814 annotations: ToolAnnotations::default(),
3815 metadata: MetadataMap::new(),
3816 },
3817 }
3818 }
3819 }
3820
3821 #[async_trait]
3822 impl Tool for HiddenTool {
3823 fn spec(&self) -> &ToolSpec {
3824 &self.spec
3825 }
3826
3827 fn current_spec(&self) -> Option<ToolSpec> {
3828 None
3829 }
3830
3831 async fn invoke(
3832 &self,
3833 request: ToolRequest,
3834 _ctx: &mut ToolContext<'_>,
3835 ) -> Result<ToolResult, ToolError> {
3836 Ok(ToolResult {
3837 result: ToolResultPart {
3838 call_id: request.call_id,
3839 output: ToolOutput::Text("hidden".into()),
3840 is_error: false,
3841 metadata: MetadataMap::new(),
3842 },
3843 duration: None,
3844 metadata: MetadataMap::new(),
3845 })
3846 }
3847 }
3848
3849 #[test]
3850 fn hidden_tools_are_omitted_from_specs_and_capabilities() {
3851 let registry = ToolRegistry::new().with(HiddenTool::new());
3852
3853 assert!(registry.specs().is_empty());
3854
3855 let provider = ToolCapabilityProvider::from_registry(
3856 ®istry,
3857 Arc::new(AllowAllPermissionChecker),
3858 Arc::new(()),
3859 );
3860 assert!(provider.invocables().is_empty());
3861 }
3862
3863 struct AllowAllPermissionChecker;
3864
3865 impl PermissionChecker for AllowAllPermissionChecker {
3866 fn evaluate(&self, _request: &dyn PermissionRequest) -> PermissionDecision {
3867 PermissionDecision::Allow
3868 }
3869 }
3870
3871 #[derive(Clone)]
3874 struct PanickingSpecTool {
3875 spec: ToolSpec,
3876 }
3877
3878 impl PanickingSpecTool {
3879 fn new(name: &str) -> Self {
3880 Self {
3881 spec: ToolSpec {
3882 name: ToolName::new(name),
3883 description: "panics on current_spec".into(),
3884 input_schema: json!({"type": "object"}),
3885 output_schema: None,
3886 annotations: ToolAnnotations::default(),
3887 metadata: MetadataMap::new(),
3888 },
3889 }
3890 }
3891 }
3892
3893 #[async_trait]
3894 impl Tool for PanickingSpecTool {
3895 fn spec(&self) -> &ToolSpec {
3896 &self.spec
3897 }
3898
3899 fn current_spec(&self) -> Option<ToolSpec> {
3900 panic!("PanickingSpecTool::current_spec");
3901 }
3902
3903 async fn invoke(
3904 &self,
3905 request: ToolRequest,
3906 _ctx: &mut ToolContext<'_>,
3907 ) -> Result<ToolResult, ToolError> {
3908 Ok(ToolResult {
3909 result: ToolResultPart {
3910 call_id: request.call_id,
3911 output: ToolOutput::Text("never".into()),
3912 is_error: false,
3913 metadata: MetadataMap::new(),
3914 },
3915 duration: None,
3916 metadata: MetadataMap::new(),
3917 })
3918 }
3919 }
3920
3921 #[test]
3932 fn catalog_recovers_from_panicked_writer() {
3933 let (writer, reader) = dynamic_catalog("test");
3934
3935 writer.upsert(Arc::new(PanickingSpecTool::new("boom")));
3939 let _ = reader.drain_catalog_events();
3940
3941 let panic_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
3946 writer.replace_all(vec![
3947 Arc::new(PanickingSpecTool::new("boom")) as Arc<dyn Tool>
3948 ]);
3949 }));
3950 assert!(
3951 panic_result.is_err(),
3952 "PanickingSpecTool::current_spec must propagate"
3953 );
3954
3955 assert!(
3959 reader.get(&ToolName::new("boom")).is_some(),
3960 "catalog still readable after poisoning panic"
3961 );
3962
3963 assert!(writer.remove(&ToolName::new("boom")));
3966
3967 writer.upsert(Arc::new(HiddenTool::new()));
3971 assert!(
3972 reader.get(&ToolName::new("hidden")).is_some(),
3973 "catalog usable for further writes + reads"
3974 );
3975 }
3976
3977 #[derive(Clone)]
3978 struct EchoTool {
3979 spec: ToolSpec,
3980 }
3981
3982 impl EchoTool {
3983 fn new(name: &str) -> Self {
3984 Self {
3985 spec: ToolSpec {
3986 name: ToolName::new(name),
3987 description: format!("echo {name}"),
3988 input_schema: json!({"type": "object"}),
3989 output_schema: None,
3990 annotations: ToolAnnotations::default(),
3991 metadata: MetadataMap::new(),
3992 },
3993 }
3994 }
3995 }
3996
3997 #[async_trait]
3998 impl Tool for EchoTool {
3999 fn spec(&self) -> &ToolSpec {
4000 &self.spec
4001 }
4002
4003 async fn invoke(
4004 &self,
4005 request: ToolRequest,
4006 _ctx: &mut ToolContext<'_>,
4007 ) -> Result<ToolResult, ToolError> {
4008 Ok(ToolResult::new(ToolResultPart::success(
4009 request.call_id,
4010 ToolOutput::text(request.tool_name.0.clone()),
4011 )))
4012 }
4013 }
4014
4015 fn registry_with(names: &[&str]) -> ToolRegistry {
4016 names.iter().fold(ToolRegistry::new(), |reg, name| {
4017 reg.with(EchoTool::new(name))
4018 })
4019 }
4020
4021 #[test]
4022 fn prefixed_rewrites_specs_and_resolves_lookups() {
4023 let source = registry_with(&["get_temp", "get_humidity"]).prefixed("weather");
4024 let names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4025 assert_eq!(names, vec!["weather_get_humidity", "weather_get_temp"]);
4026
4027 assert!(source.get(&ToolName::new("weather_get_temp")).is_some());
4028 assert!(
4029 source.get(&ToolName::new("get_temp")).is_none(),
4030 "original name must not resolve when prefixed"
4031 );
4032 assert!(source.get(&ToolName::new("unknown")).is_none());
4033 }
4034
4035 #[tokio::test]
4036 async fn prefixed_invoke_sees_inner_name_on_request() {
4037 let source = registry_with(&["get_temp"]).prefixed("weather");
4038 let tool = source.get(&ToolName::new("weather_get_temp")).unwrap();
4039
4040 assert_eq!(tool.spec().name.0, "weather_get_temp");
4042
4043 let owned = OwnedToolContext {
4045 session_id: SessionId::new("s"),
4046 turn_id: TurnId::new("t"),
4047 metadata: MetadataMap::new(),
4048 permissions: Arc::new(AllowAllPermissions),
4049 resources: Arc::new(()),
4050 cancellation: None,
4051 execution_scope: None,
4052 approved_request: None,
4053 };
4054 let mut ctx = owned.borrowed();
4055 let request = ToolRequest {
4056 call_id: ToolCallId::new("c"),
4057 tool_name: ToolName::new("weather_get_temp"),
4058 input: json!({}),
4059 session_id: SessionId::new("s"),
4060 turn_id: TurnId::new("t"),
4061 metadata: MetadataMap::new(),
4062 };
4063 let result = tool.invoke(request, &mut ctx).await.unwrap();
4064 match result.result.output {
4065 ToolOutput::Text(text) => assert_eq!(text, "get_temp"),
4066 other => panic!("unexpected output: {other:?}"),
4067 }
4068 }
4069
4070 #[derive(Clone)]
4071 struct StaticOutputTool {
4072 spec: ToolSpec,
4073 output: ToolOutput,
4074 }
4075
4076 impl StaticOutputTool {
4077 fn new(name: &str, output: ToolOutput) -> Self {
4078 Self {
4079 spec: ToolSpec::new(name, format!("static {name}"), json!({"type": "object"})),
4080 output,
4081 }
4082 }
4083
4084 fn with_output_limit(mut self, limit: ToolOutputLimit) -> Self {
4085 self.spec = self.spec.with_output_limit(limit);
4086 self
4087 }
4088 }
4089
4090 #[async_trait]
4091 impl Tool for StaticOutputTool {
4092 fn spec(&self) -> &ToolSpec {
4093 &self.spec
4094 }
4095
4096 async fn invoke(
4097 &self,
4098 request: ToolRequest,
4099 _ctx: &mut ToolContext<'_>,
4100 ) -> Result<ToolResult, ToolError> {
4101 Ok(ToolResult::new(ToolResultPart::success(
4102 request.call_id,
4103 self.output.clone(),
4104 )))
4105 }
4106 }
4107
4108 struct ApprovedContextTool {
4109 spec: ToolSpec,
4110 }
4111
4112 impl ApprovedContextTool {
4113 fn new() -> Self {
4114 Self {
4115 spec: ToolSpec::new(
4116 "approved_context",
4117 "approved context",
4118 json!({"type": "object"}),
4119 ),
4120 }
4121 }
4122 }
4123
4124 #[async_trait]
4125 impl Tool for ApprovedContextTool {
4126 fn spec(&self) -> &ToolSpec {
4127 &self.spec
4128 }
4129
4130 async fn invoke(
4131 &self,
4132 request: ToolRequest,
4133 ctx: &mut ToolContext<'_>,
4134 ) -> Result<ToolResult, ToolError> {
4135 Ok(ToolResult::new(ToolResultPart::success(
4136 request.call_id,
4137 ToolOutput::structured(json!({
4138 "approved": ctx.approved_request.is_some()
4139 })),
4140 )))
4141 }
4142 }
4143
4144 struct ScopeChildTool {
4145 spec: ToolSpec,
4146 }
4147
4148 impl ScopeChildTool {
4149 fn new() -> Self {
4150 Self {
4151 spec: ToolSpec::new("scope_child", "scope child", json!({"type": "object"})),
4152 }
4153 }
4154 }
4155
4156 #[async_trait]
4157 impl Tool for ScopeChildTool {
4158 fn spec(&self) -> &ToolSpec {
4159 &self.spec
4160 }
4161
4162 async fn invoke(
4163 &self,
4164 request: ToolRequest,
4165 _ctx: &mut ToolContext<'_>,
4166 ) -> Result<ToolResult, ToolError> {
4167 Ok(ToolResult::new(ToolResultPart::success(
4168 request.call_id,
4169 ToolOutput::structured(json!({ "child": request.input })),
4170 )))
4171 }
4172 }
4173
4174 struct ScopeParentTool {
4175 spec: ToolSpec,
4176 }
4177
4178 impl ScopeParentTool {
4179 fn new() -> Self {
4180 Self {
4181 spec: ToolSpec::new("scope_parent", "scope parent", json!({"type": "object"})),
4182 }
4183 }
4184 }
4185
4186 #[async_trait]
4187 impl Tool for ScopeParentTool {
4188 fn spec(&self) -> &ToolSpec {
4189 &self.spec
4190 }
4191
4192 async fn invoke(
4193 &self,
4194 request: ToolRequest,
4195 _ctx: &mut ToolContext<'_>,
4196 ) -> Result<ToolResult, ToolError> {
4197 Ok(ToolResult::new(ToolResultPart::success(
4198 request.call_id,
4199 ToolOutput::text("unused"),
4200 )))
4201 }
4202
4203 async fn invoke_outcome(
4204 &self,
4205 request: ToolRequest,
4206 ctx: &mut ToolContext<'_>,
4207 ) -> ToolExecutionOutcome {
4208 let Some(scope) = ctx.execution_scope.clone() else {
4209 return ToolExecutionOutcome::Failed(ToolError::Internal(
4210 "missing execution scope".into(),
4211 ));
4212 };
4213 let child = ToolRequest::new(
4214 "child-call",
4215 "scope_child",
4216 request.input.clone(),
4217 request.session_id.clone(),
4218 request.turn_id.clone(),
4219 );
4220 match scope.execute_child(child).await {
4221 ToolExecutionOutcome::Completed(child_result) => {
4222 ToolExecutionOutcome::Completed(ToolResult::new(ToolResultPart::success(
4223 request.call_id,
4224 child_result.result.output,
4225 )))
4226 }
4227 other => other,
4228 }
4229 }
4230 }
4231
4232 fn test_context() -> OwnedToolContext {
4233 OwnedToolContext {
4234 session_id: SessionId::new("s"),
4235 turn_id: TurnId::new("t"),
4236 metadata: MetadataMap::new(),
4237 permissions: Arc::new(AllowAllPermissions),
4238 resources: Arc::new(()),
4239 cancellation: None,
4240 execution_scope: None,
4241 approved_request: None,
4242 }
4243 }
4244
4245 fn test_context_with_scope(executor: Arc<dyn ToolExecutor>) -> OwnedToolContext {
4246 let session_id = SessionId::new("s");
4247 let turn_id = TurnId::new("t");
4248 let metadata = MetadataMap::new();
4249 let permissions: Arc<dyn PermissionChecker> = Arc::new(AllowAllPermissions);
4250 let resources: Arc<dyn ToolResources> = Arc::new(());
4251 let scope = ToolExecutionScope {
4252 executor,
4253 session_id: session_id.clone(),
4254 turn_id: turn_id.clone(),
4255 permissions: permissions.clone(),
4256 resources: resources.clone(),
4257 cancellation: None,
4258 };
4259 OwnedToolContext {
4260 session_id,
4261 turn_id,
4262 metadata,
4263 permissions,
4264 resources,
4265 cancellation: None,
4266 execution_scope: Some(scope),
4267 approved_request: None,
4268 }
4269 }
4270
4271 #[tokio::test]
4272 async fn default_invoke_outcome_wraps_invoke_success() {
4273 let executor = BasicToolExecutor::from_registry(ToolRegistry::new().with(
4274 StaticOutputTool::new("plain", ToolOutput::structured(json!({"ok": true}))),
4275 ));
4276 let outcome = executor
4277 .execute_owned(
4278 ToolRequest::new("call", "plain", json!({}), "s", "t"),
4279 test_context(),
4280 )
4281 .await;
4282
4283 let ToolExecutionOutcome::Completed(result) = outcome else {
4284 panic!("expected completed outcome, got {outcome:?}");
4285 };
4286 assert_eq!(
4287 result.result.output,
4288 ToolOutput::structured(json!({"ok": true}))
4289 );
4290 }
4291
4292 #[tokio::test]
4293 async fn execute_approved_passes_approval_context_to_tool() {
4294 let executor =
4295 BasicToolExecutor::from_registry(ToolRegistry::new().with(ApprovedContextTool::new()));
4296 let approval = ApprovalRequest {
4297 task_id: None,
4298 call_id: Some(ToolCallId::new("call")),
4299 id: ApprovalId::new("approval"),
4300 request_kind: "test.approval".into(),
4301 reason: ApprovalReason::PolicyRequiresConfirmation,
4302 summary: "approve".into(),
4303 metadata: MetadataMap::new(),
4304 };
4305 let outcome = executor
4306 .execute_approved_owned(
4307 ToolRequest::new("call", "approved_context", json!({}), "s", "t"),
4308 &approval,
4309 test_context(),
4310 )
4311 .await;
4312
4313 let ToolExecutionOutcome::Completed(result) = outcome else {
4314 panic!("expected completed outcome, got {outcome:?}");
4315 };
4316 assert_eq!(
4317 result.result.output,
4318 ToolOutput::structured(json!({"approved": true}))
4319 );
4320 }
4321
4322 #[tokio::test]
4323 async fn execution_scope_invokes_child_through_executor() {
4324 let executor: Arc<dyn ToolExecutor> = Arc::new(BasicToolExecutor::from_registry(
4325 ToolRegistry::new()
4326 .with(ScopeParentTool::new())
4327 .with(ScopeChildTool::new()),
4328 ));
4329 let outcome = executor
4330 .execute_owned(
4331 ToolRequest::new("parent-call", "scope_parent", json!({"value": 3}), "s", "t"),
4332 test_context_with_scope(executor.clone()),
4333 )
4334 .await;
4335
4336 let ToolExecutionOutcome::Completed(result) = outcome else {
4337 panic!("expected completed outcome, got {outcome:?}");
4338 };
4339 assert_eq!(
4340 result.result.output,
4341 ToolOutput::structured(json!({ "child": { "value": 3 } }))
4342 );
4343 }
4344
4345 #[tokio::test]
4346 async fn executor_stores_oversized_output_using_tool_metadata_limit() {
4347 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4348 let strategy = ConfigurableToolOutputTruncationStrategy::new(store.clone());
4349 let tool = StaticOutputTool::new("big", ToolOutput::text("x".repeat(500)))
4350 .with_output_limit(ToolOutputLimit::store_for_readback(300));
4351 let executor = BasicToolExecutor::from_registry(ToolRegistry::new().with(tool))
4352 .with_output_truncation_strategy(strategy);
4353
4354 let outcome = executor
4355 .execute_owned(
4356 ToolRequest::new(
4357 "call",
4358 "big",
4359 json!({}),
4360 SessionId::new("s"),
4361 TurnId::new("t"),
4362 ),
4363 test_context(),
4364 )
4365 .await;
4366
4367 let ToolExecutionOutcome::Completed(result) = outcome else {
4368 panic!("expected completed outcome, got {outcome:?}");
4369 };
4370 let ToolOutput::Structured(envelope) = result.result.output else {
4371 panic!("expected truncation envelope");
4372 };
4373 assert_eq!(envelope["truncated"], true);
4374 assert_eq!(envelope["read_tool"], TOOL_RESULT_READ_TOOL_NAME);
4375 let id = envelope["tool_result_id"].as_str().expect("tool_result_id");
4376
4377 let slice = store
4378 .read(&ToolOutputArtifactId(id.to_string()), 0, 50)
4379 .await
4380 .expect("read artifact");
4381 assert_eq!(slice.content, "x".repeat(50));
4382 assert_eq!(slice.next_offset, 50);
4383 assert!(!slice.eof);
4384 }
4385
4386 #[tokio::test]
4387 async fn tool_result_read_enforces_explicit_max_read_size() {
4388 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4389 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4390 let request = ToolRequest::new(
4391 "call",
4392 "big",
4393 json!({}),
4394 SessionId::new("s"),
4395 TurnId::new("t"),
4396 );
4397 let ctx = ToolOutputTruncationContext::from((&request, spec));
4398 let artifact = store
4399 .put(&ctx, "abcdef".to_string(), 6)
4400 .await
4401 .expect("store artifact");
4402 let tool = ToolResultReadTool::new(store, 4);
4403 let owned_ctx = test_context();
4404 let mut tool_ctx = owned_ctx.borrowed();
4405
4406 let err = tool
4407 .invoke(
4408 ToolRequest::new(
4409 "read-call",
4410 TOOL_RESULT_READ_TOOL_NAME,
4411 json!({"id": artifact.id.0, "offset": 0, "limit": 5}),
4412 SessionId::new("s"),
4413 TurnId::new("t"),
4414 ),
4415 &mut tool_ctx,
4416 )
4417 .await
4418 .expect_err("read past max must fail");
4419 match err {
4420 ToolError::InvalidInput(message) => assert!(message.contains("exceeds maximum")),
4421 other => panic!("expected InvalidInput, got {other:?}"),
4422 }
4423 }
4424
4425 #[tokio::test]
4426 async fn tool_result_read_rejects_zero_limit() {
4427 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4428 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4429 let request = ToolRequest::new(
4430 "call",
4431 "big",
4432 json!({}),
4433 SessionId::new("s"),
4434 TurnId::new("t"),
4435 );
4436 let ctx = ToolOutputTruncationContext::from((&request, spec));
4437 let artifact = store
4438 .put(&ctx, "abcdef".to_string(), 6)
4439 .await
4440 .expect("store artifact");
4441 let tool = ToolResultReadTool::new(store, 4);
4442 let owned_ctx = test_context();
4443 let mut tool_ctx = owned_ctx.borrowed();
4444
4445 let err = tool
4446 .invoke(
4447 ToolRequest::new(
4448 "read-call",
4449 TOOL_RESULT_READ_TOOL_NAME,
4450 json!({"id": artifact.id.0, "offset": 0, "limit": 0}),
4451 SessionId::new("s"),
4452 TurnId::new("t"),
4453 ),
4454 &mut tool_ctx,
4455 )
4456 .await
4457 .expect_err("zero limit must fail");
4458 match err {
4459 ToolError::InvalidInput(message) => assert!(message.contains("greater than 0")),
4460 other => panic!("expected InvalidInput, got {other:?}"),
4461 }
4462 }
4463
4464 #[tokio::test]
4465 async fn tool_result_read_executor_allows_full_content_limit_with_envelope() {
4466 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4467 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4468 let request = ToolRequest::new(
4469 "call",
4470 "big",
4471 json!({}),
4472 SessionId::new("s"),
4473 TurnId::new("t"),
4474 );
4475 let ctx = ToolOutputTruncationContext::from((&request, spec));
4476 let artifact = store
4477 .put(&ctx, "abcd".to_string(), 4)
4478 .await
4479 .expect("store artifact");
4480 let executor = BasicToolExecutor::from_registry(
4481 ToolRegistry::new().with(ToolResultReadTool::new(store.clone(), 4)),
4482 )
4483 .with_output_truncation_strategy(ConfigurableToolOutputTruncationStrategy::new(store));
4484
4485 let outcome = executor
4486 .execute_owned(
4487 ToolRequest::new(
4488 "read-call",
4489 TOOL_RESULT_READ_TOOL_NAME,
4490 json!({"id": artifact.id.0, "offset": 0, "limit": 4}),
4491 SessionId::new("s"),
4492 TurnId::new("t"),
4493 ),
4494 test_context(),
4495 )
4496 .await;
4497
4498 let ToolExecutionOutcome::Completed(result) = outcome else {
4499 panic!("expected completed outcome, got {outcome:?}");
4500 };
4501 let ToolOutput::Structured(output) = result.result.output else {
4502 panic!("expected structured readback output");
4503 };
4504 assert_eq!(output["content"], "abcd");
4505 assert_eq!(output["eof"], true);
4506 }
4507
4508 #[tokio::test]
4509 async fn tool_result_read_executor_allows_json_escaped_full_content_limit() {
4510 let store = Arc::new(InMemoryToolOutputArtifactStore::new());
4511 let spec = ToolSpec::new("big", "big output", json!({"type": "object"}));
4512 let request = ToolRequest::new(
4513 "call",
4514 "big",
4515 json!({}),
4516 SessionId::new("s"),
4517 TurnId::new("t"),
4518 );
4519 let ctx = ToolOutputTruncationContext::from((&request, spec));
4520 let content = "\0".repeat(4);
4521 let artifact = store
4522 .put(&ctx, content.clone(), content.len())
4523 .await
4524 .expect("store artifact");
4525 let executor = BasicToolExecutor::from_registry(
4526 ToolRegistry::new().with(ToolResultReadTool::new(store.clone(), 4)),
4527 )
4528 .with_output_truncation_strategy(ConfigurableToolOutputTruncationStrategy::new(store));
4529
4530 let outcome = executor
4531 .execute_owned(
4532 ToolRequest::new(
4533 "read-call",
4534 TOOL_RESULT_READ_TOOL_NAME,
4535 json!({"id": artifact.id.0, "offset": 0, "limit": 4}),
4536 SessionId::new("s"),
4537 TurnId::new("t"),
4538 ),
4539 test_context(),
4540 )
4541 .await;
4542
4543 let ToolExecutionOutcome::Completed(result) = outcome else {
4544 panic!("expected completed outcome, got {outcome:?}");
4545 };
4546 let ToolOutput::Structured(output) = result.result.output else {
4547 panic!("expected structured readback output");
4548 };
4549 assert_eq!(output["content"], content);
4550 assert_eq!(output["eof"], true);
4551 }
4552
4553 #[test]
4554 fn inline_clip_respects_limit_when_marker_exceeds_budget() {
4555 let clipped = clip_string_with_marker("abcdef", 8, 1000);
4556
4557 assert!(clipped.len() <= 8);
4558 assert!(clipped.is_char_boundary(clipped.len()));
4559 }
4560
4561 #[test]
4562 fn filtered_hides_tools_rejected_by_predicate() {
4563 let source = registry_with(&["safe", "danger_drop", "danger_delete"])
4564 .filtered(|name| !name.0.starts_with("danger_"));
4565 let names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4566 assert_eq!(names, vec!["safe"]);
4567
4568 assert!(source.get(&ToolName::new("safe")).is_some());
4569 assert!(source.get(&ToolName::new("danger_drop")).is_none());
4570 }
4571
4572 #[test]
4573 fn renamed_remaps_specs_and_lookups() {
4574 let source = registry_with(&["legacy_name", "passthrough"])
4575 .renamed([(ToolName::new("legacy_name"), ToolName::new("modern_name"))]);
4576 let mut names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4577 names.sort();
4578 assert_eq!(names, vec!["modern_name", "passthrough"]);
4579
4580 assert!(source.get(&ToolName::new("modern_name")).is_some());
4581 assert!(
4582 source.get(&ToolName::new("legacy_name")).is_none(),
4583 "original name is hidden after renaming"
4584 );
4585 assert!(source.get(&ToolName::new("passthrough")).is_some());
4586 }
4587
4588 #[cfg(feature = "schemars")]
4589 mod schemars_helpers {
4590 use super::*;
4591 use schemars::JsonSchema;
4592 use serde::Deserialize;
4593
4594 #[derive(JsonSchema, Deserialize)]
4595 #[allow(dead_code)]
4596 struct WeatherInput {
4597 location: String,
4599 #[serde(default)]
4601 celsius: bool,
4602 }
4603
4604 #[test]
4605 fn schema_for_emits_object_schema_with_typed_fields() {
4606 let schema = schema_for::<WeatherInput>();
4607 let obj = schema.as_object().expect("schema is a JSON object");
4608 assert_eq!(
4609 obj.get("type").and_then(|v| v.as_str()),
4610 Some("object"),
4611 "root type should be object"
4612 );
4613 let properties = obj
4614 .get("properties")
4615 .and_then(|v| v.as_object())
4616 .expect("properties block");
4617 assert!(properties.contains_key("location"));
4618 assert!(properties.contains_key("celsius"));
4619 }
4620
4621 #[test]
4622 fn tool_spec_for_carries_schema_name_and_description() {
4623 let spec = tool_spec_for::<WeatherInput>("get_weather", "Fetch current weather");
4624 assert_eq!(spec.name.0, "get_weather");
4625 assert_eq!(spec.description, "Fetch current weather");
4626 assert!(spec.input_schema.is_object());
4627 }
4628 }
4629
4630 #[test]
4631 fn transforms_compose_via_chained_methods() {
4632 let source = registry_with(&["read_file", "write_file", "delete_file"])
4633 .filtered(|name| name.0 != "delete_file")
4634 .prefixed("fs");
4635 let mut names: Vec<_> = source.specs().into_iter().map(|s| s.name.0).collect();
4636 names.sort();
4637 assert_eq!(names, vec!["fs_read_file", "fs_write_file"]);
4638 }
4639}