okengine 0.5.0 → 0.6.0

package/src/drivers/ai-openai-compatible.ts

@@ -1,8 +1,10 @@
 /**
  * `openai-compatible` AI driver — thin HTTP client for OpenAI-shaped APIs.
  *
- * Covers OpenAI, vLLM, Groq, Together, LM Studio, and most self-hosted
- * servers (unified-theory §16). Injectable `fetch` for tests. Never a
+ * One protocol driver for OpenAI, Groq, Together, OpenRouter, vLLM, LM Studio,
+ * Ollama `/v1`, and other chat/completions endpoints. Configure via
+ * `baseUrl` + `apiKey` + `model` (+ optional `headers`). Native Ollama stays
+ * on the separate `ollama` driver. Injectable `fetch` for tests. Never a
  * production default — prod must declare.
  */
 
@@ -14,23 +16,67 @@ import type {
   AiEmbedResult,
   AiModelClient,
   AiOpenOptions,
+  AiStreamChunk,
+  AiToolCall,
 } from "./ai-types.ts";
 
-const DEFAULT_BASE = "https://api.openai.com/v1";
+/** Default OpenAI cloud base — apiKey is required for this origin. */
+export const OPENAI_COMPAT_DEFAULT_BASE = "https://api.openai.com/v1";
+
+/**
+ * Normalize a chat/embeddings base URL (strip trailing slash).
+ *
+ * @param raw - Base URL
+ */
+export function normalizeOpenaiCompatibleBaseUrl(raw: string): string {
+  return raw.trim().replace(/\/$/, "");
+}
+
+/**
+ * Whether this base is the OpenAI cloud default (key required).
+ *
+ * @param baseUrl - Normalized base
+ */
+export function isOpenaiCloudBase(baseUrl: string): boolean {
+  return normalizeOpenaiCompatibleBaseUrl(baseUrl) === OPENAI_COMPAT_DEFAULT_BASE;
+}
+
+/**
+ * Build request headers: content-type, optional Bearer, optional extras.
+ *
+ * @param apiKey - Bearer token when present
+ * @param extra - Caller headers
+ */
+export function openaiCompatibleHeaders(
+  apiKey: string | undefined,
+  extra?: Readonly<Record<string, string>>,
+): Record<string, string> {
+  const headers: Record<string, string> = {
+    "content-type": "application/json",
+    ...(extra ?? {}),
+  };
+  if (apiKey) headers.Authorization = `Bearer ${apiKey}`;
+  return headers;
+}
 
 /**
  * Open an OpenAI-compatible chat / embeddings client.
  *
- * @param options - API key / model / base URL / injectable fetch
+ * apiKey is required for the default OpenAI cloud base. Custom `baseUrl`
+ * (LM Studio, Ollama `/v1`, Groq, …) may omit the key — Authorization is
+ * then omitted. HTTP failures always throw (never silent mock fallback).
+ *
+ * @param options - API key / model / base URL / headers / injectable fetch
  */
 export async function openOpenaiCompatible(options: AiOpenOptions = {}): Promise<AiModelClient> {
+  const baseUrl = normalizeOpenaiCompatibleBaseUrl(options.baseUrl ?? OPENAI_COMPAT_DEFAULT_BASE);
   const apiKey = options.apiKey ?? process.env.OPENAI_API_KEY;
-  if (!apiKey) {
+  if (isOpenaiCloudBase(baseUrl) && !apiKey) {
     throw new Error("openai-compatible: apiKey is required (or OPENAI_API_KEY)");
   }
   const model = options.model ?? "gpt-4o-mini";
-  const baseUrl = (options.baseUrl ?? DEFAULT_BASE).replace(/\/$/, "");
   const fetchFn = options.fetch ?? globalThis.fetch;
+  const extraHeaders = options.headers;
 
   return {
     driverId: "openai-compatible",
@@ -39,52 +85,82 @@ export async function openOpenaiCompatible(options: AiOpenOptions = {}): Promise
       const resolvedModel = opts.model ?? model;
       const body: Record<string, unknown> = {
         model: resolvedModel,
-        messages: opts.messages.map((m) => ({
-          role: m.role,
-          content: m.content,
-          ...(m.name !== undefined ? { name: m.name } : {}),
-        })),
+        messages: opts.messages.map(serializeMessage),
       };
       if (opts.temperature !== undefined) body.temperature = opts.temperature;
       if (opts.maxTokens !== undefined) body.max_tokens = opts.maxTokens;
       if (opts.responseFormat !== undefined) {
         body.response_format = opts.responseFormat;
       }
+      if (opts.tools !== undefined && opts.tools.length > 0) {
+        body.tools = opts.tools.map((t) => ({
+          type: "function",
+          function: {
+            name: t.name,
+            ...(t.description !== undefined ? { description: t.description } : {}),
+            ...(t.parameters !== undefined ? { parameters: t.parameters } : {}),
+          },
+        }));
+      }
 
       const res = await fetchFn(`${baseUrl}/chat/completions`, {
         method: "POST",
-        headers: {
-          "content-type": "application/json",
-          Authorization: `Bearer ${apiKey}`,
-        },
+        headers: openaiCompatibleHeaders(apiKey, extraHeaders),
         body: JSON.stringify(body),
+        ...(opts.signal !== undefined ? { signal: opts.signal } : {}),
       });
       const raw = (await res.json().catch(() => ({}))) as OpenAiChatResponse;
       if (!res.ok) {
         const msg = raw.error?.message ?? `openai-compatible HTTP ${res.status}`;
         throw new Error(`openai-compatible: ${msg}`);
       }
-      const text = raw.choices?.[0]?.message?.content ?? "";
+      const message = raw.choices?.[0]?.message;
+      const text = message?.content ?? "";
+      const toolCalls = parseToolCalls(message?.tool_calls);
       return {
         text,
         raw,
         model: raw.model ?? resolvedModel,
         driverId: "openai-compatible",
+        ...(toolCalls !== undefined ? { toolCalls } : {}),
         usage: {
           inputTokens: raw.usage?.prompt_tokens,
           outputTokens: raw.usage?.completion_tokens,
         },
       };
     },
+    async *stream(opts: AiCompleteOptions): AsyncIterable<AiStreamChunk> {
+      const resolvedModel = opts.model ?? model;
+      const body: Record<string, unknown> = {
+        model: resolvedModel,
+        messages: opts.messages.map(serializeMessage),
+        stream: true,
+      };
+      if (opts.temperature !== undefined) body.temperature = opts.temperature;
+      if (opts.maxTokens !== undefined) body.max_tokens = opts.maxTokens;
+
+      const res = await fetchFn(`${baseUrl}/chat/completions`, {
+        method: "POST",
+        headers: openaiCompatibleHeaders(apiKey, extraHeaders),
+        body: JSON.stringify(body),
+        ...(opts.signal !== undefined ? { signal: opts.signal } : {}),
+      });
+      if (!res.ok) {
+        const raw = (await res.json().catch(() => ({}))) as OpenAiChatResponse;
+        const msg = raw.error?.message ?? `openai-compatible HTTP ${res.status}`;
+        throw new Error(`openai-compatible: ${msg}`);
+      }
+      if (!res.body) {
+        throw new Error("openai-compatible: stream response has no body");
+      }
+      yield* readOpenaiSse(res.body, opts.signal);
+    },
     async embed(opts: AiEmbedOptions): Promise<AiEmbedResult> {
       const resolvedModel = opts.model ?? model;
       const input = opts.input;
       const res = await fetchFn(`${baseUrl}/embeddings`, {
         method: "POST",
-        headers: {
-          "content-type": "application/json",
-          Authorization: `Bearer ${apiKey}`,
-        },
+        headers: openaiCompatibleHeaders(apiKey, extraHeaders),
         body: JSON.stringify({ model: resolvedModel, input }),
       });
       const raw = (await res.json().catch(() => ({}))) as OpenAiEmbedResponse;
@@ -111,10 +187,124 @@ export const openaiCompatibleAiDriver: AiDriver = {
   open: openOpenaiCompatible,
 };
 
+function serializeMessage(m: AiCompleteOptions["messages"][number]): Record<string, unknown> {
+  const out: Record<string, unknown> = {
+    role: m.role,
+    content: m.content,
+  };
+  if (m.name !== undefined) out.name = m.name;
+  if (m.toolCallId !== undefined) out.tool_call_id = m.toolCallId;
+  if (m.toolCalls !== undefined && m.toolCalls.length > 0) {
+    out.tool_calls = m.toolCalls.map((tc) => ({
+      id: tc.id,
+      type: "function",
+      function: {
+        name: tc.name,
+        arguments:
+          typeof tc.arguments === "string" ? tc.arguments : JSON.stringify(tc.arguments ?? {}),
+      },
+    }));
+  }
+  return out;
+}
+
+function parseToolCalls(
+  raw: readonly OpenAiToolCall[] | undefined,
+): readonly AiToolCall[] | undefined {
+  if (!raw || raw.length === 0) return undefined;
+  return raw.map((tc, i) => {
+    const name = tc.function?.name ?? "";
+    const argStr = tc.function?.arguments ?? "{}";
+    let args: unknown = argStr;
+    try {
+      args = JSON.parse(argStr) as unknown;
+    } catch {
+      args = { _raw: argStr };
+    }
+    return {
+      id: tc.id ?? `call_${i}`,
+      name,
+      arguments: args,
+    };
+  });
+}
+
+/**
+ * Parse OpenAI SSE chat.completion.chunk stream.
+ *
+ * @param body - Response body
+ * @param signal - Optional abort
+ */
+async function* readOpenaiSse(
+  body: ReadableStream<Uint8Array>,
+  signal?: AbortSignal,
+): AsyncGenerator<AiStreamChunk> {
+  const reader = body.getReader();
+  const decoder = new TextDecoder();
+  let buffer = "";
+  try {
+    while (true) {
+      if (signal?.aborted) {
+        throw abortAsError(signal.reason);
+      }
+      const { done, value } = await reader.read();
+      if (done) break;
+      buffer += decoder.decode(value, { stream: true });
+      const lines = buffer.split("\n");
+      buffer = lines.pop() ?? "";
+      for (const line of lines) {
+        const trimmed = line.trim();
+        if (!trimmed.startsWith("data:")) continue;
+        const data = trimmed.slice(5).trim();
+        if (data === "[DONE]") {
+          yield { text: "", done: true };
+          return;
+        }
+        try {
+          const chunk = JSON.parse(data) as {
+            choices?: readonly {
+              delta?: { content?: string | null };
+              finish_reason?: string | null;
+            }[];
+          };
+          const delta = chunk.choices?.[0]?.delta?.content;
+          if (typeof delta === "string" && delta.length > 0) {
+            yield { text: delta };
+          }
+          if (chunk.choices?.[0]?.finish_reason) {
+            yield { text: "", done: true };
+            return;
+          }
+        } catch {
+          // ignore malformed SSE lines
+        }
+      }
+    }
+    yield { text: "", done: true };
+  } finally {
+    reader.releaseLock();
+  }
+}
+
+function abortAsError(reason?: unknown): Error {
+  if (reason instanceof Error) return reason;
+  const err = new Error(reason !== undefined ? String(reason) : "This operation was aborted");
+  err.name = "AbortError";
+  return err;
+}
+
+interface OpenAiToolCall {
+  readonly id?: string;
+  readonly function?: { readonly name?: string; readonly arguments?: string };
+}
+
 interface OpenAiChatResponse {
   readonly model?: string;
   readonly choices?: readonly {
-    readonly message?: { readonly content?: string | null };
+    readonly message?: {
+      readonly content?: string | null;
+      readonly tool_calls?: readonly OpenAiToolCall[];
+    };
   }[];
   readonly usage?: {
     readonly prompt_tokens?: number;

package/src/drivers/ai-providers.test.ts

@@ -4,8 +4,16 @@
 
 import { describe, expect, test } from "bun:test";
 import { anthropicAiDriver, openAnthropic } from "./ai-anthropic.ts";
-import { openaiCompatibleAiDriver, openOpenaiCompatible } from "./ai-openai-compatible.ts";
 import { mockAiDriver } from "./ai-mock.ts";
+import {
+  normalizeOllamaBaseUrl,
+  ollamaAiDriver,
+  OLLAMA_DEFAULT_MODEL,
+  openOllama,
+  OllamaUnavailableError,
+  resolveOllamaModel,
+} from "./ai-ollama.ts";
+import { openaiCompatibleAiDriver, openOpenaiCompatible } from "./ai-openai-compatible.ts";
 
 describe("anthropic driver", () => {
   test("id is anthropic; mock remains the only default elsewhere", () => {
@@ -144,7 +152,7 @@ describe("openai-compatible driver", () => {
     expect(paths.some((p) => p.endsWith("/embeddings"))).toBe(true);
   });
 
-  test("requires apiKey", async () => {
+  test("requires apiKey for default OpenAI cloud base", async () => {
     const prev = process.env.OPENAI_API_KEY;
     delete process.env.OPENAI_API_KEY;
     try {
@@ -153,4 +161,173 @@ describe("openai-compatible driver", () => {
       if (prev !== undefined) process.env.OPENAI_API_KEY = prev;
     }
   });
+
+  test("custom baseUrl may omit apiKey (local / self-hosted)", async () => {
+    const prev = process.env.OPENAI_API_KEY;
+    delete process.env.OPENAI_API_KEY;
+    const headersSeen: Array<Record<string, string> | undefined> = [];
+    try {
+      const fetchFn: typeof fetch = Object.assign(
+        async (_input: string | URL | Request, init?: RequestInit) => {
+          headersSeen.push(init?.headers as Record<string, string> | undefined);
+          return new Response(
+            JSON.stringify({
+              model: "local",
+              choices: [{ message: { content: "ok" } }],
+            }),
+            { status: 200 },
+          );
+        },
+        { preconnect: () => {} },
+      ) as typeof fetch;
+      const client = await openOpenaiCompatible({
+        baseUrl: "http://127.0.0.1:1234/v1",
+        model: "local",
+        fetch: fetchFn,
+      });
+      const result = await client.complete({
+        messages: [{ role: "user", content: "hi" }],
+      });
+      expect(result.text).toBe("ok");
+      expect(headersSeen[0]?.Authorization).toBeUndefined();
+    } finally {
+      if (prev !== undefined) process.env.OPENAI_API_KEY = prev;
+    }
+  });
+
+  test("same driver serves Groq-shaped baseUrl + key + OpenRouter headers", async () => {
+    const calls: Array<{ url: string; headers: Record<string, string> }> = [];
+    const fetchFn: typeof fetch = Object.assign(
+      async (input: string | URL | Request, init?: RequestInit) => {
+        calls.push({
+          url: String(input),
+          headers: init?.headers as Record<string, string>,
+        });
+        return new Response(
+          JSON.stringify({
+            model: "llama-3.1-8b-instant",
+            choices: [{ message: { content: "groq-ok" } }],
+          }),
+          { status: 200 },
+        );
+      },
+      { preconnect: () => {} },
+    ) as typeof fetch;
+
+    const groq = await openOpenaiCompatible({
+      apiKey: "gsk-test",
+      baseUrl: "https://api.groq.com/openai/v1",
+      model: "llama-3.1-8b-instant",
+      fetch: fetchFn,
+    });
+    expect((await groq.complete({ messages: [{ role: "user", content: "x" }] })).text).toBe(
+      "groq-ok",
+    );
+    expect(calls[0]!.url).toBe("https://api.groq.com/openai/v1/chat/completions");
+    expect(calls[0]!.headers.Authorization).toBe("Bearer gsk-test");
+
+    calls.length = 0;
+    const openrouter = await openOpenaiCompatible({
+      apiKey: "or-test",
+      baseUrl: "https://openrouter.ai/api/v1",
+      model: "meta-llama/llama-3.1-8b-instruct",
+      headers: {
+        "HTTP-Referer": "https://example.com",
+        "X-Title": "oke-test",
+      },
+      fetch: fetchFn,
+    });
+    expect((await openrouter.complete({ messages: [{ role: "user", content: "x" }] })).text).toBe(
+      "groq-ok",
+    );
+    expect(calls[0]!.headers["HTTP-Referer"]).toBe("https://example.com");
+    expect(calls[0]!.headers["X-Title"]).toBe("oke-test");
+    expect(calls[0]!.headers.Authorization).toBe("Bearer or-test");
+  });
+
+  test("HTTP error fails loud (no silent mock fallback)", async () => {
+    const fetchFn = Object.assign(
+      async () =>
+        new Response(JSON.stringify({ error: { message: "quota exceeded" } }), {
+          status: 429,
+        }),
+      { preconnect: () => {} },
+    ) as typeof fetch;
+    const client = await openOpenaiCompatible({
+      apiKey: "sk-test",
+      baseUrl: "https://api.together.xyz/v1",
+      fetch: fetchFn,
+    });
+    await expect(client.complete({ messages: [{ role: "user", content: "x" }] })).rejects.toThrow(
+      "quota exceeded",
+    );
+  });
+});
+
+describe("ollama driver", () => {
+  test("id is ollama; documented default model is qwen3.5:9b (overridable)", () => {
+    expect(ollamaAiDriver.id).toBe("ollama");
+    expect(OLLAMA_DEFAULT_MODEL).toBe("qwen3.5:9b");
+    expect(resolveOllamaModel({ model: "llama3.2:1b" })).toBe("llama3.2:1b");
+    expect(normalizeOllamaBaseUrl("localhost:11434")).toBe("http://localhost:11434");
+    expect(normalizeOllamaBaseUrl("http://127.0.0.1:11434/v1")).toBe("http://127.0.0.1:11434");
+  });
+
+  test("complete via injectable fetch — native /api/chat, any model name", async () => {
+    const calls: Array<{ url: string; body: Record<string, unknown> }> = [];
+    const fetchFn: typeof fetch = Object.assign(
+      async (input: string | URL | Request, init?: RequestInit) => {
+        const url = String(input);
+        if (url.endsWith("/api/tags")) {
+          return new Response(JSON.stringify({ models: [] }), { status: 200 });
+        }
+        const body = JSON.parse(String(init?.body)) as Record<string, unknown>;
+        calls.push({ url, body });
+        return new Response(
+          JSON.stringify({
+            model: "llama3.2:1b",
+            message: { role: "assistant", content: "pong" },
+            prompt_eval_count: 4,
+            eval_count: 1,
+          }),
+          { status: 200 },
+        );
+      },
+      { preconnect: () => {} },
+    ) as typeof fetch;
+
+    const client = await openOllama({
+      model: "llama3.2:1b",
+      baseUrl: "http://ollama.test:11434",
+      fetch: fetchFn,
+    });
+    const result = await client.complete({
+      messages: [{ role: "user", content: "ping" }],
+      temperature: 0,
+      maxTokens: 16,
+    });
+
+    expect(result.driverId).toBe("ollama");
+    expect(result.text).toBe("pong");
+    expect(result.usage?.inputTokens).toBe(4);
+    expect(calls).toHaveLength(1);
+    expect(calls[0]!.url).toBe("http://ollama.test:11434/api/chat");
+    expect(calls[0]!.body.model).toBe("llama3.2:1b");
+    expect(calls[0]!.body.stream).toBe(false);
+    expect(calls[0]!.body.think).toBe(false);
+    expect(calls[0]!.body.options).toEqual({ temperature: 0, num_predict: 16 });
+  });
+
+  test("unreachable server throws OllamaUnavailableError (no silent fallback)", async () => {
+    const fetchFn: typeof fetch = Object.assign(
+      async () => {
+        throw new TypeError("connection refused");
+      },
+      { preconnect: () => {} },
+    ) as typeof fetch;
+
+    await expect(
+      openOllama({ baseUrl: "http://127.0.0.1:9", fetch: fetchFn }),
+    ).rejects.toBeInstanceOf(OllamaUnavailableError);
+  });
 });

package/src/drivers/ai-stream.test.ts

@@ -0,0 +1,195 @@
+/**
+ * Real streaming + ambient AbortSignal cancellation (no second mechanism).
+ */
+
+import { describe, expect, test } from "bun:test";
+import { withAbortSignal } from "../kernel/abort-scope.ts";
+import { createFx, createFxContext } from "../kernel/fx.ts";
+import { ai, createAiRuntime } from "../elements/ai.ts";
+import { openOpenaiCompatible } from "./ai-openai-compatible.ts";
+import { openOllama } from "./ai-ollama.ts";
+import { mockAiDriver } from "./ai-mock.ts";
+
+describe("openai-compatible SSE stream", () => {
+  test("parses data: chunks and honours abort", async () => {
+    const sse =
+      'data: {"choices":[{"delta":{"content":"Hel"}}]}\n\n' +
+      'data: {"choices":[{"delta":{"content":"lo"}}]}\n\n' +
+      "data: [DONE]\n\n";
+    let aborted = false;
+    const fetchFn: typeof fetch = Object.assign(
+      async (_input: string | URL | Request, init?: RequestInit) => {
+        const signal = init?.signal;
+        if (signal) {
+          signal.addEventListener("abort", () => {
+            aborted = true;
+          });
+        }
+        return new Response(sse, {
+          status: 200,
+          headers: { "content-type": "text/event-stream" },
+        });
+      },
+      { preconnect: () => {} },
+    ) as typeof fetch;
+
+    const client = await openOpenaiCompatible({
+      apiKey: "sk-test",
+      baseUrl: "https://example.test/v1",
+      fetch: fetchFn,
+    });
+    const parts: string[] = [];
+    for await (const chunk of client.stream!({
+      messages: [{ role: "user", content: "hi" }],
+    })) {
+      if (chunk.text) parts.push(chunk.text);
+    }
+    expect(parts.join("")).toBe("Hello");
+
+    const ac = new AbortController();
+    const iter = client.stream!({
+      messages: [{ role: "user", content: "hi" }],
+      signal: ac.signal,
+    })[Symbol.asyncIterator]();
+    await iter.next();
+    ac.abort();
+    // Consumer abort should mark the signal; driver checks between reads.
+    expect(ac.signal.aborted).toBe(true);
+    void aborted;
+  });
+});
+
+describe("ollama NDJSON stream", () => {
+  test("yields message.content deltas", async () => {
+    const ndjson =
+      JSON.stringify({ message: { content: "A" }, done: false }) +
+      "\n" +
+      JSON.stringify({ message: { content: "B" }, done: true }) +
+      "\n";
+    const fetchFn: typeof fetch = Object.assign(
+      async (input: string | URL | Request) => {
+        const url = String(input);
+        if (url.endsWith("/api/tags")) {
+          return new Response(JSON.stringify({ models: [] }), { status: 200 });
+        }
+        return new Response(ndjson, { status: 200 });
+      },
+      { preconnect: () => {} },
+    ) as typeof fetch;
+
+    const client = await openOllama({
+      model: "llama3.2:1b",
+      baseUrl: "http://ollama.test:11434",
+      fetch: fetchFn,
+    });
+    const parts: string[] = [];
+    for await (const chunk of client.stream!({
+      messages: [{ role: "user", content: "x" }],
+    })) {
+      if (chunk.text) parts.push(chunk.text);
+    }
+    expect(parts.join("")).toBe("AB");
+  });
+});
+
+describe("fx.stream uses ambient AbortSignal", () => {
+  test("race abort cancels in-flight stream fetch", async () => {
+    let sawAbort = false;
+    const smart = ai.model("smart");
+    const runtime = createAiRuntime({
+      models: [smart],
+      defaultDriver: {
+        id: "mock",
+        async open() {
+          return {
+            driverId: "mock" as const,
+            model: "smart",
+            async complete() {
+              return { text: "", model: "smart", driverId: "mock" as const };
+            },
+            async *stream(opts) {
+              const signal = opts.signal;
+              yield { text: "start" };
+              await new Promise<void>((resolve, reject) => {
+                const t = setTimeout(resolve, 5_000);
+                signal?.addEventListener("abort", () => {
+                  sawAbort = true;
+                  clearTimeout(t);
+                  const err = new Error("aborted");
+                  err.name = "AbortError";
+                  reject(err);
+                });
+              });
+              yield { text: "never" };
+            },
+          };
+        },
+      },
+    });
+
+    const fx = createFx({
+      flow: "stream-host",
+      effects: { asks: ["smart"] },
+      aiRuntime: runtime,
+    });
+
+    await expect(
+      fx.race([
+        async () => {
+          for await (const _c of fx.stream("smart", { prompt: "long" })) {
+            // consume until aborted by the race sibling
+          }
+          return "stream";
+        },
+        async () => {
+          await new Promise((r) => setTimeout(r, 20));
+          return "winner";
+        },
+      ]),
+    ).resolves.toBe("winner");
+
+    // Give the abort listener a tick.
+    await new Promise((r) => setTimeout(r, 30));
+    expect(sawAbort).toBe(true);
+  });
+
+  test("mock driver streams through fx.stream", async () => {
+    const smart = ai.model("smart");
+    const runtime = createAiRuntime({
+      models: [smart],
+      defaultDriver: mockAiDriver,
+      clients: {
+        smart: await mockAiDriver.open({
+          model: "smart",
+          mockResponses: { "*": "abcdef" },
+        }),
+      },
+    });
+    const { fx } = createFxContext({
+      flow: "s",
+      effects: { asks: ["smart"] },
+      aiRuntime: runtime,
+    });
+    const parts: string[] = [];
+    for await (const c of fx.stream("smart", { prompt: "hi" })) {
+      parts.push(c);
+    }
+    expect(parts.join("")).toBe("abcdef");
+  });
+
+  test("withAbortSignal aborts cooperative mock stream", async () => {
+    const client = await mockAiDriver.open({
+      mockResponses: { "*": "abcdefghij" },
+    });
+    const ac = new AbortController();
+    await withAbortSignal(ac.signal, async () => {
+      const iter = client.stream!({
+        messages: [{ role: "user", content: "x" }],
+        signal: ac.signal,
+      })[Symbol.asyncIterator]();
+      await iter.next();
+      ac.abort();
+      await expect(iter.next()).rejects.toThrow();
+    });
+  });
+});