okengine 0.13.0 → 0.14.1

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

@@ -0,0 +1,212 @@
+/**
+ * stdio MCP transport — newline-delimited JSON-RPC over a child process.
+ *
+ * Cancel = `notifications/cancelled` then kill the process.
+ */
+
+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";
+
+/** Byte writer used by {@link StdioChild.stdin}. */
+export interface StdioWriter {
+  write(data: Uint8Array): Promise<void> | void;
+  close(): Promise<void> | void;
+}
+
+/** Options for {@link createMcpStdioTransport}. */
+export interface CreateMcpStdioTransportOptions {
+  readonly command: string;
+  readonly args?: readonly string[];
+  readonly spawn?: (cmd: string, args: readonly string[]) => StdioChild;
+}
+
+/** Minimal child-process handle (Bun.spawn-compatible). */
+export interface StdioChild {
+  readonly stdin: StdioWriter | WritableStream<Uint8Array> | null;
+  readonly stdout: ReadableStream<Uint8Array> | null;
+  readonly killed?: boolean;
+  kill(): void;
+}
+
+interface Pending {
+  readonly resolve: (value: McpWireResult) => void;
+  readonly reject: (reason: unknown) => void;
+}
+
+/**
+ * Create a stdio transport. The process is spawned on first request.
+ *
+ * @param options - Command + args (never a shell string)
+ */
+export function createMcpStdioTransport(options: CreateMcpStdioTransportOptions): McpTransport {
+  let child: StdioChild | undefined;
+  let writer: StdioWriter | undefined;
+  const pending = new Map<string, Pending>();
+  let readLoop: Promise<void> | undefined;
+  const encoder = new TextEncoder();
+
+  async function ensureChild(): Promise<void> {
+    if (child) return;
+    const spawn = options.spawn ?? defaultSpawn;
+    child = spawn(options.command, options.args ?? []);
+    if (!child.stdin || !child.stdout) {
+      throw new Error("ai.mcp: stdio process is missing stdin/stdout");
+    }
+    writer = asWriter(child.stdin);
+    readLoop = consumeStdout(child.stdout, pending);
+  }
+
+  return {
+    kind: "stdio",
+    async request(request: McpTransportRequest): Promise<McpWireResult> {
+      if (request.signal?.aborted) throw abortError(request.signal.reason);
+      await ensureChild();
+      if (!writer) throw new Error("ai.mcp: stdio stdin is not writable");
+      const key = String(request.id);
+      const envelope = JSON.stringify({
+        jsonrpc: "2.0",
+        id: request.id,
+        method: request.method,
+        ...(request.params !== undefined ? { params: request.params } : {}),
+      });
+      const result = new Promise<McpWireResult>((resolve, reject) => {
+        pending.set(key, { resolve, reject });
+      });
+      const onAbort = (): void => {
+        void writeLine(
+          writer!,
+          encoder,
+          JSON.stringify({
+            jsonrpc: "2.0",
+            method: "notifications/cancelled",
+            params: { requestId: request.id },
+          }),
+        ).finally(() => {
+          child?.kill();
+          child = undefined;
+          writer = undefined;
+          rejectAll(pending, abortError(request.signal?.reason));
+        });
+      };
+      request.signal?.addEventListener("abort", onAbort, { once: true });
+      try {
+        await writeLine(writer, encoder, envelope);
+        return await result;
+      } finally {
+        request.signal?.removeEventListener("abort", onAbort);
+        pending.delete(key);
+      }
+    },
+    async close() {
+      child?.kill();
+      child = undefined;
+      try {
+        await writer?.close();
+      } catch {
+        /* already closed */
+      }
+      writer = undefined;
+      rejectAll(pending, new Error("ai.mcp: stdio transport closed"));
+      await readLoop?.catch(() => undefined);
+    },
+  };
+}
+
+function defaultSpawn(command: string, args: readonly string[]): StdioChild {
+  const proc = Bun.spawn([command, ...args], {
+    stdin: "pipe",
+    stdout: "pipe",
+    stderr: "pipe",
+  });
+  return {
+    stdin: {
+      write(data) {
+        proc.stdin.write(data);
+      },
+      close() {
+        proc.stdin.end();
+      },
+    },
+    stdout: proc.stdout,
+    kill() {
+      proc.kill();
+    },
+  };
+}
+
+function asWriter(stdin: StdioWriter | WritableStream<Uint8Array>): StdioWriter {
+  if (isWritableStream(stdin)) {
+    const streamWriter = stdin.getWriter();
+    return {
+      write: (data) => streamWriter.write(data),
+      close: () => streamWriter.close(),
+    };
+  }
+  return stdin;
+}
+
+function isWritableStream(
+  value: StdioWriter | WritableStream<Uint8Array>,
+): value is WritableStream<Uint8Array> {
+  return "getWriter" in value && typeof value.getWriter === "function";
+}
+
+async function writeLine(writer: StdioWriter, encoder: TextEncoder, line: string): Promise<void> {
+  await writer.write(encoder.encode(`${line}\n`));
+}
+
+async function consumeStdout(
+  stdout: ReadableStream<Uint8Array>,
+  pending: Map<string, Pending>,
+): Promise<void> {
+  const decoder = new TextDecoder();
+  const reader = stdout.getReader();
+  let buf = "";
+  try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      buf += decoder.decode(value, { stream: true });
+      const lines = buf.split("\n");
+      buf = lines.pop() ?? "";
+      for (const line of lines) {
+        const trimmed = line.trim();
+        if (!trimmed) continue;
+        let json: unknown;
+        try {
+          json = JSON.parse(trimmed) as unknown;
+        } catch {
+          continue;
+        }
+        const parsed = parseJsonRpcResponse(json);
+        if (!parsed) continue;
+        const key = parsed.id === null ? undefined : String(parsed.id);
+        if (key === undefined) continue;
+        const waiter = pending.get(key);
+        if (!waiter) continue;
+        pending.delete(key);
+        if (parsed.ok) {
+          waiter.resolve({ ok: true, result: parsed.result, rawBody: json });
+        } else {
+          waiter.resolve({ ok: false, error: parsed.error, rawBody: json });
+        }
+      }
+    }
+  } catch (err) {
+    rejectAll(pending, err);
+  } finally {
+    reader.releaseLock();
+    rejectAll(pending, new McpTransportError("ai.mcp: stdio process ended", { network: true }));
+  }
+}
+
+function rejectAll(pending: Map<string, Pending>, reason: unknown): void {
+  for (const waiter of pending.values()) waiter.reject(reason);
+  pending.clear();
+}

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

@@ -0,0 +1,70 @@
+/**
+ * MCP client transport contract — HTTP, stdio, and test mock.
+ */
+
+import type { McpRpcError } from "./mcp-protocol.ts";
+
+/** Outbound JSON-RPC request. */
+export interface McpTransportRequest {
+  readonly id: string | number;
+  readonly method: string;
+  readonly params?: unknown;
+  readonly headers: Readonly<Record<string, string>>;
+  readonly signal?: AbortSignal;
+  /** HTTP clients MUST abort fetch / SSE. stdio writes `notifications/cancelled` then kills. */
+  readonly cancelMode: "http" | "stdio";
+}
+
+/** Wire result (HTTP status present only for Streamable HTTP). */
+export interface McpWireResult {
+  readonly ok: boolean;
+  readonly result?: unknown;
+  readonly error?: McpRpcError;
+  readonly httpStatus?: number;
+  readonly sessionId?: string;
+  readonly rawBody?: unknown;
+}
+
+/** Bidirectional MCP transport. */
+export interface McpTransport {
+  readonly kind: "http" | "stdio" | "mock";
+  /**
+   * Send one JSON-RPC request and wait for the matching response.
+   *
+   * @param request - Envelope + headers + abort
+   */
+  request(request: McpTransportRequest): Promise<McpWireResult>;
+  /** Release the process / connection. */
+  close(): Promise<void>;
+}
+
+/**
+ * Transport / protocol failure with enough shape for era fallback.
+ */
+export class McpTransportError extends Error {
+  readonly httpStatus?: number;
+  readonly rpc?: McpRpcError;
+  readonly network?: boolean;
+  readonly timeout?: boolean;
+
+  /**
+   * @param message - Diagnostic
+   * @param extras - HTTP / RPC / network flags
+   */
+  constructor(
+    message: string,
+    extras: {
+      readonly httpStatus?: number;
+      readonly rpc?: McpRpcError;
+      readonly network?: boolean;
+      readonly timeout?: boolean;
+    } = {},
+  ) {
+    super(message);
+    this.name = "McpTransportError";
+    if (extras.httpStatus !== undefined) this.httpStatus = extras.httpStatus;
+    if (extras.rpc !== undefined) this.rpc = extras.rpc;
+    if (extras.network !== undefined) this.network = extras.network;
+    if (extras.timeout !== undefined) this.timeout = extras.timeout;
+  }
+}

package/src/elements/ai/runtime.ts

@@ -8,11 +8,26 @@
  */
 
 import type { AiDriver, AiMessage, AiModelClient, AiToolDef } from "../../drivers/ai-types.ts";
+import { currentAbortSignal, withAbortSignal } from "../../kernel/abort-scope.ts";
+import {
+  mcpCapabilityRefFromName,
+  mcpModelToolName,
+  parseMcpToolRef,
+} from "../../manifest/mcp-ref.ts";
+import { createMcpClient, type McpClient } from "./mcp-client.ts";
+import type { McpTransport } from "./mcp-transport.ts";
 import type { IndexStore } from "../../drivers/types.ts";
 import { maskRedactedDeep } from "../../kernel/redacted.ts";
 import type { GatePolicyContext } from "../gate/declare.ts";
 import type { GateRuntime } from "../gate/runtime.ts";
-import type { AiAgentDecl, AiEmbedDecl, AiModelDecl, AiPromptDecl, AiTimeout } from "./declare.ts";
+import type {
+  AiAgentDecl,
+  AiEmbedDecl,
+  AiMcpServerDecl,
+  AiModelDecl,
+  AiPromptDecl,
+  AiTimeout,
+} from "./declare.ts";
 import {
   isRetryableAiError,
   mergeAskAbortSignal,
@@ -102,7 +117,7 @@ export interface AiFallbackAttempt {
 }
 
 /** Ask outcome class — schema failure is not a provider error. */
-export type AiAskOutcome = "ok" | "provider_error" | "schema_invalid";
+export type AiAskOutcome = "ok" | "provider_error" | "schema_invalid" | "budget_exceeded";
 
 /** Journal entry for a nondeterministic ask. */
 export interface AiJournalEntry {
@@ -171,6 +186,11 @@ export interface CreateAiRuntimeOptions {
   readonly clients?: Readonly<Record<string, AiModelClient>>;
   /** Default driver when a model has no client (usually mock in dev). */
   readonly defaultDriver?: AiDriver;
+  /**
+   * Protocol drivers keyed by id — used when a logical model sets
+   * {@link AiModelDecl.driverId} instead of the app default.
+   */
+  readonly drivers?: Readonly<Record<string, AiDriver>>;
   /** Gate runtime for agent tool calls. */
   readonly gates?: GateRuntime;
   /**
@@ -208,6 +228,12 @@ export interface CreateAiRuntimeOptions {
    * Nondeterministic contract.
    */
   readonly forceJournal?: boolean;
+  /** Declared external MCP servers. */
+  readonly mcpServers?: readonly AiMcpServerDecl[];
+  /** Resolve a vault secret by contract name (MCP bearer). */
+  readonly resolveSecret?: (name: string) => string | Promise<string>;
+  /** Injected MCP transports keyed by server name (tests). */
+  readonly mcpTransports?: Readonly<Record<string, McpTransport>>;
 }
 
 /** Ask options. */
@@ -242,6 +268,8 @@ export interface AiStreamOptions {
   readonly prompt?: string;
   readonly data?: unknown;
   readonly signal?: AbortSignal;
+  /** Ordered fallback models after the primary `stream(model)` name. */
+  readonly via?: readonly string[];
 }
 
 /** AI runtime surface. */
@@ -291,6 +319,14 @@ export interface AiRuntime {
    * @param options - Prompt / data / signal
    */
   stream(model: string, options?: AiStreamOptions): AsyncIterable<string>;
+  /**
+   * Call an allowlisted MCP tool (`mcp:<server>/<tool>`).
+   *
+   * @param ref - Capability ref
+   * @param input - Tool arguments
+   * @param signal - Cancel
+   */
+  callMcp(ref: string, input: unknown, signal?: AbortSignal): Promise<unknown>;
   /**
    * Embed text into the configured index store.
    *
@@ -341,6 +377,11 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
   for (const m of options.models ?? []) models.set(m.name, m);
 
   const clients = new Map<string, AiModelClient>(Object.entries(options.clients ?? {}));
+  const mcpClient: McpClient = createMcpClient({
+    servers: options.mcpServers,
+    ...(options.resolveSecret !== undefined ? { resolveSecret: options.resolveSecret } : {}),
+    ...(options.mcpTransports !== undefined ? { transports: options.mcpTransports } : {}),
+  });
   const denials: AgentDenial[] = [];
   const agentRuns: AgentRunRecord[] = [];
   const journal: AiJournalEntry[] = [];
@@ -351,11 +392,18 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
   async function clientFor(name: string): Promise<AiModelClient> {
     const existing = clients.get(name);
     if (existing) return existing;
-    if (!options.defaultDriver) {
-      throw new Error(`ai: no client for model "${name}" and no defaultDriver`);
-    }
     const model = models.get(name);
-    const opened = await options.defaultDriver.open({
+    const driver =
+      (model?.driverId !== undefined ? options.drivers?.[model.driverId] : undefined) ??
+      options.defaultDriver;
+    if (!driver) {
+      throw new Error(
+        model?.driverId
+          ? `ai: no driver "${model.driverId}" for model "${name}" and no defaultDriver`
+          : `ai: no client for model "${name}" and no defaultDriver`,
+      );
+    }
+    const opened = await driver.open({
       model: model?.model ?? name,
       ...(model?.baseUrl !== undefined ? { baseUrl: model.baseUrl } : {}),
       ...(model?.apiKey !== undefined ? { apiKey: model.apiKey } : {}),
@@ -373,12 +421,26 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
     return options.effectsForFlow?.(tool) ?? [];
   }
 
-  function toolDefsFor(toolNames: readonly string[]): AiToolDef[] {
-    return toolNames.map((name) => ({
-      name,
-      description: `Flow tool: ${name}`,
-      parameters: options.toolSchemaForFlow?.(name) ?? { type: "object", properties: {} },
-    }));
+  async function toolDefsFor(toolNames: readonly string[]): Promise<AiToolDef[]> {
+    const defs: AiToolDef[] = [];
+    for (const name of toolNames) {
+      const mcp = parseMcpToolRef(name);
+      if (mcp) {
+        const listed = await mcpClient.listedTool(mcp);
+        defs.push({
+          name: mcpModelToolName(mcp.server, mcp.tool),
+          description: listed?.description ?? `MCP tool: ${mcp.server}/${mcp.tool}`,
+          parameters: listed?.inputSchema ?? { type: "object", properties: {} },
+        });
+        continue;
+      }
+      defs.push({
+        name,
+        description: `Flow tool: ${name}`,
+        parameters: options.toolSchemaForFlow?.(name) ?? { type: "object", properties: {} },
+      });
+    }
+    return defs;
   }
 
   async function dispatchTool(opts: {
@@ -392,6 +454,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
     readonly meta?: GatePolicyContext["meta"];
     readonly trail: AgentToolStep[];
     readonly runDenials: AgentDenial[];
+    readonly signal?: AbortSignal;
   }): Promise<unknown> {
     const {
       tool,
@@ -404,10 +467,12 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
       meta,
       trail,
       runDenials,
+      signal,
     } = opts;
-    const effects = effectsFor(tool);
+    const capability = mcpCapabilityRefFromName(tool) ?? tool;
+    const effects = effectsFor(capability);
 
-    if (!allowedTools.has(tool)) {
+    if (!allowedTools.has(tool) && !allowedTools.has(capability)) {
       const denial: AgentDenial = {
         agent: agentLabel,
         tool,
@@ -421,7 +486,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
       throw new Error(`ai: model requested unknown tool "${tool}"`);
     }
 
-    const requiredGates = options.gatesForFlow?.(tool) ?? [];
+    const requiredGates = options.gatesForFlow?.(capability) ?? [];
     if (requiredGates.length > 0 && options.gates) {
       const ctx: GatePolicyContext = {
         auth: auth ?? { userId: null, scopes: new Set() },
@@ -460,8 +525,9 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
       return { error: denial.reason, denied: true };
     }
 
-    const output = await invoke(tool, args);
-    trail.push({ tool, status: "ok", effects, at: now() });
+    const invokeSignal = signal ?? currentAbortSignal();
+    const output = await withAbortSignal(invokeSignal, () => invoke(capability, args));
+    trail.push({ tool: capability, status: "ok", effects, at: now() });
     return output;
   }
 
@@ -491,8 +557,12 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
     readonly outputTokens?: number;
   }> {
     const messages = [...opts.messages];
-    const defs = toolDefsFor(opts.tools);
+    const defs = await toolDefsFor(opts.tools);
     const allowed = new Set(opts.tools);
+    for (const name of opts.tools) {
+      const mcp = parseMcpToolRef(name);
+      if (mcp) allowed.add(mcpModelToolName(mcp.server, mcp.tool));
+    }
     const trail: AgentToolStep[] = [];
     const runDenials: AgentDenial[] = [];
     let steps = 0;
@@ -551,6 +621,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
           meta: opts.meta,
           trail,
           runDenials,
+          ...(opts.signal !== undefined ? { signal: opts.signal } : {}),
         });
         lastToolResult = toolResult;
         messages.push({
@@ -594,7 +665,10 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
       const version = decl.version;
       const started = now();
       const tools = opts?.tools ?? [];
-      const signal = mergeAskAbortSignal(resolveTimeoutMs(opts?.timeout ?? decl.timeout));
+      const signal = mergeAskAbortSignal(
+        resolveTimeoutMs(opts?.timeout ?? decl.timeout),
+        currentAbortSignal(),
+      );
 
       // Replay from journal when input matches (nondeterministic contract)
       if (journalingForced && tools.length === 0) {
@@ -628,6 +702,28 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
         });
       };
 
+      const assertAskBudget = (spent: number): void => {
+        const cap = decl.budget?.maxCostPerCall;
+        if (cap === undefined || spent <= cap) return;
+        const message = `ai: prompt "${prompt}" exceeded maxCostPerCall ${cap}`;
+        if (journalingForced) {
+          pushJournal({
+            prompt,
+            ...(version !== undefined ? { version } : {}),
+            input,
+            output: { error: message },
+            attempts: [...attempts],
+            outcome: "budget_exceeded",
+            cost: spent,
+            latencyMs: Math.max(0, now() - started),
+            at: now(),
+          });
+        }
+        const err = new Error(message);
+        err.name = "AiBudgetExceededError";
+        throw err;
+      };
+
       for (const modelName of via) {
         let sameModelTries = 0;
         let advance = true;
@@ -656,6 +752,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
               attemptCost = loop.cost;
               totalCost += attemptCost;
               addUsageTokens(totalTokens, loop);
+              assertAskBudget(totalCost);
               if (loop.denials.length > 0 && loop.trail.every((t) => t.status === "denied")) {
                 throw new Error(
                   `ai: all tool calls denied for prompt "${prompt}": ${loop.denials[0]?.reason}`,
@@ -671,6 +768,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
               attemptCost = result.usage?.cost ?? 0;
               totalCost += attemptCost;
               addUsageTokens(totalTokens, result.usage);
+              assertAskBudget(totalCost);
               // Prefer assistant text — `raw` is often the transport envelope
               // (OpenAI chat.completion object), which must not shadow the content.
               raw =
@@ -815,7 +913,14 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
         auth: runOpts.auth,
         operator: runOpts.operator,
         meta: runOpts.meta,
+        signal: currentAbortSignal(),
       });
+      const runCap = decl.budget?.maxCostPerRun;
+      if (runCap !== undefined && loop.cost > runCap) {
+        const err = new Error(`ai: agent "${agent}" exceeded maxCostPerRun ${runCap}`);
+        err.name = "AiBudgetExceededError";
+        throw err;
+      }
 
       const record: AgentRunRecord = {
         id: `agent-run-${++runSeq}`,
@@ -842,21 +947,46 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
     },
 
     async *stream(model, streamOpts) {
-      const client = await clientFor(model);
-      if (!client.stream) {
-        throw new Error(`ai: model "${model}" (driver ${client.driverId}) does not support stream`);
-      }
+      const via = [model, ...(streamOpts?.via ?? []).filter((name) => name !== model)];
       const content =
         streamOpts?.data !== undefined
           ? promptContentFromInput(streamOpts.data)
           : promptContentFromInput(streamOpts?.prompt ?? "");
-      for await (const chunk of client.stream({
-        model: wireModel(model, client),
-        messages: [{ role: "user", content }],
-        signal: streamOpts?.signal,
-      })) {
-        if (chunk.text) yield chunk.text;
+      let lastError: unknown;
+      for (const modelName of via) {
+        let yielded = false;
+        try {
+          const client = await clientFor(modelName);
+          if (!client.stream) {
+            throw new Error(
+              `ai: model "${modelName}" (driver ${client.driverId}) does not support stream`,
+            );
+          }
+          for await (const chunk of client.stream({
+            model: wireModel(modelName, client),
+            messages: [{ role: "user", content }],
+            signal: streamOpts?.signal,
+          })) {
+            if (chunk.text) {
+              yielded = true;
+              yield chunk.text;
+            }
+          }
+          return;
+        } catch (err) {
+          lastError = err;
+          if (yielded || !isRetryableAiError(err)) {
+            throw err instanceof Error ? err : new Error(String(err));
+          }
+        }
       }
+      throw lastError instanceof Error
+        ? lastError
+        : new Error(`ai: all models failed to stream: ${String(lastError)}`);
+    },
+
+    async callMcp(ref, input, signal) {
+      return mcpClient.call(ref, input, signal ?? currentAbortSignal());
     },
 
     async embed(embedName, id, text) {