mcp_host/server/

profile.rs

//! Session profiles for role-based capability configuration
//!
//! Profiles allow pre-defining sets of tools, resources, and prompts
//! that can be applied to sessions based on user roles (admin, vip, guest, etc.)
//!
//! # Example
//!
//! ```rust,ignore
//! // Create an admin profile with extra tools
//! let mut admin = SessionProfile::new("admin");
//! admin.add_tool(Arc::new(DeleteDatabaseTool));
//! admin.add_tool(Arc::new(SystemRebootTool));
//!
//! // Create a VIP profile with upgraded implementations
//! let mut vip = SessionProfile::new("vip");
//! vip.override_tool("export", Arc::new(UnlimitedExportTool));
//!
//! // Create a guest profile with restricted access
//! let mut guest = SessionProfile::new("guest");
//! guest.hide_tool("delete");
//! guest.hide_tool("admin_panel");
//!
//! // Apply to session
//! session.apply_profile(&admin);
//! ```

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use crate::registry::prompts::Prompt;
use crate::registry::resources::Resource;
use crate::registry::tools::Tool;

/// Session profile for role-based capability configuration
///
/// A profile defines a set of modifications to apply to a session:
/// - Extra tools/resources/prompts to add
/// - Overrides for existing tools (same name, different implementation)
/// - Tools/resources/prompts to hide
/// - Aliases for ergonomic name mappings
#[derive(Default)]
pub struct SessionProfile {
    /// Profile name (e.g., "admin", "vip", "guest")
    pub name: String,

    // Tool configuration
    /// Extra tools to add to the session
    pub tool_extras: Vec<Arc<dyn Tool>>,
    /// Tool overrides (same name, different implementation)
    pub tool_overrides: HashMap<String, Arc<dyn Tool>>,
    /// Tools to hide from the session
    pub tool_hidden: HashSet<String>,
    /// Tool aliases (alias -> actual name)
    pub tool_aliases: HashMap<String, String>,

    // Resource configuration
    /// Extra resources to add to the session
    pub resource_extras: Vec<Arc<dyn Resource>>,
    /// Resource overrides (same URI, different implementation)
    pub resource_overrides: HashMap<String, Arc<dyn Resource>>,
    /// Resources to hide from the session
    pub resource_hidden: HashSet<String>,

    // Prompt configuration
    /// Extra prompts to add to the session
    pub prompt_extras: Vec<Arc<dyn Prompt>>,
    /// Prompt overrides (same name, different implementation)
    pub prompt_overrides: HashMap<String, Arc<dyn Prompt>>,
    /// Prompts to hide from the session
    pub prompt_hidden: HashSet<String>,
}

impl SessionProfile {
    /// Create a new empty profile with the given name
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            ..Default::default()
        }
    }

    /// Create an admin profile (typically has extra management tools)
    pub fn admin() -> Self {
        Self::new("admin")
    }

    /// Create a VIP profile (typically has upgraded implementations)
    pub fn vip() -> Self {
        Self::new("vip")
    }

    /// Create a readonly profile (hides write operations)
    pub fn readonly() -> Self {
        Self::new("readonly")
    }

    /// Create a guest profile (minimal access)
    pub fn guest() -> Self {
        Self::new("guest")
    }

    // Tool methods

    /// Add an extra tool to the profile
    pub fn add_tool(&mut self, tool: Arc<dyn Tool>) -> &mut Self {
        self.tool_extras.push(tool);
        self
    }

    /// Override a tool (same name, different implementation)
    pub fn override_tool(&mut self, name: impl Into<String>, tool: Arc<dyn Tool>) -> &mut Self {
        self.tool_overrides.insert(name.into(), tool);
        self
    }

    /// Hide a tool from sessions with this profile
    pub fn hide_tool(&mut self, name: impl Into<String>) -> &mut Self {
        self.tool_hidden.insert(name.into());
        self
    }

    /// Hide multiple tools
    pub fn hide_tools(&mut self, names: impl IntoIterator<Item = impl Into<String>>) -> &mut Self {
        for name in names {
            self.tool_hidden.insert(name.into());
        }
        self
    }

    /// Add a tool alias (alias -> target)
    pub fn alias_tool(&mut self, alias: impl Into<String>, target: impl Into<String>) -> &mut Self {
        self.tool_aliases.insert(alias.into(), target.into());
        self
    }

    // Resource methods

    /// Add an extra resource to the profile
    pub fn add_resource(&mut self, resource: Arc<dyn Resource>) -> &mut Self {
        self.resource_extras.push(resource);
        self
    }

    /// Override a resource (same URI, different implementation)
    pub fn override_resource(
        &mut self,
        uri: impl Into<String>,
        resource: Arc<dyn Resource>,
    ) -> &mut Self {
        self.resource_overrides.insert(uri.into(), resource);
        self
    }

    /// Hide a resource from sessions with this profile
    pub fn hide_resource(&mut self, uri: impl Into<String>) -> &mut Self {
        self.resource_hidden.insert(uri.into());
        self
    }

    /// Hide multiple resources
    pub fn hide_resources(
        &mut self,
        uris: impl IntoIterator<Item = impl Into<String>>,
    ) -> &mut Self {
        for uri in uris {
            self.resource_hidden.insert(uri.into());
        }
        self
    }

    // Prompt methods

    /// Add an extra prompt to the profile
    pub fn add_prompt(&mut self, prompt: Arc<dyn Prompt>) -> &mut Self {
        self.prompt_extras.push(prompt);
        self
    }

    /// Override a prompt (same name, different implementation)
    pub fn override_prompt(
        &mut self,
        name: impl Into<String>,
        prompt: Arc<dyn Prompt>,
    ) -> &mut Self {
        self.prompt_overrides.insert(name.into(), prompt);
        self
    }

    /// Hide a prompt from sessions with this profile
    pub fn hide_prompt(&mut self, name: impl Into<String>) -> &mut Self {
        self.prompt_hidden.insert(name.into());
        self
    }

    /// Hide multiple prompts
    pub fn hide_prompts(
        &mut self,
        names: impl IntoIterator<Item = impl Into<String>>,
    ) -> &mut Self {
        for name in names {
            self.prompt_hidden.insert(name.into());
        }
        self
    }

    /// Merge another profile into this one
    ///
    /// Later profiles take precedence for overrides and hidden items
    pub fn merge(&mut self, other: &SessionProfile) -> &mut Self {
        // Tools
        self.tool_extras.extend(other.tool_extras.iter().cloned());
        self.tool_overrides.extend(
            other
                .tool_overrides
                .iter()
                .map(|(k, v)| (k.clone(), v.clone())),
        );
        self.tool_hidden.extend(other.tool_hidden.iter().cloned());
        self.tool_aliases.extend(
            other
                .tool_aliases
                .iter()
                .map(|(k, v)| (k.clone(), v.clone())),
        );

        // Resources
        self.resource_extras
            .extend(other.resource_extras.iter().cloned());
        self.resource_overrides.extend(
            other
                .resource_overrides
                .iter()
                .map(|(k, v)| (k.clone(), v.clone())),
        );
        self.resource_hidden
            .extend(other.resource_hidden.iter().cloned());

        // Prompts
        self.prompt_extras
            .extend(other.prompt_extras.iter().cloned());
        self.prompt_overrides.extend(
            other
                .prompt_overrides
                .iter()
                .map(|(k, v)| (k.clone(), v.clone())),
        );
        self.prompt_hidden
            .extend(other.prompt_hidden.iter().cloned());

        self
    }
}

impl std::fmt::Debug for SessionProfile {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SessionProfile")
            .field("name", &self.name)
            .field("tool_extras_count", &self.tool_extras.len())
            .field(
                "tool_overrides",
                &self.tool_overrides.keys().collect::<Vec<_>>(),
            )
            .field("tool_hidden", &self.tool_hidden)
            .field("tool_aliases", &self.tool_aliases)
            .field("resource_extras_count", &self.resource_extras.len())
            .field(
                "resource_overrides",
                &self.resource_overrides.keys().collect::<Vec<_>>(),
            )
            .field("resource_hidden", &self.resource_hidden)
            .field("prompt_extras_count", &self.prompt_extras.len())
            .field(
                "prompt_overrides",
                &self.prompt_overrides.keys().collect::<Vec<_>>(),
            )
            .field("prompt_hidden", &self.prompt_hidden)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_profile_creation() {
        let profile = SessionProfile::new("test");
        assert_eq!(profile.name, "test");
        assert!(profile.tool_extras.is_empty());
        assert!(profile.tool_overrides.is_empty());
        assert!(profile.tool_hidden.is_empty());
    }

    #[test]
    fn test_profile_presets() {
        let admin = SessionProfile::admin();
        assert_eq!(admin.name, "admin");

        let vip = SessionProfile::vip();
        assert_eq!(vip.name, "vip");

        let readonly = SessionProfile::readonly();
        assert_eq!(readonly.name, "readonly");

        let guest = SessionProfile::guest();
        assert_eq!(guest.name, "guest");
    }

    #[test]
    fn test_hide_tools() {
        let mut profile = SessionProfile::new("test");
        profile.hide_tool("delete");
        profile.hide_tools(["admin_panel", "system_reboot"]);

        assert!(profile.tool_hidden.contains("delete"));
        assert!(profile.tool_hidden.contains("admin_panel"));
        assert!(profile.tool_hidden.contains("system_reboot"));
        assert_eq!(profile.tool_hidden.len(), 3);
    }

    #[test]
    fn test_tool_aliases() {
        let mut profile = SessionProfile::new("test");
        profile.alias_tool("save", "git_commit");
        profile.alias_tool("deploy", "kubernetes_apply");

        assert_eq!(
            profile.tool_aliases.get("save"),
            Some(&"git_commit".to_string())
        );
        assert_eq!(
            profile.tool_aliases.get("deploy"),
            Some(&"kubernetes_apply".to_string())
        );
    }

    #[test]
    fn test_profile_merge() {
        let mut base = SessionProfile::new("base");
        base.hide_tool("tool1");
        base.alias_tool("a", "b");

        let mut extra = SessionProfile::new("extra");
        extra.hide_tool("tool2");
        extra.alias_tool("c", "d");

        base.merge(&extra);

        assert!(base.tool_hidden.contains("tool1"));
        assert!(base.tool_hidden.contains("tool2"));
        assert_eq!(base.tool_aliases.get("a"), Some(&"b".to_string()));
        assert_eq!(base.tool_aliases.get("c"), Some(&"d".to_string()));
    }

    #[test]
    fn test_builder_pattern() {
        let mut profile = SessionProfile::new("test");
        profile
            .hide_tool("dangerous")
            .alias_tool("save", "commit")
            .hide_resource("secret://data");

        assert!(profile.tool_hidden.contains("dangerous"));
        assert_eq!(
            profile.tool_aliases.get("save"),
            Some(&"commit".to_string())
        );
        assert!(profile.resource_hidden.contains("secret://data"));
    }
}