capo_agent/permissions/

extension.rs

#![cfg_attr(test, allow(clippy::expect_used, clippy::unwrap_used))]

use std::path::PathBuf;
use std::sync::Arc;

use async_trait::async_trait;
use motosan_agent_loop::core::decision::ToolDecision;
use motosan_agent_loop::core::ext_error::ExtError;
use motosan_agent_loop::core::extension::Extension;
use motosan_agent_loop::core::hook_ctx::HookCtx;
use motosan_agent_loop::llm::ToolCallItem;
use motosan_agent_tool::ToolResult;
use serde_json::Value;
use tokio::sync::{mpsc, oneshot, RwLock};

use super::policy::Policy;
use super::session_cache::SessionCache;
use super::Decision;
use crate::events::UiEvent;

/// How `PermissionExtension` resolves a tool call that steps 1-3 of
/// `decide()` left unanswered.
///
/// Tag-only enum (v0.10): no longer carries `ui_tx` inline. The
/// extension holds `ui_tx: Option<mpsc::Sender<UiEvent>>` separately;
/// the strategy tells `decide()` what to DO with that sender (or whether
/// to short-circuit).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PromptStrategy {
    /// Interactive: emit UiEvent::PermissionRequested, await user decision via the resolver oneshot.
    Prompt,
    /// Non-interactive (`--json` without `--dangerously-allow-all`): deny anything that would prompt.
    HeadlessDeny,
    /// v0.10 `bypass` mode: skip the prompt, return Allowed (after hard-block + allowlist).
    AllowAll,
    /// v0.10 `accept-edits` mode: auto-allow write/edit (after hard-block); bash falls through to `Prompt` or `HeadlessDeny` depending on whether ui_tx is Some.
    AcceptEdits,
}

pub struct PermissionExtension {
    policy: Arc<Policy>,
    cache: Arc<SessionCache>,
    project_root: PathBuf,
    /// v0.10: ui_tx is held here regardless of mode. When strategy is
    /// Prompt or (AcceptEdits + bash request), the extension sends
    /// `UiEvent::PermissionRequested` here. When strategy is AllowAll
    /// or HeadlessDeny, ui_tx is unused for this decision.
    ui_tx: Option<mpsc::Sender<UiEvent>>,
    /// v0.10: strategy is mutable at runtime via `set_strategy`. RwLock
    /// allows the existing &self method signature on Extension trait
    /// while supporting writes from Command::SetPermissionMode.
    strategy: Arc<RwLock<PromptStrategy>>,
}

impl PermissionExtension {
    pub fn new(
        policy: Arc<Policy>,
        cache: Arc<SessionCache>,
        project_root: PathBuf,
        ui_tx: mpsc::Sender<UiEvent>,
    ) -> Self {
        Self {
            policy,
            cache,
            project_root,
            ui_tx: Some(ui_tx),
            strategy: Arc::new(RwLock::new(PromptStrategy::Prompt)),
        }
    }

    pub fn headless(policy: Arc<Policy>, cache: Arc<SessionCache>, project_root: PathBuf) -> Self {
        Self {
            policy,
            cache,
            project_root,
            ui_tx: None,
            strategy: Arc::new(RwLock::new(PromptStrategy::HeadlessDeny)),
        }
    }

    /// v0.10 `accept-edits` mode: auto-allow write/edit; prompt bash via ui_tx.
    pub fn accept_edits(
        policy: Arc<Policy>,
        cache: Arc<SessionCache>,
        project_root: PathBuf,
        ui_tx: mpsc::Sender<UiEvent>,
    ) -> Self {
        Self {
            policy,
            cache,
            project_root,
            ui_tx: Some(ui_tx),
            strategy: Arc::new(RwLock::new(PromptStrategy::AcceptEdits)),
        }
    }

    /// v0.10 `bypass` mode: auto-allow everything except hard-blocked paths.
    pub fn allow_all(policy: Arc<Policy>, cache: Arc<SessionCache>, project_root: PathBuf) -> Self {
        Self {
            policy,
            cache,
            project_root,
            ui_tx: None,
            strategy: Arc::new(RwLock::new(PromptStrategy::AllowAll)),
        }
    }

    /// v0.10: runtime mutation of the strategy. Called by
    /// `App::set_permission_mode` after a `Command::SetPermissionMode`.
    pub async fn set_strategy(&self, strategy: PromptStrategy) {
        *self.strategy.write().await = strategy;
    }

    /// Read the current strategy. Used by `decide()` and test fixtures.
    pub async fn current_strategy(&self) -> PromptStrategy {
        *self.strategy.read().await
    }

    /// v0.10: construct with an externally-owned strategy lock so rebuilt
    /// sessions keep observing `App::set_permission_mode` changes.
    pub(crate) fn with_strategy_handle(
        policy: Arc<Policy>,
        cache: Arc<SessionCache>,
        project_root: PathBuf,
        ui_tx: Option<mpsc::Sender<UiEvent>>,
        strategy: Arc<RwLock<PromptStrategy>>,
    ) -> Self {
        Self {
            policy,
            cache,
            project_root,
            ui_tx,
            strategy,
        }
    }

    async fn decide(&self, tool_name: &str, args: &Value) -> Decision {
        // 1. Hard-block: `write`/`edit` into a `blocked_paths` match can't be
        //    prompted past.
        if matches!(tool_name, "write" | "edit") {
            if let Some(path) = args.get("path").and_then(|v| v.as_str()) {
                let abs = if std::path::Path::new(path).is_absolute() {
                    PathBuf::from(path)
                } else {
                    self.project_root.join(path)
                };
                let blocked = path_is_builtin_hard_blocked(&abs, &self.project_root)
                    || match tool_name {
                        "edit" => self.policy.edit_is_blocked(&abs, &self.project_root),
                        _ => self.policy.write_is_blocked(&abs, &self.project_root),
                    };
                if blocked {
                    return Decision::Denied(format!("{} is in a blocked path", abs.display()));
                }
            }
        }

        // 2. v0.10: mode-based fast-path. Runs AFTER hard-block, BEFORE allowlist.
        let strategy = *self.strategy.read().await;
        match (strategy, tool_name) {
            (PromptStrategy::AllowAll, _) => return Decision::Allowed,
            (PromptStrategy::AcceptEdits, "write" | "edit") => return Decision::Allowed,
            _ => {}
        }

        // 3. Persistent allowlist.
        let policy_allowed = match tool_name {
            "bash" => args
                .get("command")
                .and_then(|v| v.as_str())
                .map(|c| self.policy.bash_is_allowed(c))
                .unwrap_or(false),
            "write" | "edit" => args
                .get("path")
                .and_then(|v| v.as_str())
                .map(|p| {
                    let abs = std::path::PathBuf::from(p);
                    let abs = if abs.is_absolute() {
                        abs
                    } else {
                        self.project_root.join(&abs)
                    };
                    match tool_name {
                        "edit" => self.policy.edit_is_allowed(&abs, &self.project_root),
                        _ => self.policy.write_is_allowed(&abs, &self.project_root),
                    }
                })
                .unwrap_or(false),
            "read" | "grep" | "find" | "ls" => return Decision::Allowed,
            other if other.contains("__") => {
                let mut parts = other.splitn(2, "__");
                let server = parts.next().unwrap_or("");
                let tool = parts.next().unwrap_or("");
                self.policy.mcp_auto_allow(server, tool)
            }
            _ => false,
        };
        if policy_allowed {
            return Decision::Allowed;
        }

        // 4. Session cache (read-only here — Phase F's
        //    Command::ResolvePermission consumer is what writes to it when
        //    the user picks "S" in the modal. For Phase E, the cache will
        //    only ever be populated externally; this branch returns hits
        //    that were inserted by tests or by future code).
        let cache_key = SessionCache::key(tool_name, args);
        if let Some(cached) = self.cache.get(&cache_key) {
            return cached;
        }

        // 5. Resolve per the prompt strategy.
        let strategy = *self.strategy.read().await;
        match strategy {
            PromptStrategy::HeadlessDeny => {
                Decision::Denied("non-interactive: tool requires approval".into())
            }
            PromptStrategy::Prompt | PromptStrategy::AcceptEdits => {
                let Some(ui_tx) = &self.ui_tx else {
                    return Decision::Denied(
                        "no ui_tx attached to PermissionExtension; cannot prompt".into(),
                    );
                };
                let (resolver_tx, resolver_rx) = oneshot::channel::<Decision>();
                if ui_tx
                    .send(UiEvent::PermissionRequested {
                        tool: tool_name.to_string(),
                        args: args.clone(),
                        resolver: resolver_tx,
                    })
                    .await
                    .is_err()
                {
                    return Decision::Denied("no UI channel to prompt".into());
                }
                resolver_rx
                    .await
                    .unwrap_or(Decision::Denied("prompt cancelled".into()))
            }
            PromptStrategy::AllowAll => Decision::Allowed,
        }
    }
}

fn path_is_builtin_hard_blocked(path: &std::path::Path, project_root: &std::path::Path) -> bool {
    let rel = path.strip_prefix(project_root).unwrap_or(path);
    rel.components().any(|component| {
        let std::path::Component::Normal(name) = component else {
            return false;
        };
        let Some(name) = name.to_str() else {
            return false;
        };
        name == ".git"
            || name == ".ssh"
            || name == "node_modules"
            || name == "target"
            || name.starts_with(".env")
    })
}

#[async_trait]
impl Extension for PermissionExtension {
    fn name(&self) -> &'static str {
        "capo-permissions"
    }

    async fn intercept_tool_call(
        &mut self,
        call: ToolCallItem,
        _ctx: &mut HookCtx<'_>,
    ) -> Result<ToolDecision, ExtError> {
        match self.decide(&call.name, &call.args).await {
            Decision::Allowed => Ok(ToolDecision::Proceed(call)),
            Decision::Denied(reason) => Ok(ToolDecision::ShortCircuit(ToolResult::error(format!(
                "Permission denied: {reason}"
            )))),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use tokio::sync::mpsc;

    use super::*;

    #[tokio::test]
    async fn session_cache_short_circuits_prompt() {
        let policy = Arc::new(Policy::default());
        let cache = Arc::new(SessionCache::new());
        let args = serde_json::json!({"command": "curl https://example.com"});
        cache.insert(SessionCache::key("bash", &args), Decision::Allowed);

        let (ui_tx, mut ui_rx) = mpsc::channel::<UiEvent>(4);
        let ext = PermissionExtension::new(
            Arc::clone(&policy),
            Arc::clone(&cache),
            std::env::current_dir().unwrap_or_default(),
            ui_tx,
        );

        let decision = ext.decide("bash", &args).await;
        assert!(matches!(decision, Decision::Allowed));
        assert!(ui_rx.try_recv().is_err());
    }

    #[tokio::test]
    async fn grep_find_ls_are_auto_allowed() {
        let policy = Arc::new(Policy::default());
        let cache = Arc::new(SessionCache::new());
        let (ui_tx, mut ui_rx) = mpsc::channel::<UiEvent>(4);
        let ext = PermissionExtension::new(
            Arc::clone(&policy),
            Arc::clone(&cache),
            std::env::current_dir().unwrap_or_default(),
            ui_tx,
        );

        for tool in ["grep", "find", "ls"] {
            let decision = ext.decide(tool, &serde_json::json!({})).await;
            assert!(
                matches!(decision, Decision::Allowed),
                "{tool} not auto-allowed"
            );
        }
        assert!(ui_rx.try_recv().is_err());
    }

    #[tokio::test]
    async fn accept_edits_mode_auto_allows_write_and_edit() {
        let dir = tempfile::tempdir().expect("tempdir");
        let policy = Arc::new(crate::permissions::Policy::default());
        let cache = Arc::new(SessionCache::new());
        let (tx, _rx) = tokio::sync::mpsc::channel::<UiEvent>(16);
        let ext = PermissionExtension::accept_edits(
            Arc::clone(&policy),
            Arc::clone(&cache),
            dir.path().to_path_buf(),
            tx,
        );

        // write to a NON-blocked path: should auto-allow.
        let args = serde_json::json!({ "path": "src/foo.rs", "content": "..." });
        let decision = ext.decide("write", &args).await;
        assert!(
            matches!(decision, Decision::Allowed),
            "accept_edits should auto-allow write; got {decision:?}"
        );

        // edit similarly.
        let args =
            serde_json::json!({ "path": "src/foo.rs", "old_string": "x", "new_string": "y" });
        let decision = ext.decide("edit", &args).await;
        assert!(matches!(decision, Decision::Allowed));
    }

    #[tokio::test]
    async fn accept_edits_mode_enforces_hard_blocked_paths_for_write() {
        use crate::permissions::Policy;
        let dir = tempfile::tempdir().expect("tempdir");
        let policy = Arc::new(Policy::default());
        let cache = Arc::new(SessionCache::new());
        let (tx, _rx) = tokio::sync::mpsc::channel::<UiEvent>(16);
        let ext = PermissionExtension::accept_edits(
            Arc::clone(&policy),
            Arc::clone(&cache),
            dir.path().to_path_buf(),
            tx,
        );

        // write to .git/config: must be denied even in accept-edits.
        let args = serde_json::json!({ "path": ".git/config", "content": "x" });
        let decision = ext.decide("write", &args).await;
        assert!(
            matches!(decision, Decision::Denied(_)),
            "accept_edits must still enforce hard-blocked paths; got {decision:?}"
        );
    }

    #[tokio::test]
    async fn allow_all_mode_auto_allows_bash_too() {
        let dir = tempfile::tempdir().expect("tempdir");
        let policy = Arc::new(crate::permissions::Policy::default());
        let cache = Arc::new(SessionCache::new());
        let ext = PermissionExtension::allow_all(
            Arc::clone(&policy),
            Arc::clone(&cache),
            dir.path().to_path_buf(),
        );

        let args = serde_json::json!({ "command": "ls -la" });
        let decision = ext.decide("bash", &args).await;
        assert!(matches!(decision, Decision::Allowed));
    }

    #[tokio::test]
    async fn allow_all_mode_still_enforces_hard_blocked_write() {
        let dir = tempfile::tempdir().expect("tempdir");
        let policy = Arc::new(crate::permissions::Policy::default());
        let cache = Arc::new(SessionCache::new());
        let ext = PermissionExtension::allow_all(
            Arc::clone(&policy),
            Arc::clone(&cache),
            dir.path().to_path_buf(),
        );

        let args = serde_json::json!({ "path": ".env", "content": "x" });
        let decision = ext.decide("write", &args).await;
        assert!(
            matches!(decision, Decision::Denied(_)),
            "allow_all must still enforce hard-blocked .env*; got {decision:?}"
        );
    }

    #[tokio::test]
    async fn set_strategy_runtime_switches_mode() {
        let dir = tempfile::tempdir().expect("tempdir");
        let policy = Arc::new(crate::permissions::Policy::default());
        let cache = Arc::new(SessionCache::new());

        // Start in headless-deny mode.
        let ext = PermissionExtension::headless(
            Arc::clone(&policy),
            Arc::clone(&cache),
            dir.path().to_path_buf(),
        );
        let args = serde_json::json!({ "command": "ls" });
        assert!(matches!(
            ext.decide("bash", &args).await,
            Decision::Denied(_)
        ));

        // Mutate to allow-all at runtime.
        ext.set_strategy(PromptStrategy::AllowAll).await;
        assert!(matches!(ext.decide("bash", &args).await, Decision::Allowed));

        // Mutate back to headless-deny.
        ext.set_strategy(PromptStrategy::HeadlessDeny).await;
        assert!(matches!(
            ext.decide("bash", &args).await,
            Decision::Denied(_)
        ));
    }

    #[tokio::test]
    async fn headless_denies_a_would_prompt_tool_but_keeps_auto_allows() {
        let policy = Arc::new(Policy::default());
        let cache = Arc::new(SessionCache::new());
        let ext = PermissionExtension::headless(
            Arc::clone(&policy),
            Arc::clone(&cache),
            std::env::current_dir().unwrap_or_default(),
        );
        // A bash command not on the allowlist would prompt interactively —
        // headless mode denies it instead of hanging.
        let denied = ext
            .decide("bash", &serde_json::json!({"command": "curl https://x"}))
            .await;
        assert!(matches!(denied, Decision::Denied(_)));
        // Read-only tools are still auto-allowed (step 3 unchanged).
        let allowed = ext.decide("read", &serde_json::json!({})).await;
        assert!(matches!(allowed, Decision::Allowed));
    }
}