| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232 |
- use std::collections::BTreeMap;
- #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
- pub enum PermissionMode {
- ReadOnly,
- WorkspaceWrite,
- DangerFullAccess,
- Prompt,
- Allow,
- }
- impl PermissionMode {
- #[must_use]
- pub fn as_str(self) -> &'static str {
- match self {
- Self::ReadOnly => "read-only",
- Self::WorkspaceWrite => "workspace-write",
- Self::DangerFullAccess => "danger-full-access",
- Self::Prompt => "prompt",
- Self::Allow => "allow",
- }
- }
- }
- #[derive(Debug, Clone, PartialEq, Eq)]
- pub struct PermissionRequest {
- pub tool_name: String,
- pub input: String,
- pub current_mode: PermissionMode,
- pub required_mode: PermissionMode,
- }
- #[derive(Debug, Clone, PartialEq, Eq)]
- pub enum PermissionPromptDecision {
- Allow,
- Deny { reason: String },
- }
- pub trait PermissionPrompter {
- fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision;
- }
- #[derive(Debug, Clone, PartialEq, Eq)]
- pub enum PermissionOutcome {
- Allow,
- Deny { reason: String },
- }
- #[derive(Debug, Clone, PartialEq, Eq)]
- pub struct PermissionPolicy {
- active_mode: PermissionMode,
- tool_requirements: BTreeMap<String, PermissionMode>,
- }
- impl PermissionPolicy {
- #[must_use]
- pub fn new(active_mode: PermissionMode) -> Self {
- Self {
- active_mode,
- tool_requirements: BTreeMap::new(),
- }
- }
- #[must_use]
- pub fn with_tool_requirement(
- mut self,
- tool_name: impl Into<String>,
- required_mode: PermissionMode,
- ) -> Self {
- self.tool_requirements
- .insert(tool_name.into(), required_mode);
- self
- }
- #[must_use]
- pub fn active_mode(&self) -> PermissionMode {
- self.active_mode
- }
- #[must_use]
- pub fn required_mode_for(&self, tool_name: &str) -> PermissionMode {
- self.tool_requirements
- .get(tool_name)
- .copied()
- .unwrap_or(PermissionMode::DangerFullAccess)
- }
- #[must_use]
- pub fn authorize(
- &self,
- tool_name: &str,
- input: &str,
- mut prompter: Option<&mut dyn PermissionPrompter>,
- ) -> PermissionOutcome {
- let current_mode = self.active_mode();
- let required_mode = self.required_mode_for(tool_name);
- if current_mode == PermissionMode::Allow || current_mode >= required_mode {
- return PermissionOutcome::Allow;
- }
- let request = PermissionRequest {
- tool_name: tool_name.to_string(),
- input: input.to_string(),
- current_mode,
- required_mode,
- };
- if current_mode == PermissionMode::Prompt
- || (current_mode == PermissionMode::WorkspaceWrite
- && required_mode == PermissionMode::DangerFullAccess)
- {
- return match prompter.as_mut() {
- Some(prompter) => match prompter.decide(&request) {
- PermissionPromptDecision::Allow => PermissionOutcome::Allow,
- PermissionPromptDecision::Deny { reason } => PermissionOutcome::Deny { reason },
- },
- None => PermissionOutcome::Deny {
- reason: format!(
- "tool '{tool_name}' requires approval to escalate from {} to {}",
- current_mode.as_str(),
- required_mode.as_str()
- ),
- },
- };
- }
- PermissionOutcome::Deny {
- reason: format!(
- "tool '{tool_name}' requires {} permission; current mode is {}",
- required_mode.as_str(),
- current_mode.as_str()
- ),
- }
- }
- }
- #[cfg(test)]
- mod tests {
- use super::{
- PermissionMode, PermissionOutcome, PermissionPolicy, PermissionPromptDecision,
- PermissionPrompter, PermissionRequest,
- };
- struct RecordingPrompter {
- seen: Vec<PermissionRequest>,
- allow: bool,
- }
- impl PermissionPrompter for RecordingPrompter {
- fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision {
- self.seen.push(request.clone());
- if self.allow {
- PermissionPromptDecision::Allow
- } else {
- PermissionPromptDecision::Deny {
- reason: "not now".to_string(),
- }
- }
- }
- }
- #[test]
- fn allows_tools_when_active_mode_meets_requirement() {
- let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
- .with_tool_requirement("read_file", PermissionMode::ReadOnly)
- .with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
- assert_eq!(
- policy.authorize("read_file", "{}", None),
- PermissionOutcome::Allow
- );
- assert_eq!(
- policy.authorize("write_file", "{}", None),
- PermissionOutcome::Allow
- );
- }
- #[test]
- fn denies_read_only_escalations_without_prompt() {
- let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
- .with_tool_requirement("write_file", PermissionMode::WorkspaceWrite)
- .with_tool_requirement("bash", PermissionMode::DangerFullAccess);
- assert!(matches!(
- policy.authorize("write_file", "{}", None),
- PermissionOutcome::Deny { reason } if reason.contains("requires workspace-write permission")
- ));
- assert!(matches!(
- policy.authorize("bash", "{}", None),
- PermissionOutcome::Deny { reason } if reason.contains("requires danger-full-access permission")
- ));
- }
- #[test]
- fn prompts_for_workspace_write_to_danger_full_access_escalation() {
- let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
- .with_tool_requirement("bash", PermissionMode::DangerFullAccess);
- let mut prompter = RecordingPrompter {
- seen: Vec::new(),
- allow: true,
- };
- let outcome = policy.authorize("bash", "echo hi", Some(&mut prompter));
- assert_eq!(outcome, PermissionOutcome::Allow);
- assert_eq!(prompter.seen.len(), 1);
- assert_eq!(prompter.seen[0].tool_name, "bash");
- assert_eq!(
- prompter.seen[0].current_mode,
- PermissionMode::WorkspaceWrite
- );
- assert_eq!(
- prompter.seen[0].required_mode,
- PermissionMode::DangerFullAccess
- );
- }
- #[test]
- fn honors_prompt_rejection_reason() {
- let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
- .with_tool_requirement("bash", PermissionMode::DangerFullAccess);
- let mut prompter = RecordingPrompter {
- seen: Vec::new(),
- allow: false,
- };
- assert!(matches!(
- policy.authorize("bash", "echo hi", Some(&mut prompter)),
- PermissionOutcome::Deny { reason } if reason == "not now"
- ));
- }
- }
|