import json import logging import os from datetime import datetime from typing import Any import aiofiles from mcp.server.fastmcp import FastMCP from starlette.requests import Request from starlette.responses import JSONResponse, Response from .core import GLOBAL_TOOL_REGISTRY, get_tool_schema logger = logging.getLogger(__name__) def register_custom_routes(mcp: FastMCP) -> None: """Register custom HTTP routes and discovery tools on the FastMCP instance.""" # Discovery endpoint for Agents (Phase 4.6.19) @mcp.tool() async def list_tools() -> str: """Dynamically list all registered MCP tools using official FastMCP API.""" try: tools = await mcp.list_tools() formatted_tools = [] for tool in tools: formatted_tools.append( { "type": "function", "function": { "name": tool.name, "description": tool.description or "", "parameters": get_tool_schema(tool), }, } ) return json.dumps( { "success": True, "tools": formatted_tools, "count": len(formatted_tools), "timestamp": datetime.now().isoformat(), }, indent=2, ) except Exception as e: logger.error(f"Failed to list tools: {e}") return json.dumps({"success": False, "error": str(e)}, indent=2) # Direct RPC Bridge for Agents (Phase 4.6.19) @mcp.custom_route("/rpc", methods=["POST"]) async def mcp_rpc_handler(request: Request) -> Response: """Handle standard JSON-RPC requests via POST.""" global GLOBAL_TOOL_REGISTRY try: body = await request.json() method_name = body.get("method") params = body.get("params", {}) # 1. Discovery handling (Fast path using Worker-local Registry) if method_name == "list_tools": agent_type = request.headers.get("X-Agent-Type", "anonymous") tools_to_return = [] # FAIL-SAFE A: Try disk cache first tools_cache = "/tmp/mcp_tools.json" if os.path.exists(tools_cache): try: async with aiofiles.open(tools_cache) as f: cached = json.loads(await f.read()) if cached: tools_to_return = cached logger.info( f"✅ [PID {os.getpid()}] Discovery: Serving {len(cached)} tools from disk cache for {agent_type}" ) except Exception: pass # FAIL-SAFE B: If registry is empty, try live fetch if not tools_to_return: if not GLOBAL_TOOL_REGISTRY: logger.warning(f"⚠️ [PID {os.getpid()}] Registry empty, triggering fail-safe live fetch") raw_tools = mcp._tool_manager.list_tools() GLOBAL_TOOL_REGISTRY = [ { "type": "function", "function": { "name": t.name, "description": t.description or "", "parameters": get_tool_schema(t), }, } for t in raw_tools ] tools_to_return = GLOBAL_TOOL_REGISTRY # Apply RBAC Filtering (Dynamic Tool Exposing) via RBACService from src.server.services.rbac_service import RBACService rbac = RBACService() restricted_tools = await rbac.get_restricted_mcp_tools(agent_type) if restricted_tools: filtered_tools = [] for t in tools_to_return: if isinstance(t, dict): func = t.get("function") if isinstance(func, dict) and func.get("name") not in restricted_tools: filtered_tools.append(t) tools_to_return = filtered_tools logger.info( f"🔒 [PID {os.getpid()}] RBAC applied for {agent_type}. Returning {len(tools_to_return)} tools." ) return JSONResponse({"jsonrpc": "2.0", "result": tools_to_return, "id": body.get("id")}) # 2. Tool execution via official API logger.info(f"RPC Call: [PID {os.getpid()}] Executing tool '{method_name}' via official call_tool API") # Apply RBAC Enforcement on Execution via RBACService agent_type = request.headers.get("X-Agent-Type", "anonymous") from src.server.services.rbac_service import RBACService rbac = RBACService() restricted_tools = await rbac.get_restricted_mcp_tools(agent_type) if method_name in restricted_tools: logger.warning(f"🚫 [PID {os.getpid()}] RBAC Violation: {agent_type} attempted to call {method_name}") return JSONResponse( { "error": { "code": -32003, "message": f"RBAC Violation: Tool '{method_name}' is restricted for agent '{agent_type}'", } }, status_code=403, ) try: raw_result = await mcp.call_tool(method_name, params) except Exception as tool_err: return JSONResponse( {"error": {"code": -32603, "message": f"Tool execution failed: {str(tool_err)}"}}, status_code=500 ) # 3. Physical Serialization & Unwrapping processed_result: Any = [] if isinstance(raw_result, list): for item in raw_result: if hasattr(item, "model_dump"): dump = item.model_dump() processed_result.append(dump.get("text", "") if dump.get("type") == "text" else dump) else: processed_result.append(str(item)) else: processed_result = str(raw_result) return JSONResponse({"jsonrpc": "2.0", "result": processed_result, "id": body.get("id")}) except Exception as e: logger.error(f"RPC Bridge Failed: {e}", exc_info=True) return JSONResponse({"error": {"code": -32603, "message": str(e)}}, status_code=500) # Custom route for session tracking @mcp.custom_route("/sessions", methods=["GET"]) async def get_sessions(request: Request) -> Response: """Get active session count from FastMCP internal state.""" active_sessions = 0 try: if hasattr(mcp, "session_manager"): if hasattr(mcp.session_manager, "_server_instances"): active_sessions = len(mcp.session_manager._server_instances) except Exception as e: logger.error(f"Failed to get active sessions: {e}") return JSONResponse({"active_sessions": active_sessions}) # Custom route for health check (Phase 5.2.0 Task 5.2.0.10) @mcp.custom_route("/health", methods=["GET"]) async def get_health(request: Request) -> Response: """Get MCP health status.""" return JSONResponse({"status": "healthy"})