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

//! In-memory MCP tool policy registry.

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

use domain::AgentToolPolicy;
use domain::AgentToolPolicyStore;

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

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