use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
use basis::approval::{
ApprovalAnswer, ApprovalDecision, ApprovalRequest, Approver, RuntimeError,
ToolAuthorizationDecision, ToolAuthorizationRequest, ToolAuthorizer, is_consequential,
};
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ApprovalPolicy {
Always,
#[default]
Prompt,
Never,
}
impl ApprovalPolicy {
pub const fn type_tag(self) -> &'static str {
match self {
Self::Always => "always",
Self::Prompt => "prompt",
Self::Never => "never",
}
}
pub fn from_type_tag(tag: &str) -> Option<Self> {
match tag {
"always" => Some(Self::Always),
"prompt" => Some(Self::Prompt),
"never" => Some(Self::Never),
_ => None,
}
}
}
#[derive(Clone)]
pub struct SessionApproval {
inner: Arc<Mutex<State>>,
}
struct State {
current: ApprovalPolicy,
remembered: HashMap<String, bool>,
}
impl SessionApproval {
pub fn new(initial: ApprovalPolicy) -> Self {
Self {
inner: Arc::new(Mutex::new(State {
current: initial,
remembered: HashMap::new(),
})),
}
}
pub fn current(&self) -> ApprovalPolicy {
self.lock().current
}
pub fn set(&self, policy: ApprovalPolicy) {
let mut state = self.lock();
state.current = policy;
state.remembered.clear();
}
fn remember(&self, tool_name: &str, allow: bool) {
self.lock().remembered.insert(tool_name.to_string(), allow);
}
fn remembered(&self, tool_name: &str) -> Option<bool> {
self.lock().remembered.get(tool_name).copied()
}
fn lock(&self) -> std::sync::MutexGuard<'_, State> {
self.inner
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
}
fn read_only_refusal(tool_name: &str) -> String {
format!(
"{tool_name} changes state outside this process, and this session is set to refuse that"
)
}
#[derive(Clone)]
pub struct PolicyGate {
approval: SessionApproval,
}
impl PolicyGate {
pub fn new(approval: SessionApproval) -> Self {
Self { approval }
}
}
#[async_trait::async_trait]
impl ToolAuthorizer for PolicyGate {
async fn authorize(
&self,
request: &ToolAuthorizationRequest,
) -> Result<ToolAuthorizationDecision, RuntimeError> {
if !is_consequential(request.preview.side_effect_level) {
return Ok(ToolAuthorizationDecision::allow());
}
Ok(match self.approval.current() {
ApprovalPolicy::Never => {
ToolAuthorizationDecision::deny(read_only_refusal(&request.tool_name))
}
ApprovalPolicy::Always | ApprovalPolicy::Prompt => {
ToolAuthorizationDecision::prompt(format!(
"{} wants to run and can change state outside this process",
request.tool_name
))
}
})
}
}
pub struct PolicyApprover<A> {
approval: SessionApproval,
inner: A,
}
impl<A> PolicyApprover<A> {
pub fn new(approval: SessionApproval, inner: A) -> Self {
Self { approval, inner }
}
}
#[async_trait::async_trait]
impl<A: Approver> Approver for PolicyApprover<A> {
async fn approve(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
match self.approval.current() {
ApprovalPolicy::Always => ApprovalDecision::Allow.into(),
ApprovalPolicy::Never => ApprovalAnswer::new(ApprovalDecision::Deny)
.because(read_only_refusal(&request.tool_name)),
ApprovalPolicy::Prompt => self.ask(request).await,
}
}
}
impl<A: Approver> PolicyApprover<A> {
async fn ask(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
if let Some(allow) = self.approval.remembered(&request.tool_name) {
return if allow {
ApprovalDecision::Allow.into()
} else {
ApprovalAnswer::new(ApprovalDecision::Deny).because(format!(
"{} was refused earlier in this session, and that answer still stands",
request.tool_name
))
};
}
let answer = self.inner.approve(request).await;
match answer.decision {
ApprovalDecision::AllowForSession => {
self.approval.remember(&request.tool_name, true);
ApprovalAnswer {
decision: ApprovalDecision::Allow,
..answer
}
}
ApprovalDecision::DenyForSession => {
self.approval.remember(&request.tool_name, false);
ApprovalAnswer {
decision: ApprovalDecision::Deny,
..answer
}
}
_ => answer,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use basis::ToolSideEffectLevel;
use mentra::tool::{
ToolApprovalCategory, ToolAuthorizationOutcome, ToolAuthorizationPreview, ToolCapability,
ToolDurability, ToolExecutionCategory,
};
use serde_json::json;
use std::sync::atomic::{AtomicUsize, Ordering};
struct Counting {
asked: Arc<AtomicUsize>,
answer: ApprovalDecision,
}
#[async_trait::async_trait]
impl Approver for Counting {
async fn approve(&mut self, _request: &ApprovalRequest) -> ApprovalAnswer {
self.asked.fetch_add(1, Ordering::SeqCst);
self.answer.into()
}
}
fn request(tool_name: &str) -> ApprovalRequest {
ApprovalRequest {
request_id: "r1".to_string(),
tool_call_id: "c1".to_string(),
tool_name: tool_name.to_string(),
description: "wants to write".to_string(),
input: json!({}),
side_effect_level: Some(ToolSideEffectLevel::LocalState),
}
}
fn gate(
initial: ApprovalPolicy,
answer: ApprovalDecision,
) -> (SessionApproval, PolicyApprover<Counting>, Arc<AtomicUsize>) {
let approval = SessionApproval::new(initial);
let asked = Arc::new(AtomicUsize::new(0));
let approver = PolicyApprover::new(
approval.clone(),
Counting {
asked: Arc::clone(&asked),
answer,
},
);
(approval, approver, asked)
}
#[test]
fn the_tags_round_trip() {
assert_eq!(ApprovalPolicy::default(), ApprovalPolicy::Prompt);
for policy in [
ApprovalPolicy::Always,
ApprovalPolicy::Prompt,
ApprovalPolicy::Never,
] {
let tag = policy.type_tag();
assert_eq!(ApprovalPolicy::from_type_tag(tag), Some(policy));
assert_eq!(serde_json::to_value(policy).unwrap(), json!(tag));
assert_eq!(
serde_json::from_value::<ApprovalPolicy>(json!(tag)).unwrap(),
policy
);
}
assert_eq!(ApprovalPolicy::from_type_tag("architect"), None);
}
#[tokio::test]
async fn allow_and_refuse_answer_without_asking() {
for (mode, expected) in [
(ApprovalPolicy::Always, ApprovalDecision::Allow),
(ApprovalPolicy::Never, ApprovalDecision::Deny),
] {
let (_modes, mut approver, asked) = gate(mode, ApprovalDecision::Allow);
assert_eq!(approver.approve(&request("shell")).await.decision, expected);
assert_eq!(
asked.load(Ordering::SeqCst),
0,
"{mode:?} has nothing to ask about"
);
}
}
#[tokio::test]
async fn a_read_only_session_says_so_when_it_refuses() {
let (_approval, mut approver, _asked) =
gate(ApprovalPolicy::Never, ApprovalDecision::Allow);
assert_eq!(
approver.approve(&request("shell")).await.reason.as_deref(),
Some(
"shell changes state outside this process, \
and this session is set to refuse that"
)
);
}
#[tokio::test]
async fn asking_puts_the_request_to_the_client() {
let (_approval, mut approver, asked) =
gate(ApprovalPolicy::Prompt, ApprovalDecision::Allow);
assert_eq!(
approver.approve(&request("shell")).await.decision,
ApprovalDecision::Allow
);
assert_eq!(asked.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn an_answer_for_the_session_is_not_asked_twice() {
let (_approval, mut approver, asked) =
gate(ApprovalPolicy::Prompt, ApprovalDecision::AllowForSession);
assert_eq!(
approver.approve(&request("shell")).await.decision,
ApprovalDecision::Allow
);
assert_eq!(
approver.approve(&request("shell")).await.decision,
ApprovalDecision::Allow
);
assert_eq!(asked.load(Ordering::SeqCst), 1);
assert_eq!(
approver.approve(&request("files")).await.decision,
ApprovalDecision::Allow
);
assert_eq!(asked.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn changing_mode_forgets_what_was_allowed_for_the_session() {
let (approval, mut approver, _asked) =
gate(ApprovalPolicy::Prompt, ApprovalDecision::AllowForSession);
approver.approve(&request("shell")).await;
approval.set(ApprovalPolicy::Never);
assert_eq!(
approver.approve(&request("shell")).await.decision,
ApprovalDecision::Deny,
"a stale allow must not survive the mode that replaced it"
);
approval.set(ApprovalPolicy::Prompt);
assert_eq!(
approver.approve(&request("shell")).await.decision,
ApprovalDecision::Allow,
"moving back to prompt must ask again rather than reviving a stale rule"
);
}
fn authorization(tool_name: &str, level: ToolSideEffectLevel) -> ToolAuthorizationRequest {
ToolAuthorizationRequest {
agent_id: "a1".to_string(),
agent_name: "test".to_string(),
model: "m".to_string(),
history_len: 1,
tool_call_id: "tc-1".to_string(),
tool_name: tool_name.to_string(),
preview: ToolAuthorizationPreview {
working_directory: std::path::PathBuf::from("/repo"),
capabilities: vec![ToolCapability::FilesystemWrite],
side_effect_level: level,
durability: ToolDurability::Ephemeral,
execution_category: ToolExecutionCategory::default(),
approval_category: ToolApprovalCategory::default(),
raw_input: json!({}),
structured_input: json!({}),
},
}
}
async fn gated(
approval: &SessionApproval,
level: ToolSideEffectLevel,
) -> ToolAuthorizationDecision {
PolicyGate::new(approval.clone())
.authorize(&authorization("shell", level))
.await
.expect("authorization does not error")
}
#[tokio::test]
async fn a_read_only_session_refuses_at_the_gate_rather_than_surfacing() {
let approval = SessionApproval::new(ApprovalPolicy::Never);
let decision = gated(&approval, ToolSideEffectLevel::Process).await;
assert_eq!(decision.outcome, ToolAuthorizationOutcome::Deny);
assert_eq!(
decision.reason.as_deref(),
Some(
"shell changes state outside this process, \
and this session is set to refuse that"
),
"and it refuses in the same words the approver would have used"
);
}
#[tokio::test]
async fn every_other_policy_still_surfaces_the_call() {
for policy in [ApprovalPolicy::Always, ApprovalPolicy::Prompt] {
let approval = SessionApproval::new(policy);
assert_eq!(
gated(&approval, ToolSideEffectLevel::Process).await.outcome,
ToolAuthorizationOutcome::Prompt,
"{policy:?} decides per turn, and a remembered rule may answer first"
);
}
}
#[tokio::test]
async fn a_read_is_allowed_whatever_the_policy_is() {
for policy in [
ApprovalPolicy::Always,
ApprovalPolicy::Prompt,
ApprovalPolicy::Never,
] {
let approval = SessionApproval::new(policy);
assert_eq!(
gated(&approval, ToolSideEffectLevel::None).await.outcome,
ToolAuthorizationOutcome::Allow,
"{policy:?} has nothing to ask about for a read"
);
}
}
#[tokio::test]
async fn the_gate_follows_the_session_it_was_installed_on() {
let approval = SessionApproval::new(ApprovalPolicy::Prompt);
let gate = PolicyGate::new(approval.clone());
let request = authorization("shell", ToolSideEffectLevel::Process);
assert_eq!(
gate.authorize(&request).await.expect("no error").outcome,
ToolAuthorizationOutcome::Prompt
);
approval.set(ApprovalPolicy::Never);
assert_eq!(
gate.authorize(&request).await.expect("no error").outcome,
ToolAuthorizationOutcome::Deny,
"the same gate must read the mode the session is on now"
);
approval.set(ApprovalPolicy::Always);
assert_eq!(
gate.authorize(&request).await.expect("no error").outcome,
ToolAuthorizationOutcome::Prompt,
"and must stop refusing when the session moves back"
);
}
#[test]
fn a_gate_waits_as_long_as_it_takes_and_offers_no_way_not_to() {
assert_eq!(
PolicyGate::new(SessionApproval::new(ApprovalPolicy::Prompt)).timeout(),
None
);
}
#[tokio::test]
async fn a_refusal_for_the_session_is_also_remembered() {
let (_approval, mut approver, asked) =
gate(ApprovalPolicy::Prompt, ApprovalDecision::DenyForSession);
assert_eq!(
approver.approve(&request("shell")).await.decision,
ApprovalDecision::Deny
);
let repeated = approver.approve(&request("shell")).await;
assert_eq!(repeated.decision, ApprovalDecision::Deny);
assert_eq!(
repeated.reason.as_deref(),
Some("shell was refused earlier in this session, and that answer still stands")
);
assert_eq!(asked.load(Ordering::SeqCst), 1);
}
}