1use 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
10pub type ToolFuture<'a> = Pin<Box<dyn Future<Output = Result<(), TaskError>> + Send + 'a>>;
12
13#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct ToolSpec {
16 pub name: String,
18 pub description: String,
20 pub required_permission: &'static str,
22 pub resource_id: String,
24}
25
26pub 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 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 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 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(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}