203 lines
8.4 KiB
Python
203 lines
8.4 KiB
Python
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)
|