IdeaSDK/crates/infrastructure/src/orchestrator/mcp/policy.rs

//! MCP tool policy helpers and in-memory live policy registry.

use std::collections::HashMap;
use std::sync::RwLock;

use domain::ports::McpToolPermissionStore;
use domain::{AgentId, AgentToolPolicy, AgentToolPolicyStore, McpToolPolicy, Project, StoreError};

use super::tools::{self, ToolDef};

/// Stores per-requester MCP tool policies for live assistant sessions.
#[derive(Default)]
pub struct ToolPolicyRegistry {
    policies: RwLock<HashMap<String, AgentToolPolicy>>,
}

impl ToolPolicyRegistry {
    /// Builds an empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets or replaces the policy for `requester`.
    pub fn set(&self, requester: impl Into<String>, policy: AgentToolPolicy) {
        self.policies
            .write()
            .unwrap()
            .insert(requester.into(), policy);
    }

    /// Returns the policy for `requester`, if any.
    #[must_use]
    pub fn get(&self, requester: &str) -> Option<AgentToolPolicy> {
        self.policies.read().unwrap().get(requester).cloned()
    }

    /// Clears the policy for `requester`.
    pub fn clear(&self, requester: &str) {
        self.policies.write().unwrap().remove(requester);
    }
}

impl AgentToolPolicyStore for ToolPolicyRegistry {
    fn set_policy(&self, requester: String, policy: AgentToolPolicy) {
        self.set(requester, policy);
    }

    fn get_policy(&self, requester: &str) -> Option<AgentToolPolicy> {
        self.get(requester)
    }

    fn clear_policy(&self, requester: &str) {
        self.clear(requester);
    }
}

/// Effective MCP tool surface resolved for one requester.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EffectiveToolSurface {
    /// Durable MCP policy when one applies.
    pub durable_policy: Option<McpToolPolicy>,
    /// Whether the read-only fallback was used because the requester could not be
    /// resolved to an agent policy.
    pub used_read_only_fallback: bool,
}

impl EffectiveToolSurface {
    /// Returns whether `tool` is allowed by the durable policy, when present.
    #[must_use]
    pub fn permits_durable(&self, tool: &str) -> bool {
        self.durable_policy
            .as_ref()
            .map_or(true, |policy| policy.permits(tool))
    }
}

/// Resolves the durable MCP tool policy for the given requester.
///
/// This is the shared source of truth for the surfaces that expose IdeA tools to
/// agents. Agent UUID requesters use `.ideai/mcp-tool-permissions.json`;
/// anonymous/legacy requesters and non-agent requesters without an ephemeral
/// policy fall back to the canonical read-only policy. Non-agent requesters with
/// an ephemeral policy (ticket assistants) are intentionally governed only by
/// that narrower session policy.
///
/// # Errors
/// [`StoreError`] on store load failure or invalid persisted policy.
pub async fn resolve_effective_tool_surface(
    store: Option<&dyn McpToolPermissionStore>,
    project: &Project,
    requester: &str,
    ephemeral_policy: Option<&AgentToolPolicy>,
) -> Result<EffectiveToolSurface, StoreError> {
    let Some(store) = store else {
        return Ok(EffectiveToolSurface {
            durable_policy: None,
            used_read_only_fallback: false,
        });
    };

    let known_tools = tools::classified_tool_names();
    if let Some(agent_id) = requester_agent_id(requester) {
        let doc = store.load_mcp_tool_permissions(project).await?;
        let policy = doc
            .effective_policy(agent_id, tools::READ_ONLY_TOOLS, &known_tools)
            .map_err(|err| StoreError::Invalid(err.to_string()))?;
        return Ok(EffectiveToolSurface {
            durable_policy: Some(policy),
            used_read_only_fallback: false,
        });
    }

    if requester.is_empty() || requester == "mcp" || ephemeral_policy.is_none() {
        application::diag!(
            "[mcp] unresolved requester `{}` uses read-only tool fallback",
            if requester.is_empty() {
                "mcp"
            } else {
                requester
            },
        );
        let policy = McpToolPolicy::read_only(tools::READ_ONLY_TOOLS, &known_tools)
            .map_err(|err| StoreError::Invalid(err.to_string()))?;
        return Ok(EffectiveToolSurface {
            durable_policy: Some(policy),
            used_read_only_fallback: true,
        });
    }

    Ok(EffectiveToolSurface {
        durable_policy: None,
        used_read_only_fallback: false,
    })
}

/// Filters the canonical catalogue by the effective ephemeral + durable policies.
#[must_use]
pub fn effective_tool_catalogue(
    ephemeral_policy: Option<&AgentToolPolicy>,
    surface: &EffectiveToolSurface,
) -> Vec<ToolDef> {
    tools::catalogue()
        .into_iter()
        .filter(|tool| {
            ephemeral_policy.map_or(true, |policy| policy.permits(tool.name))
                && surface.permits_durable(tool.name)
        })
        .collect()
}

fn requester_agent_id(requester: &str) -> Option<AgentId> {
    uuid::Uuid::parse_str(requester)
        .ok()
        .map(AgentId::from_uuid)
}

#[cfg(test)]
mod tests {
    use domain::IssueRef;

    use super::*;

    #[test]
    fn set_get_clear_roundtrips_by_requester() {
        let registry = ToolPolicyRegistry::new();
        let policy = AgentToolPolicy::new(
            vec!["idea_ticket_read".to_owned()],
            Some("#7".parse::<IssueRef>().unwrap()),
            true,
        );
        registry.set("assistant-1", policy.clone());
        assert_eq!(registry.get("assistant-1"), Some(policy));
        assert_eq!(registry.get("assistant-2"), None);
        registry.clear("assistant-1");
        assert_eq!(registry.get("assistant-1"), None);
    }
}