mcp-agent-selection/src/agent_mcp_server/server.py

#!/usr/bin/env python3
"""
Production-ready FastMCP server for Claude Agent recommendations.

Built using the patterns from our working test_agents.py prototype.
Provides intelligent agent recommendations with project roots functionality.
"""

import asyncio
import json
import os
import yaml
import logging
from pathlib import Path
from typing import List, Dict, Optional, Any
from dataclasses import dataclass, asdict

from fastmcp import FastMCP
from pydantic import BaseModel, Field

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Initialize FastMCP server
mcp = FastMCP("Claude Agent MCP Server")

# ===== Data Models =====

@dataclass
class AgentInfo:
    """
    Agent information structure supporting both individual and composed agents.
    
    Composed Agent Architecture:
    - FLAT STRUCTURE MAINTAINED: All agents (parent + sub-agents) exist in root agents dict
    - Claude Code sees all agents as individual entities in list_agents()
    - Parent agents marked as agent_type="composed" with sub_agents list
    - Sub-agents marked as agent_type="sub" with parent_agent reference
    - Individual agents remain agent_type="individual"
    
    Example flat structure:
    agents = {
        "🧪-testing-integration-expert": AgentInfo(agent_type="composed", sub_agents=[...]),
        "🐍-python-testing-framework-expert": AgentInfo(agent_type="sub", parent_agent="..."),
        "🌐-html-report-generation-expert": AgentInfo(agent_type="sub", parent_agent="..."),
        "🚄-fastapi-expert": AgentInfo(agent_type="individual")
    }
    """
    name: str
    emoji: str
    description: str
    tools: List[str]
    content: str
    file_path: str
    # NEW: Hierarchy support while maintaining flat structure
    parent_agent: Optional[str] = None      # For sub-agents: references parent agent name
    sub_agents: List[str] = None           # For composed agents: list of sub-agent names
    agent_type: str = "individual"         # "individual", "composed", "sub"
    
    def __post_init__(self):
        """Initialize sub_agents as empty list if None"""
        if self.sub_agents is None:
            self.sub_agents = []

@dataclass
class ProjectRoots:
    """Project roots configuration"""
    directories: List[str]
    base_path: str
    description: str = ""

class SetProjectRootsRequest(BaseModel):
    """Request model for setting project roots"""
    directories: List[str] = Field(..., description="List of directories to focus on")
    base_path: str = Field(..., description="Base path of the project")
    description: str = Field("", description="Optional description of the project focus")

class RecommendAgentsRequest(BaseModel):
    """Request model for agent recommendations"""
    task: str = Field(..., description="Description of the task or problem")
    project_context: str = Field("", description="Additional project context")
    limit: int = Field(5, ge=1, le=10, description="Maximum number of recommendations")

class GetAgentContentRequest(BaseModel):
    """Request model for getting agent content"""
    agent_name: str = Field(..., description="Name of the agent to retrieve")

class ListAgentsRequest(BaseModel):
    """Request model for listing agents"""
    search: Optional[str] = Field(None, description="Optional search term")
    tools_filter: Optional[List[str]] = Field(None, description="Filter by required tools")

# ===== Agent Library Service =====

class AgentLibrary:
    """
    Production agent library with async patterns and comprehensive error handling
    """
    
    def __init__(self, templates_path: str | Path):
        self.templates_path = Path(templates_path)
        self.agents: Dict[str, AgentInfo] = {}
        self.roots: Optional[ProjectRoots] = None
        self._initialized = False
        
    async def initialize(self):
        """Initialize the agent library by loading all agents"""
        if self._initialized:
            return
            
        try:
            await self._load_agents()
            self._initialized = True
            logger.info(f"AgentLibrary initialized with {len(self.agents)} agents")
        except Exception as e:
            logger.error(f"Failed to initialize AgentLibrary: {e}")
            raise
    
    async def _load_agents(self):
        """
        Load all agent templates from the directory, supporting both individual and composed agents.
        
        Composed Agent Loading Strategy:
        1. Load individual .md files as before (individual agents)
        2. Scan for directories with matching .md files (composed agents)
        3. For composed agents: load parent + all sub-agents into FLAT agents dict
        4. Maintain parent-child relationships via metadata
        5. Claude Code sees all agents in flat list_agents() output
        """
        if not self.templates_path.exists():
            logger.warning(f"Templates path not found: {self.templates_path}")
            return
        
        loaded_count = 0
        failed_count = 0
        composed_agents = 0
        sub_agents = 0
        
        # STEP 1: Load individual .md files (existing logic)
        for file_path in self.templates_path.glob("*.md"):
            try:
                # Check if this is part of a composed agent (has matching directory)
                potential_dir = self.templates_path / file_path.stem
                if potential_dir.exists() and potential_dir.is_dir():
                    # This will be handled in STEP 2 as a composed agent
                    continue
                    
                # Load as individual agent
                agent = await self._load_single_agent(file_path, "individual")
                if agent:
                    self.agents[agent.name] = agent
                    loaded_count += 1
                        
            except Exception as e:
                logger.error(f"Error loading individual agent {file_path}: {e}")
                failed_count += 1
        
        # STEP 2: Load composed agents (NEW)
        for dir_path in self.templates_path.iterdir():
            if not dir_path.is_dir():
                continue
                
            parent_md = self.templates_path / f"{dir_path.name}.md"
            if not parent_md.exists():
                continue  # Skip directories without matching .md file
                
            try:
                # Load parent agent as "composed"
                parent_agent = await self._load_single_agent(parent_md, "composed")
                if not parent_agent:
                    continue
                    
                # Load all sub-agents from directory
                sub_agent_names = []
                for sub_file in dir_path.glob("*.md"):
                    try:
                        sub_agent = await self._load_single_agent(sub_file, "sub", parent_agent.name)
                        if sub_agent:
                            self.agents[sub_agent.name] = sub_agent  # ADD TO FLAT STRUCTURE
                            sub_agent_names.append(sub_agent.name)
                            sub_agents += 1
                    except Exception as e:
                        logger.error(f"Error loading sub-agent {sub_file}: {e}")
                        failed_count += 1
                
                # Update parent with sub-agent references
                parent_agent.sub_agents = sub_agent_names
                
                # Add parent to flat structure
                self.agents[parent_agent.name] = parent_agent
                loaded_count += 1
                composed_agents += 1
                
            except Exception as e:
                logger.error(f"Error loading composed agent {dir_path}: {e}")
                failed_count += 1
        
        logger.info(f"Loaded {loaded_count} agents successfully ({composed_agents} composed, {sub_agents} sub-agents), {failed_count} failed")
        
    async def _load_single_agent(
        self, 
        file_path: Path, 
        agent_type: str, 
        parent_agent: Optional[str] = None
    ) -> Optional[AgentInfo]:
        """
        Load a single agent from a markdown file.
        
        Args:
            file_path: Path to the .md file
            agent_type: "individual", "composed", or "sub"  
            parent_agent: Name of parent agent (for sub-agents only)
            
        Returns:
            AgentInfo object or None if loading failed
        """
        try:
            content = await asyncio.to_thread(file_path.read_text, encoding='utf-8')
            
            # Extract YAML frontmatter (existing logic)
            if content.startswith("---\n"):
                parts = content.split("---\n", 2)
                if len(parts) >= 3:
                    frontmatter = yaml.safe_load(parts[1])
                    body = parts[2]
                    
                    agent = AgentInfo(
                        name=frontmatter.get("name", file_path.stem),
                        emoji=frontmatter.get("emoji", "🤖"),
                        description=frontmatter.get("description", ""),
                        tools=frontmatter.get("tools", []),
                        content=body.strip(),
                        file_path=str(file_path),
                        agent_type=agent_type,
                        parent_agent=parent_agent
                    )
                    
                    return agent
            
            # If no YAML frontmatter, create basic agent info
            agent = AgentInfo(
                name=file_path.stem,
                emoji="🤖",
                description="",
                tools=[],
                content=content.strip(),
                file_path=str(file_path),
                agent_type=agent_type,
                parent_agent=parent_agent
            )
            
            return agent
                        
        except Exception as e:
            logger.error(f"Error loading agent from {file_path}: {e}")
            return None
    
    async def set_roots(self, directories: List[str], base_path: str, description: str = "") -> Dict[str, Any]:
        """Set project roots for focused analysis"""
        try:
            self.roots = ProjectRoots(directories, base_path, description)
            logger.info(f"Set project roots: {directories} in {base_path}")
            return {
                "success": True,
                "message": f"Set project roots: {directories} in {base_path}",
                "roots": asdict(self.roots)
            }
        except Exception as e:
            logger.error(f"Error setting roots: {e}")
            raise
    
    async def get_current_roots(self) -> Dict[str, Any]:
        """Get current project roots"""
        if self.roots:
            return {
                "configured": True,
                "roots": asdict(self.roots)
            }
        return {
            "configured": False,
            "message": "No project roots configured"
        }
    
    async def clear_project_roots(self) -> Dict[str, Any]:
        """Clear project roots"""
        try:
            self.roots = None
            logger.info("Cleared project roots")
            return {
                "success": True,
                "message": "Project roots cleared"
            }
        except Exception as e:
            logger.error(f"Error clearing roots: {e}")
            raise
    
    async def recommend_agents(self, task: str, project_context: str = "", limit: int = 5) -> List[Dict[str, Any]]:
        """
        Recommend agents based on task description with hierarchical intelligence.
        
        Hierarchical Recommendation Strategy:
        1. First, find matching composed agents and their sub-agents
        2. Prioritize specialized sub-agents over general composed agents
        3. Include composed agents as fallbacks for domain overview
        4. Maintain flat recommendation list for Claude Code compatibility
        """
        try:
            recommendations = []
            task_lower = task.lower()
            project_lower = project_context.lower()
            combined_context = f"{task_lower} {project_lower}"
            
            # Enhanced keywords to agent mapping with hierarchical awareness
            keywords_mapping = {
                "python testing": {
                    # Prioritize specialized sub-agent over composed parent
                    "agents": ["python-testing-framework-expert", "testing-integration-expert"],
                    "confidence": 0.95
                },
                "html report": {
                    "agents": ["html-report-generation-expert", "testing-integration-expert"],
                    "confidence": 0.95
                },
                "testing framework": {
                    "agents": ["testing-integration-expert", "python-testing-framework-expert"],
                    "confidence": 0.9
                },
                "python": {
                    "agents": ["🔮-python-mcp-expert", "testing-integration-expert"],
                    "confidence": 0.9
                },
                "testing": {
                    # Check for sub-agents first, fallback to composed
                    "agents": ["testing-integration-expert"],
                    "confidence": 0.8
                },
                "fastapi": {
                    "agents": ["🚄-fastapi-expert", "🔮-python-mcp-expert"],
                    "confidence": 0.9
                },
                "docker": {
                    "agents": ["🐳-docker-infrastructure-expert"],
                    "confidence": 0.8
                },
                "security": {
                    "agents": ["🔒-security-audit-expert"],
                    "confidence": 0.8
                },
                "documentation": {
                    "agents": ["📖-readme-expert", "📝-documentation-expert"],
                    "confidence": 0.7
                },
                "mcp": {
                    "agents": ["🔮-python-mcp-expert"],
                    "confidence": 0.9
                },
                "subagent": {
                    "agents": ["🎭-subagent-expert"],
                    "confidence": 0.8
                },
                "database": {
                    "agents": ["💾-database-expert", "🔮-python-mcp-expert"],
                    "confidence": 0.7
                },
                "frontend": {
                    "agents": ["🎨-frontend-expert", "⚡-javascript-expert"],
                    "confidence": 0.7
                },
                "performance": {
                    "agents": ["⚡-performance-expert"],
                    "confidence": 0.7
                }
            }
            
            # Track agents already recommended to avoid duplicates
            recommended_agents = set()
            
            # STEP 1: Find matching agents (including sub-agents)
            for keyword, config in keywords_mapping.items():
                if keyword in combined_context:
                    for agent_name in config["agents"]:
                        if agent_name in self.agents and agent_name not in recommended_agents:
                            agent = self.agents[agent_name]
                            
                            # Boost confidence for specialized sub-agents
                            confidence = config["confidence"]
                            if agent.agent_type == "sub":
                                confidence += 0.05  # Sub-agents get slight boost for specialization
                                
                            recommendations.append({
                                "name": agent.name,
                                "emoji": agent.emoji,
                                "description": agent.description,
                                "confidence": confidence,
                                "reason": f"Matches '{keyword}' - {self._get_agent_type_description(agent)}",
                                "tools": agent.tools,
                                "agent_type": agent.agent_type,
                                "parent_agent": agent.parent_agent
                            })
                            recommended_agents.add(agent_name)
            
            # STEP 2: Add related sub-agents for composed matches
            # If we recommended a composed agent, also suggest its most relevant sub-agents
            for rec in recommendations.copy():  # Copy to avoid modification during iteration
                agent = self.agents[rec["name"]]
                if agent.agent_type == "composed" and agent.sub_agents:
                    for sub_agent_name in agent.sub_agents:
                        if sub_agent_name not in recommended_agents and len(recommendations) < limit * 2:
                            sub_agent = self.agents.get(sub_agent_name)
                            if sub_agent:
                                recommendations.append({
                                    "name": sub_agent.name,
                                    "emoji": sub_agent.emoji,
                                    "description": sub_agent.description,
                                    "confidence": rec["confidence"] - 0.1,  # Slightly lower than parent
                                    "reason": f"Specialist under {agent.name}",
                                    "tools": sub_agent.tools,
                                    "agent_type": sub_agent.agent_type,
                                    "parent_agent": sub_agent.parent_agent
                                })
                                recommended_agents.add(sub_agent_name)
            
            # STEP 3: If no specific matches, suggest general purpose agents
            if not recommendations:
                general_agents = ["🎭-subagent-expert", "🔮-python-mcp-expert", "📖-readme-expert"]
                for agent_name in general_agents:
                    if agent_name in self.agents:
                        agent = self.agents[agent_name]
                        recommendations.append({
                            "name": agent.name,
                            "emoji": agent.emoji,
                            "description": agent.description,
                            "confidence": 0.6,
                            "reason": "General purpose recommendation for task planning",
                            "tools": agent.tools,
                            "agent_type": agent.agent_type,
                            "parent_agent": agent.parent_agent
                        })
                        break
            
            # STEP 4: Sort by confidence (sub-agents should naturally rank higher)
            recommendations.sort(key=lambda x: x["confidence"], reverse=True)
            result = recommendations[:limit]
            
            logger.info(f"Generated {len(result)} hierarchical recommendations for task: {task[:50]}...")
            return result
            
        except Exception as e:
            logger.error(f"Error in recommend_agents: {e}")
            raise
            
    def _get_agent_type_description(self, agent: AgentInfo) -> str:
        """Get human-readable description of agent type"""
        if agent.agent_type == "sub":
            return f"specialized expert (part of {agent.parent_agent})"
        elif agent.agent_type == "composed":
            return f"framework architect with {len(agent.sub_agents)} specialists"
        else:
            return "individual expert"
    
    async def get_agent_content(self, agent_name: str) -> Dict[str, Any]:
        """
        Get the full content of a specific agent with dynamic sub-agent references.
        
        Content Enhancement Strategy:
        - Composed agents: Add "My Specialists" section with sub-agent references
        - Sub-agents: Add "Part of [parent] expertise" context note  
        - Individual agents: Content unchanged
        - All agents: Include project roots context if available
        """
        try:
            if agent_name not in self.agents:
                return {
                    "success": False,
                    "error": f"Agent '{agent_name}' not found"
                }
            
            agent = self.agents[agent_name]
            
            # Start with base content
            enhanced_content = agent.content
            
            # Add hierarchical context based on agent type
            hierarchy_note = self._generate_hierarchy_context(agent)
            if hierarchy_note:
                enhanced_content += hierarchy_note
            
            # Include project roots context if available
            context_note = ""
            if self.roots:
                context_note = f"\n\n## Current Project Context\n"
                context_note += f"**Base Path:** {self.roots.base_path}\n"
                context_note += f"**Focus Directories:** {', '.join(self.roots.directories)}\n"
                if self.roots.description:
                    context_note += f"**Description:** {self.roots.description}\n"
            
            return {
                "success": True,
                "agent": {
                    "name": agent.name,
                    "emoji": agent.emoji,
                    "description": agent.description,
                    "tools": agent.tools,
                    "file_path": agent.file_path,
                    "agent_type": agent.agent_type,
                    "parent_agent": agent.parent_agent,
                    "sub_agents": agent.sub_agents
                },
                "content": enhanced_content + context_note,
                "has_project_context": self.roots is not None,
                "has_hierarchy": agent.agent_type != "individual"
            }
            
        except Exception as e:
            logger.error(f"Error getting agent content: {e}")
            raise
            
    def _generate_hierarchy_context(self, agent: AgentInfo) -> str:
        """Generate hierarchical context notes for agent content"""
        if agent.agent_type == "composed" and agent.sub_agents:
            # Add sub-agent references for composed agents
            context = f"\n\n## My Specialist Team\n"
            context += "I coordinate with these specialized experts for complex tasks:\n\n"
            
            for sub_agent_name in agent.sub_agents:
                if sub_agent_name in self.agents:
                    sub_agent = self.agents[sub_agent_name]
                    context += f"- **{sub_agent.emoji} {sub_agent.name}**: {sub_agent.description}\n"
            
            context += f"\n💡 *Tip: For specific tasks, ask directly for one of my specialists by name.*"
            return context
            
        elif agent.agent_type == "sub" and agent.parent_agent:
            # Add parent context for sub-agents
            parent = self.agents.get(agent.parent_agent)
            if parent:
                context = f"\n\n## Part of {parent.emoji} {parent.name} Framework\n"
                context += f"I'm a specialized expert within the broader {parent.name} architecture. "
                context += f"For overview and coordination, consult **{parent.name}**.\n"
                
                # List sibling specialists
                siblings = [name for name in parent.sub_agents if name != agent.name]
                if siblings:
                    context += f"\n**Related specialists:**\n"
                    for sibling_name in siblings:
                        if sibling_name in self.agents:
                            sibling = self.agents[sibling_name]
                            context += f"- {sibling.emoji} {sibling.name}\n"
                
                return context
        
        return ""
    
    async def list_agents(
        self, 
        search: Optional[str] = None,
        tools_filter: Optional[List[str]] = None
    ) -> List[Dict[str, Any]]:
        """List all available agents with optional filtering"""
        try:
            agents_list = []
            
            for agent in self.agents.values():
                # Apply search filter
                if search:
                    search_lower = search.lower()
                    if not (
                        search_lower in agent.name.lower() or
                        search_lower in agent.description.lower() or
                        any(search_lower in tool.lower() for tool in agent.tools)
                    ):
                        continue
                
                # Apply tools filter
                if tools_filter:
                    if not any(tool in agent.tools for tool in tools_filter):
                        continue
                
                agents_list.append({
                    "name": agent.name,
                    "emoji": agent.emoji,
                    "description": agent.description,
                    "tools": agent.tools,
                    "file_path": agent.file_path
                })
            
            # Sort by name
            agents_list.sort(key=lambda x: x["name"])
            
            logger.info(f"Listed {len(agents_list)} agents (filtered from {len(self.agents)} total)")
            return agents_list
            
        except Exception as e:
            logger.error(f"Error listing agents: {e}")
            raise
    
    async def get_server_stats(self) -> Dict[str, Any]:
        """Get comprehensive server statistics"""
        try:
            tools_count = {}
            emojis = set()
            total_content_length = 0
            
            for agent in self.agents.values():
                # Count tools
                for tool in agent.tools:
                    tools_count[tool] = tools_count.get(tool, 0) + 1
                
                # Collect unique emojis
                emojis.add(agent.emoji)
                
                # Count content length
                total_content_length += len(agent.content)
            
            return {
                "server_info": {
                    "name": "Claude Agent MCP Server",
                    "version": "1.0.0",
                    "initialized": self._initialized
                },
                "agents": {
                    "total_count": len(self.agents),
                    "unique_emojis": len(emojis),
                    "total_content_lines": total_content_length // 80,  # Approximate lines
                    "average_tools_per_agent": len(sum([a.tools for a in self.agents.values()], [])) / len(self.agents) if self.agents else 0
                },
                "project_roots": {
                    "configured": self.roots is not None,
                    "details": asdict(self.roots) if self.roots else None
                },
                "tools_distribution": dict(sorted(tools_count.items(), key=lambda x: x[1], reverse=True)[:10]),
                "most_common_emojis": list(emojis)[:10]
            }
            
        except Exception as e:
            logger.error(f"Error getting server stats: {e}")
            raise

# ===== Global Service Instance =====

# Initialize with environment variable or default to local agent templates
templates_path = os.getenv(
    "AGENT_TEMPLATES_PATH", 
    str(Path(__file__).parent.parent.parent / "agent_templates")
)
agent_library = AgentLibrary(templates_path)

# ===== FastMCP Tool Definitions =====

@mcp.tool()
async def set_project_roots(request: SetProjectRootsRequest) -> Dict[str, Any]:
    """
    Configure project roots to focus agent recommendations on specific directories.
    
    Use this when you want the agent selection system to understand your project structure
    and provide more targeted recommendations. For example, if you're working on a FastAPI
    project in 'src/api/' and 'tests/', set those as roots to get more relevant agent
    suggestions for your specific context.
    
    Args:
        directories: List of directory paths to focus on (e.g., ["src/api", "tests"])
        base_path: The root path of your project (e.g., "/path/to/project")
        description: Optional description of what you're working on
        
    Returns:
        Success confirmation with the configured roots information
    """
    try:
        return await agent_library.set_roots(
            request.directories,
            request.base_path,
            request.description
        )
    except Exception as e:
        logger.error(f"Error in set_project_roots: {e}")
        return {
            "success": False,
            "error": str(e)
        }

@mcp.tool()
async def get_current_roots() -> Dict[str, Any]:
    """
    Get the currently configured project roots to see what context is being used.
    
    Call this to check if project roots are configured and what directories are being
    focused on. Useful to verify your project context before asking for agent
    recommendations.
    
    Returns:
        If configured: {"configured": true, "roots": {"directories": [...], "base_path": "...", "description": "..."}}
        If not configured: {"configured": false, "message": "No project roots configured"}
    """
    try:
        return await agent_library.get_current_roots()
    except Exception as e:
        logger.error(f"Error in get_current_roots: {e}")
        return {
            "success": False,
            "error": str(e)
        }

@mcp.tool()
async def clear_project_roots() -> Dict[str, Any]:
    """
    Clear the currently configured project roots to return to general recommendations.
    
    Use this when you want to reset the context and get general agent recommendations
    that aren't focused on any specific project structure. After calling this,
    recommend_agents will provide broader suggestions.
    
    Returns:
        Success confirmation that roots have been cleared
    """
    try:
        return await agent_library.clear_project_roots()
    except Exception as e:
        logger.error(f"Error in clear_project_roots: {e}")
        return {
            "success": False,
            "error": str(e)
        }

@mcp.tool()
async def recommend_agents(request: RecommendAgentsRequest) -> List[Dict[str, Any]]:
    """
    Get intelligent agent recommendations based on your task description.
    
    This is the main tool for finding the right Claude Code agent for your needs.
    Provide a description of what you want to accomplish, and it will analyze your
    task and suggest the most relevant specialists. Uses project context if roots
    are configured for even better recommendations.
    
    Args:
        task: Description of what you want to accomplish (e.g., "Help with Python FastAPI development")
        project_context: Additional context about your project (optional)
        limit: Maximum number of recommendations to return (default: 3)
        
    Returns:
        List of recommended agents with confidence scores, descriptions, and reasoning
    """
    try:
        return await agent_library.recommend_agents(
            request.task,
            request.project_context,
            request.limit
        )
    except Exception as e:
        logger.error(f"Error in recommend_agents: {e}")
        raise

@mcp.tool()
async def get_agent_content(request: GetAgentContentRequest) -> Dict[str, Any]:
    """
    Retrieve the complete content and instructions for a specific agent.
    
    Use this after getting recommendations to fetch the full agent template with
    all their expertise, patterns, examples, and guidance. The content includes
    everything needed to use that agent effectively, and may include project-specific
    context if roots are configured.
    
    Args:
        agent_name: The exact name of the agent (e.g., "🚄-fastapi-expert")
        
    Returns:
        Full agent content including expertise areas, code examples, best practices,
        and usage guidance
    """
    try:
        return await agent_library.get_agent_content(request.agent_name)
    except Exception as e:
        logger.error(f"Error in get_agent_content: {e}")
        raise

@mcp.tool()
async def list_agents(request: ListAgentsRequest) -> List[Dict[str, Any]]:
    """
    Browse the complete catalog of available agents with optional filtering.
    
    Use this to discover all available Claude Code agents or search for specific
    capabilities. Perfect for exploring what agents exist when you're not sure
    which one to use. Returns a comprehensive list with names, descriptions, and tools.
    
    Args:
        search: Optional search term to filter agents (searches names and descriptions)
        tools_filter: Optional list of required tools to filter by (e.g., ["Bash", "Edit"])
        
    Returns:
        List of agent dictionaries with name, emoji, description, tools, and file_path
        Each agent entry includes full metadata for evaluation
    """
    try:
        return await agent_library.list_agents(
            request.search,
            request.tools_filter
        )
    except Exception as e:
        logger.error(f"Error in list_agents: {e}")
        raise

@mcp.tool()
async def server_stats() -> Dict[str, Any]:
    """
    Get comprehensive statistics about the agent library and server status.
    
    Provides high-level overview of the entire agent ecosystem including counts,
    tool distribution, hierarchy information, and current project configuration.
    Useful for understanding system capabilities and health monitoring.
    
    Args:
        None required
        
    Returns:
        Dictionary with:
        - total_agents: Total number of loaded agents
        - agent_types: Breakdown by individual/composed/sub types
        - most_common_tools: Tools used across multiple agents
        - project_roots: Current project configuration (if any)
        - server_status: Loading status and any errors
    """
    try:
        return await agent_library.get_server_stats()
    except Exception as e:
        logger.error(f"Error in server_stats: {e}")
        raise

# ===== NEW: Hierarchy Navigation Tools =====

class GetSubAgentsRequest(BaseModel):
    """Request model for getting sub-agents"""
    agent_name: str = Field(..., description="Name of the composed agent")

@mcp.tool()
async def get_sub_agents(request: GetSubAgentsRequest) -> Dict[str, Any]:
    """
    Get all specialist sub-agents for a composed agent.
    
    When you know a composed agent exists but want to see its specialized
    sub-components, use this tool. Perfect for discovering specific expertise
    within a broader domain (e.g., getting Python testing specialists from
    a general testing framework agent).
    
    Args:
        agent_name: Name of the composed agent to explore (must be type "composed")
        
    Returns:
        Dictionary with:
        - success: Boolean indicating if operation succeeded
        - parent_agent: Details about the composed agent
        - sub_agents: List of specialist agents with full metadata
        - total_specialists: Count of available specialists
        - error: Error message if agent not found or not composed
    """
    try:
        if request.agent_name not in agent_library.agents:
            return {
                "success": False,
                "error": f"Agent '{request.agent_name}' not found"
            }
            
        agent = agent_library.agents[request.agent_name]
        
        if agent.agent_type != "composed":
            return {
                "success": False,
                "error": f"Agent '{request.agent_name}' is not a composed agent (type: {agent.agent_type})"
            }
            
        sub_agents = []
        for sub_name in agent.sub_agents:
            if sub_name in agent_library.agents:
                sub_agent = agent_library.agents[sub_name]
                sub_agents.append({
                    "name": sub_agent.name,
                    "emoji": sub_agent.emoji,
                    "description": sub_agent.description,
                    "tools": sub_agent.tools,
                    "file_path": sub_agent.file_path
                })
        
        return {
            "success": True,
            "parent_agent": {
                "name": agent.name,
                "emoji": agent.emoji,
                "description": agent.description
            },
            "sub_agents": sub_agents,
            "total_specialists": len(sub_agents)
        }
        
    except Exception as e:
        logger.error(f"Error in get_sub_agents: {e}")
        raise

@mcp.tool()
async def get_agent_hierarchy() -> Dict[str, Any]:
    """
    Get the complete agent hierarchy showing parent-child relationships.
    
    Provides a comprehensive organizational map of all agents in the system.
    Shows how individual agents, composed agents, and their specialists relate
    to each other. Essential for understanding the full scope of available
    expertise and finding related agents.
    
    Args:
        None required
        
    Returns:
        Dictionary with:
        - individual_agents: List of standalone agents
        - composed_agents: List of parent agents with their sub-agents
        - total_agents: Total count of all agents
        - hierarchy_stats: Breakdown of agent types and counts
        Each agent includes name, emoji, description, and relationship info
    """
    try:
        hierarchy = {
            "individual_agents": [],
            "composed_agents": [],
            "total_agents": len(agent_library.agents),
            "hierarchy_stats": {
                "composed_count": 0,
                "sub_agents_count": 0,
                "individual_count": 0
            }
        }
        
        # Organize agents by type
        for agent in agent_library.agents.values():
            if agent.agent_type == "individual":
                hierarchy["individual_agents"].append({
                    "name": agent.name,
                    "emoji": agent.emoji,
                    "description": agent.description
                })
                hierarchy["hierarchy_stats"]["individual_count"] += 1
                
            elif agent.agent_type == "composed":
                sub_agent_details = []
                for sub_name in agent.sub_agents:
                    if sub_name in agent_library.agents:
                        sub_agent = agent_library.agents[sub_name]
                        sub_agent_details.append({
                            "name": sub_agent.name,
                            "emoji": sub_agent.emoji,
                            "description": sub_agent.description
                        })
                
                hierarchy["composed_agents"].append({
                    "name": agent.name,
                    "emoji": agent.emoji,
                    "description": agent.description,
                    "sub_agents": sub_agent_details,
                    "specialist_count": len(sub_agent_details)
                })
                hierarchy["hierarchy_stats"]["composed_count"] += 1
                hierarchy["hierarchy_stats"]["sub_agents_count"] += len(sub_agent_details)
        
        return hierarchy
        
    except Exception as e:
        logger.error(f"Error in get_agent_hierarchy: {e}")
        raise

class GetParentAgentRequest(BaseModel):
    """Request model for getting parent agent"""
    agent_name: str = Field(..., description="Name of the sub-agent")

@mcp.tool()
async def get_parent_agent(request: GetParentAgentRequest) -> Dict[str, Any]:
    """
    Get the parent composed agent for a specialist sub-agent.
    
    When working with a specialized sub-agent, use this to understand its
    broader context and discover related specialists. Shows the parent framework
    and sibling agents that work in the same domain.
    
    Args:
        agent_name: Name of the sub-agent to look up (must be type "sub")
        
    Returns:
        Dictionary with:
        - success: Boolean indicating if operation succeeded
        - parent_agent: Details about the composed parent agent
        - sub_agent: Details about the requested specialist
        - sibling_specialists: Other related specialists in same domain
        - total_siblings: Count of related specialists
        - error: Error message if agent not found or not a sub-agent
    """
    try:
        if request.agent_name not in agent_library.agents:
            return {
                "success": False,
                "error": f"Agent '{request.agent_name}' not found"
            }
            
        agent = agent_library.agents[request.agent_name]
        
        if agent.agent_type != "sub":
            return {
                "success": False,
                "error": f"Agent '{request.agent_name}' is not a sub-agent (type: {agent.agent_type})"
            }
            
        if not agent.parent_agent or agent.parent_agent not in agent_library.agents:
            return {
                "success": False,
                "error": f"Parent agent not found for '{request.agent_name}'"
            }
            
        parent = agent_library.agents[agent.parent_agent]
        
        # Get sibling specialists
        siblings = []
        for sibling_name in parent.sub_agents:
            if sibling_name != agent.name and sibling_name in agent_library.agents:
                sibling = agent_library.agents[sibling_name]
                siblings.append({
                    "name": sibling.name,
                    "emoji": sibling.emoji,
                    "description": sibling.description
                })
        
        return {
            "success": True,
            "sub_agent": {
                "name": agent.name,
                "emoji": agent.emoji,
                "description": agent.description
            },
            "parent_agent": {
                "name": parent.name,
                "emoji": parent.emoji,
                "description": parent.description,
                "tools": parent.tools
            },
            "sibling_specialists": siblings,
            "total_siblings": len(siblings)
        }
        
    except Exception as e:
        logger.error(f"Error in get_parent_agent: {e}")
        raise

# ===== Server Lifecycle =====

async def initialize_server():
    """Initialize the MCP server"""
    try:
        logger.info("Initializing Claude Agent MCP Server...")
        await agent_library.initialize()
        logger.info("Server initialization completed successfully")
    except Exception as e:
        logger.error(f"Server initialization failed: {e}")
        raise

async def main():
    """Main entry point for the MCP server"""
    try:
        # Initialize services
        await initialize_server()
        
        # Start the MCP server with stdio transport
        logger.info("Starting MCP server with stdio transport...")
        await mcp.run(transport="stdio")
        
    except Exception as e:
        logger.error(f"Server error: {e}")
        raise
    finally:
        logger.info("MCP server shutting down")

def run_server():
    """Entry point for the MCP server"""
    asyncio.run(main())

if __name__ == "__main__":
    run_server()