diff --git a/claude_desktop/rootfs/usr/local/bin/headroom-mcp-proxy.py b/claude_desktop/rootfs/usr/local/bin/headroom-mcp-proxy.py new file mode 100644 index 0000000000..87164b3f48 --- /dev/null +++ b/claude_desktop/rootfs/usr/local/bin/headroom-mcp-proxy.py @@ -0,0 +1,142 @@ +#!/lsiopy/bin/python3 +"""Run Headroom MCP with compression delegated to the single proxy backend. + +Upstream Headroom MCP normally performs ``headroom_compress`` in its own +process. That imports the compression pipeline and can load a second copy of the +Kompress ONNX model in addition to the HTTP proxy. This adapter preserves the +upstream MCP protocol and retrieve/stats implementations, but replaces only its +local compression method with a call to the proxy's loopback-only +``/v1/compress`` endpoint. + +The proxy gate starts the heavy backend on this first request and later unloads +it after the configured idle timeout. The MCP process therefore remains a +lightweight protocol bridge and never imports ``headroom.compress``. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import sys +from types import MethodType +from typing import Any + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Headroom MCP single-runtime adapter") + parser.add_argument( + "--proxy-url", + default=os.environ.get("HEADROOM_PROXY_URL", "http://127.0.0.1:8787"), + ) + parser.add_argument("--transport", default="stdio") + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8788) + parser.add_argument("--path", default="/mcp") + parser.add_argument("--debug", action="store_true") + parser.add_argument("--direct", action="store_true") + args, unknown = parser.parse_known_args() + if unknown: + print(f"headroom-mcp-proxy: ignoring unsupported arguments: {unknown}", file=sys.stderr) + if args.transport.lower() != "stdio": + parser.error("only stdio transport is supported by the add-on adapter") + return args + + +def _number(data: dict[str, Any], key: str) -> int: + value = data.get(key, 0) + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def make_proxy_compressor(proxy_url: str): + endpoint = f"{proxy_url.rstrip('/')}/v1/compress" + model = os.environ.get("HEADROOM_MCP_MODEL", "claude-sonnet-4-5-20250929") + + def compress_via_proxy(_self, content: str) -> dict[str, Any]: + # Imported here, not at module import, so MCP startup stays small. httpx is + # already a dependency of Headroom's MCP server for retrieve/stats. + import httpx + + response = httpx.post( + endpoint, + json={ + "messages": [{"role": "tool", "content": content}], + "model": model, + }, + timeout=httpx.Timeout(180.0, connect=75.0), + ) + response.raise_for_status() + data = response.json() + if not isinstance(data, dict): + raise RuntimeError("Headroom proxy returned a non-object compression response") + + messages = data.get("messages") + compressed: Any = content + if isinstance(messages, list) and messages: + last = messages[-1] + if isinstance(last, dict) and "content" in last: + compressed = last["content"] + if not isinstance(compressed, str): + compressed = json.dumps(compressed, ensure_ascii=False) + + hashes = data.get("ccr_hashes") + hash_key = next((item for item in hashes if isinstance(item, str)), None) if isinstance(hashes, list) else None + before = _number(data, "tokens_before") + after = _number(data, "tokens_after") + saved = _number(data, "tokens_saved") + if saved <= 0: + saved = max(0, before - after) + savings_percent = round(saved / before * 100, 1) if before > 0 else 0.0 + transforms = data.get("transforms_applied") + if not isinstance(transforms, list): + transforms = [] + + note = "Compression was executed by the shared Headroom proxy runtime." + if hash_key: + note += ( + f" Original stored with hash={hash_key}. " + "Use mcp__headroom__headroom_retrieve to recover it." + ) + return { + "compressed": compressed, + "hash": hash_key, + "original_tokens": before, + "compressed_tokens": after, + "tokens_saved": saved, + "savings_percent": savings_percent, + "transforms": transforms, + "note": note, + } + + return compress_via_proxy + + +async def run() -> None: + args = parse_args() + + # Defense in depth. The adapter never calls the local compressor, but keep + # the MCP process explicitly unable to initialize Kompress if upstream code + # changes or an unrelated import probes the compression pipeline. + os.environ.pop("HF_HOME", None) + os.environ["HEADROOM_DISABLE_KOMPRESS"] = "1" + os.environ["HEADROOM_PROXY_URL"] = args.proxy_url + + from headroom.ccr.mcp_server import HeadroomMCPServer + + server = HeadroomMCPServer(proxy_url=args.proxy_url, check_proxy=True) + server._compress_content = MethodType(make_proxy_compressor(args.proxy_url), server) + try: + await server.run_stdio() + finally: + await server.cleanup() + + +if __name__ == "__main__": + try: + asyncio.run(run()) + except KeyboardInterrupt: + pass