capo-agent 0.1.0

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};

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

pub struct PermissionExtension {
    policy: Arc<Policy>,
    cache: Arc<SessionCache>,
    project_root: PathBuf,
    /// Where to emit `UiEvent::PermissionRequested` when a prompt is needed.
    /// Cloning an mpsc::Sender is cheap — extensions are registered as
    /// `Box<dyn Extension>` and hooks take `&mut self`, so we keep the sender
    /// inside the extension and clone per prompt.
    ui_tx: mpsc::Sender<UiEvent>,
}

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,
        }
    }

    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 = 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. 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" => 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;
        }

        // 3. 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;
        }

        // 4. Prompt user via UiEvent.
        let (resolver_tx, resolver_rx) = oneshot::channel::<Decision>();
        if self
            .ui_tx
            .send(UiEvent::PermissionRequested {
                tool: tool_name.to_string(),
                args: args.clone(),
                resolver: resolver_tx,
            })
            .await
            .is_err()
        {
            // No UI attached — fail closed.
            return Decision::Denied("no UI channel to prompt".into());
        }
        resolver_rx
            .await
            .unwrap_or(Decision::Denied("prompt cancelled".into()))
    }
}

#[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());
    }
}