import json import logging import contextvars from fastapi import APIRouter, Depends, Request, HTTPException from fastapi.responses import Response, JSONResponse from mcp.server import Server from mcp.server.sse import SseServerTransport from mcp.types import Tool, TextContent from app.auth import get_current_agent from app.scopes import has_scope_access, ALL_SCOPES from app.models import get_rule, get_all_rules, get_all_ports, get_port, set_rule, get_agent_by_key, update_key_last_used logger = logging.getLogger(__name__) router = APIRouter() mcp_server = Server("context-hub") current_agent_var: contextvars.ContextVar[str] = contextvars.ContextVar("current_agent", default="") @mcp_server.list_tools() async def handle_list_tools() -> list[Tool]: return [ Tool( name="get_rules", description="Get rules for a specific scope if authorized", inputSchema={ "type": "object", "properties": { "scope": {"type": "string", "description": "Scope name (infra, llm, nyora, perso, tt)"} }, "required": ["scope"] } ), Tool( name="get_ports", description="Get the full ports registry", inputSchema={ "type": "object", "properties": {} } ), Tool( name="search_rules", description="Search rules across all authorized scopes", inputSchema={ "type": "object", "properties": { "query": {"type": "string", "description": "Search term"} }, "required": ["query"] } ), Tool( name="get_agent_config", description="Get the compiled configuration for a specific agent based on its scopes", inputSchema={ "type": "object", "properties": { "name": {"type": "string", "description": "Agent name (e.g., GEMINI, HERMES_TT)"} }, "required": ["name"] } ), Tool( name="check_port", description="Check if a port is free or occupied", inputSchema={ "type": "object", "properties": { "port": {"type": "integer", "description": "Port number"} }, "required": ["port"] } ), Tool( name="get_my_context", description="Get all scopes authorized for this agent — call at START of session", inputSchema={"type": "object", "properties": {"agent": {"type": "string"}}, "required": ["agent"]} ), Tool( name="update_rule", description="Merge-update a scope with new key-value data — call at END of session", inputSchema={ "type": "object", "properties": { "scope": {"type": "string", "description": "Scope to update (infra/llm/nyora/perso/tt)"}, "agent": {"type": "string", "description": "Agent name for scope access check"}, "data": {"type": "object", "description": "Key-value pairs to merge into the scope"} }, "required": ["scope", "agent", "data"] } ) ] @mcp_server.call_tool() async def handle_call_tool(name: str, arguments: dict) -> list[TextContent]: # Autorisation liee a l'identite reelle de la connexion SSE (current_agent_var), # jamais a un champ "agent"/"name" declare par le client dans les arguments. caller_agent = current_agent_var.get() if name == "get_rules": scope = arguments.get("scope") if not has_scope_access(caller_agent, scope): return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden on scope {scope}"}))] rule = await get_rule(scope) if not rule: return [TextContent(type="text", text=json.dumps({"error": "Scope not found"}))] return [TextContent(type="text", text=json.dumps(rule))] elif name == "get_ports": if not has_scope_access(caller_agent, "infra"): return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden: requires infra scope"}))] ports = await get_all_ports() return [TextContent(type="text", text=json.dumps(ports))] elif name == "search_rules": query = arguments.get("query", "").lower() all_rules = await get_all_rules() results = [] for scope, content_ in all_rules.items(): if not has_scope_access(caller_agent, scope): continue if query in json.dumps(content_).lower(): results.append({"scope": scope, "content": content_}) return [TextContent(type="text", text=json.dumps(results))] elif name == "get_agent_config": all_rules = await get_all_rules() config = {} for scope in ALL_SCOPES: if has_scope_access(caller_agent, scope): config[scope] = all_rules.get(scope, {}) return [TextContent(type="text", text=json.dumps(config))] elif name == "check_port": if not has_scope_access(caller_agent, "infra"): return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden: requires infra scope"}))] port = arguments.get("port") port_info = await get_port(port) if port_info: return [TextContent(type="text", text=json.dumps({"status": "occupied", "info": port_info}))] else: return [TextContent(type="text", text=json.dumps({"status": "free", "port": port}))] elif name == "get_my_context": all_rules = await get_all_rules() result = {} for scope in ALL_SCOPES: if has_scope_access(caller_agent, scope): result[scope] = all_rules.get(scope, {}) return [TextContent(type="text", text=json.dumps({"agent": caller_agent, "scopes": result}))] elif name == "update_rule": scope = arguments.get("scope") data = arguments.get("data", {}) if not has_scope_access(caller_agent, scope): return [TextContent(type="text", text=json.dumps({"error": f"Agent {caller_agent} forbidden on scope {scope}"}))] if not isinstance(data, dict): return [TextContent(type="text", text=json.dumps({"error": "data must be a JSON object"}))] existing = await get_rule(scope) or {} existing.update(data) await set_rule(scope, existing) return [TextContent(type="text", text=json.dumps({"status": "ok", "scope": scope, "updated_keys": list(data.keys())}))] return [TextContent(type="text", text=json.dumps({"error": "Unknown tool"}))] # FastMCP / SSE Integration # connect_sse() ET handle_post_message() envoient chacun leur reponse ASGI # completement par eux-memes (via le send() qu'on leur passe). Les faire # passer par des routes FastAPI classiques (Depends + return Response) fait # que FastAPI tente un second envoi une fois le SDK termine -> RuntimeError # uvicorn ("Unexpected ASGI message 'http.response.start' sent, after response # already completed"). Un seul montage ASGI brut gere GET (ouverture SSE) et # POST (messages) sans jamais repasser par le wrapping FastAPI ; l'auth # X-API-Key est donc verifiee ici a la main plutot que via Depends. sse_transport = SseServerTransport("/messages") async def _mcp_asgi_app(scope, receive, send): if scope["type"] != "http": return headers = dict(scope.get("headers") or []) api_key = headers.get(b"x-api-key", b"").decode() agent = await get_agent_by_key(api_key) if api_key else None if not agent: response = JSONResponse({"detail": "Invalid or missing API Key"}, status_code=401) await response(scope, receive, send) return await update_key_last_used(api_key) if scope["method"] == "POST": await sse_transport.handle_post_message(scope, receive, send) return async with sse_transport.connect_sse(scope, receive, send) as (read_stream, write_stream): token = current_agent_var.set(agent) try: await mcp_server.run(read_stream, write_stream, mcp_server.create_initialization_options()) finally: current_agent_var.reset(token) router.mount("/mcp", _mcp_asgi_app)