agent-relay-runner 0.22.0 → 0.24.0

package/src/relay-mcp-proxy.ts

@@ -0,0 +1,383 @@
+import { errMessage, isRecord } from "agent-relay-sdk";
+import { logger } from "./logger";
+
+// Loose fetch signature so tests can inject a plain async stub without Bun's `preconnect`
+// member; the real global `fetch` satisfies it.
+export type FetchLike = (input: string | URL | Request, init?: RequestInit) => Promise<Response>;
+
+// Stage 2 of #213/#215 — the Runner becomes the MCP endpoint the agent connects to,
+// fronting the relay. The agent's MCP client points at this localhost server instead of
+// directly at the relay, so the Runner owns the relay connection, reconnect/backoff, and a
+// durable buffer. A relay restart/crash becomes invisible to the agent.
+//
+// This is a TRANSPARENT JSON-RPC pass-through to the relay's `/api/mcp` (which stays the
+// sole enforcement authority). It intervenes in exactly four places:
+//
+//   1. `initialize`   — forward, then advertise `capabilities.tools.listChanged: true`,
+//                       the live-tool-set capability the relay's static endpoint can't.
+//   2. `tools/list`   — forward (the relay already scope-filters per token), then context-
+//                       NARROW: hide tools that don't apply to this agent's workspace mode/
+//                       state (e.g. `relay_workspace_*` for a non-worktree agent). Strictly
+//                       a subset — never widens. Served from a last-known cache during a blip.
+//   3. `tools/call`   — forward with the runner's LIVE token; for bufferable writes during a
+//                       relay outage, enqueue durably + return a synthetic "queued" ok so the
+//                       call is never lost. Mutating spawn/shutdown are forwarded, never local.
+//   4. `GET` (SSE)    — the server→client notification channel; emits
+//                       `notifications/tools/list_changed` when a transition changes the
+//                       narrowed set, so the agent's tool menu updates mid-session.
+//
+// Narrow-never-widen is the safety invariant: filtering `tools/list` is UX/token-efficiency,
+// NEVER enforcement. The proxy can only ever surface a subset of what the relay's scope filter
+// already returned, so a proxy bug can never become an auth bypass — the relay owns the lock at
+// `tools/call`, the proxy owns the menu.
+
+const PROXY_PATH = "/mcp";
+// Relay-side failures we treat as "relay down" (buffer/serve-cache), as opposed to a real 4xx
+// rejection that must be surfaced to the agent verbatim.
+const GATEWAY_STATUSES = new Set([502, 503, 504]);
+const SSE_KEEPALIVE_MS = 25_000;
+
+// The write tools whose loss during a relay outage is unacceptable and whose result the agent
+// does not need synchronously — safe to queue durably and replay on reconnect. Reads, claims
+// (409 contention), spawn/shutdown (need a real ack) are deliberately NOT bufferable.
+export const DEFAULT_BUFFERABLE_TOOLS = new Set<string>([
+  "relay_send_message",
+  "relay_reply",
+  "relay_workspace_ready",
+]);
+
+// Tools surfaced ONLY to an agent that owns a live isolated worktree. For any other agent the
+// proxy narrows them out of `tools/list` even though the token scope permits them (the relay
+// would return them). This is the context the coarse token scope can't express (#214/#215).
+const WORKTREE_ONLY_TOOLS = new Set<string>([
+  "relay_workspace_status",
+  "relay_workspace_ready",
+  "relay_workspace_deps",
+  "relay_workspace_list",
+  "relay_workspace_claim",
+  "relay_workspace_release",
+  "relay_workspace_land",
+]);
+
+export interface ProxyContext {
+  // The agent owns a live (non-terminal) isolated git worktree → workspace tools apply.
+  isolatedWorktree: boolean;
+}
+
+export interface BufferedToolCall {
+  tool: string;
+  arguments: Record<string, unknown>;
+  idempotencyKey: string;
+}
+
+export interface RelayMcpProxyOptions {
+  // The relay's MCP endpoint, e.g. http://localhost:4850/api/mcp.
+  relayMcpEndpoint: string;
+  // The runner's LIVE relay token (read on every forward so rotation is invisible to the agent).
+  getToken(): string | undefined;
+  // The bearer the agent must present to this localhost proxy (a per-session secret the runner
+  // mints and injects into the agent env). Decouples the agent from the rotating relay token.
+  authSecret: string;
+  // Persist a bufferable write durably for replay on reconnect (wired to the runner outbox).
+  enqueueBuffered(call: BufferedToolCall): void;
+  initialContext?: ProxyContext;
+  bufferableTools?: Set<string>;
+  // Test seam.
+  fetchImpl?: FetchLike;
+}
+
+interface JsonRpcMessage {
+  jsonrpc?: string;
+  id?: string | number | null;
+  method?: string;
+  params?: unknown;
+}
+
+interface SseClient {
+  controller: ReadableStreamDefaultController<Uint8Array>;
+  keepalive: ReturnType<typeof setInterval>;
+}
+
+export class RelayMcpProxy {
+  private readonly relayMcpEndpoint: string;
+  private readonly getToken: () => string | undefined;
+  private readonly authSecret: string;
+  private readonly enqueueBuffered: (call: BufferedToolCall) => void;
+  private readonly bufferableTools: Set<string>;
+  private readonly fetchImpl: FetchLike;
+  private readonly encoder = new TextEncoder();
+
+  private context: ProxyContext;
+  private server?: ReturnType<typeof Bun.serve>;
+  private readonly sseClients = new Set<SseClient>();
+  // Last successful relay tools/list — narrowed and served when the relay is briefly down so a
+  // read still works (reads serve from last-known where safe).
+  private lastRelayTools: Array<Record<string, unknown>> = [];
+  private lastNarrowedNames = "";
+
+  constructor(options: RelayMcpProxyOptions) {
+    this.relayMcpEndpoint = options.relayMcpEndpoint;
+    this.getToken = options.getToken;
+    this.authSecret = options.authSecret;
+    this.enqueueBuffered = options.enqueueBuffered;
+    this.bufferableTools = options.bufferableTools ?? DEFAULT_BUFFERABLE_TOOLS;
+    this.fetchImpl = options.fetchImpl ?? fetch;
+    this.context = options.initialContext ?? { isolatedWorktree: false };
+  }
+
+  start(): { url: string; port: number } {
+    const self = this;
+    this.server = Bun.serve({
+      hostname: "127.0.0.1",
+      port: 0,
+      // SSE streams are long-lived; disable Bun's idle timeout and keep them alive with pings.
+      idleTimeout: 0,
+      fetch(req) {
+        return self.handle(req);
+      },
+    });
+    const port = this.server.port;
+    if (port === undefined) throw new Error("relay MCP proxy did not bind a port");
+    return { url: `http://127.0.0.1:${port}${PROXY_PATH}`, port };
+  }
+
+  stop(): void {
+    for (const client of this.sseClients) {
+      clearInterval(client.keepalive);
+      try { client.controller.close(); } catch { /* already closed */ }
+    }
+    this.sseClients.clear();
+    this.server?.stop(true);
+    this.server = undefined;
+  }
+
+  // The runner calls this on a workspace mode/state transition (active→ready→merged→terminal,
+  // or shared↔worktree). If it changes which tools the agent can see, emit list_changed so the
+  // agent's menu updates mid-session instead of staying frozen until reconnect.
+  setContext(context: ProxyContext): void {
+    this.context = context;
+    this.maybeEmitListChanged();
+  }
+
+  // The runner calls this after re-minting its runtime token (scope may have changed — e.g. a
+  // profile change grants/revokes command:spawn). Re-fetch the relay's now-differently-scoped
+  // tool list with the live token and emit list_changed if the visible set changed. This is the
+  // "token scope transition" path; setContext covers the workspace mode/state path. Best-effort —
+  // a failed refresh keeps the last-known list (the next tools/list refreshes it anyway).
+  async refreshTools(): Promise<void> {
+    const relay = await this.forward({ method: "tools/list", id: 0 }).catch(() => null);
+    const tools = relay && isRecord(relay.body?.result) && Array.isArray((relay.body!.result as Record<string, unknown>).tools)
+      ? ((relay.body!.result as Record<string, unknown>).tools as Array<Record<string, unknown>>)
+      : null;
+    if (!tools) return;
+    this.lastRelayTools = tools;
+    this.maybeEmitListChanged();
+  }
+
+  private async handle(req: Request): Promise<Response> {
+    const url = new URL(req.url);
+    if (url.pathname !== PROXY_PATH) return new Response("not found", { status: 404 });
+    if (!this.authorized(req)) {
+      return Response.json(jsonRpcError(null, -32001, "proxy auth required"), { status: 401 });
+    }
+    // GET → open the server→client SSE notification channel (streamable-HTTP transport).
+    if (req.method === "GET") return this.openSse();
+    if (req.method === "DELETE") return new Response(null, { status: 204 });
+    if (req.method !== "POST") return new Response("method not allowed", { status: 405 });
+    return this.handleRpc(req);
+  }
+
+  private authorized(req: Request): boolean {
+    const header = req.headers.get("authorization") ?? "";
+    const bearer = header.startsWith("Bearer ") ? header.slice(7) : "";
+    return bearer === this.authSecret;
+  }
+
+  private async handleRpc(req: Request): Promise<Response> {
+    let message: JsonRpcMessage;
+    try {
+      const body = await req.json();
+      if (!body || typeof body !== "object" || Array.isArray(body)) {
+        return Response.json(jsonRpcError(null, -32600, "JSON-RPC body must be an object"));
+      }
+      message = body as JsonRpcMessage;
+    } catch {
+      return Response.json(jsonRpcError(null, -32700, "invalid JSON-RPC body"));
+    }
+    const id = message.id ?? null;
+    const method = message.method;
+
+    if (method === "initialize") return this.handleInitialize(id, message);
+    if (method === "tools/list") return this.handleToolsList(id, message);
+    if (method === "tools/call") return this.handleToolsCall(id, message);
+    // Everything else (notifications/initialized, ping, …) forwards verbatim.
+    return this.forwardRaw(message);
+  }
+
+  private async handleInitialize(id: string | number | null, message: JsonRpcMessage): Promise<Response> {
+    const relay = await this.forward({ ...message, method: "initialize", id }).catch((error) => {
+      logger.warn("mcp-proxy", `initialize forward failed: ${errMessage(error)}`);
+      return null;
+    });
+    // Degrade gracefully if the relay is momentarily down at connect: still hand the agent a
+    // usable initialize result advertising the proxy's capabilities.
+    const result: Record<string, unknown> = isRecord(relay?.body?.result) ? { ...(relay!.body!.result as Record<string, unknown>) } : {
+      protocolVersion: "2024-11-05",
+      serverInfo: { name: "agent-relay", title: "Agent Relay (via runner)", version: "proxy" },
+    };
+    const caps = isRecord(result.capabilities) ? { ...result.capabilities } : {};
+    // The capability the relay's static endpoint doesn't advertise: live tool sets.
+    caps.tools = { ...(isRecord(caps.tools) ? caps.tools : {}), listChanged: true };
+    result.capabilities = caps;
+    return Response.json(jsonRpcResult(id, result));
+  }
+
+  private async handleToolsList(id: string | number | null, message: JsonRpcMessage): Promise<Response> {
+    const relay = await this.forward({ ...message, method: "tools/list", id }).catch(() => null);
+    const tools = relay && isRecord(relay.body?.result) && Array.isArray((relay.body!.result as Record<string, unknown>).tools)
+      ? ((relay.body!.result as Record<string, unknown>).tools as Array<Record<string, unknown>>)
+      : null;
+    if (tools) this.lastRelayTools = tools; // refresh the last-known base list
+    // If the relay is down and we have no cache yet, surface an empty list rather than erroring —
+    // the agent can still operate (writes buffer; tools/list refreshes on the next call).
+    const base = tools ?? this.lastRelayTools;
+    const narrowed = this.narrow(base);
+    this.lastNarrowedNames = toolNames(narrowed);
+    return Response.json(jsonRpcResult(id, { tools: narrowed }));
+  }
+
+  // Strict subset of the relay's already-scope-filtered list. Only removes — never adds. The one
+  // narrowing rule today: workspace tools apply only to an agent that owns a live worktree.
+  private narrow(tools: Array<Record<string, unknown>>): Array<Record<string, unknown>> {
+    return tools.filter((tool) => {
+      const name = typeof tool.name === "string" ? tool.name : "";
+      if (WORKTREE_ONLY_TOOLS.has(name) && !this.context.isolatedWorktree) return false;
+      return true;
+    });
+  }
+
+  private async handleToolsCall(id: string | number | null, message: JsonRpcMessage): Promise<Response> {
+    const params = isRecord(message.params) ? message.params : {};
+    const toolName = typeof params.name === "string" ? params.name : "";
+    const args = isRecord(params.arguments) ? params.arguments : {};
+
+    const relay = await this.forward(message).catch((error) => {
+      // A thrown fetch = transport failure (relay unreachable / DNS / connection refused).
+      return { ok: false, status: 0, body: null, transportError: errMessage(error) } as ForwardResult;
+    });
+
+    const relayDown = !relay.ok && (relay.status === 0 || GATEWAY_STATUSES.has(relay.status));
+    if (relayDown && this.bufferableTools.has(toolName)) {
+      // Durably queue the write and tell the agent it's safely accepted. It replays on reconnect.
+      const idempotencyKey = typeof args.idempotencyKey === "string" && args.idempotencyKey
+        ? args.idempotencyKey
+        : `mcp-${toolName}-${crypto.randomUUID()}`;
+      this.enqueueBuffered({ tool: toolName, arguments: { ...args, idempotencyKey }, idempotencyKey });
+      logger.info("mcp-proxy", `relay unreachable — buffered ${toolName} (idempotencyKey=${idempotencyKey})`);
+      return Response.json(jsonRpcResult(id, toolResult({
+        queued: true,
+        tool: toolName,
+        idempotencyKey,
+        note: "Relay was unreachable; your call was queued durably and will be delivered automatically when the relay comes back. It is not lost.",
+      })));
+    }
+
+    if (relay.body) return Response.json(relay.body);
+    // Relay down and not bufferable: a real error the agent must see (and can retry).
+    return Response.json(jsonRpcError(id, -32002, `relay unreachable: ${relay.transportError ?? `status ${relay.status}`}`));
+  }
+
+  // Verbatim forward for methods with no proxy-specific handling. Returns the relay's response
+  // body+status unchanged (or a JSON-RPC error if the relay is down).
+  private async forwardRaw(message: JsonRpcMessage): Promise<Response> {
+    const relay = await this.forward(message).catch((error) => {
+      return { ok: false, status: 0, body: null, transportError: errMessage(error) } as ForwardResult;
+    });
+    if (relay.body) return Response.json(relay.body, { status: relay.status || 200 });
+    // Notifications (no id) tolerate a down relay — ack locally; requests get an error.
+    if (message.id === undefined || message.id === null) return new Response(null, { status: 202 });
+    return Response.json(jsonRpcError(message.id, -32002, `relay unreachable: ${relay.transportError ?? `status ${relay.status}`}`));
+  }
+
+  // POST the JSON-RPC message to the relay with the runner's LIVE token. The agent's incoming
+  // bearer is the proxy secret; we substitute the real relay credential here so the agent never
+  // holds it and token rotation is invisible. The body is re-serialized from the already-parsed
+  // message (the request stream was consumed during dispatch and can't be re-read).
+  private async forward(message: JsonRpcMessage): Promise<ForwardResult> {
+    const token = this.getToken();
+    const headers: Record<string, string> = { "content-type": "application/json" };
+    if (token) headers.authorization = `Bearer ${token}`;
+    const payload: Record<string, unknown> = { jsonrpc: "2.0", method: message.method };
+    if (message.id !== undefined) payload.id = message.id;
+    if (message.params !== undefined) payload.params = message.params;
+    const response = await this.fetchImpl(this.relayMcpEndpoint, { method: "POST", headers, body: JSON.stringify(payload) });
+    const text = await response.text();
+    let parsed: Record<string, unknown> | null = null;
+    if (text) { try { parsed = JSON.parse(text); } catch { parsed = null; } }
+    return { ok: response.ok, status: response.status, body: parsed };
+  }
+
+  private openSse(): Response {
+    const self = this;
+    let client: SseClient;
+    const stream = new ReadableStream<Uint8Array>({
+      start(controller) {
+        controller.enqueue(self.encoder.encode(": connected\n\n"));
+        const keepalive = setInterval(() => {
+          try { controller.enqueue(self.encoder.encode(": keepalive\n\n")); } catch { /* closed */ }
+        }, SSE_KEEPALIVE_MS);
+        keepalive.unref?.();
+        client = { controller, keepalive };
+        self.sseClients.add(client);
+      },
+      cancel() {
+        if (client) {
+          clearInterval(client.keepalive);
+          self.sseClients.delete(client);
+        }
+      },
+    });
+    return new Response(stream, {
+      headers: { "content-type": "text/event-stream", "cache-control": "no-cache", connection: "keep-alive" },
+    });
+  }
+
+  private maybeEmitListChanged(): void {
+    const narrowed = this.narrow(this.lastRelayTools);
+    const names = toolNames(narrowed);
+    if (names === this.lastNarrowedNames) return; // no visible change → no notification
+    this.lastNarrowedNames = names;
+    const frame = this.encoder.encode(`event: message\ndata: ${JSON.stringify({ jsonrpc: "2.0", method: "notifications/tools/list_changed" })}\n\n`);
+    for (const client of this.sseClients) {
+      try { client.controller.enqueue(frame); } catch { /* dropped client; cancel() cleans up */ }
+    }
+    logger.debug("mcp-proxy", `tools/list_changed emitted to ${this.sseClients.size} client(s)`);
+  }
+
+  // Test/observability hooks.
+  sseClientCount(): number { return this.sseClients.size; }
+  narrowedToolNames(): string[] { return this.narrow(this.lastRelayTools).map((t) => String(t.name)); }
+}
+
+interface ForwardResult {
+  ok: boolean;
+  status: number;
+  body: Record<string, unknown> | null;
+  transportError?: string;
+}
+
+function toolNames(tools: Array<Record<string, unknown>>): string {
+  return tools.map((t) => String(t.name)).sort().join(",");
+}
+
+function jsonRpcResult(id: string | number | null, result: unknown): Record<string, unknown> {
+  return { jsonrpc: "2.0", id, result };
+}
+
+function jsonRpcError(id: string | number | null, code: number, message: string): Record<string, unknown> {
+  return { jsonrpc: "2.0", id, error: { code, message } };
+}
+
+function toolResult(result: unknown): Record<string, unknown> {
+  return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }], structuredContent: result };
+}

package/src/relay-mcp.ts

@@ -26,14 +26,17 @@ export function relayMcpEndpoint(relayUrl: string): string {
 // Claude: additive `--mcp-config` JSON (NOT --strict-mcp-config, which would clobber
 // the user's own servers). HTTP transport, token via env-var expansion so it never
 // hits argv. Returns the full ["--mcp-config", "<json>"] arg pair.
-export function relayMcpClaudeConfigArg(relayUrl: string): string[] {
+//
+// `endpoint` overrides the target URL: the runner passes its local MCP proxy URL (Stage 2,
+// #215) so the agent connects to the Runner, not the relay. Omitted → the direct relay endpoint.
+export function relayMcpClaudeConfigArg(relayUrl: string, endpoint?: string): string[] {
   return [
     "--mcp-config",
     JSON.stringify({
       mcpServers: {
         [RELAY_MCP_SERVER_NAME]: {
           type: "http",
-          url: relayMcpEndpoint(relayUrl),
+          url: endpoint ?? relayMcpEndpoint(relayUrl),
           headers: { Authorization: `Bearer \${${RELAY_MCP_TOKEN_ENV}}` },
         },
       },
@@ -43,11 +46,12 @@ export function relayMcpClaudeConfigArg(relayUrl: string): string[] {
 
 // Codex: `-c mcp_servers.<name>.*` overrides. `bearer_token_env_var` tells Codex to
 // read the token from the env var itself → transport resolves to streamable_http.
-export function relayMcpCodexConfigArgs(relayUrl: string): string[] {
+// `endpoint` overrides the target URL (runner-local proxy, Stage 2 #215) — see above.
+export function relayMcpCodexConfigArgs(relayUrl: string, endpoint?: string): string[] {
   const key = `mcp_servers.${RELAY_MCP_SERVER_NAME}`;
   return [
     "-c",
-    `${key}.url=${tomlString(relayMcpEndpoint(relayUrl))}`,
+    `${key}.url=${tomlString(endpoint ?? relayMcpEndpoint(relayUrl))}`,
     "-c",
     `${key}.bearer_token_env_var=${tomlString(RELAY_MCP_TOKEN_ENV)}`,
   ];