use std::time::Duration;
use async_trait::async_trait;
use mentra::{
error::RuntimeError,
tool::{
ToolAuthorizationDecision, ToolAuthorizationRequest, ToolAuthorizer, ToolSideEffectLevel,
},
};
use serde_json::Value;
#[derive(Debug, Clone, PartialEq)]
pub struct ApprovalRequest {
pub request_id: String,
pub tool_call_id: String,
pub tool_name: String,
pub description: String,
pub input: Value,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ApprovalDecision {
Allow,
#[default]
Deny,
AllowForSession,
DenyForSession,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ApprovalAnswer {
pub decision: ApprovalDecision,
pub reason: Option<String>,
}
impl ApprovalAnswer {
pub fn new(decision: ApprovalDecision) -> Self {
Self {
decision,
reason: None,
}
}
pub fn because(self, reason: impl Into<String>) -> Self {
Self {
reason: Some(reason.into()),
..self
}
}
}
impl From<ApprovalDecision> for ApprovalAnswer {
fn from(decision: ApprovalDecision) -> Self {
Self::new(decision)
}
}
#[async_trait]
pub trait Approver: Send + 'static {
async fn approve(&mut self, request: &ApprovalRequest) -> ApprovalAnswer;
}
#[async_trait]
impl<A: Approver + ?Sized> Approver for Box<A> {
async fn approve(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
(**self).approve(request).await
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct AllowAll;
#[async_trait]
impl Approver for AllowAll {
async fn approve(&mut self, _request: &ApprovalRequest) -> ApprovalAnswer {
ApprovalDecision::Allow.into()
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct DenyAll;
#[async_trait]
impl Approver for DenyAll {
async fn approve(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
ApprovalAnswer::new(ApprovalDecision::Deny).because(format!(
"{} changes state outside this process, which this run does not allow",
request.tool_name
))
}
}
pub fn is_consequential(level: ToolSideEffectLevel) -> bool {
!matches!(level, ToolSideEffectLevel::None)
}
#[derive(Debug, Default, Clone, Copy)]
pub struct ApprovalGate {
timeout: Option<Duration>,
}
impl ApprovalGate {
pub fn new() -> Self {
Self {
timeout: None,
}
}
pub fn with_timeout(self, timeout: Duration) -> Self {
Self {
timeout: Some(timeout),
}
}
}
#[async_trait]
impl ToolAuthorizer for ApprovalGate {
async fn authorize(
&self,
request: &ToolAuthorizationRequest,
) -> Result<ToolAuthorizationDecision, RuntimeError> {
if !is_consequential(request.preview.side_effect_level) {
return Ok(ToolAuthorizationDecision::allow());
}
Ok(ToolAuthorizationDecision::prompt(format!(
"{} wants to run and can change state outside this process",
request.tool_name
)))
}
fn timeout(&self) -> Option<Duration> {
self.timeout
}
}
#[cfg(test)]
mod tests {
use super::*;
use mentra::tool::{
ToolApprovalCategory, ToolAuthorizationOutcome, ToolAuthorizationPreview, ToolCapability,
ToolDurability, ToolExecutionCategory,
};
use serde_json::json;
use std::path::PathBuf;
fn request(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: name.to_string(),
preview: ToolAuthorizationPreview {
working_directory: 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 outcome(level: ToolSideEffectLevel) -> ToolAuthorizationOutcome {
ApprovalGate::new()
.authorize(&request("shell", level))
.await
.expect("authorization does not error")
.outcome
}
fn approval_request() -> ApprovalRequest {
ApprovalRequest {
request_id: "r".to_string(),
tool_call_id: "t".to_string(),
tool_name: "shell".to_string(),
description: "d".to_string(),
input: json!({}),
}
}
#[test]
fn only_side_effects_are_consequential() {
assert!(!is_consequential(ToolSideEffectLevel::None));
assert!(is_consequential(ToolSideEffectLevel::LocalState));
assert!(is_consequential(ToolSideEffectLevel::Process));
assert!(is_consequential(ToolSideEffectLevel::External));
}
#[tokio::test]
async fn a_read_only_call_is_never_worth_asking_about() {
assert_eq!(
outcome(ToolSideEffectLevel::None).await,
ToolAuthorizationOutcome::Allow,
"prompting for reads trains people to approve without reading"
);
}
#[tokio::test]
async fn every_other_call_is_put_to_the_approver() {
for level in [
ToolSideEffectLevel::LocalState,
ToolSideEffectLevel::Process,
ToolSideEffectLevel::External,
] {
assert_eq!(
outcome(level).await,
ToolAuthorizationOutcome::Prompt,
"{level:?} changes something outside this process"
);
}
}
#[tokio::test]
async fn the_request_says_which_tool_wants_to_run() {
let decision = ApprovalGate::new()
.authorize(&request("files", ToolSideEffectLevel::LocalState))
.await
.expect("no error");
let reason = decision.reason.expect("a prompt must say what it is about");
assert!(reason.contains("files"), "{reason}");
}
#[test]
fn a_gate_waits_as_long_as_it_takes_unless_told_otherwise() {
assert_eq!(ApprovalGate::new().timeout(), None);
assert_eq!(
ApprovalGate::new()
.with_timeout(Duration::from_secs(60))
.timeout(),
Some(Duration::from_secs(60))
);
}
#[tokio::test]
async fn the_trivial_approvers_answer_as_named() {
let request = approval_request();
assert_eq!(
AllowAll.approve(&request).await.decision,
ApprovalDecision::Allow
);
assert_eq!(
DenyAll.approve(&request).await.decision,
ApprovalDecision::Deny
);
}
#[tokio::test]
async fn a_blanket_refusal_tells_the_model_why_it_was_refused() {
let reason = DenyAll
.approve(&approval_request())
.await
.reason
.expect("a refusal the model can act on must explain itself");
assert_eq!(
reason,
"shell changes state outside this process, which this run does not allow"
);
}
#[tokio::test]
async fn a_boxed_approver_answers_exactly_as_the_one_inside() {
let mut chosen: Box<dyn Approver> = Box::new(DenyAll);
let answer = chosen.approve(&approval_request()).await;
assert_eq!(answer.decision, ApprovalDecision::Deny);
assert!(
answer.reason.is_some(),
"the reason must survive the indirection too"
);
}
#[test]
fn an_unanswered_request_is_a_refusal() {
assert_eq!(ApprovalDecision::default(), ApprovalDecision::Deny);
assert_eq!(
ApprovalAnswer::default(),
ApprovalAnswer::new(ApprovalDecision::Deny)
);
}
#[test]
fn a_reason_rides_along_without_changing_the_decision() {
let answer = ApprovalAnswer::from(ApprovalDecision::DenyForSession).because("no writes");
assert_eq!(answer.decision, ApprovalDecision::DenyForSession);
assert_eq!(answer.reason.as_deref(), Some("no writes"));
}
}