From 96699f2a17815eaa36cced6476bb1d9a77f9e39b Mon Sep 17 00:00:00 2001 From: pi Date: Thu, 2 Jul 2026 13:09:29 +0200 Subject: [PATCH] feat(pi-bridge): external MemPalace transport via MEMPALACE_REMOTE_URL Let the pi<->mempalace bridge connect to a shared MemPalace over HTTP instead of always spawning a local mempalace-mcp: - Extract IMcpClient; rename McpClient -> StdioMcpClient (ctor command, arg-less start()). - Add RemoteMcpClient (vendored from pi-extensions/mcp-loader.ts): streamable-HTTP with AbortController timeouts, protocolVersion pinned 2024-11-05, alive/ensureAlive. mempalace-mcp --transport http is sessionless JSON-RPC today; session/SSE/404 branches retained for a future streamable-HTTP server. - createClient() selects transport from MEMPALACE_REMOTE_URL; MEMPALACE_REMOTE_TOKEN -> Authorization: Bearer. Lifecycle automation (wake-up, /mempalace-diary) unchanged. - scripts/check-mcp-client-sync.sh: drift guard vs canonical mcp-loader.ts. - README: document local-vs-external transport. Typechecks clean (strict); both transports smoke-tested against live mempalace-mcp. --- extensions/pi/README.md | 37 +++- extensions/pi/mempalace.ts | 340 +++++++++++++++++++++++++++++-- scripts/check-mcp-client-sync.sh | 77 +++++++ 3 files changed, 437 insertions(+), 17 deletions(-) create mode 100755 scripts/check-mcp-client-sync.sh diff --git a/extensions/pi/README.md b/extensions/pi/README.md index e2462e5..26ed4ec 100644 --- a/extensions/pi/README.md +++ b/extensions/pi/README.md @@ -15,6 +15,7 @@ dependencies (~300 MB). **Jump to:** - [What it does](#what-it-does) +- [Transport: local vs external](#transport-local-vs-external) - [The `Type.Unsafe` gotcha](#the-typeunsafe-gotcha) - [Deploying pi with mempalace on a new machine](#deploying-pi-with-mempalace-on-a-new-machine) - [Fail-soft, identity, debugging](#fail-soft) @@ -23,9 +24,12 @@ dependencies (~300 MB). ## What it does -1. **Spawns `mempalace-mcp`** as a subprocess and does the MCP stdio - JSON-RPC handshake (`initialize` + `notifications/initialized` + - `tools/list`). +1. **Connects to MemPalace** and does the MCP handshake (`initialize` + + `notifications/initialized` + `tools/list`). By default it **spawns + `mempalace-mcp`** as a local stdio subprocess (`StdioMcpClient`); if + `$MEMPALACE_REMOTE_URL` is set it instead talks to a shared MemPalace over + HTTP (`RemoteMcpClient`) and spawns no local process — see + [Transport](#transport-local-vs-external). 2. **Registers each MCP tool** as a pi tool with its real `inputSchema` passed through via `Type.Unsafe(...)` (see gotcha below). 3. **Wake-up auto-injection** (`before_agent_start`, one-shot per fresh @@ -39,6 +43,33 @@ dependencies (~300 MB). because pi sessions are typically short/tactical and `session_shutdown` fires too late to drive another LLM turn. +## Transport: local vs external + +The bridge speaks the same MCP protocol over two interchangeable transports, +chosen at load time: + +- **Local (default)** — spawns `mempalace-mcp` as a stdio subprocess; the palace + lives wherever that process opens it (default `~/.mempalace`). This is the + hardened path with per-request timeouts and respawn/self-heal (below). +- **External** — set `MEMPALACE_REMOTE_URL` to a MemPalace HTTP endpoint (e.g. + `http://mempalace.lan:8765/mcp`) and the bridge connects over HTTP instead, + spawning no local process. Use this to share **one** palace across several + harnesses/containers (pi + opencode + native). `MEMPALACE_REMOTE_TOKEN`, if + set, is sent as `Authorization: Bearer `. + + Serve such an endpoint with `mempalace-mcp --transport http --host 0.0.0.0 + --port 8765` (the `pi-devbox` / `opencode-devbox` repos ship a + `docker-compose.mempalace.yml` for exactly this). Note: that HTTP transport is + currently sessionless and **unauthenticated** — keep it on a trusted network + or behind a reverse proxy that enforces the bearer token. + + Implementation note: the HTTP client (`RemoteMcpClient`) is **vendored** from + [`pi-extensions`](https://gitea.jordbo.se/joakimp/pi-extensions)' + `mcp-loader.ts`. A `MCP-STREAMABLE-HTTP-CLIENT-SYNC` token keeps the two + copies from drifting — [`scripts/check-mcp-client-sync.sh`](../../scripts/check-mcp-client-sync.sh) + fails if they diverge (it skips gracefully when the `pi-extensions` checkout + isn't present). + ## Fail-soft If `mempalace-mcp` can't be spawned (PATH missing, binary crashes at diff --git a/extensions/pi/mempalace.ts b/extensions/pi/mempalace.ts index 40d7f62..76dfa18 100644 --- a/extensions/pi/mempalace.ts +++ b/extensions/pi/mempalace.ts @@ -1,9 +1,20 @@ /** * MemPalace ↔ pi bridge. * - * Spawns the `mempalace-mcp` MCP stdio server as a subprocess, performs the - * MCP `initialize` handshake, lists available tools, and registers each - * one as a pi tool that proxies to `tools/call`. + * Registers every MemPalace MCP tool as a pi tool that proxies to `tools/call`. + * Two interchangeable transports, selected at load time: + * + * - LOCAL (default): spawn the `mempalace-mcp` stdio server as a subprocess + * (StdioMcpClient) — hardened with per-request timeouts + generation-tracked + * respawn/self-heal for slow cold-opens (see below). + * - EXTERNAL: connect to a shared MemPalace over HTTP (RemoteMcpClient) when + * $MEMPALACE_REMOTE_URL is set — e.g. one palace serving pi + opencode + + * native. Optional bearer auth via $MEMPALACE_REMOTE_TOKEN. No local + * `mempalace-mcp` process is spawned in this mode. + * + * Either way the client performs the MCP `initialize` handshake, lists tools, + * and the extension body below is transport-agnostic (depends only on the + * IMcpClient interface). * * Lifecycle automation (per ~/.agents/skills/mempalace/SKILL.md): * - Wake-up (auto): on first user prompt of a fresh session, inject @@ -68,6 +79,21 @@ interface Pending { timer: ReturnType | null; } +// Transport-agnostic client contract. The extension body depends only on this, +// so LOCAL (StdioMcpClient) and EXTERNAL (RemoteMcpClient) are drop-in swaps. +// Each implementation owns its own reliability model — stdio uses timeout + +// kill + respawn; http uses fetch + AbortController timeout (+ 404 re-init if +// a future streamable-HTTP server issues sessions). +interface IMcpClient { + tools: McpTool[]; + readonly alive: boolean; + onExit: (() => void) | null; + start(): Promise; + callTool(name: string, args: Record): Promise; + ensureAlive(): Promise; + stop(): void | Promise; +} + const num = (envVal: string | undefined, fallback: number): number => { const n = envVal !== undefined ? Number(envVal) : Number.NaN; return Number.isFinite(n) && n >= 0 ? n : fallback; @@ -79,7 +105,7 @@ const sleep = (ms: number): Promise => if (typeof t.unref === "function") t.unref(); }); -class McpClient { +class StdioMcpClient implements IMcpClient { private proc: ChildProcessWithoutNullStreams | null = null; private nextId = 1; private pending = new Map(); @@ -102,25 +128,28 @@ class McpClient { // they were attached for and no-op if a newer server has since taken over // (prevents a stale OLD-proc 'exit' from clobbering a freshly respawned one). private gen = 0; - // Spawn args remembered so a respawn can reuse them. - private command = "mempalace-mcp"; - private args: string[] = []; + // Spawn args remembered so a respawn can reuse them (set via constructor). + private command: string; + private args: string[]; // Fired when the child process dies (exit or stall-kill). Lets the // extension flip `available` so later tool calls fail fast. public onExit: (() => void) | null = null; + constructor(command = "mempalace-mcp", args: string[] = []) { + this.command = command; + this.args = args; + } + /** True only when a server has completed init and not since died. */ get alive(): boolean { return this.healthy; } - async start(command: string, args: string[] = []): Promise { + async start(): Promise { if (this.ready) return this.ready; - this.command = command; - this.args = args; this.ready = (async () => { const myGen = ++this.gen; - const child = spawn(command, args, { stdio: ["pipe", "pipe", "pipe"] }); + const child = spawn(this.command, this.args, { stdio: ["pipe", "pipe", "pipe"] }); this.proc = child; child.on("error", (err) => this.handleDeath(myGen, err)); child.on("exit", (code) => @@ -242,7 +271,7 @@ class McpClient { // runs while not healthy — no useful in-flight start() can exist. this.ready = null; try { - await this.start(this.command, this.args); + await this.start(); } catch { // start() rejected (e.g. respawn cold-open also stalled and was // killed). Loop until the budget is exhausted. @@ -328,8 +357,291 @@ class McpClient { } } +// ─────────────────────────────────────────────────────────────────────── +// RemoteMcpClient — EXTERNAL transport (streamable-HTTP / sessionless JSON-RPC). +// +// VENDORED from pi-extensions/extensions/mcp-loader.ts (class RemoteMcpClient). +// Kept as a copy rather than a shared import because the two repos ship +// independently; the canonical implementation lives in mcp-loader.ts. Port +// protocol fixes in both directions. +// +// MCP-STREAMABLE-HTTP-CLIENT-SYNC: v1 +// This token must match the canonical block's. Bump it in BOTH files whenever +// the streamable-HTTP protocol handling changes; scripts/check-mcp-client-sync.sh +// fails when they diverge (skips gracefully if pi-extensions isn't checked out). +// +// Deltas from the canonical copy: +// • protocolVersion pinned to "2024-11-05" to match mempalace-mcp's stdio +// handshake (the server echoes whatever it is sent, so cosmetic today, +// but keeps both transports consistent). +// • per-request AbortController timeout honouring MEMPALACE_MCP_TIMEOUT_MS / +// MEMPALACE_MCP_INIT_TIMEOUT_MS, mirroring StdioMcpClient's timeout ethos. +// • alive / ensureAlive / onExit to satisfy IMcpClient. +// +// NOTE: mempalace-mcp --transport http is a SESSIONLESS, stateless JSON-RPC +// server (no Mcp-Session-Id, always application/json, Connection: close), so +// the session-id / SSE / 404-reinit branches below are never exercised against +// it today. They are retained so this client also works unchanged if mempalace +// later moves to a full streamable-HTTP (FastMCP) transport. +// ─────────────────────────────────────────────────────────────────────── +const REMOTE_PROTOCOL_VERSION = "2024-11-05"; +const REMOTE_CLIENT_INFO = { name: "pi-mempalace-ext", version: "0.1.0" }; + +class RemoteMcpClient implements IMcpClient { + private url: string; + private extraHeaders: Record; + private sessionId: string | null = null; + private nextId = 1; + public tools: McpTool[] = []; + // Unused for HTTP (no child process to die) — present to satisfy IMcpClient. + public onExit: (() => void) | null = null; + + private requestTimeoutMs = num(process.env.MEMPALACE_MCP_TIMEOUT_MS, 60_000); + private initTimeoutMs = num(process.env.MEMPALACE_MCP_INIT_TIMEOUT_MS, 300_000); + private healthy = false; + private reviving: Promise | null = null; + + constructor(url: string, headers?: Record) { + this.url = url; + this.extraHeaders = headers ?? {}; + } + + get alive(): boolean { + return this.healthy; + } + + async start(): Promise { + await this.request( + "initialize", + { protocolVersion: REMOTE_PROTOCOL_VERSION, capabilities: {}, clientInfo: REMOTE_CLIENT_INFO }, + { timeoutMs: this.initTimeoutMs }, + ); + await this.notify("notifications/initialized", {}); + const listed = await this.request("tools/list", {}, { timeoutMs: this.initTimeoutMs }); + this.tools = (listed?.tools as McpTool[]) ?? []; + this.healthy = true; + } + + async callTool(name: string, args: Record): Promise { + return this.request("tools/call", { name, arguments: args }); + } + + /** + * Re-establish connectivity (and re-list tools) after a fetch failure. + * For the stateless server this is just a fresh initialize+tools/list; + * concurrent callers share one in-flight attempt. + */ + async ensureAlive(): Promise { + if (this.healthy) return true; + if (this.reviving) return this.reviving; + this.reviving = (async () => { + try { + this.sessionId = null; + await this.start(); + } catch { + // stays unhealthy + } + return this.healthy; + })(); + try { + return await this.reviving; + } finally { + this.reviving = null; + } + } + + async stop(): Promise { + this.healthy = false; + if (!this.sessionId) return; + try { + await fetch(this.url, { method: "DELETE", headers: this.buildHeaders() }); + } catch {} + } + + private buildHeaders(): Record { + const h: Record = { + "Content-Type": "application/json", + Accept: "application/json, text/event-stream", + "MCP-Protocol-Version": REMOTE_PROTOCOL_VERSION, + ...this.extraHeaders, + }; + if (this.sessionId) h["Mcp-Session-Id"] = this.sessionId; + return h; + } + + private signalFor(timeoutMs: number): AbortSignal | undefined { + return timeoutMs > 0 ? AbortSignal.timeout(timeoutMs) : undefined; + } + + private async reinitialize(): Promise { + await this.request( + "initialize", + { protocolVersion: REMOTE_PROTOCOL_VERSION, capabilities: {}, clientInfo: REMOTE_CLIENT_INFO }, + { timeoutMs: this.initTimeoutMs, allowReinitOn404: false }, + ); + await this.notify("notifications/initialized", {}); + } + + private async request( + method: string, + params: unknown, + opts: { timeoutMs?: number; allowReinitOn404?: boolean } = {}, + ): Promise { + const timeoutMs = opts.timeoutMs ?? this.requestTimeoutMs; + const allowReinitOn404 = opts.allowReinitOn404 ?? true; + const id = this.nextId++; + const body = JSON.stringify({ jsonrpc: "2.0", id, method, params }); + const sessionAtStart = this.sessionId; + let res: Response; + try { + res = await fetch(this.url, { + method: "POST", + headers: this.buildHeaders(), + body, + signal: this.signalFor(timeoutMs), + }); + } catch (err) { + // Network failure / timeout (AbortError) — mark unreachable so the next + // tool call triggers ensureAlive(). + this.healthy = false; + const e = err as Error; + const reason = + e.name === "AbortError" || e.name === "TimeoutError" + ? `timed out after ${timeoutMs}ms` + : e.message; + throw new Error(`mempalace remote request '${method}' failed: ${reason}`); + } + + const sid = res.headers.get("Mcp-Session-Id") ?? res.headers.get("mcp-session-id"); + if (sid) this.sessionId = sid; + + // 404 + we sent a session id → server forgot us. Drop id, re-init, retry once. + if (res.status === 404 && sessionAtStart && allowReinitOn404) { + try { + await res.arrayBuffer(); + } catch {} + this.sessionId = null; + await this.reinitialize(); + return this.request(method, params, { timeoutMs, allowReinitOn404: false }); + } + + if (!res.ok) { + let detail = ""; + try { + detail = (await res.text()).slice(0, 200); + } catch {} + throw new Error(`mempalace remote: HTTP ${res.status} on ${method}${detail ? `: ${detail}` : ""}`); + } + + const ct = (res.headers.get("content-type") ?? "").toLowerCase(); + if (ct.includes("text/event-stream")) { + return await this.readSseForResponse(res, id); + } + if (ct.includes("application/json")) { + const json: any = await res.json(); + if (json.error) throw new Error(json.error.message ?? "MCP error"); + return json.result; + } + if (res.status === 202) return undefined; + throw new Error(`mempalace remote: unexpected content-type "${ct}" on ${method}`); + } + + private async notify(method: string, params: unknown): Promise { + const body = JSON.stringify({ jsonrpc: "2.0", method, params }); + let res: Response; + try { + res = await fetch(this.url, { + method: "POST", + headers: this.buildHeaders(), + body, + signal: this.signalFor(this.requestTimeoutMs), + }); + } catch (err) { + throw new Error(`mempalace remote: fetch failed on notify ${method}: ${(err as Error).message}`); + } + try { + await res.arrayBuffer(); + } catch {} + if (!res.ok && res.status !== 202) { + throw new Error(`mempalace remote: HTTP ${res.status} on notify ${method}`); + } + } + + private async readSseForResponse(res: Response, expectedId: number): Promise { + if (!res.body) throw new Error("mempalace remote: SSE response had no body"); + const reader = res.body.getReader(); + const decoder = new TextDecoder(); + let buf = ""; + let dataLines: string[] = []; + + const tryDispatch = (): { matched: boolean; result?: any } => { + if (dataLines.length === 0) return { matched: false }; + const data = dataLines.join("\n"); + dataLines = []; + let msg: any; + try { + msg = JSON.parse(data); + } catch { + return { matched: false }; + } + if (typeof msg.id === "number" && msg.id === expectedId) { + if (msg.error) throw new Error(msg.error.message ?? "MCP error"); + return { matched: true, result: msg.result }; + } + return { matched: false }; + }; + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + buf += decoder.decode(value, { stream: true }); + let nl: number; + while ((nl = buf.indexOf("\n")) !== -1) { + const rawLine = buf.slice(0, nl); + buf = buf.slice(nl + 1); + const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine; + if (line === "") { + const out = tryDispatch(); + if (out.matched) return out.result; + } else if (line.startsWith(":")) { + // SSE comment / keepalive — ignore + } else if (line.startsWith("data:")) { + dataLines.push(line.slice(5).replace(/^ /, "")); + } + } + } + const out = tryDispatch(); + if (out.matched) return out.result; + } finally { + try { + reader.cancel(); + } catch {} + } + throw new Error(`mempalace remote: SSE stream closed without response for id=${expectedId}`); + } +} + +/** + * Pick the transport from the environment: MEMPALACE_REMOTE_URL selects the + * EXTERNAL (HTTP) client; unset falls back to the LOCAL stdio server. + */ +function createClient(): IMcpClient { + const remoteUrl = process.env.MEMPALACE_REMOTE_URL?.trim(); + if (remoteUrl) { + return new RemoteMcpClient(remoteUrl, authHeaders()); + } + return new StdioMcpClient("mempalace-mcp"); +} + +/** Bearer auth header for the remote transport, if MEMPALACE_REMOTE_TOKEN is set. */ +function authHeaders(): Record | undefined { + const token = process.env.MEMPALACE_REMOTE_TOKEN?.trim(); + return token ? { Authorization: `Bearer ${token}` } : undefined; +} + export default async function mempalaceExtension(pi: ExtensionAPI) { - const client = new McpClient(); + const client = createClient(); let available = false; const agentName = process.env.MEMPALACE_AGENT_NAME ?? "pi"; @@ -343,7 +655,7 @@ export default async function mempalaceExtension(pi: ExtensionAPI) { // call will attempt a bounded respawn via client.ensureAlive(). available = false; }; - await client.start("mempalace-mcp"); + await client.start(); available = true; } catch (err) { // First cold-open stalled/crashed. Give the bounded self-heal a chance diff --git a/scripts/check-mcp-client-sync.sh b/scripts/check-mcp-client-sync.sh new file mode 100755 index 0000000..2d2a4a4 --- /dev/null +++ b/scripts/check-mcp-client-sync.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# check-mcp-client-sync.sh — drift guard for the vendored streamable-HTTP MCP client. +# +# The RemoteMcpClient in extensions/pi/mempalace.ts is a VENDORED copy of the +# canonical implementation in pi-extensions/extensions/mcp-loader.ts (the two +# repos ship independently, so a shared import isn't practical for one small, +# protocol-frozen class — see the annotated header in mempalace.ts). +# +# Both blocks carry a matching sync token: +# MCP-STREAMABLE-HTTP-CLIENT-SYNC: vN +# Bump it in BOTH files whenever the streamable-HTTP protocol handling changes. +# This guard fails when the tokens diverge, so a protocol fix in one file can't +# silently skip the other. +# +# It is a LOCAL developer guard: it needs the sibling pi-extensions checkout to +# compare against. If that isn't found it SKIPS (exit 0) rather than failing — +# so it's safe to wire into CI that only checks out this repo. +# +# Point it at a specific checkout with: PI_EXTENSIONS_DIR=/path/to/pi-extensions +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +VENDORED="$REPO_ROOT/extensions/pi/mempalace.ts" + +TOKEN_RE='MCP-STREAMABLE-HTTP-CLIENT-SYNC: v[0-9]+' +extract_token() { grep -oE "$TOKEN_RE" "$1" 2>/dev/null | head -1 | grep -oE 'v[0-9]+' || true; } + +if [[ ! -f "$VENDORED" ]]; then + echo "ERROR: vendored file not found: $VENDORED" >&2 + exit 1 +fi + +VENDORED_TOKEN="$(extract_token "$VENDORED")" +if [[ -z "$VENDORED_TOKEN" ]]; then + echo "FAIL: no '$TOKEN_RE' token in $VENDORED" >&2 + echo " The vendored RemoteMcpClient must carry a sync token." >&2 + exit 1 +fi + +# Locate the canonical mcp-loader.ts across common dev layouts. +CANDIDATES=( + "${PI_EXTENSIONS_DIR:-}/extensions/mcp-loader.ts" + "$REPO_ROOT/../pi-extensions/extensions/mcp-loader.ts" + "/workspace/pi-extensions/extensions/mcp-loader.ts" + "$HOME/src/pi-extensions/extensions/mcp-loader.ts" + "$HOME/src/src_local/pi-extensions/extensions/mcp-loader.ts" + "/opt/pi-extensions/extensions/mcp-loader.ts" +) +CANONICAL="" +for c in "${CANDIDATES[@]}"; do + [[ -n "$c" && -f "$c" ]] && { CANONICAL="$c"; break; } +done + +if [[ -z "$CANONICAL" ]]; then + echo "SKIP: canonical pi-extensions/extensions/mcp-loader.ts not found — cannot compare." + echo " (vendored token: $VENDORED_TOKEN). Set PI_EXTENSIONS_DIR to enable the check." + exit 0 +fi + +CANONICAL_TOKEN="$(extract_token "$CANONICAL")" +if [[ -z "$CANONICAL_TOKEN" ]]; then + echo "FAIL: no '$TOKEN_RE' token in canonical $CANONICAL" >&2 + exit 1 +fi + +if [[ "$VENDORED_TOKEN" != "$CANONICAL_TOKEN" ]]; then + echo "FAIL: MCP client sync token drift." >&2 + echo " vendored ($VENDORED): $VENDORED_TOKEN" >&2 + echo " canonical ($CANONICAL): $CANONICAL_TOKEN" >&2 + echo " A streamable-HTTP protocol change landed in one file but not the other." >&2 + echo " Reconcile RemoteMcpClient in both, then bump the token in BOTH to match." >&2 + exit 1 +fi + +echo "OK: MCP client sync token matches ($VENDORED_TOKEN)." +echo " vendored: $VENDORED" +echo " canonical: $CANONICAL"