1use serde::{Deserialize, Serialize};
4use std::collections::{BTreeMap, BTreeSet};
5
6#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
7pub enum Capability {
8 RepoRead,
9 RepoWriteSrc,
10 RepoWriteTests,
11 RepoWriteDocs,
12 RepoWriteCi,
13 RepoWriteLockfile,
14 SecretsRead,
15 NetworkGithub,
16 NetworkNpm,
17 ShellSafe,
18 ShellPrivileged,
19 Mcp { server: String, capability: String },
20 Custom(String),
21}
22
23impl Capability {
24 pub fn parse(value: &str) -> Self {
25 match value {
26 "repo.read" => Self::RepoRead,
27 "repo.write.src" | "repo.write" => Self::RepoWriteSrc,
28 "repo.write.tests" => Self::RepoWriteTests,
29 "repo.write.docs" => Self::RepoWriteDocs,
30 "repo.write.ci" => Self::RepoWriteCi,
31 "repo.write.lockfile" => Self::RepoWriteLockfile,
32 "secrets.read" => Self::SecretsRead,
33 "network.github" => Self::NetworkGithub,
34 "network.npm" | "network.package" => Self::NetworkNpm,
35 "shell.safe" | "shell.exec" => Self::ShellSafe,
36 "shell.privileged" | "shell.bash" => Self::ShellPrivileged,
37 other if other.starts_with("mcp.") => {
38 let mut parts = other.splitn(3, '.');
39 let _ = parts.next();
40 let server = parts.next().unwrap_or_default().to_string();
41 let capability = parts.next().unwrap_or("*").to_string();
42 Self::Mcp { server, capability }
43 }
44 other => Self::Custom(other.to_string()),
45 }
46 }
47
48 pub fn as_key(&self) -> String {
49 match self {
50 Self::RepoRead => "repo.read".to_string(),
51 Self::RepoWriteSrc => "repo.write.src".to_string(),
52 Self::RepoWriteTests => "repo.write.tests".to_string(),
53 Self::RepoWriteDocs => "repo.write.docs".to_string(),
54 Self::RepoWriteCi => "repo.write.ci".to_string(),
55 Self::RepoWriteLockfile => "repo.write.lockfile".to_string(),
56 Self::SecretsRead => "secrets.read".to_string(),
57 Self::NetworkGithub => "network.github".to_string(),
58 Self::NetworkNpm => "network.npm".to_string(),
59 Self::ShellSafe => "shell.safe".to_string(),
60 Self::ShellPrivileged => "shell.privileged".to_string(),
61 Self::Mcp { server, capability } => format!("mcp.{server}.{capability}"),
62 Self::Custom(value) => value.clone(),
63 }
64 }
65
66 pub fn is_guarded(&self) -> bool {
67 matches!(
68 self,
69 Self::SecretsRead
70 | Self::ShellPrivileged
71 | Self::RepoWriteCi
72 | Self::RepoWriteLockfile
73 | Self::NetworkGithub
74 | Self::NetworkNpm
75 )
76 }
77}
78
79#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
80pub enum CapabilityScope {
81 Session,
82 Turn(String),
83 Branch(String),
84 Subagent(String),
85 SingleCall(String),
86}
87
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
89pub enum CapabilityDecision {
90 Requested,
91 Granted,
92 Denied,
93 Consumed,
94 Expired,
95 Violated,
96}
97
98#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
99pub struct CapabilityGrant {
100 pub capability: Capability,
101 pub scope: CapabilityScope,
102 pub justification: String,
103 pub expires_at_ms: Option<u64>,
104 pub guarded: bool,
105}
106
107#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
108pub struct CapabilityLedgerEntry {
109 pub capability: Capability,
110 pub scope: CapabilityScope,
111 pub decision: CapabilityDecision,
112 pub at_ms: u64,
113 pub justification: String,
114}
115
116#[derive(Debug, Clone, Default, Serialize, Deserialize)]
117pub struct CapabilityLedger {
118 grants: BTreeMap<String, CapabilityGrant>,
119 consumed: BTreeSet<String>,
120 entries: Vec<CapabilityLedgerEntry>,
121}
122
123impl CapabilityLedger {
124 pub fn entries(&self) -> &[CapabilityLedgerEntry] {
125 &self.entries
126 }
127
128 pub fn request(&mut self, capability: Capability, scope: CapabilityScope, at_ms: u64) {
129 self.entries.push(CapabilityLedgerEntry {
130 capability,
131 scope,
132 decision: CapabilityDecision::Requested,
133 at_ms,
134 justification: String::new(),
135 });
136 }
137
138 pub fn grant(
139 &mut self,
140 capability: Capability,
141 scope: CapabilityScope,
142 justification: impl Into<String>,
143 expires_at_ms: Option<u64>,
144 at_ms: u64,
145 explicit_user_approval: bool,
146 ) -> bool {
147 if capability.is_guarded() && !explicit_user_approval {
148 self.entries.push(CapabilityLedgerEntry {
149 capability,
150 scope,
151 decision: CapabilityDecision::Denied,
152 at_ms,
153 justification: "guarded capability requires explicit approval".to_string(),
154 });
155 return false;
156 }
157
158 let grant = CapabilityGrant {
159 guarded: capability.is_guarded(),
160 capability: capability.clone(),
161 scope: scope.clone(),
162 justification: justification.into(),
163 expires_at_ms,
164 };
165 self.grants
166 .insert(grant_key(&capability, &scope), grant.clone());
167 self.entries.push(CapabilityLedgerEntry {
168 capability,
169 scope,
170 decision: CapabilityDecision::Granted,
171 at_ms,
172 justification: grant.justification,
173 });
174 true
175 }
176
177 pub fn consume(
178 &mut self,
179 capability: &Capability,
180 scope: &CapabilityScope,
181 call_id: &str,
182 at_ms: u64,
183 ) -> bool {
184 let key = grant_key(capability, scope);
185 let Some(grant) = self.grants.get(&key).cloned() else {
186 self.entries.push(CapabilityLedgerEntry {
187 capability: capability.clone(),
188 scope: scope.clone(),
189 decision: CapabilityDecision::Violated,
190 at_ms,
191 justification: "capability was not granted".to_string(),
192 });
193 return false;
194 };
195 if grant
196 .expires_at_ms
197 .map(|expires| at_ms > expires)
198 .unwrap_or(false)
199 {
200 self.entries.push(CapabilityLedgerEntry {
201 capability: capability.clone(),
202 scope: scope.clone(),
203 decision: CapabilityDecision::Expired,
204 at_ms,
205 justification: "capability grant expired".to_string(),
206 });
207 return false;
208 }
209 self.consumed.insert(format!("{key}:{call_id}"));
210 self.entries.push(CapabilityLedgerEntry {
211 capability: capability.clone(),
212 scope: scope.clone(),
213 decision: CapabilityDecision::Consumed,
214 at_ms,
215 justification: call_id.to_string(),
216 });
217 true
218 }
219}
220
221pub fn capabilities_from_tool_metadata(values: &[String]) -> Vec<Capability> {
222 values
223 .iter()
224 .map(|value| Capability::parse(value))
225 .collect()
226}
227
228fn grant_key(capability: &Capability, scope: &CapabilityScope) -> String {
229 format!("{}::{scope:?}", capability.as_key())
230}
231
232#[cfg(test)]
233mod tests {
234 use super::*;
235
236 #[test]
237 fn guarded_capability_requires_explicit_approval() {
238 let mut ledger = CapabilityLedger::default();
239
240 let granted = ledger.grant(
241 Capability::ShellPrivileged,
242 CapabilityScope::Session,
243 "run shell",
244 None,
245 1,
246 false,
247 );
248
249 assert!(!granted);
250 assert_eq!(ledger.entries()[0].decision, CapabilityDecision::Denied);
251 }
252
253 #[test]
254 fn granted_capability_can_be_consumed_before_expiry() {
255 let mut ledger = CapabilityLedger::default();
256 assert!(ledger.grant(
257 Capability::RepoRead,
258 CapabilityScope::Turn("t1".to_string()),
259 "read repo",
260 Some(10),
261 1,
262 false,
263 ));
264
265 assert!(ledger.consume(
266 &Capability::RepoRead,
267 &CapabilityScope::Turn("t1".to_string()),
268 "call-1",
269 5,
270 ));
271 }
272
273 #[test]
274 fn expired_capability_is_not_consumed() {
275 let mut ledger = CapabilityLedger::default();
276 assert!(ledger.grant(
277 Capability::RepoRead,
278 CapabilityScope::Session,
279 "read repo",
280 Some(10),
281 1,
282 false,
283 ));
284
285 assert!(!ledger.consume(&Capability::RepoRead, &CapabilityScope::Session, "call", 11));
286 assert_eq!(
287 ledger.entries().last().unwrap().decision,
288 CapabilityDecision::Expired
289 );
290 }
291}