agentkit_tools_core/
lib.rs

1use std::any::Any;
2use std::collections::{BTreeMap, BTreeSet};
3use std::fmt;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6use std::sync::atomic::{AtomicU64, Ordering};
7use std::time::Duration;
8
9use agentkit_capabilities::{
10    CapabilityContext, CapabilityError, CapabilityName, CapabilityProvider, Invocable,
11    InvocableOutput, InvocableRequest, InvocableResult, InvocableSpec, PromptProvider,
12    ResourceProvider,
13};
14use agentkit_core::{
15    ApprovalId, Item, ItemKind, MetadataMap, Part, SessionId, ToolCallId, ToolOutput,
16    ToolResultPart, TurnCancellation, TurnId,
17};
18use async_trait::async_trait;
19use serde::{Deserialize, Serialize};
20use serde_json::Value;
21use thiserror::Error;
22
23#[derive(Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
24pub struct ToolName(pub String);
25
26impl ToolName {
27    pub fn new(value: impl Into<String>) -> Self {
28        Self(value.into())
29    }
30}
31
32impl fmt::Display for ToolName {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        self.0.fmt(f)
35    }
36}
37
38impl From<&str> for ToolName {
39    fn from(value: &str) -> Self {
40        Self::new(value)
41    }
42}
43
44#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
45pub struct ToolAnnotations {
46    pub read_only_hint: bool,
47    pub destructive_hint: bool,
48    pub idempotent_hint: bool,
49    pub needs_approval_hint: bool,
50    pub supports_streaming_hint: bool,
51}
52
53#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
54pub struct ToolSpec {
55    pub name: ToolName,
56    pub description: String,
57    pub input_schema: Value,
58    pub annotations: ToolAnnotations,
59    pub metadata: MetadataMap,
60}
61
62#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
63pub struct ToolRequest {
64    pub call_id: ToolCallId,
65    pub tool_name: ToolName,
66    pub input: Value,
67    pub session_id: SessionId,
68    pub turn_id: TurnId,
69    pub metadata: MetadataMap,
70}
71
72#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
73pub struct ToolResult {
74    pub result: ToolResultPart,
75    pub duration: Option<Duration>,
76    pub metadata: MetadataMap,
77}
78
79pub trait ToolResources: Send + Sync {
80    fn as_any(&self) -> &dyn Any;
81}
82
83impl ToolResources for () {
84    fn as_any(&self) -> &dyn Any {
85        self
86    }
87}
88
89pub struct ToolContext<'a> {
90    pub capability: CapabilityContext<'a>,
91    pub permissions: &'a dyn PermissionChecker,
92    pub resources: &'a dyn ToolResources,
93    pub cancellation: Option<TurnCancellation>,
94}
95
96pub trait PermissionRequest: Send + Sync {
97    fn kind(&self) -> &'static str;
98    fn summary(&self) -> String;
99    fn metadata(&self) -> &MetadataMap;
100    fn as_any(&self) -> &dyn Any;
101}
102
103#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
104pub enum PermissionCode {
105    PathNotAllowed,
106    CommandNotAllowed,
107    NetworkNotAllowed,
108    ServerNotTrusted,
109    AuthScopeNotAllowed,
110    CustomPolicyDenied,
111    UnknownRequest,
112}
113
114#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
115pub struct PermissionDenial {
116    pub code: PermissionCode,
117    pub message: String,
118    pub metadata: MetadataMap,
119}
120
121#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
122pub enum ApprovalReason {
123    PolicyRequiresConfirmation,
124    EscalatedRisk,
125    UnknownTarget,
126    SensitivePath,
127    SensitiveCommand,
128    SensitiveServer,
129    SensitiveAuthScope,
130}
131
132#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
133pub struct ApprovalRequest {
134    pub id: ApprovalId,
135    pub request_kind: String,
136    pub reason: ApprovalReason,
137    pub summary: String,
138    pub metadata: MetadataMap,
139}
140
141#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
142pub enum ApprovalDecision {
143    Approve,
144    Deny { reason: Option<String> },
145}
146
147#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
148pub struct AuthRequest {
149    pub id: String,
150    pub provider: String,
151    pub operation: AuthOperation,
152    pub challenge: MetadataMap,
153}
154
155#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
156pub enum AuthOperation {
157    ToolCall {
158        tool_name: String,
159        input: Value,
160        call_id: Option<ToolCallId>,
161        session_id: Option<SessionId>,
162        turn_id: Option<TurnId>,
163        metadata: MetadataMap,
164    },
165    McpConnect {
166        server_id: String,
167        metadata: MetadataMap,
168    },
169    McpToolCall {
170        server_id: String,
171        tool_name: String,
172        input: Value,
173        metadata: MetadataMap,
174    },
175    McpResourceRead {
176        server_id: String,
177        resource_id: String,
178        metadata: MetadataMap,
179    },
180    McpPromptGet {
181        server_id: String,
182        prompt_id: String,
183        args: Value,
184        metadata: MetadataMap,
185    },
186    Custom {
187        kind: String,
188        payload: Value,
189        metadata: MetadataMap,
190    },
191}
192
193impl AuthOperation {
194    pub fn server_id(&self) -> Option<&str> {
195        match self {
196            Self::McpConnect { server_id, .. }
197            | Self::McpToolCall { server_id, .. }
198            | Self::McpResourceRead { server_id, .. }
199            | Self::McpPromptGet { server_id, .. } => Some(server_id.as_str()),
200            Self::ToolCall { metadata, .. } | Self::Custom { metadata, .. } => {
201                metadata.get("server_id").and_then(Value::as_str)
202            }
203        }
204    }
205}
206
207#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
208pub enum AuthResolution {
209    Provided {
210        request: AuthRequest,
211        credentials: MetadataMap,
212    },
213    Cancelled {
214        request: AuthRequest,
215    },
216}
217
218impl AuthResolution {
219    pub fn request(&self) -> &AuthRequest {
220        match self {
221            Self::Provided { request, .. } | Self::Cancelled { request } => request,
222        }
223    }
224}
225
226impl AuthRequest {
227    pub fn server_id(&self) -> Option<&str> {
228        self.operation.server_id()
229    }
230}
231
232#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
233pub enum ToolInterruption {
234    ApprovalRequired(ApprovalRequest),
235    AuthRequired(AuthRequest),
236}
237
238#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
239pub enum PermissionDecision {
240    Allow,
241    Deny(PermissionDenial),
242    RequireApproval(ApprovalRequest),
243}
244
245pub trait PermissionChecker: Send + Sync {
246    fn evaluate(&self, request: &dyn PermissionRequest) -> PermissionDecision;
247}
248
249#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
250pub enum PolicyMatch {
251    NoOpinion,
252    Allow,
253    Deny(PermissionDenial),
254    RequireApproval(ApprovalRequest),
255}
256
257pub trait PermissionPolicy: Send + Sync {
258    fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch;
259}
260
261pub struct CompositePermissionChecker {
262    policies: Vec<Box<dyn PermissionPolicy>>,
263    fallback: PermissionDecision,
264}
265
266impl CompositePermissionChecker {
267    pub fn new(fallback: PermissionDecision) -> Self {
268        Self {
269            policies: Vec::new(),
270            fallback,
271        }
272    }
273
274    pub fn with_policy(mut self, policy: impl PermissionPolicy + 'static) -> Self {
275        self.policies.push(Box::new(policy));
276        self
277    }
278}
279
280impl PermissionChecker for CompositePermissionChecker {
281    fn evaluate(&self, request: &dyn PermissionRequest) -> PermissionDecision {
282        let mut saw_allow = false;
283        let mut approval = None;
284
285        for policy in &self.policies {
286            match policy.evaluate(request) {
287                PolicyMatch::NoOpinion => {}
288                PolicyMatch::Allow => saw_allow = true,
289                PolicyMatch::Deny(denial) => return PermissionDecision::Deny(denial),
290                PolicyMatch::RequireApproval(req) => approval = Some(req),
291            }
292        }
293
294        if let Some(req) = approval {
295            PermissionDecision::RequireApproval(req)
296        } else if saw_allow {
297            PermissionDecision::Allow
298        } else {
299            self.fallback.clone()
300        }
301    }
302}
303
304#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
305pub struct ShellPermissionRequest {
306    pub executable: String,
307    pub argv: Vec<String>,
308    pub cwd: Option<PathBuf>,
309    pub env_keys: Vec<String>,
310    pub metadata: MetadataMap,
311}
312
313impl PermissionRequest for ShellPermissionRequest {
314    fn kind(&self) -> &'static str {
315        "shell.command"
316    }
317
318    fn summary(&self) -> String {
319        if self.argv.is_empty() {
320            self.executable.clone()
321        } else {
322            format!("{} {}", self.executable, self.argv.join(" "))
323        }
324    }
325
326    fn metadata(&self) -> &MetadataMap {
327        &self.metadata
328    }
329
330    fn as_any(&self) -> &dyn Any {
331        self
332    }
333}
334
335#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
336pub enum FileSystemPermissionRequest {
337    Read {
338        path: PathBuf,
339        metadata: MetadataMap,
340    },
341    Write {
342        path: PathBuf,
343        metadata: MetadataMap,
344    },
345    Edit {
346        path: PathBuf,
347        metadata: MetadataMap,
348    },
349    Delete {
350        path: PathBuf,
351        metadata: MetadataMap,
352    },
353    Move {
354        from: PathBuf,
355        to: PathBuf,
356        metadata: MetadataMap,
357    },
358    List {
359        path: PathBuf,
360        metadata: MetadataMap,
361    },
362    CreateDir {
363        path: PathBuf,
364        metadata: MetadataMap,
365    },
366}
367
368impl FileSystemPermissionRequest {
369    fn metadata_map(&self) -> &MetadataMap {
370        match self {
371            Self::Read { metadata, .. }
372            | Self::Write { metadata, .. }
373            | Self::Edit { metadata, .. }
374            | Self::Delete { metadata, .. }
375            | Self::Move { metadata, .. }
376            | Self::List { metadata, .. }
377            | Self::CreateDir { metadata, .. } => metadata,
378        }
379    }
380}
381
382impl PermissionRequest for FileSystemPermissionRequest {
383    fn kind(&self) -> &'static str {
384        match self {
385            Self::Read { .. } => "filesystem.read",
386            Self::Write { .. } => "filesystem.write",
387            Self::Edit { .. } => "filesystem.edit",
388            Self::Delete { .. } => "filesystem.delete",
389            Self::Move { .. } => "filesystem.move",
390            Self::List { .. } => "filesystem.list",
391            Self::CreateDir { .. } => "filesystem.mkdir",
392        }
393    }
394
395    fn summary(&self) -> String {
396        match self {
397            Self::Read { path, .. } => format!("Read {}", path.display()),
398            Self::Write { path, .. } => format!("Write {}", path.display()),
399            Self::Edit { path, .. } => format!("Edit {}", path.display()),
400            Self::Delete { path, .. } => format!("Delete {}", path.display()),
401            Self::Move { from, to, .. } => {
402                format!("Move {} to {}", from.display(), to.display())
403            }
404            Self::List { path, .. } => format!("List {}", path.display()),
405            Self::CreateDir { path, .. } => format!("Create directory {}", path.display()),
406        }
407    }
408
409    fn metadata(&self) -> &MetadataMap {
410        self.metadata_map()
411    }
412
413    fn as_any(&self) -> &dyn Any {
414        self
415    }
416}
417
418#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
419pub enum McpPermissionRequest {
420    Connect {
421        server_id: String,
422        metadata: MetadataMap,
423    },
424    InvokeTool {
425        server_id: String,
426        tool_name: String,
427        metadata: MetadataMap,
428    },
429    ReadResource {
430        server_id: String,
431        resource_id: String,
432        metadata: MetadataMap,
433    },
434    FetchPrompt {
435        server_id: String,
436        prompt_id: String,
437        metadata: MetadataMap,
438    },
439    UseAuthScope {
440        server_id: String,
441        scope: String,
442        metadata: MetadataMap,
443    },
444}
445
446impl McpPermissionRequest {
447    fn metadata_map(&self) -> &MetadataMap {
448        match self {
449            Self::Connect { metadata, .. }
450            | Self::InvokeTool { metadata, .. }
451            | Self::ReadResource { metadata, .. }
452            | Self::FetchPrompt { metadata, .. }
453            | Self::UseAuthScope { metadata, .. } => metadata,
454        }
455    }
456}
457
458impl PermissionRequest for McpPermissionRequest {
459    fn kind(&self) -> &'static str {
460        match self {
461            Self::Connect { .. } => "mcp.connect",
462            Self::InvokeTool { .. } => "mcp.invoke_tool",
463            Self::ReadResource { .. } => "mcp.read_resource",
464            Self::FetchPrompt { .. } => "mcp.fetch_prompt",
465            Self::UseAuthScope { .. } => "mcp.use_auth_scope",
466        }
467    }
468
469    fn summary(&self) -> String {
470        match self {
471            Self::Connect { server_id, .. } => format!("Connect MCP server {server_id}"),
472            Self::InvokeTool {
473                server_id,
474                tool_name,
475                ..
476            } => format!("Invoke MCP tool {server_id}.{tool_name}"),
477            Self::ReadResource {
478                server_id,
479                resource_id,
480                ..
481            } => format!("Read MCP resource {server_id}:{resource_id}"),
482            Self::FetchPrompt {
483                server_id,
484                prompt_id,
485                ..
486            } => format!("Fetch MCP prompt {server_id}:{prompt_id}"),
487            Self::UseAuthScope {
488                server_id, scope, ..
489            } => format!("Use MCP auth scope {server_id}:{scope}"),
490        }
491    }
492
493    fn metadata(&self) -> &MetadataMap {
494        self.metadata_map()
495    }
496
497    fn as_any(&self) -> &dyn Any {
498        self
499    }
500}
501
502pub struct CustomKindPolicy {
503    allowed_kinds: BTreeSet<String>,
504    denied_kinds: BTreeSet<String>,
505    require_approval_by_default: bool,
506}
507
508impl CustomKindPolicy {
509    pub fn new(require_approval_by_default: bool) -> Self {
510        Self {
511            allowed_kinds: BTreeSet::new(),
512            denied_kinds: BTreeSet::new(),
513            require_approval_by_default,
514        }
515    }
516
517    pub fn allow_kind(mut self, kind: impl Into<String>) -> Self {
518        self.allowed_kinds.insert(kind.into());
519        self
520    }
521
522    pub fn deny_kind(mut self, kind: impl Into<String>) -> Self {
523        self.denied_kinds.insert(kind.into());
524        self
525    }
526}
527
528impl PermissionPolicy for CustomKindPolicy {
529    fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
530        let kind = request.kind();
531        if !kind.starts_with("custom.") {
532            return PolicyMatch::NoOpinion;
533        }
534        if self.denied_kinds.contains(kind) {
535            return PolicyMatch::Deny(PermissionDenial {
536                code: PermissionCode::CustomPolicyDenied,
537                message: format!("custom permission kind {kind} is denied"),
538                metadata: request.metadata().clone(),
539            });
540        }
541        if self.allowed_kinds.contains(kind) {
542            return PolicyMatch::Allow;
543        }
544        if self.require_approval_by_default {
545            PolicyMatch::RequireApproval(ApprovalRequest {
546                id: ApprovalId::new(format!("approval:{kind}")),
547                request_kind: kind.to_string(),
548                reason: ApprovalReason::PolicyRequiresConfirmation,
549                summary: request.summary(),
550                metadata: request.metadata().clone(),
551            })
552        } else {
553            PolicyMatch::NoOpinion
554        }
555    }
556}
557
558pub struct PathPolicy {
559    allowed_roots: Vec<PathBuf>,
560    protected_roots: Vec<PathBuf>,
561    require_approval_outside_allowed: bool,
562}
563
564impl PathPolicy {
565    pub fn new() -> Self {
566        Self {
567            allowed_roots: Vec::new(),
568            protected_roots: Vec::new(),
569            require_approval_outside_allowed: true,
570        }
571    }
572
573    pub fn allow_root(mut self, root: impl Into<PathBuf>) -> Self {
574        self.allowed_roots.push(root.into());
575        self
576    }
577
578    pub fn protect_root(mut self, root: impl Into<PathBuf>) -> Self {
579        self.protected_roots.push(root.into());
580        self
581    }
582
583    pub fn require_approval_outside_allowed(mut self, value: bool) -> Self {
584        self.require_approval_outside_allowed = value;
585        self
586    }
587}
588
589impl Default for PathPolicy {
590    fn default() -> Self {
591        Self::new()
592    }
593}
594
595impl PermissionPolicy for PathPolicy {
596    fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
597        let Some(fs) = request
598            .as_any()
599            .downcast_ref::<FileSystemPermissionRequest>()
600        else {
601            return PolicyMatch::NoOpinion;
602        };
603
604        let candidate_paths: Vec<&Path> = match fs {
605            FileSystemPermissionRequest::Move { from, to, .. } => {
606                vec![from.as_path(), to.as_path()]
607            }
608            FileSystemPermissionRequest::Read { path, .. }
609            | FileSystemPermissionRequest::Write { path, .. }
610            | FileSystemPermissionRequest::Edit { path, .. }
611            | FileSystemPermissionRequest::Delete { path, .. }
612            | FileSystemPermissionRequest::List { path, .. }
613            | FileSystemPermissionRequest::CreateDir { path, .. } => vec![path.as_path()],
614        };
615
616        if candidate_paths.iter().any(|path| {
617            self.protected_roots
618                .iter()
619                .any(|root| path.starts_with(root))
620        }) {
621            return PolicyMatch::Deny(PermissionDenial {
622                code: PermissionCode::PathNotAllowed,
623                message: format!("path access denied for {}", fs.summary()),
624                metadata: fs.metadata().clone(),
625            });
626        }
627
628        if self.allowed_roots.is_empty() {
629            return PolicyMatch::NoOpinion;
630        }
631
632        let all_allowed = candidate_paths
633            .iter()
634            .all(|path| self.allowed_roots.iter().any(|root| path.starts_with(root)));
635
636        if all_allowed {
637            PolicyMatch::Allow
638        } else if self.require_approval_outside_allowed {
639            PolicyMatch::RequireApproval(ApprovalRequest {
640                id: ApprovalId::new(format!("approval:{}", fs.kind())),
641                request_kind: fs.kind().to_string(),
642                reason: ApprovalReason::SensitivePath,
643                summary: fs.summary(),
644                metadata: fs.metadata().clone(),
645            })
646        } else {
647            PolicyMatch::Deny(PermissionDenial {
648                code: PermissionCode::PathNotAllowed,
649                message: format!("path outside allowed roots for {}", fs.summary()),
650                metadata: fs.metadata().clone(),
651            })
652        }
653    }
654}
655
656pub struct CommandPolicy {
657    allowed_executables: BTreeSet<String>,
658    denied_executables: BTreeSet<String>,
659    allowed_cwds: Vec<PathBuf>,
660    denied_env_keys: BTreeSet<String>,
661    require_approval_for_unknown: bool,
662}
663
664impl CommandPolicy {
665    pub fn new() -> Self {
666        Self {
667            allowed_executables: BTreeSet::new(),
668            denied_executables: BTreeSet::new(),
669            allowed_cwds: Vec::new(),
670            denied_env_keys: BTreeSet::new(),
671            require_approval_for_unknown: true,
672        }
673    }
674
675    pub fn allow_executable(mut self, executable: impl Into<String>) -> Self {
676        self.allowed_executables.insert(executable.into());
677        self
678    }
679
680    pub fn deny_executable(mut self, executable: impl Into<String>) -> Self {
681        self.denied_executables.insert(executable.into());
682        self
683    }
684
685    pub fn allow_cwd(mut self, cwd: impl Into<PathBuf>) -> Self {
686        self.allowed_cwds.push(cwd.into());
687        self
688    }
689
690    pub fn deny_env_key(mut self, key: impl Into<String>) -> Self {
691        self.denied_env_keys.insert(key.into());
692        self
693    }
694
695    pub fn require_approval_for_unknown(mut self, value: bool) -> Self {
696        self.require_approval_for_unknown = value;
697        self
698    }
699}
700
701impl Default for CommandPolicy {
702    fn default() -> Self {
703        Self::new()
704    }
705}
706
707impl PermissionPolicy for CommandPolicy {
708    fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
709        let Some(shell) = request.as_any().downcast_ref::<ShellPermissionRequest>() else {
710            return PolicyMatch::NoOpinion;
711        };
712
713        if self.denied_executables.contains(&shell.executable)
714            || shell
715                .env_keys
716                .iter()
717                .any(|key| self.denied_env_keys.contains(key))
718        {
719            return PolicyMatch::Deny(PermissionDenial {
720                code: PermissionCode::CommandNotAllowed,
721                message: format!("command denied for {}", shell.summary()),
722                metadata: shell.metadata().clone(),
723            });
724        }
725
726        if let Some(cwd) = &shell.cwd
727            && !self.allowed_cwds.is_empty()
728            && !self.allowed_cwds.iter().any(|root| cwd.starts_with(root))
729        {
730            return PolicyMatch::RequireApproval(ApprovalRequest {
731                id: ApprovalId::new("approval:shell.cwd"),
732                request_kind: shell.kind().to_string(),
733                reason: ApprovalReason::SensitiveCommand,
734                summary: shell.summary(),
735                metadata: shell.metadata().clone(),
736            });
737        }
738
739        if self.allowed_executables.is_empty()
740            || self.allowed_executables.contains(&shell.executable)
741        {
742            PolicyMatch::Allow
743        } else if self.require_approval_for_unknown {
744            PolicyMatch::RequireApproval(ApprovalRequest {
745                id: ApprovalId::new("approval:shell.command"),
746                request_kind: shell.kind().to_string(),
747                reason: ApprovalReason::SensitiveCommand,
748                summary: shell.summary(),
749                metadata: shell.metadata().clone(),
750            })
751        } else {
752            PolicyMatch::Deny(PermissionDenial {
753                code: PermissionCode::CommandNotAllowed,
754                message: format!("executable {} is not allowed", shell.executable),
755                metadata: shell.metadata().clone(),
756            })
757        }
758    }
759}
760
761pub struct McpServerPolicy {
762    trusted_servers: BTreeSet<String>,
763    allowed_auth_scopes: BTreeSet<String>,
764    require_approval_for_untrusted: bool,
765}
766
767impl McpServerPolicy {
768    pub fn new() -> Self {
769        Self {
770            trusted_servers: BTreeSet::new(),
771            allowed_auth_scopes: BTreeSet::new(),
772            require_approval_for_untrusted: true,
773        }
774    }
775
776    pub fn trust_server(mut self, server_id: impl Into<String>) -> Self {
777        self.trusted_servers.insert(server_id.into());
778        self
779    }
780
781    pub fn allow_auth_scope(mut self, scope: impl Into<String>) -> Self {
782        self.allowed_auth_scopes.insert(scope.into());
783        self
784    }
785}
786
787impl Default for McpServerPolicy {
788    fn default() -> Self {
789        Self::new()
790    }
791}
792
793impl PermissionPolicy for McpServerPolicy {
794    fn evaluate(&self, request: &dyn PermissionRequest) -> PolicyMatch {
795        let Some(mcp) = request.as_any().downcast_ref::<McpPermissionRequest>() else {
796            return PolicyMatch::NoOpinion;
797        };
798
799        let server_id = match mcp {
800            McpPermissionRequest::Connect { server_id, .. }
801            | McpPermissionRequest::InvokeTool { server_id, .. }
802            | McpPermissionRequest::ReadResource { server_id, .. }
803            | McpPermissionRequest::FetchPrompt { server_id, .. }
804            | McpPermissionRequest::UseAuthScope { server_id, .. } => server_id,
805        };
806
807        if !self.trusted_servers.is_empty() && !self.trusted_servers.contains(server_id) {
808            return if self.require_approval_for_untrusted {
809                PolicyMatch::RequireApproval(ApprovalRequest {
810                    id: ApprovalId::new(format!("approval:mcp:{server_id}")),
811                    request_kind: mcp.kind().to_string(),
812                    reason: ApprovalReason::SensitiveServer,
813                    summary: mcp.summary(),
814                    metadata: mcp.metadata().clone(),
815                })
816            } else {
817                PolicyMatch::Deny(PermissionDenial {
818                    code: PermissionCode::ServerNotTrusted,
819                    message: format!("MCP server {server_id} is not trusted"),
820                    metadata: mcp.metadata().clone(),
821                })
822            };
823        }
824
825        if let McpPermissionRequest::UseAuthScope { scope, .. } = mcp
826            && !self.allowed_auth_scopes.is_empty()
827            && !self.allowed_auth_scopes.contains(scope)
828        {
829            return PolicyMatch::Deny(PermissionDenial {
830                code: PermissionCode::AuthScopeNotAllowed,
831                message: format!("MCP auth scope {scope} is not allowed"),
832                metadata: mcp.metadata().clone(),
833            });
834        }
835
836        PolicyMatch::Allow
837    }
838}
839
840#[async_trait]
841pub trait Tool: Send + Sync {
842    fn spec(&self) -> &ToolSpec;
843
844    fn proposed_requests(
845        &self,
846        _request: &ToolRequest,
847    ) -> Result<Vec<Box<dyn PermissionRequest>>, ToolError> {
848        Ok(Vec::new())
849    }
850
851    async fn invoke(
852        &self,
853        request: ToolRequest,
854        ctx: &mut ToolContext<'_>,
855    ) -> Result<ToolResult, ToolError>;
856}
857
858#[derive(Clone, Default)]
859pub struct ToolRegistry {
860    tools: BTreeMap<ToolName, Arc<dyn Tool>>,
861}
862
863impl ToolRegistry {
864    pub fn new() -> Self {
865        Self::default()
866    }
867
868    pub fn register<T>(&mut self, tool: T) -> &mut Self
869    where
870        T: Tool + 'static,
871    {
872        self.tools.insert(tool.spec().name.clone(), Arc::new(tool));
873        self
874    }
875
876    pub fn with<T>(mut self, tool: T) -> Self
877    where
878        T: Tool + 'static,
879    {
880        self.register(tool);
881        self
882    }
883
884    pub fn register_arc(&mut self, tool: Arc<dyn Tool>) -> &mut Self {
885        self.tools.insert(tool.spec().name.clone(), tool);
886        self
887    }
888
889    pub fn get(&self, name: &ToolName) -> Option<Arc<dyn Tool>> {
890        self.tools.get(name).cloned()
891    }
892
893    pub fn tools(&self) -> Vec<Arc<dyn Tool>> {
894        self.tools.values().cloned().collect()
895    }
896
897    pub fn specs(&self) -> Vec<ToolSpec> {
898        self.tools
899            .values()
900            .map(|tool| tool.spec().clone())
901            .collect()
902    }
903}
904
905impl ToolSpec {
906    pub fn as_invocable_spec(&self) -> InvocableSpec {
907        InvocableSpec {
908            name: CapabilityName::new(self.name.0.clone()),
909            description: self.description.clone(),
910            input_schema: self.input_schema.clone(),
911            metadata: self.metadata.clone(),
912        }
913    }
914}
915
916pub struct ToolInvocableAdapter {
917    spec: InvocableSpec,
918    tool: Arc<dyn Tool>,
919    permissions: Arc<dyn PermissionChecker>,
920    resources: Arc<dyn ToolResources>,
921    next_call_id: AtomicU64,
922}
923
924impl ToolInvocableAdapter {
925    pub fn new(
926        tool: Arc<dyn Tool>,
927        permissions: Arc<dyn PermissionChecker>,
928        resources: Arc<dyn ToolResources>,
929    ) -> Self {
930        let spec = tool.spec().as_invocable_spec();
931        Self {
932            spec,
933            tool,
934            permissions,
935            resources,
936            next_call_id: AtomicU64::new(1),
937        }
938    }
939}
940
941#[async_trait]
942impl Invocable for ToolInvocableAdapter {
943    fn spec(&self) -> &InvocableSpec {
944        &self.spec
945    }
946
947    async fn invoke(
948        &self,
949        request: InvocableRequest,
950        ctx: &mut CapabilityContext<'_>,
951    ) -> Result<InvocableResult, CapabilityError> {
952        let tool_request = ToolRequest {
953            call_id: ToolCallId::new(format!(
954                "tool-call-{}",
955                self.next_call_id.fetch_add(1, Ordering::Relaxed)
956            )),
957            tool_name: self.tool.spec().name.clone(),
958            input: request.input,
959            session_id: ctx
960                .session_id
961                .cloned()
962                .unwrap_or_else(|| SessionId::new("capability-session")),
963            turn_id: ctx
964                .turn_id
965                .cloned()
966                .unwrap_or_else(|| TurnId::new("capability-turn")),
967            metadata: request.metadata,
968        };
969
970        for permission_request in self
971            .tool
972            .proposed_requests(&tool_request)
973            .map_err(|error| CapabilityError::InvalidInput(error.to_string()))?
974        {
975            match self.permissions.evaluate(permission_request.as_ref()) {
976                PermissionDecision::Allow => {}
977                PermissionDecision::Deny(denial) => {
978                    return Err(CapabilityError::ExecutionFailed(format!(
979                        "tool permission denied: {denial:?}"
980                    )));
981                }
982                PermissionDecision::RequireApproval(req) => {
983                    return Err(CapabilityError::Unavailable(format!(
984                        "tool invocation requires approval: {}",
985                        req.summary
986                    )));
987                }
988            }
989        }
990
991        let mut tool_ctx = ToolContext {
992            capability: CapabilityContext {
993                session_id: ctx.session_id,
994                turn_id: ctx.turn_id,
995                metadata: ctx.metadata,
996            },
997            permissions: self.permissions.as_ref(),
998            resources: self.resources.as_ref(),
999            cancellation: None,
1000        };
1001
1002        let result = self
1003            .tool
1004            .invoke(tool_request, &mut tool_ctx)
1005            .await
1006            .map_err(|error| CapabilityError::ExecutionFailed(error.to_string()))?;
1007
1008        Ok(InvocableResult {
1009            output: match result.result.output {
1010                ToolOutput::Text(text) => InvocableOutput::Text(text),
1011                ToolOutput::Structured(value) => InvocableOutput::Structured(value),
1012                ToolOutput::Parts(parts) => InvocableOutput::Items(vec![Item {
1013                    id: None,
1014                    kind: ItemKind::Tool,
1015                    parts,
1016                    metadata: MetadataMap::new(),
1017                }]),
1018                ToolOutput::Files(files) => {
1019                    let parts = files.into_iter().map(Part::File).collect();
1020                    InvocableOutput::Items(vec![Item {
1021                        id: None,
1022                        kind: ItemKind::Tool,
1023                        parts,
1024                        metadata: MetadataMap::new(),
1025                    }])
1026                }
1027            },
1028            metadata: result.metadata,
1029        })
1030    }
1031}
1032
1033pub struct ToolCapabilityProvider {
1034    invocables: Vec<Arc<dyn Invocable>>,
1035}
1036
1037impl ToolCapabilityProvider {
1038    pub fn from_registry(
1039        registry: &ToolRegistry,
1040        permissions: Arc<dyn PermissionChecker>,
1041        resources: Arc<dyn ToolResources>,
1042    ) -> Self {
1043        let invocables = registry
1044            .tools()
1045            .into_iter()
1046            .map(|tool| {
1047                Arc::new(ToolInvocableAdapter::new(
1048                    tool,
1049                    permissions.clone(),
1050                    resources.clone(),
1051                )) as Arc<dyn Invocable>
1052            })
1053            .collect();
1054
1055        Self { invocables }
1056    }
1057}
1058
1059impl CapabilityProvider for ToolCapabilityProvider {
1060    fn invocables(&self) -> Vec<Arc<dyn Invocable>> {
1061        self.invocables.clone()
1062    }
1063
1064    fn resources(&self) -> Vec<Arc<dyn ResourceProvider>> {
1065        Vec::new()
1066    }
1067
1068    fn prompts(&self) -> Vec<Arc<dyn PromptProvider>> {
1069        Vec::new()
1070    }
1071}
1072
1073#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1074pub enum ToolExecutionOutcome {
1075    Completed(ToolResult),
1076    Interrupted(ToolInterruption),
1077    Failed(ToolError),
1078}
1079
1080#[async_trait]
1081pub trait ToolExecutor: Send + Sync {
1082    async fn execute(
1083        &self,
1084        request: ToolRequest,
1085        ctx: &mut ToolContext<'_>,
1086    ) -> ToolExecutionOutcome;
1087
1088    async fn execute_approved(
1089        &self,
1090        request: ToolRequest,
1091        approved_request: &ApprovalRequest,
1092        ctx: &mut ToolContext<'_>,
1093    ) -> ToolExecutionOutcome {
1094        let _ = approved_request;
1095        self.execute(request, ctx).await
1096    }
1097}
1098
1099pub struct BasicToolExecutor {
1100    registry: ToolRegistry,
1101}
1102
1103impl BasicToolExecutor {
1104    pub fn new(registry: ToolRegistry) -> Self {
1105        Self { registry }
1106    }
1107
1108    pub fn specs(&self) -> Vec<ToolSpec> {
1109        self.registry.specs()
1110    }
1111
1112    async fn execute_inner(
1113        &self,
1114        request: ToolRequest,
1115        approved_request_id: Option<&ApprovalId>,
1116        ctx: &mut ToolContext<'_>,
1117    ) -> ToolExecutionOutcome {
1118        let Some(tool) = self.registry.get(&request.tool_name) else {
1119            return ToolExecutionOutcome::Failed(ToolError::NotFound(request.tool_name));
1120        };
1121
1122        match tool.proposed_requests(&request) {
1123            Ok(requests) => {
1124                for permission_request in requests {
1125                    match ctx.permissions.evaluate(permission_request.as_ref()) {
1126                        PermissionDecision::Allow => {}
1127                        PermissionDecision::Deny(denial) => {
1128                            return ToolExecutionOutcome::Failed(ToolError::PermissionDenied(
1129                                denial,
1130                            ));
1131                        }
1132                        PermissionDecision::RequireApproval(req) => {
1133                            if approved_request_id != Some(&req.id) {
1134                                return ToolExecutionOutcome::Interrupted(
1135                                    ToolInterruption::ApprovalRequired(req),
1136                                );
1137                            }
1138                        }
1139                    }
1140                }
1141            }
1142            Err(error) => return ToolExecutionOutcome::Failed(error),
1143        }
1144
1145        match tool.invoke(request, ctx).await {
1146            Ok(result) => ToolExecutionOutcome::Completed(result),
1147            Err(ToolError::AuthRequired(request)) => {
1148                ToolExecutionOutcome::Interrupted(ToolInterruption::AuthRequired(*request))
1149            }
1150            Err(error) => ToolExecutionOutcome::Failed(error),
1151        }
1152    }
1153}
1154
1155#[async_trait]
1156impl ToolExecutor for BasicToolExecutor {
1157    async fn execute(
1158        &self,
1159        request: ToolRequest,
1160        ctx: &mut ToolContext<'_>,
1161    ) -> ToolExecutionOutcome {
1162        self.execute_inner(request, None, ctx).await
1163    }
1164
1165    async fn execute_approved(
1166        &self,
1167        request: ToolRequest,
1168        approved_request: &ApprovalRequest,
1169        ctx: &mut ToolContext<'_>,
1170    ) -> ToolExecutionOutcome {
1171        self.execute_inner(request, Some(&approved_request.id), ctx)
1172            .await
1173    }
1174}
1175
1176#[derive(Debug, Error, Clone, PartialEq, Serialize, Deserialize)]
1177pub enum ToolError {
1178    #[error("tool not found: {0}")]
1179    NotFound(ToolName),
1180    #[error("invalid tool input: {0}")]
1181    InvalidInput(String),
1182    #[error("tool permission denied: {0:?}")]
1183    PermissionDenied(PermissionDenial),
1184    #[error("tool execution failed: {0}")]
1185    ExecutionFailed(String),
1186    #[error("tool auth required: {0:?}")]
1187    AuthRequired(Box<AuthRequest>),
1188    #[error("tool unavailable: {0}")]
1189    Unavailable(String),
1190    #[error("tool execution cancelled")]
1191    Cancelled,
1192    #[error("internal tool error: {0}")]
1193    Internal(String),
1194}
1195
1196impl ToolError {
1197    pub fn permission_denied(denial: PermissionDenial) -> Self {
1198        Self::PermissionDenied(denial)
1199    }
1200}
1201
1202impl From<PermissionDenial> for ToolError {
1203    fn from(value: PermissionDenial) -> Self {
1204        Self::permission_denied(value)
1205    }
1206}
1207
1208#[cfg(test)]
1209mod tests {
1210    use super::*;
1211
1212    #[test]
1213    fn command_policy_can_deny_unknown_executables_without_approval() {
1214        let policy = CommandPolicy::new()
1215            .allow_executable("pwd")
1216            .require_approval_for_unknown(false);
1217        let request = ShellPermissionRequest {
1218            executable: "rm".into(),
1219            argv: vec!["-rf".into(), "/tmp/demo".into()],
1220            cwd: None,
1221            env_keys: Vec::new(),
1222            metadata: MetadataMap::new(),
1223        };
1224
1225        match policy.evaluate(&request) {
1226            PolicyMatch::Deny(denial) => {
1227                assert_eq!(denial.code, PermissionCode::CommandNotAllowed);
1228            }
1229            other => panic!("unexpected policy match: {other:?}"),
1230        }
1231    }
1232}
agentkit_tools_core/lib.rs

agentkit_tools_core/
lib.rs
