"""Companion stdio MCP-style server over activation_drift_monitor public APIs.""" from __future__ import annotations import argparse import json import sys from typing import Any, Callable from .core import ActivationDriftMonitor MONITOR = ActivationDriftMonitor() def _json_ok(result: Any, request_id: Any = None) -> dict[str, Any]: return {"ok": True, "id": request_id, "result": result} def _json_err(message: str, request_id: Any = None) -> dict[str, Any]: return {"ok": False, "id": request_id, "error": message} def _tool_profile_reference(args: dict[str, Any]) -> dict[str, Any]: layers = args.get("layers") if not isinstance(layers, dict): raise ValueError("'layers' must be a mapping of layer name to activation values") profiled = MONITOR.profile_reference_batch(layers) return {"profiled_layers": profiled, "state": MONITOR.export_state()} def _tool_evaluate_batch(args: dict[str, Any]) -> dict[str, Any]: layers = args.get("layers") if not isinstance(layers, dict): raise ValueError("'layers' must be a mapping of layer name to activation values") batch = MONITOR.evaluate_batch(layers) return batch.as_dict() def _tool_reset_monitor(args: dict[str, Any]) -> dict[str, Any]: layer_name = args.get("layer_name") if layer_name is not None and not isinstance(layer_name, str): raise ValueError("'layer_name' must be a string or null") MONITOR.reset(layer_name) return {"state": MONITOR.export_state()} def _tool_get_state(args: dict[str, Any]) -> dict[str, Any]: _ = args return MONITOR.export_state() TOOL_REGISTRY: dict[str, Callable[[dict[str, Any]], dict[str, Any]]] = { "profile_reference": _tool_profile_reference, "evaluate_batch": _tool_evaluate_batch, "reset_monitor": _tool_reset_monitor, "get_state": _tool_get_state, } def describe_surface() -> dict[str, Any]: return { "server_name": "activation-drift-monitor", "transport": "stdio", "tools": [ { "name": "profile_reference", "description": "Profile baseline activations for one or more layers.", "input_schema": { "type": "object", "properties": { "layers": { "type": "object", "description": "Mapping layer name -> activation sequence" } }, "required": ["layers"] }, }, { "name": "evaluate_batch", "description": "Evaluate one batch and return per-layer/global drift alerts.", "input_schema": { "type": "object", "properties": { "layers": { "type": "object", "description": "Mapping layer name -> activation sequence" } }, "required": ["layers"] }, }, { "name": "reset_monitor", "description": "Reset all state or one layer.", "input_schema": { "type": "object", "properties": { "layer_name": {"type": ["string", "null"]} } }, }, { "name": "get_state", "description": "Return serialized monitor state.", "input_schema": {"type": "object", "properties": {}}, }, ], "resources": [ { "uri": "monitor://state", "description": "Current monitor state snapshot", } ], "api_mapping": { "profile_reference": "ActivationDriftMonitor.profile_reference_batch", "evaluate_batch": "ActivationDriftMonitor.evaluate_batch", "reset_monitor": "ActivationDriftMonitor.reset", "get_state": "ActivationDriftMonitor.export_state", }, } def _handle_request(request: dict[str, Any]) -> dict[str, Any]: request_id = request.get("id") method = request.get("method") params = request.get("params") or {} try: if method == "list_tools": return _json_ok({"tools": describe_surface()["tools"]}, request_id) if method == "list_resources": return _json_ok({"resources": describe_surface()["resources"]}, request_id) if method == "read_resource": uri = params.get("uri") if uri != "monitor://state": raise ValueError("unknown resource URI") return _json_ok({"uri": uri, "contents": MONITOR.export_state()}, request_id) if method == "call_tool": tool_name = params.get("name") tool_args = params.get("arguments") or {} if tool_name not in TOOL_REGISTRY: raise ValueError(f"unknown tool: {tool_name}") result = TOOL_REGISTRY[tool_name](tool_args) return _json_ok(result, request_id) raise ValueError(f"unsupported method: {method}") except Exception as exc: return _json_err(str(exc), request_id) def run_stdio_server() -> int: """Read JSON lines from stdin and write JSON line responses to stdout.""" for raw in sys.stdin: line = raw.strip() if not line: continue try: request = json.loads(line) except json.JSONDecodeError as exc: response = _json_err(f"invalid JSON request: {exc}") else: response = _handle_request(request) sys.stdout.write(json.dumps(response) + "\n") sys.stdout.flush() return 0 def _run_fastmcp_if_available() -> bool: """Try FastMCP runtime when installed; return False when unavailable.""" try: from fastmcp import FastMCP except Exception: return False app = FastMCP("activation-drift-monitor") @app.tool def profile_reference(layers: dict[str, list[float]]) -> dict[str, Any]: return _tool_profile_reference({"layers": layers}) @app.tool def evaluate_batch(layers: dict[str, list[float]]) -> dict[str, Any]: return _tool_evaluate_batch({"layers": layers}) @app.tool def reset_monitor(layer_name: str | None = None) -> dict[str, Any]: return _tool_reset_monitor({"layer_name": layer_name}) @app.tool def get_state() -> dict[str, Any]: return _tool_get_state({}) app.run(transport="stdio") return True def _run_single_tool(tool: str, args_json: str) -> int: if tool not in TOOL_REGISTRY: sys.stderr.write(f"unknown tool: {tool}\n") return 2 arguments = json.loads(args_json) if args_json else {} result = TOOL_REGISTRY[tool](arguments) sys.stdout.write(json.dumps(result, indent=2) + "\n") return 0 def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description="activation-drift-monitor companion stdio MCP server" ) parser.add_argument( "--mode", choices=["stdio", "describe", "tool"], default="stdio", help="Run as stdio server, print surface metadata, or run one tool call", ) parser.add_argument("--tool", default="", help="Tool name for --mode tool") parser.add_argument( "--args", default="{}", help="JSON arguments object for --mode tool", ) parser.add_argument( "--prefer-fastmcp", action="store_true", help="Try FastMCP runtime before fallback stdio JSON-lines mode", ) args = parser.parse_args(argv) if args.mode == "describe": sys.stdout.write(json.dumps(describe_surface(), indent=2) + "\n") return 0 if args.mode == "tool": return _run_single_tool(args.tool, args.args) if args.prefer_fastmcp and _run_fastmcp_if_available(): return 0 return run_stdio_server() if __name__ == "__main__": raise SystemExit(main())