Skip to main content

typesec_agent/
tool.rs

1//! Capability-bound tool wrappers for agent and MCP-style tool execution.
2
3use std::marker::PhantomData;
4use std::{future::Future, pin::Pin};
5
6use typesec_core::{Capability, Permission, Resource, typestate::Authenticated};
7
8use crate::{SecureAgent, executor::TaskError};
9
10/// Boxed future returned by protected tool handlers.
11pub type ToolFuture<'a> = Pin<Box<dyn Future<Output = Result<(), TaskError>> + Send + 'a>>;
12
13/// Metadata describing the authorization boundary for a protected tool.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct ToolSpec {
16    /// Tool name exposed to an agent or MCP client.
17    pub name: String,
18    /// Human-readable description.
19    pub description: String,
20    /// Permission required to invoke this tool.
21    pub required_permission: &'static str,
22    /// Resource identifier the permission applies to.
23    pub resource_id: String,
24}
25
26/// A tool that cannot run unless the caller supplies a matching capability.
27pub struct ProtectedTool<P, R, F>
28where
29    P: Permission,
30    R: Resource,
31{
32    spec: ToolSpec,
33    resource: R,
34    action: F,
35    _permission: PhantomData<fn() -> P>,
36}
37
38impl<P, R, F> ProtectedTool<P, R, F>
39where
40    P: Permission,
41    R: Resource,
42{
43    /// Create a new protected tool.
44    pub fn new(
45        name: impl Into<String>,
46        description: impl Into<String>,
47        resource: R,
48        action: F,
49    ) -> Self {
50        let resource_id = resource.resource_id().to_string();
51        Self {
52            spec: ToolSpec {
53                name: name.into(),
54                description: description.into(),
55                required_permission: P::name(),
56                resource_id,
57            },
58            resource,
59            action,
60            _permission: PhantomData,
61        }
62    }
63
64    /// Return this tool's authorization metadata.
65    pub fn spec(&self) -> &ToolSpec {
66        &self.spec
67    }
68}
69
70impl<P, R, F> ProtectedTool<P, R, F>
71where
72    P: Permission,
73    R: Resource,
74    F: for<'a> Fn(&'a R) -> ToolFuture<'a>,
75{
76    /// Invoke the tool with a typed capability.
77    pub async fn invoke(
78        &self,
79        agent: &SecureAgent<Authenticated>,
80        cap: &Capability<P, R>,
81    ) -> Result<(), TaskError> {
82        tracing::info!(
83            subject = %agent.subject(),
84            permission = %Capability::<P, R>::permission_name(),
85            resource = %cap.resource_id(),
86            tool = %self.spec.name,
87            "invoking protected tool"
88        );
89        (self.action)(&self.resource).await
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96    use std::sync::Arc;
97    use typesec_core::{
98        CanExecute,
99        policy::{PolicyEngine, PolicyResult},
100        resource::GenericResource,
101        typestate::Credentials,
102    };
103
104    struct AllowAll;
105    impl PolicyEngine for AllowAll {
106        fn check(&self, _: &str, _: &str, _: &str) -> PolicyResult {
107            PolicyResult::Allow
108        }
109    }
110
111    fn no_op_tool(_resource: &GenericResource) -> ToolFuture<'_> {
112        Box::pin(async { Ok(()) })
113    }
114
115    #[tokio::test]
116    async fn protected_tool_invokes_with_capability() {
117        let agent = SecureAgent::new(Arc::new(AllowAll))
118            .authenticate_unverified(Credentials::new("agent:test", "tok"))
119            .expect("auth ok");
120        let resource = GenericResource::new("Gmail.ListEmails", "tool");
121        let cap: Capability<CanExecute, GenericResource> =
122            agent.request_capability(&resource).await.expect("cap ok");
123        let tool = ProtectedTool::<CanExecute, _, _>::new(
124            "gmail.list",
125            "List email messages",
126            resource,
127            no_op_tool,
128        );
129
130        assert_eq!(tool.spec().required_permission, "execute");
131        tool.invoke(&agent, &cap).await.expect("tool should run");
132    }
133}