okengine 0.11.2 → 0.13.0

package/src/elements/ai/runtime.ts

@@ -22,6 +22,8 @@ import {
 import {
   AiSchemaValidationError,
   coerceModelObject,
+  promptOutJsonSchema,
+  promptResponseFormat,
   validatePromptOut,
   type AiSchemaMismatch,
 } from "./schema.ts";
@@ -32,6 +34,22 @@ const AI_SAME_MODEL_RETRY_BACKOFF_MS = 250;
 /** Default bound for tool / agent loops. */
 export const AI_DEFAULT_MAX_STEPS = 6;
 
+/**
+ * Split `name` / `name@version` the same way capability pins do.
+ *
+ * @param ref - Prompt id from `fx.ask`
+ */
+export function parsePromptRef(ref: string): {
+  readonly name: string;
+  readonly version?: number;
+} {
+  const at = ref.lastIndexOf("@");
+  if (at <= 0) return { name: ref };
+  const tail = ref.slice(at + 1);
+  if (!/^\d+$/.test(tail)) return { name: ref };
+  return { name: ref.slice(0, at), version: Number(tail) };
+}
+
 /** Recorded agent tool denial (containment proof — not an error). */
 export interface AgentDenial {
   readonly agent: string;
@@ -96,10 +114,50 @@ export interface AiJournalEntry {
   readonly outcome: AiAskOutcome;
   readonly cost: number;
   readonly latencyMs: number;
+  /** Driver-reported input tokens (omitted when the driver did not supply them). */
+  readonly inputTokens?: number;
+  /** Driver-reported output tokens (omitted when the driver did not supply them). */
+  readonly outputTokens?: number;
   readonly schemaMismatch?: AiSchemaMismatch;
   readonly at: number;
 }
 
+/** Optional token counts from a complete() / tool-loop usage bag. */
+type UsageTokens = {
+  readonly inputTokens?: number;
+  readonly outputTokens?: number;
+};
+
+/**
+ * Copy driver token fields when they are numbers — never invent counts.
+ *
+ * @param usage - Complete / loop usage
+ */
+function tokenFields(usage: UsageTokens | undefined): UsageTokens {
+  return {
+    ...(typeof usage?.inputTokens === "number" ? { inputTokens: usage.inputTokens } : {}),
+    ...(typeof usage?.outputTokens === "number" ? { outputTokens: usage.outputTokens } : {}),
+  };
+}
+
+/**
+ * Add one complete() usage onto running token totals.
+ *
+ * @param acc - Mutable totals
+ * @param usage - Driver usage
+ */
+function addUsageTokens(
+  acc: { inputTokens?: number; outputTokens?: number },
+  usage: UsageTokens | undefined,
+): void {
+  if (typeof usage?.inputTokens === "number") {
+    acc.inputTokens = (acc.inputTokens ?? 0) + usage.inputTokens;
+  }
+  if (typeof usage?.outputTokens === "number") {
+    acc.outputTokens = (acc.outputTokens ?? 0) + usage.outputTokens;
+  }
+}
+
 /** Options for {@link createAiRuntime}. */
 export interface CreateAiRuntimeOptions {
   readonly models?: readonly AiModelDecl[];
@@ -254,6 +312,19 @@ export function promptContentFromInput(input: unknown): string {
   return JSON.stringify(masked ?? {});
 }
 
+/**
+ * User message plus a JSON-only contract when `out` is declared.
+ *
+ * @param input - Ask input
+ * @param out - Prompt output schema
+ */
+function askUserContent(input: unknown, out: unknown): string {
+  const base = promptContentFromInput(input);
+  const schema = promptOutJsonSchema(out);
+  if (!schema) return base;
+  return `${base}\nReply with JSON only matching this schema: ${JSON.stringify(schema)}`;
+}
+
 /**
  * Create an AI runtime.
  *
@@ -416,6 +487,8 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
     readonly denials: AgentDenial[];
     readonly steps: number;
     readonly cost: number;
+    readonly inputTokens?: number;
+    readonly outputTokens?: number;
   }> {
     const messages = [...opts.messages];
     const defs = toolDefsFor(opts.tools);
@@ -424,6 +497,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
     const runDenials: AgentDenial[] = [];
     let steps = 0;
     let cost = 0;
+    const tokens: { inputTokens?: number; outputTokens?: number } = {};
     let lastText = "";
     let lastRaw: unknown = {};
     let lastToolResult: unknown;
@@ -438,6 +512,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
         ...(opts.signal !== undefined ? { signal: opts.signal } : {}),
       });
       cost += result.usage?.cost ?? 0;
+      addUsageTokens(tokens, result.usage);
       lastText = result.text;
       lastRaw = result.raw !== undefined ? result.raw : result.text;
 
@@ -452,6 +527,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
           denials: runDenials,
           steps,
           cost,
+          ...tokenFields(tokens),
         };
       }
 
@@ -495,6 +571,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
       denials: runDenials,
       steps,
       cost,
+      ...tokenFields(tokens),
     };
   }
 
@@ -508,8 +585,12 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
     agentRuns,
     journal,
     async ask(prompt, input, opts) {
-      const decl = prompts.get(prompt);
+      const pin = parsePromptRef(prompt);
+      const decl = prompts.get(pin.name) ?? prompts.get(prompt);
       if (!decl) throw new Error(`ai: unknown prompt "${prompt}"`);
+      if (pin.version !== undefined && decl.version !== undefined && pin.version !== decl.version) {
+        throw new Error(`ai: unknown prompt "${prompt}"`);
+      }
       const version = decl.version;
       const started = now();
       const tools = opts?.tools ?? [];
@@ -536,7 +617,16 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
       let lastError: string | undefined;
       let lastSchema: AiSchemaMismatch | undefined;
       let totalCost = 0;
-      const userContent = promptContentFromInput(input);
+      const totalTokens: { inputTokens?: number; outputTokens?: number } = {};
+      const userContent = askUserContent(input, decl.out);
+      const responseFormat = promptResponseFormat(prompt, decl.out);
+
+      const pushJournal = (entry: Omit<AiJournalEntry, "inputTokens" | "outputTokens">): void => {
+        journal.push({
+          ...entry,
+          ...tokenFields(totalTokens),
+        });
+      };
 
       for (const modelName of via) {
         let sameModelTries = 0;
@@ -557,7 +647,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
                 tools,
                 maxSteps: opts?.maxSteps ?? AI_DEFAULT_MAX_STEPS,
                 agentLabel: prompt,
-                responseFormat: decl.out,
+                ...(responseFormat !== undefined ? { responseFormat } : {}),
                 callTool: opts?.callTool,
                 ...(signal !== undefined ? { signal } : {}),
               });
@@ -565,6 +655,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
                 loop.lastToolResult !== undefined && !loop.text ? loop.lastToolResult : loop.raw;
               attemptCost = loop.cost;
               totalCost += attemptCost;
+              addUsageTokens(totalTokens, loop);
               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}`,
@@ -574,11 +665,12 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
               const result = await client.complete({
                 model: wireModel(modelName, client),
                 messages: [{ role: "user", content: userContent }],
-                responseFormat: decl.out,
+                ...(responseFormat !== undefined ? { responseFormat } : {}),
                 ...(signal !== undefined ? { signal } : {}),
               });
               attemptCost = result.usage?.cost ?? 0;
               totalCost += attemptCost;
+              addUsageTokens(totalTokens, result.usage);
               // Prefer assistant text — `raw` is often the transport envelope
               // (OpenAI chat.completion object), which must not shadow the content.
               raw =
@@ -607,7 +699,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
                 at: now(),
               });
               if (journalingForced) {
-                journal.push({
+                pushJournal({
                   prompt,
                   ...(version !== undefined ? { version } : {}),
                   input,
@@ -631,7 +723,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
                   at: now(),
                 });
                 if (journalingForced) {
-                  journal.push({
+                  pushJournal({
                     prompt,
                     ...(version !== undefined ? { version } : {}),
                     input,
@@ -662,7 +754,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
 
             if (!isRetryableAiError(err)) {
               if (journalingForced) {
-                journal.push({
+                pushJournal({
                   prompt,
                   ...(version !== undefined ? { version } : {}),
                   input,
@@ -688,7 +780,7 @@ export function createAiRuntime(options: CreateAiRuntimeOptions = {}): AiRuntime
       }
 
       if (journalingForced) {
-        journal.push({
+        pushJournal({
           prompt,
           ...(version !== undefined ? { version } : {}),
           input,

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

@@ -0,0 +1,133 @@
+/**
+ * Prompt `out` — Zod / JSON Schema / shorthand, plus chat-envelope coerce.
+ */
+
+import { afterEach, describe, expect, test } from "bun:test";
+import { z } from "zod";
+import { mockAiDriver } from "../../drivers/ai-mock.ts";
+import { ai, resetAiDecls } from "./declare.ts";
+import { createAiRuntime } from "./runtime.ts";
+
+afterEach(() => {
+  resetAiDecls();
+});
+import {
+  coerceModelObject,
+  fixtureFromJsonSchema,
+  matchOutSchema,
+  promptOutJsonSchema,
+  promptResponseFormat,
+  validatePromptOut,
+} from "./schema.ts";
+
+describe("promptOutJsonSchema", () => {
+  test("converts Zod object out", () => {
+    const schema = promptOutJsonSchema(z.object({ summary: z.string() }));
+    expect(schema?.properties).toEqual({ summary: { type: "string" } });
+    expect(schema?.required).toEqual(["summary"]);
+  });
+
+  test("passes JSON Schema through", () => {
+    const json = {
+      type: "object",
+      properties: { summary: { type: "string" } },
+      required: ["summary"],
+    };
+    expect(promptOutJsonSchema(json)).toEqual(json);
+  });
+
+  test("lifts { field: type } shorthand", () => {
+    expect(promptOutJsonSchema({ answer: "string" })).toEqual({
+      type: "object",
+      properties: { answer: { type: "string" } },
+      required: ["answer"],
+    });
+  });
+});
+
+describe("validatePromptOut", () => {
+  test("Zod out rejects a mock echo object", () => {
+    expect(() =>
+      validatePromptOut("document-summary", 1, z.object({ summary: z.string() }), {
+        ok: true,
+        echo: "x",
+      }),
+    ).toThrow(/missing \[summary\]/);
+  });
+
+  test("Zod out accepts { summary }", () => {
+    expect(
+      validatePromptOut("document-summary", 1, z.object({ summary: z.string() }), {
+        summary: "API is down",
+      }),
+    ).toEqual({ summary: "API is down" });
+  });
+});
+
+describe("coerceModelObject", () => {
+  test("unwraps chat.completion content JSON", () => {
+    expect(
+      coerceModelObject({
+        object: "chat.completion",
+        choices: [{ message: { content: '{"summary":"ok"}' } }],
+      }),
+    ).toEqual({ summary: "ok" });
+  });
+});
+
+describe("fixtureFromJsonSchema", () => {
+  test("fills required string fields", () => {
+    expect(fixtureFromJsonSchema(z.object({ summary: z.string() }))).toEqual({ summary: "ok" });
+  });
+});
+
+describe("promptResponseFormat", () => {
+  test("wraps out as json_schema", () => {
+    const rf = promptResponseFormat("document-summary", z.object({ summary: z.string() }));
+    expect(rf?.type).toBe("json_schema");
+    expect(rf?.json_schema).toMatchObject({ name: "document-summary", strict: true });
+  });
+});
+
+describe("fx.ask Zod out", () => {
+  test("mock fills declared summary instead of { ok, echo }", async () => {
+    const fast = ai.model("fast", { provider: "mock" });
+    const prompt = fast.prompt("document-summary", {
+      version: 1,
+      out: z.object({ summary: z.string() }),
+    });
+    const runtime = createAiRuntime({
+      models: [fast],
+      prompts: [prompt],
+      defaultDriver: mockAiDriver,
+    });
+    const out = await runtime.ask("document-summary", { title: "API", body: "down" });
+    expect(out.summary).toBe("ok");
+  });
+
+  test("resolves fx.ask name@version against the bare prompt id", async () => {
+    const fast = ai.model("fast", { provider: "mock" });
+    const prompt = fast.prompt("document-summary", {
+      version: 1,
+      out: z.object({ summary: z.string() }),
+    });
+    const runtime = createAiRuntime({
+      models: [fast],
+      prompts: [prompt],
+      defaultDriver: mockAiDriver,
+    });
+    const out = await runtime.ask("document-summary@1", { title: "API", body: "down" });
+    expect(out.summary).toBe("ok");
+    await expect(runtime.ask("document-summary@2", { title: "API" })).rejects.toThrow(
+      'ai: unknown prompt "document-summary@2"',
+    );
+  });
+});
+
+describe("matchOutSchema", () => {
+  test("ignores via on the payload", () => {
+    expect(
+      matchOutSchema(z.object({ summary: z.string() }), { summary: "x", via: "fast" }),
+    ).toBeNull();
+  });
+});

package/src/elements/ai/schema.ts

@@ -83,20 +83,78 @@ export function isObjectSchema(schema: unknown): schema is {
  */
 export function coerceModelObject(value: unknown): Record<string, unknown> {
   if (value && typeof value === "object" && !Array.isArray(value)) {
+    const fromEnvelope = textFromChatEnvelope(value);
+    if (fromEnvelope !== undefined) return coerceModelObject(fromEnvelope);
     return value as Record<string, unknown>;
   }
   if (typeof value === "string") {
+    const parsed = parseJsonObject(value);
+    return parsed ?? { text: value };
+  }
+  return { value };
+}
+
+/**
+ * Normalize a prompt `out` (Zod, JSON Schema, or `{ field: "string" }`
+ * shorthand) into a JSON Schema object for validation and `response_format`.
+ *
+ * @param out - Declared prompt output
+ */
+export function promptOutJsonSchema(out: unknown): Record<string, unknown> | undefined {
+  if (out == null || typeof out !== "object" || Array.isArray(out)) return undefined;
+  if (hasToJSONSchema(out)) {
     try {
-      const parsed = JSON.parse(value) as unknown;
-      if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
-        return parsed as Record<string, unknown>;
-      }
+      const json = out.toJSONSchema();
+      if (isObjectSchema(json) && json.properties) return json as Record<string, unknown>;
     } catch {
-      return { text: value };
+      return undefined;
     }
-    return { text: value };
   }
-  return { value };
+  if (isObjectSchema(out) && out.properties) return out as Record<string, unknown>;
+  const shorthand = shorthandProperties(out);
+  if (shorthand) {
+    return { type: "object", properties: shorthand, required: Object.keys(shorthand) };
+  }
+  return undefined;
+}
+
+/**
+ * Driver `response_format` for a declared `out` — JSON Schema wrapper.
+ * OpenAI-compatible wires this as `json_object` (llama.cpp granite empties
+ * `content` on `json_schema`).
+ *
+ * @param prompt - Prompt name (schema id)
+ * @param out - Declared output
+ */
+export function promptResponseFormat(
+  prompt: string,
+  out: unknown,
+): { readonly type: "json_schema"; readonly json_schema: Record<string, unknown> } | undefined {
+  const schema = promptOutJsonSchema(out);
+  if (!schema) return undefined;
+  const name = prompt.replace(/[^A-Za-z0-9_-]/g, "_").slice(0, 64) || "out";
+  return {
+    type: "json_schema",
+    json_schema: { name, schema, strict: true },
+  };
+}
+
+/**
+ * Deterministic object that satisfies a JSON Schema's required properties.
+ *
+ * @param schema - JSON Schema object
+ */
+export function fixtureFromJsonSchema(schema: unknown): Record<string, unknown> | undefined {
+  const json = promptOutJsonSchema(schema) ?? (isObjectSchema(schema) ? schema : undefined);
+  if (!json?.properties || typeof json.properties !== "object") return undefined;
+  const props = json.properties as Record<string, unknown>;
+  const required = Array.isArray(json.required) ? json.required : Object.keys(props);
+  const out: Record<string, unknown> = {};
+  for (const key of required) {
+    if (typeof key !== "string" || !(key in props)) continue;
+    out[key] = fixtureValue(props[key]);
+  }
+  return Object.keys(out).length > 0 ? out : undefined;
 }
 
 /**
@@ -109,13 +167,15 @@ export function matchOutSchema(
   schema: unknown,
   value: Record<string, unknown>,
 ): AiSchemaMismatch | null {
-  if (!isObjectSchema(schema)) return null;
+  const resolved = promptOutJsonSchema(schema) ?? schema;
+  if (!isObjectSchema(resolved)) return null;
 
-  const props = (schema.properties ?? {}) as Record<string, unknown>;
-  const required = Array.isArray(schema.required) ? [...schema.required] : Object.keys(props);
+  const props = (resolved.properties ?? {}) as Record<string, unknown>;
+  const required = Array.isArray(resolved.required) ? [...resolved.required] : Object.keys(props);
   const keys = new Set(Object.keys(value));
   const missing = required.filter((k) => !keys.has(k));
-  const extra = Object.keys(props).length > 0 ? [...keys].filter((k) => !(k in props)) : [];
+  const extra =
+    Object.keys(props).length > 0 ? [...keys].filter((k) => k !== "via" && !(k in props)) : [];
   const typeMismatches: string[] = [];
   for (const [key, propSchema] of Object.entries(props)) {
     if (!(key in value)) continue;
@@ -162,6 +222,91 @@ export function validatePromptOut(
   return object;
 }
 
+function hasToJSONSchema(value: object): value is { toJSONSchema: () => unknown } {
+  return (
+    "toJSONSchema" in value &&
+    typeof (value as { toJSONSchema?: unknown }).toJSONSchema === "function"
+  );
+}
+
+function shorthandProperties(out: object): Record<string, { type: string }> | undefined {
+  const entries = Object.entries(out);
+  if (entries.length === 0) return undefined;
+  if (hasToJSONSchema(out) || isObjectSchema(out)) return undefined;
+  const properties: Record<string, { type: string }> = {};
+  for (const [key, value] of entries) {
+    if (typeof value !== "string") return undefined;
+    properties[key] = { type: value };
+  }
+  return properties;
+}
+
+function textFromChatEnvelope(value: object): string | undefined {
+  if (!("choices" in value) || !Array.isArray((value as { choices?: unknown }).choices)) {
+    return undefined;
+  }
+  const choice = (value as { choices: unknown[] }).choices[0];
+  if (!choice || typeof choice !== "object") return undefined;
+  const message = (choice as { message?: { content?: unknown; reasoning_content?: unknown } })
+    .message;
+  if (!message) return undefined;
+  if (typeof message.content === "string" && message.content.trim().length > 0) {
+    return message.content;
+  }
+  if (
+    typeof message.reasoning_content === "string" &&
+    message.reasoning_content.trim().length > 0
+  ) {
+    return message.reasoning_content;
+  }
+  return undefined;
+}
+
+function parseJsonObject(text: string): Record<string, unknown> | undefined {
+  const trimmed = text.trim();
+  const candidates = [trimmed];
+  const fenced = /```(?:json)?\s*([\s\S]*?)```/i.exec(trimmed);
+  if (fenced?.[1]) candidates.unshift(fenced[1].trim());
+  const braced = /\{[\s\S]*\}/.exec(trimmed);
+  if (braced?.[0] && braced[0] !== trimmed) candidates.push(braced[0]);
+  for (const candidate of candidates) {
+    try {
+      const parsed: unknown = JSON.parse(candidate);
+      if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
+        return parsed as Record<string, unknown>;
+      }
+    } catch {
+      // try next candidate
+    }
+  }
+  return undefined;
+}
+
+function fixtureValue(propSchema: unknown): unknown {
+  const type =
+    propSchema &&
+    typeof propSchema === "object" &&
+    !Array.isArray(propSchema) &&
+    "type" in propSchema
+      ? String((propSchema as { type?: unknown }).type)
+      : "string";
+  switch (type) {
+    case "number":
+    case "integer":
+      return 0;
+    case "boolean":
+      return true;
+    case "array":
+      return [];
+    case "object":
+      return {};
+    case "null":
+      return null;
+    default:
+      return "ok";
+  }
+}
+
 function matchesJsonType(value: unknown, type: string): boolean {
   switch (type) {
     case "string":

package/src/elements/ai.test.ts

@@ -523,3 +523,56 @@ describe("agent tool trail carries effects; denials are not errors", () => {
     expect(runtime.denials).toHaveLength(1);
   });
 });
+
+describe("ask journal tokens", () => {
+  test("mock driver tokens reach the journal", async () => {
+    const smart = ai.model("smart", { provider: "mock" });
+    const triage = smart.prompt("ticket-triage", { version: 1 });
+    const client = await mockAiDriver.open({
+      model: "mock",
+      mockResponses: { "*": { urgency: "high" } },
+    });
+    const runtime = createAiRuntime({
+      models: [smart],
+      prompts: [triage],
+      clients: { smart: client },
+    });
+
+    await runtime.ask("ticket-triage", { subject: "hello" });
+    expect(runtime.journal).toHaveLength(1);
+    const entry = runtime.journal[0]!;
+    expect(entry.inputTokens).toBeGreaterThan(0);
+    expect(entry.outputTokens).toBeGreaterThan(0);
+    expect(entry.cost).toBe(0);
+  });
+
+  test("token-only complete result journals tokens and omits invented cost", async () => {
+    const smart = ai.model("smart", { provider: "openai-compatible" });
+    const triage = smart.prompt("ticket-triage", { version: 2 });
+    const runtime = createAiRuntime({
+      models: [smart],
+      prompts: [triage],
+      clients: {
+        smart: {
+          driverId: "openai-compatible",
+          model: "gpt-test",
+          async complete() {
+            return {
+              text: JSON.stringify({ ok: true }),
+              raw: { ok: true },
+              model: "gpt-test",
+              driverId: "openai-compatible",
+              usage: { inputTokens: 12, outputTokens: 7 },
+            };
+          },
+        },
+      },
+    });
+
+    await runtime.ask("ticket-triage", { subject: "x" });
+    const entry = runtime.journal[0]!;
+    expect(entry.inputTokens).toBe(12);
+    expect(entry.outputTokens).toBe(7);
+    expect(entry.cost).toBe(0);
+  });
+});

package/src/elements/ai.ts

@@ -37,6 +37,7 @@ export type { AiErrorFields } from "./ai/errors.ts";
 
 export {
   createAiRuntime,
+  parsePromptRef,
   AiSchemaValidationError,
   promptContentFromInput,
   AI_DEFAULT_MAX_STEPS,