okengine 0.12.0 → 0.14.1

package/src/elements/ai/mcp-client.ts

@@ -0,0 +1,362 @@
+/**
+ * MCP client session — modern `_meta` first, dual-era fallback, allowlist filter.
+ */
+
+import type { AiMcpServerDecl } from "./declare.ts";
+import { parseMcpToolRef, type ParsedMcpToolRef } from "../../manifest/mcp-ref.ts";
+import { createMcpHttpTransport } from "./mcp-http.ts";
+import { createMcpStdioTransport } from "./mcp-stdio.ts";
+import {
+  interpretToolsCallResult,
+  isModernProtocolReject,
+  MCP_CLIENT_PROTOCOL_VERSION,
+  MCP_LEGACY_PROTOCOL_VERSION,
+  mcpClientMeta,
+  mcpHeaderAnnotationsValid,
+  mcpParamHeaders,
+  parseToolsListResult,
+  type McpClientMeta,
+  type McpListedTool,
+  type McpProtocolEra,
+} from "./mcp-protocol.ts";
+import { McpTransportError, type McpTransport, type McpWireResult } from "./mcp-transport.ts";
+
+/** Options for {@link createMcpClient}. */
+export interface CreateMcpClientOptions {
+  readonly servers?: readonly AiMcpServerDecl[];
+  readonly resolveSecret?: (name: string) => string | Promise<string>;
+  readonly transports?: Readonly<Record<string, McpTransport>>;
+  readonly clientVersion?: string;
+}
+
+/** Allowlisted tool + cached schema. */
+export interface McpResolvedTool {
+  readonly server: string;
+  readonly tool: string;
+  readonly ref: string;
+  readonly listed?: McpListedTool;
+}
+
+/** MCP client used by the AI runtime. */
+export interface McpClient {
+  readonly servers: ReadonlyMap<string, AiMcpServerDecl>;
+  listedTool(ref: ParsedMcpToolRef, signal?: AbortSignal): Promise<McpListedTool | undefined>;
+  call(ref: string, input: unknown, signal?: AbortSignal): Promise<unknown>;
+  close(): Promise<void>;
+}
+
+interface ServerSession {
+  readonly decl: AiMcpServerDecl;
+  readonly transport: McpTransport;
+  readonly eraKey: string;
+  readonly cancelMode: "http" | "stdio";
+  era?: McpProtocolEra;
+  sessionId?: string;
+  listed?: readonly McpListedTool[];
+}
+
+/**
+ * Create an MCP client for declared servers.
+ *
+ * @param options - Decls + optional injected transports (tests)
+ */
+export function createMcpClient(options: CreateMcpClientOptions = {}): McpClient {
+  const meta = mcpClientMeta(options.clientVersion ?? process.env.npm_package_version ?? "0.0.0");
+  const servers = new Map<string, AiMcpServerDecl>();
+  for (const decl of options.servers ?? []) servers.set(decl.name, decl);
+  const sessions = new Map<string, ServerSession>();
+  const eraByKey = new Map<string, McpProtocolEra>();
+  let nextId = 1;
+
+  function sessionFor(decl: AiMcpServerDecl): ServerSession {
+    const existing = sessions.get(decl.name);
+    if (existing) return existing;
+    const injected = options.transports?.[decl.name];
+    const transport = injected ?? transportFor(decl, options.resolveSecret);
+    const eraKey = decl.url
+      ? httpOrigin(decl.url)
+      : `stdio:${decl.command}:${(decl.args ?? []).join("\0")}`;
+    const session: ServerSession = {
+      decl,
+      transport,
+      eraKey,
+      cancelMode: decl.url ? "http" : "stdio",
+      era: eraByKey.get(eraKey),
+    };
+    sessions.set(decl.name, session);
+    return session;
+  }
+
+  async function rpc(
+    session: ServerSession,
+    method: string,
+    params: unknown,
+    signal: AbortSignal | undefined,
+    extra: { readonly toolName?: string; readonly extraHeaders?: Record<string, string> } = {},
+  ): Promise<unknown> {
+    const id = nextId++;
+    const era = session.era ?? "modern";
+    const headers = requestHeaders(
+      method,
+      era,
+      session.sessionId,
+      extra.toolName,
+      extra.extraHeaders,
+    );
+    const wireParams = era === "modern" ? withMeta(params, meta) : params;
+    const result = await session.transport.request({
+      id,
+      method,
+      ...(wireParams !== undefined ? { params: wireParams } : {}),
+      headers,
+      ...(signal !== undefined ? { signal } : {}),
+      cancelMode: session.cancelMode,
+    });
+    if (result.sessionId) session.sessionId = result.sessionId;
+    return result;
+  }
+
+  async function ensureEra(session: ServerSession, signal?: AbortSignal): Promise<void> {
+    if (session.era) return;
+    const cached = eraByKey.get(session.eraKey);
+    if (cached) {
+      session.era = cached;
+      if (cached === "legacy") await initializeLegacy(session, signal);
+      return;
+    }
+
+    const probeMethod = session.cancelMode === "stdio" ? "server/discover" : "tools/list";
+    const probeSignal = session.cancelMode === "stdio" ? withTimeout(signal, 2_000) : signal;
+    let first: McpWireResult;
+    try {
+      first = (await rpc(session, probeMethod, {}, probeSignal)) as McpWireResult;
+    } catch (err) {
+      if (session.cancelMode === "stdio") {
+        setEra(session, "legacy");
+        await initializeLegacy(session, signal);
+        return;
+      }
+      throw err;
+    }
+    if (first.ok) {
+      setEra(session, "modern");
+      if (probeMethod === "tools/list") {
+        session.listed = filterAllowlist(session.decl, parseToolsListResult(first.result));
+      }
+      return;
+    }
+
+    const decision = classifyFirstError(first, session.cancelMode);
+    if (decision === "modern") {
+      setEra(session, "modern");
+      throw rpcFailure(first, probeMethod);
+    }
+    if (decision === "fail") {
+      throw rpcFailure(first, probeMethod);
+    }
+
+    setEra(session, "legacy");
+    await initializeLegacy(session, signal);
+  }
+
+  async function initializeLegacy(session: ServerSession, signal?: AbortSignal): Promise<void> {
+    const result = (await rpc(
+      session,
+      "initialize",
+      {
+        protocolVersion: MCP_LEGACY_PROTOCOL_VERSION,
+        capabilities: {},
+        clientInfo: meta.clientInfo,
+      },
+      signal,
+    )) as McpWireResult;
+    if (!result.ok) throw rpcFailure(result, "initialize");
+    if (result.sessionId) session.sessionId = result.sessionId;
+    // `initialized` notification is optional; HTTP clients do not need it.
+  }
+
+  function setEra(session: ServerSession, era: McpProtocolEra): void {
+    session.era = era;
+    eraByKey.set(session.eraKey, era);
+  }
+
+  async function listTools(
+    session: ServerSession,
+    signal?: AbortSignal,
+  ): Promise<readonly McpListedTool[]> {
+    if (session.listed) return session.listed;
+    await ensureEra(session, signal);
+    if (session.listed) return session.listed;
+    const result = (await rpc(session, "tools/list", {}, signal)) as McpWireResult;
+    if (!result.ok) throw rpcFailure(result, "tools/list");
+    session.listed = filterAllowlist(session.decl, parseToolsListResult(result.result));
+    return session.listed;
+  }
+
+  return {
+    servers,
+    async listedTool(ref, signal) {
+      const decl = servers.get(ref.server);
+      if (!decl) return undefined;
+      const listed = await listTools(sessionFor(decl), signal);
+      return listed.find((t) => t.name === ref.tool);
+    },
+    async call(ref, input, signal) {
+      const parsed = parseMcpToolRef(ref);
+      if (!parsed) throw new Error(`ai.mcp: invalid tool ref "${ref}"`);
+      const decl = servers.get(parsed.server);
+      if (!decl) throw new Error(`ai.mcp: unknown server "${parsed.server}"`);
+      if (!decl.tools.includes(parsed.tool)) {
+        throw new Error(`ai.mcp: tool "${parsed.tool}" is not allowlisted on "${parsed.server}"`);
+      }
+      const session = sessionFor(decl);
+      await ensureEra(session, signal);
+      const listed = await listTools(session, signal);
+      const tool = listed.find((t) => t.name === parsed.tool);
+      const args =
+        input && typeof input === "object" && !Array.isArray(input)
+          ? (input as Record<string, unknown>)
+          : {};
+      const split = mcpParamHeaders(tool?.inputSchema, args);
+      if (split === null) {
+        throw new Error(`ai.mcp: tool "${parsed.tool}" has an invalid x-mcp-header annotation`);
+      }
+      const result = (await rpc(
+        session,
+        "tools/call",
+        {
+          name: parsed.tool,
+          arguments: split.body,
+        },
+        signal,
+        {
+          toolName: parsed.tool,
+          extraHeaders: split.headers,
+        },
+      )) as McpWireResult;
+      if (!result.ok) throw rpcFailure(result, "tools/call");
+      return interpretToolsCallResult(result.result);
+    },
+    async close() {
+      await Promise.all([...sessions.values()].map((s) => s.transport.close()));
+      sessions.clear();
+    },
+  };
+}
+
+function transportFor(
+  decl: AiMcpServerDecl,
+  resolveSecret?: (name: string) => string | Promise<string>,
+): McpTransport {
+  if (decl.url) {
+    return {
+      kind: "http",
+      async request(request) {
+        const bearer = decl.auth && resolveSecret ? await resolveSecret(decl.auth) : undefined;
+        const inner = createMcpHttpTransport({
+          url: decl.url!,
+          ...(bearer !== undefined ? { bearer } : {}),
+        });
+        return inner.request(request);
+      },
+      async close() {
+        /* per-request */
+      },
+    };
+  }
+  return createMcpStdioTransport({
+    command: decl.command!,
+    ...(decl.args !== undefined ? { args: decl.args } : {}),
+  });
+}
+
+function withMeta(params: unknown, meta: McpClientMeta): unknown {
+  const base =
+    params && typeof params === "object" && !Array.isArray(params)
+      ? (params as Record<string, unknown>)
+      : {};
+  return { ...base, _meta: meta };
+}
+
+function requestHeaders(
+  method: string,
+  era: McpProtocolEra,
+  sessionId: string | undefined,
+  toolName: string | undefined,
+  extra: Record<string, string> | undefined,
+): Record<string, string> {
+  const headers: Record<string, string> = {
+    "MCP-Protocol-Version":
+      era === "modern" ? MCP_CLIENT_PROTOCOL_VERSION : MCP_LEGACY_PROTOCOL_VERSION,
+    "Mcp-Method": method,
+    ...extra,
+  };
+  if (method === "tools/call" && toolName) headers["Mcp-Name"] = toolName;
+  if (era === "legacy" && sessionId) headers["Mcp-Session-Id"] = sessionId;
+  return headers;
+}
+
+function filterAllowlist(decl: AiMcpServerDecl, listed: McpListedTool[]): McpListedTool[] {
+  const allow = new Set(decl.tools);
+  return listed.filter((t) => allow.has(t.name) && mcpHeaderAnnotationsValid(t.inputSchema));
+}
+
+function classifyFirstError(
+  result: McpWireResult,
+  mode: "http" | "stdio",
+): "modern" | "legacy" | "fail" {
+  const code = result.error?.code;
+  if (code !== undefined && isModernProtocolReject(code)) return "modern";
+  if (mode === "stdio") return "legacy";
+  const status = result.httpStatus;
+  if (status === 401 || status === 403) return "fail";
+  if (status !== undefined && status >= 500) return "fail";
+  if (status === 400) return "legacy";
+  return "fail";
+}
+
+function rpcFailure(result: McpWireResult, method: string): Error {
+  if (result.error) {
+    return new McpTransportError(`ai.mcp: ${method} failed: ${result.error.message}`, {
+      ...(result.httpStatus !== undefined ? { httpStatus: result.httpStatus } : {}),
+      rpc: result.error,
+    });
+  }
+  return new McpTransportError(`ai.mcp: ${method} failed`, {
+    ...(result.httpStatus !== undefined ? { httpStatus: result.httpStatus } : {}),
+    network: result.httpStatus === undefined,
+  });
+}
+
+function httpOrigin(url: string): string {
+  try {
+    return new URL(url).origin;
+  } catch {
+    return url;
+  }
+}
+
+function withTimeout(parent: AbortSignal | undefined, ms: number): AbortSignal {
+  const ctrl = new AbortController();
+  const timer = setTimeout(() => ctrl.abort(new Error("ai.mcp: stdio discover timeout")), ms);
+  const unlink = (): void => {
+    clearTimeout(timer);
+  };
+  if (parent) {
+    if (parent.aborted) {
+      clearTimeout(timer);
+      ctrl.abort(parent.reason);
+    } else {
+      parent.addEventListener(
+        "abort",
+        () => {
+          clearTimeout(timer);
+          ctrl.abort(parent.reason);
+        },
+        { once: true },
+      );
+    }
+  }
+  ctrl.signal.addEventListener("abort", unlink, { once: true });
+  return ctrl.signal;
+}

package/src/elements/ai/mcp-http.ts

@@ -0,0 +1,159 @@
+/**
+ * Streamable HTTP MCP transport — POST JSON-RPC, optional SSE response.
+ *
+ * Cancel = abort `fetch` / close the SSE stream. Never POST
+ * `notifications/cancelled` on HTTP.
+ */
+
+import { abortError } from "../../kernel/abort-scope.ts";
+import { parseJsonRpcResponse } from "./mcp-protocol.ts";
+import {
+  McpTransportError,
+  type McpTransport,
+  type McpTransportRequest,
+  type McpWireResult,
+} from "./mcp-transport.ts";
+
+/** Options for {@link createMcpHttpTransport}. */
+export interface CreateMcpHttpTransportOptions {
+  readonly url: string;
+  readonly bearer?: string;
+  readonly fetchImpl?: typeof fetch;
+}
+
+/**
+ * Create a Streamable HTTP transport for one MCP origin.
+ *
+ * @param options - URL + optional bearer
+ */
+export function createMcpHttpTransport(options: CreateMcpHttpTransportOptions): McpTransport {
+  const fetchImpl = options.fetchImpl ?? fetch;
+  return {
+    kind: "http",
+    async request(request: McpTransportRequest): Promise<McpWireResult> {
+      if (request.signal?.aborted) throw abortError(request.signal.reason);
+      const headers: Record<string, string> = {
+        "content-type": "application/json",
+        accept: "application/json, text/event-stream",
+        ...request.headers,
+      };
+      if (options.bearer) headers.authorization = `Bearer ${options.bearer}`;
+      const body = JSON.stringify({
+        jsonrpc: "2.0",
+        id: request.id,
+        method: request.method,
+        ...(request.params !== undefined ? { params: request.params } : {}),
+      });
+      let response: Response;
+      try {
+        response = await fetchImpl(options.url, {
+          method: "POST",
+          headers,
+          body,
+          signal: request.signal,
+        });
+      } catch (err) {
+        if (request.signal?.aborted || (err instanceof Error && err.name === "AbortError")) {
+          throw abortError(request.signal?.reason ?? err);
+        }
+        throw new McpTransportError(
+          err instanceof Error ? err.message : "ai.mcp: HTTP request failed",
+          { network: true },
+        );
+      }
+
+      const sessionId = response.headers.get("mcp-session-id") ?? undefined;
+      const ctype = response.headers.get("content-type") ?? "";
+      const rawBody = ctype.includes("text/event-stream")
+        ? await readSseJsonRpc(response, request.id, request.signal)
+        : await readJsonBody(response);
+
+      const parsed = parseJsonRpcResponse(rawBody);
+      if (parsed?.ok) {
+        return {
+          ok: true,
+          result: parsed.result,
+          httpStatus: response.status,
+          ...(sessionId !== undefined ? { sessionId } : {}),
+          rawBody,
+        };
+      }
+      if (parsed && !parsed.ok) {
+        return {
+          ok: false,
+          error: parsed.error,
+          httpStatus: response.status,
+          ...(sessionId !== undefined ? { sessionId } : {}),
+          rawBody,
+        };
+      }
+      return {
+        ok: false,
+        httpStatus: response.status,
+        ...(sessionId !== undefined ? { sessionId } : {}),
+        rawBody,
+      };
+    },
+    async close() {
+      /* HTTP is per-request */
+    },
+  };
+}
+
+async function readJsonBody(response: Response): Promise<unknown> {
+  const text = await response.text();
+  if (text.length === 0) return undefined;
+  try {
+    return JSON.parse(text) as unknown;
+  } catch {
+    return text;
+  }
+}
+
+async function readSseJsonRpc(
+  response: Response,
+  id: string | number,
+  signal?: AbortSignal,
+): Promise<unknown> {
+  let buf = "";
+  let last: unknown;
+  for await (const piece of responseTextStream(response)) {
+    if (signal?.aborted) throw abortError(signal.reason);
+    buf += piece;
+    const parts = buf.split("\n\n");
+    buf = parts.pop() ?? "";
+    for (const part of parts) {
+      const data = part
+        .split("\n")
+        .filter((line) => line.startsWith("data:"))
+        .map((line) => line.slice(5).trimStart())
+        .join("\n");
+      if (!data || data === "[DONE]") continue;
+      try {
+        const json: unknown = JSON.parse(data);
+        if (
+          json !== null &&
+          typeof json === "object" &&
+          !Array.isArray(json) &&
+          (json as { id?: unknown }).id === id
+        ) {
+          last = json;
+        }
+      } catch {
+        /* skip malformed frames */
+      }
+    }
+  }
+  if (last === undefined) {
+    throw new Error("ai.mcp: SSE stream ended without a matching response");
+  }
+  return last;
+}
+
+function responseTextStream(res: Response): AsyncIterable<string> {
+  const stream = (res as Response & { textStream?: () => AsyncIterable<string> }).textStream;
+  if (typeof stream !== "function") {
+    throw new Error("ai.mcp: Response.textStream is required (Bun >= 1.4.0)");
+  }
+  return stream.call(res);
+}

package/src/elements/ai/mcp-mock.ts

@@ -0,0 +1,134 @@
+/**
+ * In-process MCP transport for tests — no child process, no network.
+ */
+
+import { abortError, abortableSleep } from "../../kernel/abort-scope.ts";
+import type { McpListedTool, McpRpcError } from "./mcp-protocol.ts";
+import {
+  type McpTransport,
+  type McpTransportRequest,
+  type McpWireResult,
+} from "./mcp-transport.ts";
+
+/** One tool the mock server exposes. */
+export interface MockMcpTool extends McpListedTool {}
+
+/** Options for {@link createMockMcpTransport}. */
+export interface CreateMockMcpTransportOptions {
+  readonly era?: "modern" | "legacy";
+  readonly tools?: readonly MockMcpTool[];
+  /** HTTP status of the first request (default 200). */
+  readonly httpStatusOnFirst?: number;
+  /** JSON-RPC error on the first request. */
+  readonly firstError?: McpRpcError;
+  /** Unrecognized first-request body (400 without a modern error object). */
+  readonly firstBody?: unknown;
+  readonly onCall?: (name: string, args: unknown, signal?: AbortSignal) => unknown;
+  readonly delayMs?: number;
+  /** `server/discover` hangs until abort / never resolves. */
+  readonly discoverHang?: boolean;
+  readonly extraTools?: readonly MockMcpTool[];
+}
+
+/**
+ * Create a mock MCP transport.
+ *
+ * @param options - Era / tools / first-request fault
+ */
+export function createMockMcpTransport(options: CreateMockMcpTransportOptions = {}): McpTransport {
+  let requestCount = 0;
+  const era = options.era ?? "modern";
+  const tools = [...(options.tools ?? []), ...(options.extraTools ?? [])];
+
+  return {
+    kind: "mock",
+    async request(request: McpTransportRequest): Promise<McpWireResult> {
+      if (request.signal?.aborted) throw abortError(request.signal.reason);
+      requestCount += 1;
+      if (options.delayMs) {
+        await abortableSleep(options.delayMs, request.signal);
+      }
+      if (options.discoverHang && request.method === "server/discover") {
+        await hangUntilAbort(request.signal);
+      }
+
+      if (requestCount === 1 && options.httpStatusOnFirst !== undefined) {
+        if (options.firstError) {
+          return {
+            ok: false,
+            error: options.firstError,
+            httpStatus: options.httpStatusOnFirst,
+            rawBody: { jsonrpc: "2.0", id: request.id, error: options.firstError },
+          };
+        }
+        return {
+          ok: false,
+          httpStatus: options.httpStatusOnFirst,
+          rawBody: options.firstBody,
+        };
+      }
+
+      if (request.method === "initialize") {
+        if (era === "modern") {
+          return {
+            ok: false,
+            error: { code: -32022, message: "UnsupportedProtocolVersion" },
+            httpStatus: 400,
+          };
+        }
+        return {
+          ok: true,
+          result: {
+            protocolVersion: "2024-11-05",
+            capabilities: { tools: {} },
+            serverInfo: { name: "mock", version: "0" },
+          },
+          httpStatus: 200,
+          sessionId: "mock-session",
+        };
+      }
+
+      if (request.method === "server/discover" || request.method === "tools/list") {
+        return {
+          ok: true,
+          result: { tools },
+          httpStatus: 200,
+        };
+      }
+
+      if (request.method === "tools/call") {
+        const params = request.params as { name?: unknown; arguments?: unknown } | undefined;
+        const name = typeof params?.name === "string" ? params.name : "";
+        const output = options.onCall
+          ? await Promise.resolve(options.onCall(name, params?.arguments ?? {}, request.signal))
+          : { content: [{ type: "text", text: `ok:${name}` }] };
+        return { ok: true, result: output, httpStatus: 200 };
+      }
+
+      return {
+        ok: false,
+        error: { code: -32601, message: `method not found: ${request.method}` },
+        httpStatus: 200,
+      };
+    },
+    async close() {
+      /* mock */
+    },
+  };
+}
+
+function hangUntilAbort(signal?: AbortSignal): Promise<never> {
+  return new Promise((_, reject) => {
+    if (signal?.aborted) {
+      reject(abortError(signal.reason));
+      return;
+    }
+    signal?.addEventListener(
+      "abort",
+      () => {
+        reject(abortError(signal.reason));
+      },
+      { once: true },
+    );
+  });
+}