@jxsuite/ai 0.37.2 → 0.39.0

package/src/gateway/sse.ts

@@ -0,0 +1,42 @@
+/**
+ * The chat route's framing: one Server-Sent Events `data:` line per frame.
+ *
+ * @module @jxsuite/ai/gateway
+ */
+
+import type { StreamEvent } from "../streaming-client.ts";
+
+/**
+ * The headers a chat stream is answered with. `no-cache` because a replayed stream is a different
+ * turn, and `keep-alive` because the answer is long-lived by design.
+ */
+export const SSE_HEADERS: Readonly<Record<string, string>> = {
+  "Content-Type": "text/event-stream",
+  "Cache-Control": "no-cache",
+  Connection: "keep-alive",
+};
+
+/**
+ * Frame `events` as a Server-Sent Events body: `data: <JSON>` and a blank line per frame, and
+ * nothing else (no `event:`, `id:` or `retry:` field), so the JSON's own `type` is the only
+ * discriminator a reader needs.
+ *
+ * The stream is PUSHED: iteration starts when the stream is constructed, not when it is first read,
+ * and every frame is enqueued as soon as it exists. That is how the proxy has always behaved (its
+ * upstream request is on the wire before the response is returned), and this extraction keeps it.
+ * An iterable that throws errors the stream.
+ *
+ * @param {AsyncIterable<StreamEvent>} events - The frames, in wire order
+ * @returns {ReadableStream<Uint8Array>}
+ */
+export function encodeSse(events: AsyncIterable<StreamEvent>): ReadableStream<Uint8Array> {
+  const encoder = new TextEncoder();
+  return new ReadableStream<Uint8Array>({
+    async start(controller) {
+      for await (const event of events) {
+        controller.enqueue(encoder.encode(`data: ${JSON.stringify(event)}\n\n`));
+      }
+      controller.close();
+    },
+  });
+}

package/src/gateway/types.ts

@@ -0,0 +1,87 @@
+/**
+ * The gateway's contract types: what a host hands `createChatHandler`, and what the handler reports
+ * back to it.
+ *
+ * The split these types draw is the point of the gateway. The WIRE (reading the body, shaping the
+ * upstream request, normalizing the upstream stream, framing SSE, rendering a refusal) is the same
+ * for every host, so it lives here once. The POLICY (whose key a request may use, which base URLs
+ * are safe to reach, how many requests a caller may make) differs per host, so the host supplies it
+ * as `admit` and `resolveUpstream` and the gateway never second-guesses the answer.
+ *
+ * Type-only, so it is on the allowlist in `scripts/check-coverage-manifest.ts`.
+ *
+ * @module @jxsuite/ai/gateway
+ */
+
+import type { ProblemDetails } from "@jxsuite/protocol/problem";
+
+/** Where one chat request is forwarded. A host decides it per request, in `resolveUpstream`. */
+export interface Upstream {
+  /** The wire the upstream speaks. The OpenAI-compatible chat-completions API is the only one. */
+  readonly family: "openai-compat";
+  /** The provider's base URL, without the `/chat/completions` the gateway appends. */
+  readonly baseUrl: string;
+  /** Sent as a bearer token in the `Authorization` header. */
+  readonly apiKey: string;
+  /** The model a request that names none is forwarded with. */
+  readonly defaultModel: string;
+  /** True when the platform brokers the credential, rather than the user supplying one. */
+  readonly managed: boolean;
+  /**
+   * The prompt-cache affinity key for the request's conversation, when the host derived one: the
+   * gateway forwards it in whatever form this upstream understands, and to no other upstream
+   * (`upstreamCacheHints`). A host derives it with `affinityKey(scope, id)` from the client's
+   * `X-Jx-Ai-Session` header, so the raw id never leaves the host.
+   */
+  readonly sessionAffinity?: string;
+}
+
+/**
+ * A request the gateway (or the host's policy) declined, answered as `application/problem+json`
+ * before any upstream is contacted.
+ */
+export interface GatewayRefusal {
+  /** The HTTP status of the refusal. */
+  readonly status: number;
+  /** The RFC 9457 problem document the body carries. */
+  readonly problem: ProblemDetails;
+  /** A machine code for clients that key on `code` rather than on the problem type. */
+  readonly code?: string;
+}
+
+/** What a request the gateway accepted asked for, reported to `onAccepted` for audit. */
+export interface ChatAdmission {
+  /** The wire version of the accepted body. Version 1 is the only one. */
+  readonly wire: 1;
+  /** The model the request is forwarded with: the body's, or the upstream's default. */
+  readonly model: string;
+  /** How many messages the body carried, not counting the system prompt the gateway prepends. */
+  readonly messageCount: number;
+  /** The body's size in bytes, as read rather than as declared. */
+  readonly bytes: number;
+}
+
+/** A host's half of the chat route. Only `resolveUpstream` is required. */
+export interface ChatGatewayOptions<C> {
+  /** Runs first, before the body is read (a rate limit, say). A refusal ends the request. */
+  readonly admit?: (
+    request: Request,
+    context: C,
+  ) => Promise<GatewayRefusal | null> | GatewayRefusal | null;
+  /**
+   * Where the request goes, or why it goes nowhere. This is the host's policy: key provenance and
+   * any guard on a caller-supplied base URL belong here, not in the gateway.
+   */
+  readonly resolveUpstream: (
+    request: Request,
+    context: C,
+  ) => Promise<Upstream | GatewayRefusal> | Upstream | GatewayRefusal;
+  /** The largest body accepted, in bytes actually read. Reading stops once it is passed. */
+  readonly maxBodyBytes?: number;
+  /** The most messages one request may carry. */
+  readonly maxMessages?: number;
+  /** Called once per accepted request, after every check and before the upstream is contacted. */
+  readonly onAccepted?: (admission: ChatAdmission, upstream: Upstream, context: C) => void;
+  /** The fetch the upstream is called with. The global `fetch`, looked up per call, when absent. */
+  readonly fetch?: typeof fetch;
+}

package/src/gateway/upstream-error.ts

@@ -0,0 +1,54 @@
+/**
+ * The sentence inside an upstream provider's error body.
+ *
+ * @module @jxsuite/ai/gateway
+ */
+
+/**
+ * Extract a human-readable message from an upstream error body, falling back to the raw body (or
+ * `fallback` when the body is empty).
+ *
+ * OpenAI-compatible providers return `{ error: "..." }` or `{ error: { message: "..." } }`.
+ * Cloudflare's own REST API (what a BYOK base URL pointed at Workers AI actually returns) uses a
+ * different envelope, `{ success: false, errors: [{ code, message }] }`, with no top-level `error`
+ * key at all. Without this branch that shape parses as valid JSON with nothing matched, and the
+ * caller was left showing the whole raw body (e.g. the full `{"success":false,"errors":[...]}`
+ * blob) instead of the human sentence inside it.
+ *
+ * The chat stream's `error` frame reads it, and so does a host's model catalogue, which reports the
+ * same sentence as `upstreamMessage`.
+ *
+ * @param {string} rawBody - The upstream's response body, as text
+ * @param {string} fallback - What to say when the body is empty (the status text, usually)
+ * @returns {string}
+ */
+export function extractUpstreamErrorMessage(rawBody: string, fallback: string): string {
+  if (!rawBody) {
+    return fallback;
+  }
+  try {
+    const { error, errors } = JSON.parse(rawBody) as {
+      error?: string | { message?: string };
+      errors?: { code?: number; message?: string }[];
+    };
+    if (typeof error === "string") {
+      return error;
+    }
+    if (error) {
+      const { message } = error;
+      if (message) {
+        return message;
+      }
+    }
+    const [first] = errors ?? [];
+    if (first) {
+      const { message } = first;
+      if (message) {
+        return message;
+      }
+    }
+  } catch {
+    /* Not JSON: use the raw body. */
+  }
+  return rawBody;
+}

package/src/index.ts

@@ -1,8 +1,8 @@
 /**
  * Index.js — AI infrastructure for Jx Suite (`@jxsuite/ai`)
  *
- * Re-exports the core modules: streaming client abstraction, tool registry, and
- * reactive chat state management. Provider-agnostic; no Studio or Jx dependencies.
+ * Re-exports the core modules: streaming client abstraction, tool registry, and reactive chat state
+ * management. Provider-agnostic; no Studio or Jx dependencies.
  *
  * @license MIT
  */
@@ -13,6 +13,8 @@ export { createToolDefinition, createToolRegistry } from "./tools.js";
 
 export { createChatState } from "./chat-state.js";
 
+export type { ChatUsage } from "./chat-state.js";
+
 export {
   STREAM_EVENT_TYPES,
   createOpenAIStreamingClient,
@@ -20,6 +22,15 @@ export {
   createProxyStreamingClient,
 } from "./streaming-client.js";
 
+export {
+  AI_SESSION_HEADER,
+  affinityKey,
+  isAiSessionId,
+  upstreamCacheHints,
+} from "./cache-hints.js";
+
+export type { UpstreamCacheHints } from "./cache-hints.js";
+
 export type {
   AnthropicStreamingClientOptions,
   OpenAIStreamingClientOptions,
@@ -33,4 +44,5 @@ export type {
   StreamToolCallDeltaEvent,
   StreamToolCallEndEvent,
   StreamToolCallStartEvent,
+  StreamUsageEvent,
 } from "./streaming-client.js";

package/src/streaming-client.ts

@@ -1,19 +1,30 @@
 /**
  * Streaming-client.js — Provider-agnostic streaming LLM client abstraction
  *
- * Defines the StreamingClient interface, the StreamEvent union every implementation yields,
- * and concrete implementations for OpenAI and Anthropic. Designed upfront so switching
- * providers is a new implementation, not a refactor.
+ * Defines the StreamingClient interface, the StreamEvent union every implementation yields, and
+ * concrete implementations for OpenAI and Anthropic. Designed upfront so switching providers is a
+ * new implementation, not a refactor.
  *
- * The union and its seven members are exported because they are the contract a third-party Studio
+ * The union and its eight members are exported because they are the contract a third-party Studio
  * backend implements for the `ai/chat` route, not an internal detail — see the docs page below.
  *
- * @license MIT
  * @module @jxsuite/ai/streaming-client
+ * @license MIT
  * @docs extending/embedding/backend-protocol
  */
 
 import type { ProblemDetails } from "@jxsuite/protocol";
+import { AI_SESSION_HEADER, upstreamCacheHints } from "./cache-hints.ts";
+
+/* Re-exported so `./streaming-client`, the subpath a backend already imports, carries the cache
+   hints too: a Worker backend hashing the session header needs nothing else from this package. */
+export {
+  AI_SESSION_HEADER,
+  affinityKey,
+  isAiSessionId,
+  upstreamCacheHints,
+} from "./cache-hints.ts";
+export type { UpstreamCacheHints } from "./cache-hints.ts";
 
 export type StreamEvent =
   | StreamDeltaEvent
@@ -21,6 +32,7 @@ export type StreamEvent =
   | StreamToolCallStartEvent
   | StreamToolCallDeltaEvent
   | StreamToolCallEndEvent
+  | StreamUsageEvent
   | StreamDoneEvent
   | StreamErrorEvent;
 
@@ -61,6 +73,27 @@ export interface StreamToolCallEndEvent {
   id: string;
 }
 
+/**
+ * What the provider counted for the request that produced this stream, when it reported a count.
+ *
+ * Emitted at most once, immediately BEFORE `done`, because a reader stops at `done`. It is optional
+ * and additive: a backend with no figure sends none, and a reader that predates it ignores it. It
+ * exists because every budget the client keeps was an estimate of four characters per token, and
+ * the provider's own count, which the client was asking for all along (`include_usage`), was
+ * discarded before it arrived.
+ */
+export interface StreamUsageEvent {
+  type: "usage";
+  /** Tokens the provider read: system prompt, tool schemas and the whole history. */
+  inputTokens: number;
+  /** Tokens the provider generated, reasoning included. */
+  outputTokens: number;
+  /** The part of `inputTokens` served from a prompt cache, when the provider says. */
+  cachedInputTokens?: number;
+  /** The part of `outputTokens` spent on chain-of-thought, when the provider says. */
+  reasoningTokens?: number;
+}
+
 export interface StreamDoneEvent {
   type: "done";
   stopReason: string;
@@ -107,6 +140,8 @@ interface OpenAIChatRequestBody {
   tools?: object[];
   tool_choice?: string;
   parallel_tool_calls?: boolean;
+  /** OpenAI's prompt-cache routing key, sent only to OpenAI itself (see `cache-hints.ts`). */
+  prompt_cache_key?: string;
 }
 
 /** A tool-call fragment inside an OpenAI streaming `delta`. */
@@ -131,6 +166,53 @@ interface OpenAIStreamChunk {
     };
     finish_reason?: string | null;
   }[];
+  /** The count `stream_options.include_usage` asks for — normally a final chunk with no choices. */
+  usage?: OpenAIUsage | null;
+}
+
+/** OpenAI's usage object, as chat-completions streams it. */
+export interface OpenAIUsage {
+  prompt_tokens?: number;
+  completion_tokens?: number;
+  prompt_tokens_details?: { cached_tokens?: number } | null;
+  completion_tokens_details?: { reasoning_tokens?: number } | null;
+}
+
+/**
+ * Normalize OpenAI's usage object into a {@link StreamUsageEvent}, or null when it carries no count.
+ * The two detail figures are included only when the provider sent them: absent means "not said",
+ * which is not the same as zero.
+ */
+export function usageEventFromOpenAI(usage?: OpenAIUsage | null): StreamUsageEvent | null {
+  if (!usage || typeof usage.prompt_tokens !== "number") {
+    return null;
+  }
+  const event: StreamUsageEvent = {
+    type: "usage",
+    inputTokens: usage.prompt_tokens,
+    outputTokens: usage.completion_tokens ?? 0,
+  };
+  const cached = usage.prompt_tokens_details?.cached_tokens;
+  if (typeof cached === "number") {
+    event.cachedInputTokens = cached;
+  }
+  const reasoning = usage.completion_tokens_details?.reasoning_tokens;
+  if (typeof reasoning === "number") {
+    event.reasoningTokens = reasoning;
+  }
+  return event;
+}
+
+/**
+ * The stop reasons the client reports, from an OpenAI `finish_reason`. Anything else (a content
+ * filter, a provider's own spelling) is not a stop the loop acts on, so it maps to null and the
+ * stream keeps reading.
+ */
+function stopReasonFromOpenAI(reason: string | null | undefined): string | null {
+  if (reason === "tool_calls" || reason === "length" || reason === "stop") {
+    return reason;
+  }
+  return null;
 }
 
 /** Error response body — the proxy's flat `{ error: "..." }` or OpenAI's `{ error: { message } }`. */
@@ -199,6 +281,7 @@ export const STREAM_EVENT_TYPES = {
   TOOL_CALL_START: "tool_call_start",
   TOOL_CALL_DELTA: "tool_call_delta",
   TOOL_CALL_END: "tool_call_end",
+  USAGE: "usage",
   DONE: "done",
   ERROR: "error",
 } as const;
@@ -235,6 +318,9 @@ export const STREAM_EVENT_TYPES = {
  * @param {number} [opts.temperature] - Sampling temperature, forwarded only when defined. Omit for
  *   reasoning models (GPT-5.x, o-series) that reject a custom temperature; set 0 for
  *   near-deterministic eval runs.
+ * @param {string} [opts.sessionAffinity] - A stable key for the conversation, forwarded as the
+ *   prompt-cache hint the provider understands (`prompt_cache_key` to OpenAI, `x-session-affinity`
+ *   to Workers AI) and to no other host. See `upstreamCacheHints`.
  * @returns {StreamingClient}
  */
 export interface OpenAIStreamingClientOptions {
@@ -242,6 +328,7 @@ export interface OpenAIStreamingClientOptions {
   apiKey: string;
   model?: string;
   temperature?: number | undefined;
+  sessionAffinity?: string | undefined;
 }
 
 export function createOpenAIStreamingClient({
@@ -249,7 +336,11 @@ export function createOpenAIStreamingClient({
   apiKey,
   model = "gpt-4o",
   temperature,
+  sessionAffinity,
 }: OpenAIStreamingClientOptions): StreamingClient {
+  // Fixed for the client's life, so it is classified once rather than per request.
+  const hints = upstreamCacheHints(baseUrl, sessionAffinity);
+
   /**
    * @param {object[]} messages
    * @param {object[]} tools
@@ -283,6 +374,7 @@ export function createOpenAIStreamingClient({
       body.tool_choice = "auto";
       body.parallel_tool_calls = true;
     }
+    Object.assign(body, hints.body);
 
     let response;
     try {
@@ -291,6 +383,7 @@ export function createOpenAIStreamingClient({
         headers: {
           "Content-Type": "application/json",
           Authorization: `Bearer ${apiKey}`,
+          ...hints.headers,
         },
         body: JSON.stringify(body),
         signal,
@@ -332,8 +425,37 @@ export function createOpenAIStreamingClient({
     const decoder = new TextDecoder();
     let buffer = "";
 
-    /** @type {Map<string, { id: string; name: string; args: string }>} */
     const pendingToolCalls = new Map<number, { id: string; name: string; args: string }>();
+    /* The finish reason is REMEMBERED rather than acted on. With `include_usage`, the count arrives
+       in a chunk of its own after the one carrying `finish_reason`, so returning on the finish
+       dropped the count on every request. The stream ends at `[DONE]` or when the body closes. */
+    let stopReason: string | null = null;
+    let usage: StreamUsageEvent | null = null;
+
+    /**
+     * Close the tool calls still open. Idempotent, so a finish chunk and `[DONE]` can both call it.
+     *
+     * @yields {StreamEvent} A `tool_call_end` per open call.
+     */
+    function* closePendingToolCalls(): Generator<StreamEvent> {
+      for (const tc of pendingToolCalls.values()) {
+        yield { type: "tool_call_end", id: tc.id };
+      }
+      pendingToolCalls.clear();
+    }
+
+    /**
+     * The stream's last frames: any open calls, the count when there is one, then `done`.
+     *
+     * @yields {StreamEvent} The closing frames, in order.
+     */
+    function* finish(): Generator<StreamEvent> {
+      yield* closePendingToolCalls();
+      if (usage) {
+        yield usage;
+      }
+      yield { type: "done", stopReason: stopReason ?? "stop" };
+    }
 
     try {
       while (true) {
@@ -355,12 +477,7 @@ export function createOpenAIStreamingClient({
 
           const dataStr = trimmed.slice(6);
           if (dataStr === "[DONE]") {
-            // Emit any pending tool call ends before done
-            for (const tc of pendingToolCalls.values()) {
-              yield { type: "tool_call_end", id: tc.id };
-            }
-            pendingToolCalls.clear();
-            yield { type: "done", stopReason: "stop" };
+            yield* finish();
             return;
           }
 
@@ -371,29 +488,29 @@ export function createOpenAIStreamingClient({
             continue; // Skip unparseable chunks
           }
 
+          // Usually a final chunk with no choices, but a provider may put it beside the finish.
+          usage = usageEventFromOpenAI(parsed.usage) ?? usage;
+
           const choice = parsed.choices?.[0];
           if (!choice) {
             continue;
           }
 
           const { delta } = choice;
-          if (!delta) {
-            continue;
-          }
 
           // Text content
-          if (delta.content) {
+          if (delta?.content) {
             yield { type: "delta", content: delta.content };
           }
 
           // Chain-of-thought, under either of the two names providers give it.
-          const reasoning = delta.reasoning_content ?? delta.reasoning;
+          const reasoning = delta?.reasoning_content ?? delta?.reasoning;
           if (typeof reasoning === "string" && reasoning) {
             yield { type: "reasoning", content: reasoning };
           }
 
           // Tool calls
-          if (delta.tool_calls) {
+          if (delta?.tool_calls) {
             for (const tc of delta.tool_calls) {
               const existing = pendingToolCalls.get(tc.index);
 
@@ -423,38 +540,17 @@ export function createOpenAIStreamingClient({
             }
           }
 
-          // Finish reason
-          if (choice.finish_reason === "tool_calls") {
-            // Emit pending tool call ends
-            for (const tc of pendingToolCalls.values()) {
-              yield { type: "tool_call_end", id: tc.id };
-            }
-            pendingToolCalls.clear();
-            yield { type: "done", stopReason: "tool_calls" };
-            return;
-          }
-
-          if (choice.finish_reason === "stop" || choice.finish_reason === "length") {
-            // Emit any pending tool call ends
-            for (const tc of pendingToolCalls.values()) {
-              yield { type: "tool_call_end", id: tc.id };
-            }
-            pendingToolCalls.clear();
-            yield {
-              type: "done",
-              stopReason: choice.finish_reason === "length" ? "length" : "stop",
-            };
-            return;
+          // Finish reason: close the calls now, report the reason at the end of the stream.
+          const finished = stopReasonFromOpenAI(choice.finish_reason);
+          if (finished) {
+            yield* closePendingToolCalls();
+            stopReason = finished;
           }
         }
       }
 
-      // Stream ended without explicit finish_reason
-      for (const tc of pendingToolCalls.values()) {
-        yield { type: "tool_call_end", id: tc.id };
-      }
-      pendingToolCalls.clear();
-      yield { type: "done", stopReason: "stop" };
+      // Stream ended without `[DONE]`
+      yield* finish();
     } catch (error) {
       void reader.cancel();
       if ((error as Error).name === "AbortError") {
@@ -513,17 +609,26 @@ export function createAnthropicStreamingClient(
  * `@jxsuite/server/ai-api`). No API key handling here — the proxy owns provider credentials.
  *
  * @param {object} opts
- * @param {string} opts.chatUrl - URL to POST `{ messages, tools, systemPrompt, model }` to
+ * @param {string | (() => string | Promise<string>)} opts.chatUrl - URL to POST `{ messages, tools,
+ *   systemPrompt, model }` to, or a function that answers it. A function is called once, inside the
+ *   first stream, so a caller can hand over its turn's signal before it waits on anything: a Stop
+ *   during a lookup over IPC then takes effect once it answers, and the stream ends cancelled with
+ *   nothing sent. The outcome is kept for the client's later streams. A rejection propagates out of
+ *   the stream, unless the signal was aborted by then, which ends it cancelled.
  * @param {string} [opts.model] - Default model if not specified per-request
  * @param {string} [opts.apiKey] - Optional client-supplied key, sent as the `X-Api-Key` header
  * @param {string} [opts.baseUrl] - Optional OpenAI-compatible base URL, sent as `X-Api-Base-URL`
+ * @param {string} [opts.sessionId] - The conversation's id, sent as the `X-Jx-Ai-Session` header so
+ *   the backend can derive a prompt-cache affinity key from it. The raw id goes to the backend
+ *   only; a provider sees a hash of it, scoped to that backend (`affinityKey`).
  * @returns {StreamingClient}
  */
 export interface ProxyStreamingClientOptions {
-  chatUrl: string;
+  chatUrl: string | (() => string | Promise<string>);
   model?: string;
   apiKey?: string | undefined;
   baseUrl?: string | undefined;
+  sessionId?: string | undefined;
 }
 
 export function createProxyStreamingClient({
@@ -531,7 +636,11 @@ export function createProxyStreamingClient({
   model = "gpt-4o",
   apiKey,
   baseUrl,
+  sessionId,
 }: ProxyStreamingClientOptions): StreamingClient {
+  /** The URL, once the first stream has asked for it. */
+  let resolvedUrl: Promise<string> | null = null;
+
   /**
    * @param {object[]} messages
    * @param {object[]} tools
@@ -555,10 +664,34 @@ export function createProxyStreamingClient({
     if (baseUrl) {
       headers["X-Api-Base-URL"] = baseUrl;
     }
+    /* The conversation, so consecutive rounds can be routed to a warm prompt cache. Same-origin on
+       every platform, so the extra header costs no preflight; a backend that ignores it loses only
+       the cache hit. */
+    if (sessionId) {
+      headers[AI_SESSION_HEADER] = sessionId;
+    }
+
+    resolvedUrl ??= Promise.resolve(typeof chatUrl === "function" ? chatUrl() : chatUrl);
+    let url: string;
+    try {
+      url = await resolvedUrl;
+    } catch (error) {
+      // A lookup that failed after the turn was stopped is part of the stop, not an error.
+      if (signal?.aborted) {
+        yield { type: "done", stopReason: "cancelled" };
+        return;
+      }
+      throw error;
+    }
+    // Stopped while the URL was being resolved: nothing has been sent, and nothing will be.
+    if (signal?.aborted) {
+      yield { type: "done", stopReason: "cancelled" };
+      return;
+    }
 
     let response;
     try {
-      response = await fetch(chatUrl, {
+      response = await fetch(url, {
         method: "POST",
         headers,
         body: JSON.stringify({ messages, tools, systemPrompt, model }),