1use std::{
38 collections::HashMap,
39 sync::{Arc, Mutex},
40};
41
42use basis::approval::{
43 ApprovalAnswer, ApprovalDecision, ApprovalRequest, Approver, RuntimeError,
44 ToolAuthorizationDecision, ToolAuthorizationRequest, ToolAuthorizer, is_consequential,
45};
46
47#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
52#[serde(rename_all = "lowercase")]
53pub enum ApprovalPolicy {
54 Always,
56 #[default]
58 Prompt,
59 Never,
61}
62
63impl ApprovalPolicy {
64 pub const fn type_tag(self) -> &'static str {
67 match self {
68 Self::Always => "always",
69 Self::Prompt => "prompt",
70 Self::Never => "never",
71 }
72 }
73
74 pub fn from_type_tag(tag: &str) -> Option<Self> {
76 match tag {
77 "always" => Some(Self::Always),
78 "prompt" => Some(Self::Prompt),
79 "never" => Some(Self::Never),
80 _ => None,
81 }
82 }
83}
84
85#[derive(Clone)]
90pub struct SessionApproval {
91 inner: Arc<Mutex<State>>,
92}
93
94struct State {
95 current: ApprovalPolicy,
96 remembered: HashMap<String, bool>,
97}
98
99impl SessionApproval {
100 pub fn new(initial: ApprovalPolicy) -> Self {
102 Self {
103 inner: Arc::new(Mutex::new(State {
104 current: initial,
105 remembered: HashMap::new(),
106 })),
107 }
108 }
109
110 pub fn current(&self) -> ApprovalPolicy {
111 self.lock().current
112 }
113
114 pub fn set(&self, policy: ApprovalPolicy) {
116 let mut state = self.lock();
117 state.current = policy;
118 state.remembered.clear();
119 }
120
121 fn remember(&self, tool_name: &str, allow: bool) {
122 self.lock().remembered.insert(tool_name.to_string(), allow);
123 }
124
125 fn remembered(&self, tool_name: &str) -> Option<bool> {
126 self.lock().remembered.get(tool_name).copied()
127 }
128
129 fn lock(&self) -> std::sync::MutexGuard<'_, State> {
130 self.inner
131 .lock()
132 .unwrap_or_else(|poisoned| poisoned.into_inner())
133 }
134}
135
136fn read_only_refusal(tool_name: &str) -> String {
144 format!(
145 "{tool_name} changes state outside this process, and this session is set to refuse that"
146 )
147}
148
149#[derive(Clone)]
219pub struct PolicyGate {
220 approval: SessionApproval,
221}
222
223impl PolicyGate {
224 pub fn new(approval: SessionApproval) -> Self {
226 Self { approval }
227 }
228}
229
230#[async_trait::async_trait]
231impl ToolAuthorizer for PolicyGate {
232 async fn authorize(
233 &self,
234 request: &ToolAuthorizationRequest,
235 ) -> Result<ToolAuthorizationDecision, RuntimeError> {
236 if !is_consequential(request.preview.side_effect_level) {
237 return Ok(ToolAuthorizationDecision::allow());
238 }
239
240 Ok(match self.approval.current() {
241 ApprovalPolicy::Never => {
242 ToolAuthorizationDecision::deny(read_only_refusal(&request.tool_name))
243 }
244 ApprovalPolicy::Always | ApprovalPolicy::Prompt => {
247 ToolAuthorizationDecision::prompt(format!(
248 "{} wants to run and can change state outside this process",
249 request.tool_name
250 ))
251 }
252 })
253 }
254}
255
256pub struct PolicyApprover<A> {
259 approval: SessionApproval,
260 inner: A,
261}
262
263impl<A> PolicyApprover<A> {
264 pub fn new(approval: SessionApproval, inner: A) -> Self {
265 Self { approval, inner }
266 }
267}
268
269#[async_trait::async_trait]
270impl<A: Approver> Approver for PolicyApprover<A> {
271 async fn approve(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
272 match self.approval.current() {
273 ApprovalPolicy::Always => ApprovalDecision::Allow.into(),
274 ApprovalPolicy::Never => ApprovalAnswer::new(ApprovalDecision::Deny)
278 .because(read_only_refusal(&request.tool_name)),
279 ApprovalPolicy::Prompt => self.ask(request).await,
280 }
281 }
282}
283
284impl<A: Approver> PolicyApprover<A> {
285 async fn ask(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
286 if let Some(allow) = self.approval.remembered(&request.tool_name) {
287 return if allow {
288 ApprovalDecision::Allow.into()
289 } else {
290 ApprovalAnswer::new(ApprovalDecision::Deny).because(format!(
291 "{} was refused earlier in this session, and that answer still stands",
292 request.tool_name
293 ))
294 };
295 }
296
297 let answer = self.inner.approve(request).await;
298 match answer.decision {
299 ApprovalDecision::AllowForSession => {
300 self.approval.remember(&request.tool_name, true);
301 ApprovalAnswer {
302 decision: ApprovalDecision::Allow,
303 ..answer
304 }
305 }
306 ApprovalDecision::DenyForSession => {
307 self.approval.remember(&request.tool_name, false);
308 ApprovalAnswer {
309 decision: ApprovalDecision::Deny,
310 ..answer
311 }
312 }
313 _ => answer,
314 }
315 }
316}
317
318#[cfg(test)]
319mod tests {
320 use super::*;
321 use basis::ToolSideEffectLevel;
322 use mentra::tool::{
323 ToolApprovalCategory, ToolAuthorizationOutcome, ToolAuthorizationPreview, ToolCapability,
324 ToolDurability, ToolExecutionCategory,
325 };
326 use serde_json::json;
327 use std::sync::atomic::{AtomicUsize, Ordering};
328
329 struct Counting {
330 asked: Arc<AtomicUsize>,
331 answer: ApprovalDecision,
332 }
333
334 #[async_trait::async_trait]
335 impl Approver for Counting {
336 async fn approve(&mut self, _request: &ApprovalRequest) -> ApprovalAnswer {
337 self.asked.fetch_add(1, Ordering::SeqCst);
338 self.answer.into()
339 }
340 }
341
342 fn request(tool_name: &str) -> ApprovalRequest {
343 ApprovalRequest {
344 request_id: "r1".to_string(),
345 tool_call_id: "c1".to_string(),
346 tool_name: tool_name.to_string(),
347 description: "wants to write".to_string(),
348 input: json!({}),
349 side_effect_level: Some(ToolSideEffectLevel::LocalState),
350 }
351 }
352
353 fn gate(
354 initial: ApprovalPolicy,
355 answer: ApprovalDecision,
356 ) -> (SessionApproval, PolicyApprover<Counting>, Arc<AtomicUsize>) {
357 let approval = SessionApproval::new(initial);
358 let asked = Arc::new(AtomicUsize::new(0));
359 let approver = PolicyApprover::new(
360 approval.clone(),
361 Counting {
362 asked: Arc::clone(&asked),
363 answer,
364 },
365 );
366 (approval, approver, asked)
367 }
368
369 #[test]
370 fn the_tags_round_trip() {
371 assert_eq!(ApprovalPolicy::default(), ApprovalPolicy::Prompt);
372
373 for policy in [
374 ApprovalPolicy::Always,
375 ApprovalPolicy::Prompt,
376 ApprovalPolicy::Never,
377 ] {
378 let tag = policy.type_tag();
379 assert_eq!(ApprovalPolicy::from_type_tag(tag), Some(policy));
380 assert_eq!(serde_json::to_value(policy).unwrap(), json!(tag));
381 assert_eq!(
382 serde_json::from_value::<ApprovalPolicy>(json!(tag)).unwrap(),
383 policy
384 );
385 }
386 assert_eq!(ApprovalPolicy::from_type_tag("architect"), None);
387 }
388
389 #[tokio::test]
390 async fn allow_and_refuse_answer_without_asking() {
391 for (mode, expected) in [
392 (ApprovalPolicy::Always, ApprovalDecision::Allow),
393 (ApprovalPolicy::Never, ApprovalDecision::Deny),
394 ] {
395 let (_modes, mut approver, asked) = gate(mode, ApprovalDecision::Allow);
396
397 assert_eq!(approver.approve(&request("shell")).await.decision, expected);
398 assert_eq!(
399 asked.load(Ordering::SeqCst),
400 0,
401 "{mode:?} has nothing to ask about"
402 );
403 }
404 }
405
406 #[tokio::test]
407 async fn a_read_only_session_says_so_when_it_refuses() {
408 let (_approval, mut approver, _asked) =
409 gate(ApprovalPolicy::Never, ApprovalDecision::Allow);
410
411 assert_eq!(
412 approver.approve(&request("shell")).await.reason.as_deref(),
413 Some(
414 "shell changes state outside this process, \
415 and this session is set to refuse that"
416 )
417 );
418 }
419
420 #[tokio::test]
421 async fn asking_puts_the_request_to_the_client() {
422 let (_approval, mut approver, asked) =
423 gate(ApprovalPolicy::Prompt, ApprovalDecision::Allow);
424
425 assert_eq!(
426 approver.approve(&request("shell")).await.decision,
427 ApprovalDecision::Allow
428 );
429 assert_eq!(asked.load(Ordering::SeqCst), 1);
430 }
431
432 #[tokio::test]
433 async fn an_answer_for_the_session_is_not_asked_twice() {
434 let (_approval, mut approver, asked) =
435 gate(ApprovalPolicy::Prompt, ApprovalDecision::AllowForSession);
436
437 assert_eq!(
438 approver.approve(&request("shell")).await.decision,
439 ApprovalDecision::Allow
440 );
441 assert_eq!(
442 approver.approve(&request("shell")).await.decision,
443 ApprovalDecision::Allow
444 );
445 assert_eq!(asked.load(Ordering::SeqCst), 1);
446
447 assert_eq!(
448 approver.approve(&request("files")).await.decision,
449 ApprovalDecision::Allow
450 );
451 assert_eq!(asked.load(Ordering::SeqCst), 2);
452 }
453
454 #[tokio::test]
455 async fn changing_mode_forgets_what_was_allowed_for_the_session() {
456 let (approval, mut approver, _asked) =
457 gate(ApprovalPolicy::Prompt, ApprovalDecision::AllowForSession);
458
459 approver.approve(&request("shell")).await;
460 approval.set(ApprovalPolicy::Never);
461
462 assert_eq!(
463 approver.approve(&request("shell")).await.decision,
464 ApprovalDecision::Deny,
465 "a stale allow must not survive the mode that replaced it"
466 );
467
468 approval.set(ApprovalPolicy::Prompt);
469 assert_eq!(
470 approver.approve(&request("shell")).await.decision,
471 ApprovalDecision::Allow,
472 "moving back to prompt must ask again rather than reviving a stale rule"
473 );
474 }
475
476 fn authorization(tool_name: &str, level: ToolSideEffectLevel) -> ToolAuthorizationRequest {
477 ToolAuthorizationRequest {
478 agent_id: "a1".to_string(),
479 agent_name: "test".to_string(),
480 model: "m".to_string(),
481 history_len: 1,
482 tool_call_id: "tc-1".to_string(),
483 tool_name: tool_name.to_string(),
484 preview: ToolAuthorizationPreview {
485 working_directory: std::path::PathBuf::from("/repo"),
486 capabilities: vec![ToolCapability::FilesystemWrite],
487 side_effect_level: level,
488 durability: ToolDurability::Ephemeral,
489 execution_category: ToolExecutionCategory::default(),
490 approval_category: ToolApprovalCategory::default(),
491 raw_input: json!({}),
492 structured_input: json!({}),
493 },
494 }
495 }
496
497 async fn gated(
498 approval: &SessionApproval,
499 level: ToolSideEffectLevel,
500 ) -> ToolAuthorizationDecision {
501 PolicyGate::new(approval.clone())
502 .authorize(&authorization("shell", level))
503 .await
504 .expect("authorization does not error")
505 }
506
507 #[tokio::test]
508 async fn a_read_only_session_refuses_at_the_gate_rather_than_surfacing() {
509 let approval = SessionApproval::new(ApprovalPolicy::Never);
513 let decision = gated(&approval, ToolSideEffectLevel::Process).await;
514
515 assert_eq!(decision.outcome, ToolAuthorizationOutcome::Deny);
516 assert_eq!(
517 decision.reason.as_deref(),
518 Some(
519 "shell changes state outside this process, \
520 and this session is set to refuse that"
521 ),
522 "and it refuses in the same words the approver would have used"
523 );
524 }
525
526 #[tokio::test]
527 async fn every_other_policy_still_surfaces_the_call() {
528 for policy in [ApprovalPolicy::Always, ApprovalPolicy::Prompt] {
529 let approval = SessionApproval::new(policy);
530
531 assert_eq!(
532 gated(&approval, ToolSideEffectLevel::Process).await.outcome,
533 ToolAuthorizationOutcome::Prompt,
534 "{policy:?} decides per turn, and a remembered rule may answer first"
535 );
536 }
537 }
538
539 #[tokio::test]
540 async fn a_read_is_allowed_whatever_the_policy_is() {
541 for policy in [
544 ApprovalPolicy::Always,
545 ApprovalPolicy::Prompt,
546 ApprovalPolicy::Never,
547 ] {
548 let approval = SessionApproval::new(policy);
549
550 assert_eq!(
551 gated(&approval, ToolSideEffectLevel::None).await.outcome,
552 ToolAuthorizationOutcome::Allow,
553 "{policy:?} has nothing to ask about for a read"
554 );
555 }
556 }
557
558 #[tokio::test]
559 async fn the_gate_follows_the_session_it_was_installed_on() {
560 let approval = SessionApproval::new(ApprovalPolicy::Prompt);
564 let gate = PolicyGate::new(approval.clone());
565 let request = authorization("shell", ToolSideEffectLevel::Process);
566
567 assert_eq!(
568 gate.authorize(&request).await.expect("no error").outcome,
569 ToolAuthorizationOutcome::Prompt
570 );
571
572 approval.set(ApprovalPolicy::Never);
573 assert_eq!(
574 gate.authorize(&request).await.expect("no error").outcome,
575 ToolAuthorizationOutcome::Deny,
576 "the same gate must read the mode the session is on now"
577 );
578
579 approval.set(ApprovalPolicy::Always);
580 assert_eq!(
581 gate.authorize(&request).await.expect("no error").outcome,
582 ToolAuthorizationOutcome::Prompt,
583 "and must stop refusing when the session moves back"
584 );
585 }
586
587 #[test]
588 fn a_gate_waits_as_long_as_it_takes_and_offers_no_way_not_to() {
589 assert_eq!(
595 PolicyGate::new(SessionApproval::new(ApprovalPolicy::Prompt)).timeout(),
596 None
597 );
598 }
599
600 #[tokio::test]
601 async fn a_refusal_for_the_session_is_also_remembered() {
602 let (_approval, mut approver, asked) =
603 gate(ApprovalPolicy::Prompt, ApprovalDecision::DenyForSession);
604
605 assert_eq!(
606 approver.approve(&request("shell")).await.decision,
607 ApprovalDecision::Deny
608 );
609 let repeated = approver.approve(&request("shell")).await;
610 assert_eq!(repeated.decision, ApprovalDecision::Deny);
611 assert_eq!(
612 repeated.reason.as_deref(),
613 Some("shell was refused earlier in this session, and that answer still stands")
614 );
615 assert_eq!(asked.load(Ordering::SeqCst), 1);
616 }
617}