okengine 0.22.0 → 0.23.1

package/src/elements/ai/stream-turn.ts

@@ -0,0 +1,214 @@
+/**
+ * One model turn for an agent: stream deltas when the driver can, otherwise
+ * one `complete()` result. The caller journals the assembled result.
+ */
+
+import type {
+  AiCompleteOptions,
+  AiCompleteResult,
+  AiModelClient,
+  AiStreamChunk,
+  AiToolCall,
+} from "../../drivers/ai-types.ts";
+import type { AgentEventEmit } from "./events.ts";
+
+/** A model turn plus whether live deltas were already emitted. */
+export interface ModelTurn {
+  readonly result: AiCompleteResult;
+  /** True when text and tool-call events were emitted from stream chunks. */
+  readonly streamed: boolean;
+}
+
+interface ToolAcc {
+  id: string;
+  name: string;
+  args: string;
+  started: boolean;
+  ended: boolean;
+  buffered: string;
+}
+
+/**
+ * Run one model turn. Streams when `emit` is set and the client implements
+ * `stream`. Otherwise returns `complete()` and leaves events to the caller.
+ *
+ * @param client - Opened model client
+ * @param options - Completion options, including tools
+ * @param emit - Live AG-UI sink. Absent means do not stream.
+ * @param nextMessageId - Allocates `m-N` ids for emitted text
+ */
+export async function readModelTurn(
+  client: AiModelClient,
+  options: AiCompleteOptions,
+  emit: AgentEventEmit | undefined,
+  nextMessageId: () => string,
+): Promise<ModelTurn> {
+  if (emit && client.stream) {
+    const result = await consumeStream(client, options, emit, nextMessageId);
+    return { result, streamed: true };
+  }
+  const result = await client.complete(options);
+  return { result, streamed: false };
+}
+
+async function consumeStream(
+  client: AiModelClient,
+  options: AiCompleteOptions,
+  emit: AgentEventEmit,
+  nextMessageId: () => string,
+): Promise<AiCompleteResult> {
+  const stream = client.stream;
+  if (!stream) {
+    return client.complete(options);
+  }
+  let text = "";
+  let messageId: string | undefined;
+  let textOpen = false;
+  let emittedText = false;
+  const tools = new Map<number, ToolAcc>();
+  let usage: AiCompleteResult["usage"];
+  let raw: unknown;
+  let external: AiCompleteResult["external"];
+  let sawDone = false;
+  const closeText = (): void => {
+    if (!textOpen || messageId === undefined) return;
+    emit({ type: "TEXT_MESSAGE_END", messageId });
+    textOpen = false;
+  };
+  const endTool = (acc: ToolAcc): void => {
+    if (!acc.started || acc.ended) return;
+    emit({ type: "TOOL_CALL_END", toolCallId: acc.id });
+    acc.ended = true;
+  };
+  for await (const chunk of stream.call(client, options)) {
+    if (chunk.raw !== undefined) raw = chunk.raw;
+    if (chunk.external !== undefined) external = chunk.external;
+    if (chunk.done) sawDone = true;
+    applyChunk(chunk, {
+      emit,
+      tools,
+      closeText,
+      endTool,
+      onText(delta) {
+        if (!textOpen) {
+          messageId = nextMessageId();
+          emit({ type: "TEXT_MESSAGE_START", messageId, role: "assistant" });
+          textOpen = true;
+          emittedText = true;
+        }
+        text += delta;
+        emit({ type: "TEXT_MESSAGE_CONTENT", messageId: messageId!, delta });
+      },
+      parentMessageId: () => (emittedText ? messageId : undefined),
+      onUsage(next) {
+        usage = {
+          ...(usage?.inputTokens !== undefined || next.inputTokens !== undefined
+            ? { inputTokens: next.inputTokens ?? usage?.inputTokens }
+            : {}),
+          ...(usage?.outputTokens !== undefined || next.outputTokens !== undefined
+            ? { outputTokens: next.outputTokens ?? usage?.outputTokens }
+            : {}),
+          ...(usage?.cost !== undefined || next.cost !== undefined
+            ? { cost: next.cost ?? usage?.cost }
+            : {}),
+        };
+      },
+    });
+  }
+  closeText();
+  for (const acc of tools.values()) endTool(acc);
+  if (!sawDone) throw new Error("ai: model stream ended before a terminal chunk");
+  const toolCalls = assembleToolCalls(tools);
+  return {
+    text,
+    model: options.model ?? client.model,
+    driverId: client.driverId,
+    ...(raw !== undefined ? { raw } : {}),
+    ...(external !== undefined ? { external } : {}),
+    ...(toolCalls !== undefined ? { toolCalls } : {}),
+    ...(usage !== undefined ? { usage } : {}),
+  };
+}
+
+function applyChunk(
+  chunk: AiStreamChunk,
+  sink: {
+    emit: AgentEventEmit;
+    tools: Map<number, ToolAcc>;
+    closeText: () => void;
+    endTool: (acc: ToolAcc) => void;
+    parentMessageId: () => string | undefined;
+    onText: (delta: string) => void;
+    onUsage: (usage: NonNullable<AiStreamChunk["usage"]>) => void;
+  },
+): void {
+  if (chunk.text) sink.onText(chunk.text);
+  const call = chunk.toolCall;
+  if (call) {
+    sink.closeText();
+    let acc = sink.tools.get(call.index);
+    if (!acc) {
+      acc = {
+        id: call.id ?? `call-${call.index}`,
+        name: call.name ?? "",
+        args: "",
+        started: false,
+        ended: false,
+        buffered: "",
+      };
+      sink.tools.set(call.index, acc);
+    }
+    if (call.id) acc.id = call.id;
+    if (call.name) acc.name = call.name;
+    const delta = call.argumentsDelta ?? "";
+    if (!acc.started && acc.name) {
+      const parentMessageId = sink.parentMessageId();
+      sink.emit({
+        type: "TOOL_CALL_START",
+        toolCallId: acc.id,
+        toolCallName: acc.name,
+        ...(parentMessageId !== undefined ? { parentMessageId } : {}),
+      });
+      acc.started = true;
+      if (acc.buffered) {
+        sink.emit({ type: "TOOL_CALL_ARGS", toolCallId: acc.id, delta: acc.buffered });
+        acc.args += acc.buffered;
+        acc.buffered = "";
+      }
+    }
+    if (delta) {
+      if (acc.started) {
+        acc.args += delta;
+        sink.emit({ type: "TOOL_CALL_ARGS", toolCallId: acc.id, delta });
+      } else {
+        acc.buffered += delta;
+      }
+    }
+  }
+  if (chunk.usage) sink.onUsage(chunk.usage);
+  if (chunk.done) {
+    sink.closeText();
+    for (const acc of sink.tools.values()) sink.endTool(acc);
+  }
+}
+
+function assembleToolCalls(tools: Map<number, ToolAcc>): readonly AiToolCall[] | undefined {
+  if (tools.size === 0) return undefined;
+  return [...tools.entries()]
+    .sort((a, b) => a[0] - b[0])
+    .filter(([, acc]) => acc.name.length > 0)
+    .map(([, acc]) => ({
+      id: acc.id,
+      name: acc.name,
+      arguments: parseArgs(acc.args || acc.buffered),
+    }));
+}
+
+function parseArgs(raw: string): unknown {
+  if (!raw) return {};
+  try {
+    return JSON.parse(raw) as unknown;
+  } catch {
+    return { _raw: raw };
+  }
+}

package/src/elements/ai/stream.live.test.ts

@@ -0,0 +1,76 @@
+/**
+ * Opt-in OpenRouter agent stream.
+ *
+ * Skipped unless `OPENROUTER_API_KEY` is set. `bun run test` ignores
+ * `*.live.test.ts`. The key is never printed.
+ */
+
+import { describe, expect, test } from "bun:test";
+import { readFileSync } from "node:fs";
+import { resolve } from "node:path";
+import { openOpenaiCompatible } from "../../drivers/ai-openai-compatible.ts";
+
+function openRouterKey(): string | undefined {
+  const fromEnv = process.env["OPENROUTER_API_KEY"];
+  if (fromEnv && fromEnv.length > 0) return fromEnv;
+  try {
+    const text = readFileSync(resolve(import.meta.dir, "../../../.env.local"), "utf8");
+    for (const line of text.split("\n")) {
+      const trimmed = line.trim();
+      if (!trimmed.startsWith("OPENROUTER_API_KEY=")) continue;
+      const value = trimmed.slice("OPENROUTER_API_KEY=".length).replace(/^["']|["']$/g, "");
+      if (value.length > 0) return value;
+    }
+  } catch {
+    return undefined;
+  }
+  return undefined;
+}
+
+const apiKey = process.env["CI"] === "true" ? undefined : openRouterKey();
+
+describe.skipIf(apiKey === undefined)("openrouter agent stream live", () => {
+  test("a tool call streams split args and matches complete()", async () => {
+    const client = await openOpenaiCompatible({
+      apiKey,
+      baseUrl: "https://openrouter.ai/api/v1",
+      model: "openrouter/free",
+    });
+    const messages = [
+      {
+        role: "user" as const,
+        content:
+          "Call orders_get with id 14 and note set to exactly: alpha bravo charlie delta echo foxtrot golf hotel india juliet. No prose.",
+      },
+    ];
+    const tools = [
+      {
+        name: "orders_get",
+        description: "Look up an order",
+        parameters: {
+          type: "object",
+          properties: {
+            id: { type: "string" },
+            note: { type: "string" },
+          },
+          required: ["id", "note"],
+        },
+      },
+    ];
+    let name = "";
+    const streamedArgs: string[] = [];
+    for await (const chunk of client.stream!({ messages, tools, model: "openrouter/free" })) {
+      if (chunk.toolCall?.name) name = chunk.toolCall.name;
+      if (chunk.toolCall?.argumentsDelta) streamedArgs.push(chunk.toolCall.argumentsDelta);
+    }
+    expect(name).toBe("orders_get");
+    expect(streamedArgs.length).toBeGreaterThan(1);
+    const assembled = JSON.parse(streamedArgs.join("")) as { id?: string; note?: string };
+    expect(assembled.id).toBe("14");
+    expect(typeof assembled.note).toBe("string");
+    const completed = await client.complete({ messages, tools, model: "openrouter/free" });
+    const call = completed.toolCalls?.[0];
+    expect(call?.name).toBe(name);
+    expect(call?.arguments).toEqual(assembled);
+  }, 60_000);
+});

package/src/elements/ai/subagent.test.ts

@@ -0,0 +1,136 @@
+/**
+ * Agents used as tools — cost rolls up, depth stops, parent run id is stamped.
+ */
+
+import { describe, expect, test } from "bun:test";
+import type { AiCompleteOptions } from "../../drivers/ai-types.ts";
+import { ai, createAiRuntime } from "../ai.ts";
+import { createFxContext } from "../../kernel/fx.ts";
+
+describe("agents as tools", () => {
+  test("child cap is the min of its budget and the parent remaining", async () => {
+    let childCalls = 0;
+    const runtime = createAiRuntime({
+      models: [ai.model("smart")],
+      agents: [
+        ai.agent("lookup", {
+          model: "smart",
+          maxSteps: 4,
+          tools: ["notes.read"],
+          budget: { maxCostPerRun: 1 },
+        }),
+        ai.agent("support", {
+          model: "smart",
+          maxSteps: 2,
+          tools: ["lookup"],
+          budget: { maxCostPerRun: 0.2 },
+        }),
+      ],
+      clients: {
+        smart: {
+          driverId: "mock",
+          model: "smart",
+          async complete(opts: AiCompleteOptions) {
+            const tool = opts.tools?.[0]?.name;
+            if (tool === "lookup") {
+              return {
+                text: "",
+                raw: {},
+                model: "smart",
+                driverId: "mock",
+                usage: { cost: 0 },
+                toolCalls: [{ id: "p1", name: "lookup", arguments: { q: "a" } }],
+              };
+            }
+            childCalls += 1;
+            return {
+              text: "",
+              raw: {},
+              model: "smart",
+              driverId: "mock",
+              usage: { cost: 0.25 },
+              toolCalls: [{ id: `c${childCalls}`, name: "notes.read", arguments: {} }],
+            };
+          },
+        },
+      },
+      callFlow: async () => ({ ok: true }),
+    });
+    const result = await runtime.runAgent("support", { message: "find it" });
+    expect(childCalls).toBe(1);
+    expect(result.stopReason).toBe("budget");
+    expect(result.cost).toBe(0.25);
+  });
+
+  test("depth 3 refuses a fourth level and stamps parentRunId", async () => {
+    const runtime = createAiRuntime({
+      models: [ai.model("smart")],
+      agents: [
+        ai.agent("d", { model: "smart", maxSteps: 1, tools: [] }),
+        ai.agent("c", { model: "smart", maxSteps: 2, tools: ["d"] }),
+        ai.agent("b", { model: "smart", maxSteps: 2, tools: ["c"] }),
+        ai.agent("a", { model: "smart", maxSteps: 2, tools: ["b"] }),
+      ],
+      clients: {
+        smart: {
+          driverId: "mock",
+          model: "smart",
+          async complete(opts: AiCompleteOptions) {
+            const name = opts.tools?.[0]?.name;
+            if (!name) return { text: "leaf", raw: {}, model: "smart", driverId: "mock" };
+            return {
+              text: "",
+              raw: {},
+              model: "smart",
+              driverId: "mock",
+              toolCalls: [{ id: `call-${name}`, name, arguments: {} }],
+            };
+          },
+        },
+      },
+    });
+    await runtime.runAgent("a", { message: "go" });
+    const names = runtime.agentRuns.map((run) => run.agent);
+    expect(names).not.toContain("d");
+    expect(names).toContain("c");
+    const parent = runtime.agentRuns.find((run) => run.agent === "a");
+    const child = runtime.agentRuns.find((run) => run.agent === "b");
+    expect(child?.parentRunId).toBe(parent?.id);
+  });
+
+  test("fx.run of an agent with an agent tool records ask and call", async () => {
+    const lookup = ai.agent("lookup", { model: "smart", tools: [] });
+    const support = ai.agent("support", { model: "smart", tools: [lookup] });
+    const runtime = createAiRuntime({
+      models: [ai.model("smart")],
+      agents: [lookup, support],
+      clients: {
+        smart: {
+          driverId: "mock",
+          model: "smart",
+          async complete(opts: AiCompleteOptions) {
+            if (opts.tools?.[0]?.name === "lookup") {
+              return {
+                text: "",
+                raw: {},
+                model: "smart",
+                driverId: "mock",
+                toolCalls: [{ id: "t1", name: "lookup", arguments: {} }],
+              };
+            }
+            return { text: "ok", raw: {}, model: "smart", driverId: "mock" };
+          },
+        },
+      },
+    });
+    const { fx, ledger } = createFxContext({
+      flow: "assist",
+      effects: { asks: ["support"], calls: ["lookup"] },
+      aiRuntime: runtime,
+    });
+    await fx.run(support, { message: "hi" });
+    const resources = ledger.entries.map((entry) => `${entry.kind}:${entry.resource}`);
+    expect(resources).toContain("ask:support");
+    expect(resources).toContain("call:lookup");
+  });
+});