@capagents/forgeqe 0.0.0-stage → 1.0.0

package/src/mcp/server.ts

@@ -0,0 +1,234 @@
+import { createHash, randomBytes, timingSafeEqual } from "node:crypto"
+
+import { logEvent } from "../sessionLog.ts"
+import type { McpTool, ToolContext, ToolEvent } from "./tools.ts"
+
+type Json = Record<string, any>
+type Message = { jsonrpc: "2.0"; id?: string | number | null; method?: string; params?: Json; result?: unknown; error?: unknown }
+type Send = (message: Json) => void
+
+const PROTOCOLS = ["2025-06-18", "2025-03-26", "2024-11-05"]
+const PROGRESS_EVERY_MS = 10_000
+
+/** The MCP server name ForgeQE's tools get in the agent runtime, so they appear as forgeqe_* tools. */
+export const MCP_NAME = "forgeqe"
+/** Tool calls may take as long as a whole suite. */
+export const TOOL_TIMEOUT_MS = 6 * 60 * 60 * 1000
+
+export type McpServerOptions = {
+  tools: McpTool[]
+  /** Read on every call, so the shell's current profile and headed setting apply. */
+  context: () => ToolContext
+  instructions?: () => string
+  version?: string
+  progressMs?: number
+}
+
+/** JSON-RPC dispatch for MCP: initialize, ping, tools/list, tools/call, and cancellation. */
+export function createMcpServer(opts: McpServerOptions) {
+  const running = new Map<string | number, AbortController>()
+
+  async function callTool(params: Json, id: string | number, send: Send, outer?: AbortSignal): Promise<Json> {
+    const tool = opts.tools.find((t) => t.name === params.name)
+    if (!tool) return { content: [{ type: "text", text: `Unknown tool ${params.name}` }], isError: true }
+    const base = opts.context()
+    const controller = new AbortController()
+    const abort = () => controller.abort()
+    base.signal?.addEventListener("abort", abort, { once: true })
+    outer?.addEventListener("abort", abort, { once: true })
+    running.set(id, controller)
+    const token = params._meta?.progressToken
+    let tick = 0
+    const timer = token !== undefined ? setInterval(() => send({ jsonrpc: "2.0", method: "notifications/progress", params: { progressToken: token, progress: ++tick, message: `${tool.name} running` } }), opts.progressMs || PROGRESS_EVERY_MS) : null
+    const started = Date.now()
+    logEvent("tool.call", { tool: tool.name, arguments: params.arguments || {} })
+    const finish = (result: Json) => {
+      logEvent("tool.result", { tool: tool.name, ms: Date.now() - started, error: Boolean(result.isError), text: result.content?.[0]?.text })
+      return result
+    }
+    try {
+      const emit = base.emit
+        ? (e: ToolEvent) => {
+            // Browser run events are logged by the engine bridge as engine.event.
+            if (e.type !== "run-event") logEvent("tool.event", { tool: tool.name, ...e })
+            base.emit?.(e)
+          }
+        : undefined
+      const out = await tool.run(params.arguments || {}, { ...base, emit, signal: controller.signal })
+      if (controller.signal.aborted) return finish({ content: [{ type: "text", text: "Stopped by the user." }], isError: true })
+      return finish({ content: [{ type: "text", text: typeof out === "string" ? out : JSON.stringify(out, null, 2) }] })
+    } catch (err) {
+      return finish({ content: [{ type: "text", text: controller.signal.aborted ? "Stopped by the user." : (err as Error).message || String(err) }], isError: true })
+    } finally {
+      if (timer) clearInterval(timer)
+      running.delete(id)
+      base.signal?.removeEventListener("abort", abort)
+      outer?.removeEventListener("abort", abort)
+    }
+  }
+
+  /** The response for a request, or null for notifications and responses. */
+  async function handle(msg: Message, send: Send, signal?: AbortSignal): Promise<Json | null> {
+    if (!msg.method) return null
+    if (msg.id === undefined || msg.id === null) {
+      if (msg.method === "notifications/cancelled") running.get(msg.params?.requestId)?.abort()
+      return null
+    }
+    const ok = (result: unknown) => ({ jsonrpc: "2.0", id: msg.id, result })
+    switch (msg.method) {
+      case "initialize": {
+        const asked = String(msg.params?.protocolVersion || "")
+        return ok({
+          protocolVersion: PROTOCOLS.includes(asked) ? asked : PROTOCOLS[0],
+          capabilities: { tools: { listChanged: false } },
+          serverInfo: { name: "forgeqe", version: opts.version || "0" },
+          ...(opts.instructions ? { instructions: opts.instructions() } : {}),
+        })
+      }
+      case "ping":
+        return ok({})
+      case "tools/list":
+        return ok({ tools: opts.tools.map((t) => ({ name: t.name, description: t.description, inputSchema: t.inputSchema })) })
+      case "tools/call":
+        return ok(await callTool(msg.params || {}, msg.id, send, signal))
+      default:
+        return { jsonrpc: "2.0", id: msg.id, error: { code: -32601, message: `Method not found: ${msg.method}` } }
+    }
+  }
+
+  return { handle, abortAll: () => running.forEach((c) => c.abort()) }
+}
+
+export type McpServer = ReturnType<typeof createMcpServer>
+
+/** The bearer token check, in constant time. */
+export function tokenMatches(header: string | null, token: string): boolean {
+  const digest = (s: string) => createHash("sha256").update(s).digest()
+  return timingSafeEqual(digest(header || ""), digest(`Bearer ${token}`))
+}
+
+/**
+ * A browser page may not call the tools (DNS rebinding): requests that carry an Origin must come from this machine.
+ * MCP clients that are not browsers send no Origin.
+ */
+export function originAllowed(origin: string | null): boolean {
+  if (!origin) return true
+  try {
+    return ["127.0.0.1", "localhost", "[::1]"].includes(new URL(origin).hostname)
+  } catch {
+    return false
+  }
+}
+
+/**
+ * Streamable HTTP on 127.0.0.1 with a bearer token. Tool calls answer as an event stream so progress notifications
+ * and keep-alives flow while a browser run takes minutes; everything else answers as JSON.
+ */
+export function serveHttp(server: McpServer, token: string): { url: string; stop(): void } {
+  const http = Bun.serve({
+    hostname: "127.0.0.1",
+    port: 0,
+    idleTimeout: 0,
+    async fetch(req) {
+      if (new URL(req.url).pathname !== "/mcp") return new Response("Not found", { status: 404 })
+      if (!originAllowed(req.headers.get("origin"))) return new Response("Forbidden", { status: 403 })
+      if (!tokenMatches(req.headers.get("authorization"), token)) return new Response("Unauthorized", { status: 401 })
+      if (req.method === "DELETE") return new Response(null, { status: 200 })
+      if (req.method !== "POST") return new Response(null, { status: 405, headers: { Allow: "POST, DELETE" } })
+      let body: Message | Message[]
+      try {
+        body = await req.json()
+      } catch {
+        return Response.json({ jsonrpc: "2.0", id: null, error: { code: -32700, message: "Parse error" } }, { status: 400 })
+      }
+      const batch = Array.isArray(body) ? body : [body]
+      const requests = batch.filter((m) => m.method && m.id !== undefined && m.id !== null)
+      if (!requests.length) {
+        for (const m of batch) await server.handle(m, () => {})
+        return new Response(null, { status: 202 })
+      }
+      if (!requests.some((m) => m.method === "tools/call")) {
+        const out = (await Promise.all(batch.map((m) => server.handle(m, () => {})))).filter(Boolean)
+        return Response.json(Array.isArray(body) ? out : out[0])
+      }
+      const enc = new TextEncoder()
+      const stream = new ReadableStream({
+        async start(controller) {
+          let open = true
+          const write = (s: string) => {
+            if (!open) return
+            try {
+              controller.enqueue(enc.encode(s))
+            } catch {
+              open = false
+            }
+          }
+          const send: Send = (m) => write(`event: message\ndata: ${JSON.stringify(m)}\n\n`)
+          const keepAlive = setInterval(() => write(": keep-alive\n\n"), 15_000)
+          try {
+            await Promise.all(
+              batch.map(async (m) => {
+                const res = await server.handle(m, send, req.signal)
+                if (res) send(res)
+              }),
+            )
+          } finally {
+            clearInterval(keepAlive)
+            open = false
+            try {
+              controller.close()
+            } catch {}
+          }
+        },
+      })
+      return new Response(stream, { headers: { "Content-Type": "text/event-stream", "Cache-Control": "no-cache" } })
+    },
+  })
+  return { url: `http://127.0.0.1:${http.port}/mcp`, stop: () => http.stop(true) }
+}
+
+/**
+ * ForgeQE's tools on 127.0.0.1 behind a random bearer token, for one agent runtime session: `entry` goes under `mcp` in
+ * its config. Tool calls may take as long as a whole suite; progress notifications keep the connection busy meanwhile.
+ */
+export function localToolServer(opts: McpServerOptions & { timeoutMs: number }): { server: McpServer; entry: Record<string, unknown>; stop(): void } {
+  const server = createMcpServer(opts)
+  const token = randomBytes(24).toString("hex")
+  const http = serveHttp(server, token)
+  return {
+    server,
+    entry: { type: "remote", url: http.url, headers: { Authorization: `Bearer ${token}` }, oauth: false, enabled: true, timeout: opts.timeoutMs },
+    stop: () => http.stop(),
+  }
+}
+
+/** Newline-delimited JSON-RPC on stdin/stdout, for MCP clients that start `forgeqe mcp` themselves. */
+export async function serveStdio(server: McpServer, input: ReadableStream<Uint8Array> = Bun.stdin.stream()): Promise<void> {
+  const send: Send = (m) => process.stdout.write(`${JSON.stringify(m)}\n`)
+  const dec = new TextDecoder()
+  let buf = ""
+  const pending = new Set<Promise<void>>()
+  for await (const chunk of input) {
+    buf += dec.decode(chunk, { stream: true })
+    let nl: number
+    while ((nl = buf.indexOf("\n")) >= 0) {
+      const line = buf.slice(0, nl).trim()
+      buf = buf.slice(nl + 1)
+      if (!line) continue
+      let msg: Message
+      try {
+        msg = JSON.parse(line)
+      } catch {
+        send({ jsonrpc: "2.0", id: null, error: { code: -32700, message: "Parse error" } })
+        continue
+      }
+      const p = server.handle(msg, send).then((res) => {
+        if (res) send(res)
+      })
+      pending.add(p)
+      void p.finally(() => pending.delete(p))
+    }
+  }
+  server.abortAll()
+  await Promise.allSettled([...pending])
+}