headroom-ai 0.2.13__py3-none-any.whl

headroom/ccr/mcp_server.py

@@ -0,0 +1,319 @@
+"""CCR MCP Server - Exposes headroom_retrieve as an MCP tool.
+
+This MCP server allows LLMs to retrieve compressed content via MCP instead
+of through injected tool definitions. It connects to the Headroom proxy's
+CompressionStore to serve retrieval requests.
+
+Usage:
+    # As standalone server (stdio transport)
+    python -m headroom.ccr.mcp_server
+
+    # With custom proxy URL
+    python -m headroom.ccr.mcp_server --proxy-url http://localhost:8787
+
+    # Add to Claude Code's MCP config (~/.claude/mcp.json):
+    {
+        "mcpServers": {
+            "headroom": {
+                "command": "python",
+                "args": ["-m", "headroom.ccr.mcp_server"]
+            }
+        }
+    }
+
+When MCP is configured, the proxy will detect the tool is already present
+and skip tool injection, avoiding duplicate tools.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import logging
+import os
+from typing import Any
+
+# Try to import MCP SDK
+try:
+    from mcp.server import Server
+    from mcp.server.stdio import stdio_server
+    from mcp.types import TextContent, Tool
+
+    MCP_AVAILABLE = True
+except ImportError:
+    MCP_AVAILABLE = False
+    Server = None
+    stdio_server = None
+
+# Try to import httpx for proxy communication
+try:
+    import httpx
+
+    HTTPX_AVAILABLE = True
+except ImportError:
+    HTTPX_AVAILABLE = False
+    httpx = None  # type: ignore[assignment]
+
+from .tool_injection import CCR_TOOL_NAME
+
+logger = logging.getLogger("headroom.ccr.mcp")
+
+# Default proxy URL (can be overridden via env or args)
+DEFAULT_PROXY_URL = os.environ.get("HEADROOM_PROXY_URL", "http://127.0.0.1:8787")
+
+
+class CCRMCPServer:
+    """MCP Server that exposes headroom_retrieve tool.
+
+    This server can operate in two modes:
+    1. HTTP mode: Calls the proxy's /v1/retrieve endpoint (default)
+    2. Direct mode: Uses CompressionStore directly (same process)
+
+    HTTP mode is recommended as it ensures consistency with the proxy.
+    """
+
+    def __init__(
+        self,
+        proxy_url: str = DEFAULT_PROXY_URL,
+        direct_mode: bool = False,
+    ):
+        """Initialize CCR MCP Server.
+
+        Args:
+            proxy_url: URL of the Headroom proxy server.
+            direct_mode: If True, access CompressionStore directly instead of via HTTP.
+        """
+        self.proxy_url = proxy_url
+        self.direct_mode = direct_mode
+        self._http_client: httpx.AsyncClient | None = None
+
+        if not MCP_AVAILABLE:
+            raise ImportError("MCP SDK not installed. Install with: pip install mcp")
+
+        if not direct_mode and not HTTPX_AVAILABLE:
+            raise ImportError(
+                "httpx not installed (required for HTTP mode). Install with: pip install httpx"
+            )
+
+        self.server = Server("headroom-ccr")
+        self._setup_handlers()
+
+    def _setup_handlers(self):
+        """Set up MCP tool handlers."""
+
+        @self.server.list_tools()
+        async def list_tools() -> list[Tool]:
+            """Return available tools."""
+            return [
+                Tool(
+                    name=CCR_TOOL_NAME,
+                    description=(
+                        "Retrieve original uncompressed content that was compressed "
+                        "to save tokens. Use this when you need more data than what's "
+                        "shown in compressed tool results. The hash is provided in "
+                        "compression markers like [N items compressed... hash=abc123]."
+                    ),
+                    inputSchema={
+                        "type": "object",
+                        "properties": {
+                            "hash": {
+                                "type": "string",
+                                "description": "Hash key from the compression marker (e.g., 'abc123' from hash=abc123)",
+                            },
+                            "query": {
+                                "type": "string",
+                                "description": (
+                                    "Optional search query to filter results. "
+                                    "If provided, only returns items matching the query. "
+                                    "If omitted, returns all original items."
+                                ),
+                            },
+                        },
+                        "required": ["hash"],
+                    },
+                )
+            ]
+
+        @self.server.call_tool()
+        async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
+            """Handle tool calls."""
+            if name != CCR_TOOL_NAME:
+                return [
+                    TextContent(
+                        type="text",
+                        text=json.dumps({"error": f"Unknown tool: {name}"}),
+                    )
+                ]
+
+            hash_key = arguments.get("hash")
+            query = arguments.get("query")
+
+            if not hash_key:
+                return [
+                    TextContent(
+                        type="text",
+                        text=json.dumps({"error": "hash parameter is required"}),
+                    )
+                ]
+
+            # Retrieve content
+            try:
+                if self.direct_mode:
+                    result = await self._retrieve_direct(hash_key, query)
+                else:
+                    result = await self._retrieve_via_proxy(hash_key, query)
+
+                return [
+                    TextContent(
+                        type="text",
+                        text=json.dumps(result, indent=2),
+                    )
+                ]
+            except Exception as e:
+                logger.error(f"Retrieval failed: {e}")
+                return [
+                    TextContent(
+                        type="text",
+                        text=json.dumps({"error": str(e)}),
+                    )
+                ]
+
+    async def _retrieve_via_proxy(
+        self,
+        hash_key: str,
+        query: str | None,
+    ) -> dict[str, Any]:
+        """Retrieve content via proxy's HTTP endpoint."""
+        if self._http_client is None:
+            self._http_client = httpx.AsyncClient(timeout=30.0)
+
+        url = f"{self.proxy_url}/v1/retrieve"
+        payload = {"hash": hash_key}
+        if query:
+            payload["query"] = query
+
+        response = await self._http_client.post(url, json=payload)
+
+        if response.status_code == 404:
+            return {
+                "error": "Entry not found or expired (TTL: 5 minutes)",
+                "hash": hash_key,
+            }
+
+        response.raise_for_status()
+        result: dict[str, Any] = response.json()
+        return result
+
+    async def _retrieve_direct(
+        self,
+        hash_key: str,
+        query: str | None,
+    ) -> dict[str, Any]:
+        """Retrieve content directly from CompressionStore."""
+        from headroom.cache.compression_store import get_compression_store
+
+        store = get_compression_store()
+
+        if query:
+            results = store.search(hash_key, query)
+            return {
+                "hash": hash_key,
+                "query": query,
+                "results": results,
+                "count": len(results),
+            }
+        else:
+            entry = store.retrieve(hash_key)
+            if entry:
+                return {
+                    "hash": hash_key,
+                    "original_content": entry.original_content,
+                    "original_item_count": entry.original_item_count,
+                    "compressed_item_count": entry.compressed_item_count,
+                    "retrieval_count": entry.retrieval_count,
+                }
+            return {
+                "error": "Entry not found or expired (TTL: 5 minutes)",
+                "hash": hash_key,
+            }
+
+    async def run_stdio(self):
+        """Run the server with stdio transport."""
+        async with stdio_server() as (read_stream, write_stream):
+            logger.info(f"CCR MCP Server starting (proxy: {self.proxy_url})")
+            await self.server.run(
+                read_stream,
+                write_stream,
+                self.server.create_initialization_options(),
+            )
+
+    async def cleanup(self):
+        """Clean up resources."""
+        if self._http_client:
+            await self._http_client.aclose()
+
+
+def create_ccr_mcp_server(
+    proxy_url: str = DEFAULT_PROXY_URL,
+    direct_mode: bool = False,
+) -> CCRMCPServer:
+    """Create a CCR MCP server instance.
+
+    Args:
+        proxy_url: URL of the Headroom proxy server.
+        direct_mode: If True, access CompressionStore directly.
+
+    Returns:
+        CCRMCPServer instance.
+
+    Example:
+        ```python
+        server = create_ccr_mcp_server()
+        await server.run_stdio()
+        ```
+    """
+    return CCRMCPServer(proxy_url=proxy_url, direct_mode=direct_mode)
+
+
+async def main():
+    """Run the CCR MCP server."""
+    parser = argparse.ArgumentParser(
+        description="CCR MCP Server - Retrieve compressed content via MCP"
+    )
+    parser.add_argument(
+        "--proxy-url",
+        default=DEFAULT_PROXY_URL,
+        help=f"Headroom proxy URL (default: {DEFAULT_PROXY_URL})",
+    )
+    parser.add_argument(
+        "--direct",
+        action="store_true",
+        help="Use direct CompressionStore access instead of HTTP",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Enable debug logging",
+    )
+
+    args = parser.parse_args()
+
+    if args.debug:
+        logging.basicConfig(level=logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
+
+    server = create_ccr_mcp_server(
+        proxy_url=args.proxy_url,
+        direct_mode=args.direct,
+    )
+
+    try:
+        await server.run_stdio()
+    finally:
+        await server.cleanup()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())