Files
AI-Profile-Router/platform/mcp/platform_context_mcp.py
T

266 lines
10 KiB
Python

#!/usr/bin/env python3
"""Small read-only Athena knowledge MCP.
The normal entry point is ATHENA.md. Large historical documentation remains
available through bounded search/read tools but is never loaded automatically.
"""
from __future__ import annotations
import json
import os
import re
import sys
from pathlib import Path
from typing import Any
VERSION = "2.0.0"
REPO_ROOT = Path(os.environ.get("ATHENA_REPO_ROOT", "/knowledge/repo")).resolve()
RUNTIME_FILE = Path(os.environ.get("ATHENA_RUNTIME_FILE", "/runtime/runtime.json"))
MAX_OVERVIEW_CHARS = 14_000
MAX_READ_LINES = 160
MAX_SEARCH_RESULTS = 8
ALLOWED_SUFFIXES = {".md", ".json", ".yaml", ".yml", ".txt"}
BLOCKED_PARTS = {".git", "secrets", "private", "credentials"}
if hasattr(sys.stdin, "reconfigure"):
sys.stdin.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
TOOLS = [
{
"name": "athena_get_overview",
"description": (
"START HERE for Athena architecture or administration. Returns the compact, "
"authoritative ATHENA.md. Do not read additional platform documents unless a "
"specific unresolved question remains."
),
"inputSchema": {"type": "object", "properties": {}, "additionalProperties": False},
},
{
"name": "athena_get_current_state",
"description": "Return the compact generated runtime snapshot: active profile, containers, GPUs, source commit and recovery status.",
"inputSchema": {"type": "object", "properties": {}, "additionalProperties": False},
},
{
"name": "athena_get_external_services",
"description": "List known services outside Athena so an existing Unraid or home-network backend is reused instead of duplicated.",
"inputSchema": {"type": "object", "properties": {}, "additionalProperties": False},
},
{
"name": "athena_search_reference",
"description": (
"Search ATHENA.md and documentation for one concrete term. Returns at most eight "
"short excerpts. Use only when ATHENA.md did not answer the question."
),
"inputSchema": {
"type": "object",
"properties": {"query": {"type": "string", "minLength": 2, "maxLength": 120}},
"required": ["query"],
"additionalProperties": False,
},
},
{
"name": "athena_read_reference",
"description": (
"Read a bounded line range from one known documentation file. Missing paths are "
"reported as a normal not-found result and must not be retried by guessing."
),
"inputSchema": {
"type": "object",
"properties": {
"path": {"type": "string", "pattern": "^[A-Za-z0-9_.+/-]{1,200}$"},
"start_line": {"type": "integer", "minimum": 1, "maximum": 1000000, "default": 1},
"line_count": {"type": "integer", "minimum": 1, "maximum": MAX_READ_LINES, "default": 80},
},
"required": ["path"],
"additionalProperties": False,
},
},
]
def result_error(message: str, **details: Any) -> dict[str, Any]:
return {"ok": False, "error": message, "retry": False, **details}
def safe_path(relative: str) -> Path | None:
if not relative or relative.startswith("/"):
return None
candidate = Path(relative)
if ".." in candidate.parts or any(part.lower() in BLOCKED_PARTS for part in candidate.parts):
return None
target = (REPO_ROOT / candidate).resolve(strict=False)
try:
target.relative_to(REPO_ROOT)
except ValueError:
return None
if candidate.name != "ATHENA.md" and (not candidate.parts or candidate.parts[0] != "docs"):
return None
if target.suffix.lower() not in ALLOWED_SUFFIXES:
return None
return target
def read_text(path: Path, limit: int | None = None) -> str:
text = path.read_text(encoding="utf-8", errors="replace")
return text if limit is None else text[:limit]
def overview() -> dict[str, Any]:
path = REPO_ROOT / "ATHENA.md"
try:
content = read_text(path, MAX_OVERVIEW_CHARS)
except (OSError, PermissionError) as exc:
return result_error("ATHENA.md is unavailable", path="ATHENA.md", detail=str(exc))
return {"ok": True, "source": "ATHENA.md", "content": content, "truncated": path.stat().st_size > len(content.encode())}
def current_state() -> dict[str, Any]:
try:
value = json.loads(RUNTIME_FILE.read_text(encoding="utf-8"))
except (OSError, ValueError) as exc:
return result_error("runtime snapshot is unavailable", detail=str(exc))
containers = value.get("containers") or []
return {
"ok": True,
"generated_at": value.get("generated_at"),
"hostname": value.get("hostname"),
"active_inference_profiles": value.get("active_inference_profiles") or [],
"source_commit": value.get("source_commit"),
"gpus": value.get("gpus") or [],
"containers": containers,
"container_count": len(containers),
"recovery_kit": value.get("recovery_kit") or {"present": False},
}
def external_services() -> dict[str, Any]:
path = REPO_ROOT / "config/service-catalog.json"
try:
value = json.loads(path.read_text(encoding="utf-8"))
except (OSError, ValueError) as exc:
return result_error("service catalog is unavailable", detail=str(exc))
services = []
for item in value.get("services", []):
services.append({key: item.get(key) for key in ("id", "name", "host", "address", "port", "protocol", "purpose") if item.get(key) is not None})
return {"ok": True, "services": services, "count": len(services)}
def reference_files() -> list[Path]:
files = [REPO_ROOT / "ATHENA.md"]
docs = REPO_ROOT / "docs"
try:
files.extend(sorted(path for path in docs.glob("*.md") if path.is_file()))
except OSError:
pass
return files
def search_reference(arguments: dict[str, Any]) -> dict[str, Any]:
query = str(arguments.get("query", "")).strip()
if len(query) < 2:
return result_error("query must contain at least two characters")
pattern = re.compile(re.escape(query), re.IGNORECASE)
matches: list[dict[str, Any]] = []
for path in reference_files():
try:
lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
except OSError:
continue
for number, line in enumerate(lines, 1):
if pattern.search(line):
matches.append({
"path": str(path.relative_to(REPO_ROOT)),
"line": number,
"excerpt": line.strip()[:280],
})
if len(matches) >= MAX_SEARCH_RESULTS:
return {"ok": True, "query": query, "matches": matches, "truncated": True}
return {"ok": True, "query": query, "matches": matches, "truncated": False}
def read_reference(arguments: dict[str, Any]) -> dict[str, Any]:
relative = str(arguments.get("path", ""))
path = safe_path(relative)
if path is None:
return result_error("path is not an allowed documentation path", path=relative)
if not path.is_file():
return result_error("documentation file not found", path=relative)
start = max(1, int(arguments.get("start_line", 1)))
count = min(MAX_READ_LINES, max(1, int(arguments.get("line_count", 80))))
try:
lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
except OSError as exc:
return result_error("documentation file is unreadable", path=relative, detail=str(exc))
selected = lines[start - 1:start - 1 + count]
return {
"ok": True,
"path": relative,
"start_line": start,
"end_line": start + len(selected) - 1 if selected else start - 1,
"total_lines": len(lines),
"content": "\n".join(selected),
"truncated": start - 1 + len(selected) < len(lines),
}
def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
if name == "athena_get_overview":
return overview()
if name == "athena_get_current_state":
return current_state()
if name == "athena_get_external_services":
return external_services()
if name == "athena_search_reference":
return search_reference(arguments)
if name == "athena_read_reference":
return read_reference(arguments)
return result_error("unknown tool", tool=name)
def emit(request_id: Any, result: Any = None, error: dict[str, Any] | None = None) -> None:
message = {"jsonrpc": "2.0", "id": request_id}
message["error" if error else "result"] = error or result
sys.stdout.write(json.dumps(message, ensure_ascii=False, separators=(",", ":")) + "\n")
sys.stdout.flush()
def handle(message: dict[str, Any]) -> None:
method, request_id = message.get("method"), message.get("id")
if method == "initialize":
emit(request_id, {
"protocolVersion": message.get("params", {}).get("protocolVersion", "2024-11-05"),
"capabilities": {"tools": {"listChanged": False}},
"serverInfo": {"name": "mike-ai-platform-context", "version": VERSION},
})
elif method == "tools/list":
emit(request_id, {"tools": TOOLS})
elif method == "tools/call":
params = message.get("params") or {}
value = call_tool(str(params.get("name", "")), params.get("arguments") or {})
emit(request_id, {
"content": [{"type": "text", "text": json.dumps(value, ensure_ascii=False, separators=(",", ":"))}],
"structuredContent": value,
"isError": False,
})
elif request_id is not None:
emit(request_id, error={"code": -32601, "message": "method not found"})
def main() -> None:
for line in sys.stdin:
try:
if line.strip():
handle(json.loads(line))
except Exception as exc:
sys.stderr.write(f"MCP input error: {exc}\n")
sys.stderr.flush()
if __name__ == "__main__":
main()