|
| 1 | +"""Thin stdio front end for the MemPalace MCP server. |
| 2 | +
|
| 3 | +``mempalace-mcp`` is spawned once per agent session, and when a hub is |
| 4 | +running every one of those processes is a pure proxy: ``_dispatch_stdio_request`` |
| 5 | +forwards each JSON-RPC request over HTTP and the local storage stack is never |
| 6 | +touched. Importing :mod:`mempalace.mcp_server` to do that costs ~77 MB anyway, |
| 7 | +because chromadb (+61 MB on its own), numpy, pydantic, grpc and opentelemetry |
| 8 | +are all pulled in at module scope. A fleet of 50 agents therefore paid ~3.9 GB |
| 9 | +to hold proxies that do no work. |
| 10 | +
|
| 11 | +This module is the entry point instead. It imports only the standard library |
| 12 | +plus :mod:`mempalace.config` and :mod:`mempalace.server_registry` (~5 MB each), |
| 13 | +so a proxied session runs at roughly 22 MB. The full server is imported lazily, |
| 14 | +and only when this process actually has to serve a request itself. |
| 15 | +
|
| 16 | +The fallback is deliberately preserved: a session whose hub dies keeps working. |
| 17 | +It just stops being free at that point, so it says so — once to the log, and on |
| 18 | +the tool result itself, because the agent driving the session is the one who |
| 19 | +needs to know its memory backend changed shape underneath it. |
| 20 | +""" |
| 21 | + |
| 22 | +from __future__ import annotations |
| 23 | + |
| 24 | +import json |
| 25 | +import logging |
| 26 | +import os |
| 27 | +import sys |
| 28 | +import urllib.error |
| 29 | +import urllib.request |
| 30 | + |
| 31 | +logger = logging.getLogger(__name__) |
| 32 | + |
| 33 | +# Shared with the in-server forwarder and the CLI forwarder. |
| 34 | +_HUB_FORWARD_ENV = "MEMPALACE_HUB_FORWARD" |
| 35 | +_HUB_PROXY_TIMEOUT_S = 600.0 |
| 36 | + |
| 37 | +_DEGRADED_NOTICE = ( |
| 38 | + "MemPalace is running WITHOUT its shared hub. This session is now serving " |
| 39 | + "the palace directly, which loads the whole index into this process " |
| 40 | + "(hundreds of MB) instead of reusing the hub's. Memory tools still work. " |
| 41 | + "If several agents are running, expect memory pressure until the hub is " |
| 42 | + "back — check that the MemPalace hub process is alive." |
| 43 | +) |
| 44 | + |
| 45 | + |
| 46 | +def _truthy_env_off(name: str) -> bool: |
| 47 | + return os.environ.get(name, "").strip().lower() in {"0", "false", "no", "off"} |
| 48 | + |
| 49 | + |
| 50 | +def _is_plain_stdio_invocation(argv: list) -> bool: |
| 51 | + """True when this is an ordinary stdio session that a hub could serve. |
| 52 | +
|
| 53 | + Anything that asks for a different transport, or that configures the |
| 54 | + serving process itself, goes straight to the full server. Being wrong in |
| 55 | + this direction only costs the old startup weight; being wrong the other |
| 56 | + way would silently drop a flag, so unknown arguments count as "not plain". |
| 57 | + """ |
| 58 | + allowed_flags = {"--palace", "--collection", "--backend"} |
| 59 | + i = 0 |
| 60 | + while i < len(argv): |
| 61 | + arg = argv[i] |
| 62 | + if arg == "--transport": |
| 63 | + if i + 1 >= len(argv) or argv[i + 1] != "stdio": |
| 64 | + return False |
| 65 | + i += 2 |
| 66 | + continue |
| 67 | + if arg.startswith("--transport="): |
| 68 | + if arg.split("=", 1)[1] != "stdio": |
| 69 | + return False |
| 70 | + i += 1 |
| 71 | + continue |
| 72 | + if arg in allowed_flags: |
| 73 | + i += 2 |
| 74 | + continue |
| 75 | + if any(arg.startswith(flag + "=") for flag in allowed_flags): |
| 76 | + i += 1 |
| 77 | + continue |
| 78 | + return False |
| 79 | + return True |
| 80 | + |
| 81 | + |
| 82 | +def _palace_path(argv: list): |
| 83 | + """Resolve the palace path without importing the server.""" |
| 84 | + for i, arg in enumerate(argv): |
| 85 | + if arg == "--palace" and i + 1 < len(argv): |
| 86 | + return argv[i + 1] |
| 87 | + if arg.startswith("--palace="): |
| 88 | + return arg.split("=", 1)[1] |
| 89 | + try: |
| 90 | + from .config import MempalaceConfig |
| 91 | + |
| 92 | + return MempalaceConfig().palace_path |
| 93 | + except Exception: |
| 94 | + logger.debug("palace path unresolved; serving locally", exc_info=True) |
| 95 | + return None |
| 96 | + |
| 97 | + |
| 98 | +def _hub_target(palace_path): |
| 99 | + """Return ``(base_url, headers)`` for a live hub serving our palace, else None.""" |
| 100 | + if _truthy_env_off(_HUB_FORWARD_ENV) or not palace_path: |
| 101 | + return None |
| 102 | + try: |
| 103 | + from . import server_registry |
| 104 | + |
| 105 | + info = server_registry.read_live_serverinfo(palace_path) |
| 106 | + if not info or info.get("pid") == os.getpid(): |
| 107 | + return None |
| 108 | + base_url = server_registry.client_base_url(info) |
| 109 | + headers = {"Content-Type": "application/json"} |
| 110 | + token = server_registry.load_server_token(palace_path) |
| 111 | + except Exception: |
| 112 | + logger.debug("hub discovery failed", exc_info=True) |
| 113 | + return None |
| 114 | + if token: |
| 115 | + headers["Authorization"] = f"Bearer {token}" |
| 116 | + return base_url, headers |
| 117 | + |
| 118 | + |
| 119 | +def _forward(base_url: str, headers: dict, request: dict): |
| 120 | + """POST one JSON-RPC request to the hub; None for notifications (202).""" |
| 121 | + body = json.dumps(request, ensure_ascii=False).encode("utf-8") |
| 122 | + http_request = urllib.request.Request(f"{base_url}/mcp", data=body, headers=headers) |
| 123 | + with urllib.request.urlopen(http_request, timeout=_HUB_PROXY_TIMEOUT_S) as resp: |
| 124 | + raw = resp.read() |
| 125 | + if not raw: |
| 126 | + return None |
| 127 | + return json.loads(raw.decode("utf-8")) |
| 128 | + |
| 129 | + |
| 130 | +def _annotate_degraded(response): |
| 131 | + """Prepend the hub-is-gone notice to a tools/call result. |
| 132 | +
|
| 133 | + The driving agent only ever sees ``result.content``; a log line it cannot |
| 134 | + read is not a warning. Prepended rather than appended so it survives a |
| 135 | + client that renders only the first block, and only on tools/call, so |
| 136 | + tools/list and the handshake keep their exact shapes. |
| 137 | + """ |
| 138 | + if not isinstance(response, dict): |
| 139 | + return response |
| 140 | + result = response.get("result") |
| 141 | + if not isinstance(result, dict): |
| 142 | + return response |
| 143 | + content = result.get("content") |
| 144 | + if not isinstance(content, list): |
| 145 | + return response |
| 146 | + result["content"] = [{"type": "text", "text": f"[mempalace] {_DEGRADED_NOTICE}"}, *content] |
| 147 | + return response |
| 148 | + |
| 149 | + |
| 150 | +class _LocalServer: |
| 151 | + """Lazily-imported full server, plus the background services it expects. |
| 152 | +
|
| 153 | + Import is deferred to the first request this process has to answer itself, |
| 154 | + which is the whole point of this module: a proxied session never pays it. |
| 155 | + """ |
| 156 | + |
| 157 | + def __init__(self): |
| 158 | + self._module = None |
| 159 | + |
| 160 | + @property |
| 161 | + def loaded(self) -> bool: |
| 162 | + return self._module is not None |
| 163 | + |
| 164 | + def load(self): |
| 165 | + if self._module is None: |
| 166 | + logger.warning( |
| 167 | + "MemPalace hub unavailable; serving this session locally. " |
| 168 | + "Loading the local storage stack (this process will grow)." |
| 169 | + ) |
| 170 | + from . import mcp_server |
| 171 | + |
| 172 | + # Importing the server installs its stdio protection: os.dup2(2, 1) |
| 173 | + # plus sys.stdout = sys.stderr, so stray library prints cannot |
| 174 | + # corrupt JSON-RPC. That also redirects *our* responses to stderr — |
| 175 | + # fd 1 itself is moved, so holding a reference to the old object is |
| 176 | + # not enough. _restore_stdout undoes both levels, exactly as the |
| 177 | + # server's own stdio loop does before it starts answering. |
| 178 | + mcp_server._restore_stdout() |
| 179 | + if hasattr(sys.stdout, "reconfigure"): |
| 180 | + try: |
| 181 | + sys.stdout.reconfigure(encoding="utf-8", errors="replace") |
| 182 | + except (AttributeError, OSError): |
| 183 | + pass |
| 184 | + |
| 185 | + for start in ( |
| 186 | + mcp_server._start_idle_exit_watchdog, |
| 187 | + mcp_server._start_write_stall_watchdog, |
| 188 | + ): |
| 189 | + try: |
| 190 | + start() |
| 191 | + except Exception: |
| 192 | + logger.debug("local service %s failed to start", start, exc_info=True) |
| 193 | + self._module = mcp_server |
| 194 | + return self._module |
| 195 | + |
| 196 | + |
| 197 | +def _proxy_error(request: dict, base_url: str, exc: Exception): |
| 198 | + """Mirror the in-server proxy failure shape for a request we must not replay.""" |
| 199 | + if request.get("id") is None: |
| 200 | + return None |
| 201 | + return { |
| 202 | + "jsonrpc": "2.0", |
| 203 | + "id": request.get("id"), |
| 204 | + "error": { |
| 205 | + "code": -32000, |
| 206 | + "message": f"palace hub proxy failed: {exc}", |
| 207 | + "data": { |
| 208 | + "hub": base_url, |
| 209 | + "hint": ( |
| 210 | + "The palace hub did not complete this request. Mutating tools " |
| 211 | + "are not replayed locally — the hub may still be executing the " |
| 212 | + "call. Check the hub process, then retry." |
| 213 | + ), |
| 214 | + }, |
| 215 | + }, |
| 216 | + } |
| 217 | + |
| 218 | + |
| 219 | +def _handle(request: dict, palace_path, local: _LocalServer): |
| 220 | + """Route one request: live hub first, this process otherwise.""" |
| 221 | + target = _hub_target(palace_path) |
| 222 | + if target is not None: |
| 223 | + base_url, headers = target |
| 224 | + try: |
| 225 | + return _forward(base_url, headers, request) |
| 226 | + except (urllib.error.URLError, OSError, TimeoutError, ValueError) as exc: |
| 227 | + # Reaching the hub and getting an HTTP error means it may have run |
| 228 | + # the call; so does any mid-flight failure on a mutating tool. |
| 229 | + # Neither may be replayed here. |
| 230 | + module = local.load() |
| 231 | + if isinstance(exc, urllib.error.HTTPError) or module._request_is_mutating(request): |
| 232 | + return _proxy_error(request, base_url, exc) |
| 233 | + logger.warning("Hub at %s unreachable (%s); handling request locally", base_url, exc) |
| 234 | + return _annotate_degraded(module.handle_request(request)) |
| 235 | + module = local.load() |
| 236 | + return _annotate_degraded(module.handle_request(request)) |
| 237 | + |
| 238 | + |
| 239 | +def _run_proxy_loop(palace_path) -> None: |
| 240 | + for stream in (sys.stdin, sys.stdout): |
| 241 | + if hasattr(stream, "reconfigure"): |
| 242 | + try: |
| 243 | + stream.reconfigure(encoding="utf-8", errors="replace") |
| 244 | + except (AttributeError, OSError): |
| 245 | + pass |
| 246 | + |
| 247 | + local = _LocalServer() |
| 248 | + while True: |
| 249 | + try: |
| 250 | + line = sys.stdin.readline() |
| 251 | + except KeyboardInterrupt: |
| 252 | + break |
| 253 | + except OSError as exc: |
| 254 | + logger.info("stdin read failed (%s) -- client disconnected, shutting down", exc) |
| 255 | + break |
| 256 | + if not line: |
| 257 | + logger.info("stdin EOF -- client disconnected, shutting down") |
| 258 | + break |
| 259 | + line = line.strip() |
| 260 | + if not line: |
| 261 | + continue |
| 262 | + |
| 263 | + payload = None |
| 264 | + try: |
| 265 | + response = _handle(json.loads(line), palace_path, local) |
| 266 | + if response is not None: |
| 267 | + payload = json.dumps(response, ensure_ascii=False) |
| 268 | + except KeyboardInterrupt: |
| 269 | + break |
| 270 | + except Exception as e: |
| 271 | + logger.error(f"Server error: {e}") |
| 272 | + continue |
| 273 | + |
| 274 | + if payload is None: |
| 275 | + continue |
| 276 | + try: |
| 277 | + sys.stdout.write(payload + "\n") |
| 278 | + sys.stdout.flush() |
| 279 | + except KeyboardInterrupt: |
| 280 | + break |
| 281 | + except (BrokenPipeError, OSError) as exc: |
| 282 | + logger.info("stdout write failed (%s) -- client disconnected, shutting down", exc) |
| 283 | + break |
| 284 | + |
| 285 | + |
| 286 | +def main() -> None: |
| 287 | + """Entry point for ``mempalace-mcp``. |
| 288 | +
|
| 289 | + Delegates to the full server for anything but a plain stdio session, and |
| 290 | + for a plain stdio session with no hub to proxy to — in both cases the |
| 291 | + heavy import was going to happen regardless. |
| 292 | + """ |
| 293 | + argv = sys.argv[1:] |
| 294 | + if not _is_plain_stdio_invocation(argv): |
| 295 | + from . import mcp_server |
| 296 | + |
| 297 | + return mcp_server.main() |
| 298 | + |
| 299 | + palace_path = _palace_path(argv) |
| 300 | + if _hub_target(palace_path) is None: |
| 301 | + from . import mcp_server |
| 302 | + |
| 303 | + return mcp_server.main() |
| 304 | + |
| 305 | + logging.basicConfig(level=logging.INFO, stream=sys.stderr) |
| 306 | + _run_proxy_loop(palace_path) |
0 commit comments