@jxsuite/ai 0.37.2 → 0.39.0

package/README.md

@@ -1,15 +1,16 @@
 # @jxsuite/ai
 
-The provider-agnostic substrate Jx Studio's assistant is built on: a streaming LLM client abstraction that normalizes provider SSE into seven event shapes, a tool registry that speaks OpenAI function-calling, and a reactive chat store. It has no Jx-domain dependencies (just a type-only import of `ProblemDetails` from `@jxsuite/protocol`, plus `@vue/reactivity`), and its only platform APIs are `fetch`, `TextDecoder` and `AbortSignal`, so the same modules run in the browser and in Bun/Node.
+The provider-agnostic substrate Jx Studio's assistant is built on: a streaming LLM client abstraction that normalizes provider SSE into eight event shapes, a tool registry that speaks OpenAI function-calling, and a reactive chat store. It has no Jx-domain dependencies beyond `@jxsuite/protocol`, for the wire's problem documents (the gateway builds them at runtime; everything else imports only the `ProblemDetails` type), plus `@vue/reactivity` for the chat store. Its platform APIs are the web-standard ones (`fetch`, `Request` and `Response`, streams, `TextEncoder` and `TextDecoder`, `AbortSignal`), so the same modules run in the browser, in Bun/Node and, for the Worker-safe subpaths, in a Cloudflare Worker.
 
 Governing spec: [`specs/ai.md`](../../specs/ai.md) §1-§2 (Status: Partial). User-facing documentation for the Studio feature lives at [`docs/studio/ai.md`](../../docs/studio/ai.md).
 
-| Entrypoint                     | Exports                                                                                                                                     |
-| ------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------- |
-| `@jxsuite/ai`                  | Barrel: the three factories, `createToolDefinition`, `createToolRegistry`, `createChatState`, `STREAM_EVENT_TYPES`, and the streaming types |
-| `@jxsuite/ai/streaming-client` | `StreamingClient`, `StreamEvent` and its seven members, the three factories and their option types, `STREAM_EVENT_TYPES`                    |
-| `@jxsuite/ai/tools`            | `createToolDefinition`, `createToolRegistry`, `toolSuccess`, `toolError`, and their types                                                   |
-| `@jxsuite/ai/chat-state`       | `createChatState`, `ChatStore`, `Message`, `ToolCallRecord`, `ChatState`, `MessageRole`                                                     |
+| Entrypoint                     | Exports                                                                                                                                                                                                                                     |
+| ------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `@jxsuite/ai`                  | Barrel: the three factories, `createToolDefinition`, `createToolRegistry`, `createChatState`, the `ChatUsage` type, `STREAM_EVENT_TYPES`, the streaming types, and the prompt-cache hints (below)                                           |
+| `@jxsuite/ai/streaming-client` | `StreamingClient`, `StreamEvent` and its eight members, the three factories and their option types, `STREAM_EVENT_TYPES`, `usageEventFromOpenAI` and its `OpenAIUsage` input type, and the prompt-cache hints                               |
+| `@jxsuite/ai/tools`            | `createToolDefinition`, `createToolRegistry`, `toolSuccess`, `toolError`, the call context (`createToolContext`, `createLedger`, `createSessionFacts`, `linkCallSignal`), and their types                                                   |
+| `@jxsuite/ai/chat-state`       | `createChatState`, `ChatStore`, `ChatUsage`, `Message`, `ToolCallRecord`, `ChatState`, `MessageRole`                                                                                                                                        |
+| `@jxsuite/ai/gateway`          | `createChatHandler`, `modelsResponse`, `normalizeOpenAIStream`, `encodeSse`, `problemResponse`, `extractUpstreamErrorMessage`, the prompt-cache hints, and the `Upstream`, `GatewayRefusal`, `ChatAdmission` and `ChatGatewayOptions` types |
 
 `toolSuccess` / `toolError` are **not** re-exported by the barrel. Import them from `@jxsuite/ai/tools`. Importing the subpaths instead of the root is also what makes the package tree-shakeable.
 
@@ -22,7 +23,7 @@ streamChat(messages: object[], tools: object[], systemPrompt: string, signal: Ab
   : AsyncGenerator<StreamEvent>;
 ```
 
-Every implementation yields the same discriminated union, whose type strings are also published as `STREAM_EVENT_TYPES`: `delta {content}`, `tool_call_start {id, name}`, `tool_call_delta {id, args}`, `tool_call_end {id}`, `done {stopReason}`, `error {message, code?, problem?}`. This union is the contract a third-party Studio backend must satisfy for the `ai/chat` route (see [`packages/protocol/README.md`](../protocol/README.md)); `@jxsuite/server` emits the same format without depending on this package. `StreamEvent` and each member are exported, so an implementer can import the type instead of reconstructing it.
+Every implementation yields the same discriminated union, whose type strings are also published as `STREAM_EVENT_TYPES`: `delta {content}`, `reasoning {content}`, `tool_call_start {id, name}`, `tool_call_delta {id, args}`, `tool_call_end {id}`, `usage {inputTokens, outputTokens, cachedInputTokens?, reasoningTokens?}`, `done {stopReason}`, `error {message, code?, problem?}`. This union is the contract a third-party Studio backend must satisfy for the `ai/chat` route (see [`packages/protocol/README.md`](../protocol/README.md)); `@jxsuite/server` emits the same format through [the gateway](#the-gateway). `StreamEvent` and each member are exported, so an implementer can import the type instead of reconstructing it.
 
 ```ts
 import { createProxyStreamingClient } from "@jxsuite/ai";
@@ -38,15 +39,47 @@ for await (const event of client.streamChat(messages, tools, systemPrompt, signa
 }
 ```
 
-- **`createOpenAIStreamingClient({ baseUrl, apiKey, model = "gpt-4o", temperature })`** POSTs to `${baseUrl}/chat/completions` with `Authorization: Bearer`, `stream: true`, `stream_options: { include_usage: true }`, and `{ role: "system", content: systemPrompt }` prepended to your messages. A non-empty `tools` array also sets `tool_choice: "auto"` and `parallel_tool_calls: true`; an empty one sends no `tools` key. `temperature` is forwarded **only when defined**, because reasoning models (GPT-5.x, o-series) reject a custom temperature. That is a provider behaviour recorded in the source comment, not covered by a test here. It does its own SSE framing and reassembles fragmented tool-call argument JSON keyed by the provider's `delta.tool_calls[].index`.
+- **`createOpenAIStreamingClient({ baseUrl, apiKey, model = "gpt-4o", temperature, sessionAffinity })`** POSTs to `${baseUrl}/chat/completions` with `Authorization: Bearer`, `stream: true`, `stream_options: { include_usage: true }`, and `{ role: "system", content: systemPrompt }` prepended to your messages. A non-empty `tools` array also sets `tool_choice: "auto"` and `parallel_tool_calls: true`; an empty one sends no `tools` key. `temperature` is forwarded **only when defined**, because reasoning models (GPT-5.x, o-series) reject a custom temperature. That is a provider behaviour recorded in the source comment, not covered by a test here. It does its own SSE framing and reassembles fragmented tool-call argument JSON keyed by the provider's `delta.tool_calls[].index`.
 - **`createAnthropicStreamingClient({ baseUrl, apiKey })`** is a stub. It makes no network call and yields exactly one `error` event with `code: "NOT_IMPLEMENTED"`. Only the OpenAI-compatible path ships (`specs/ai.md` §2).
-- **`createProxyStreamingClient({ chatUrl, model = "gpt-4o", apiKey, baseUrl })`** POSTs `{ messages, tools, systemPrompt, model }` to an endpoint that already speaks the normalized format (the dev/desktop server's `/__studio/ai/chat`). It handles no credentials of its own; `apiKey` and `baseUrl` are passed through as the `X-Api-Key` and `X-Api-Base-URL` headers only when truthy, and the proxy owns the provider key. It is a pass-through, **not** a translator: it re-yields whatever parses out of each `data:` line and returns at the first `done` or `error`, so frames after `done` are dropped.
+- **`createProxyStreamingClient({ chatUrl, model = "gpt-4o", apiKey, baseUrl, sessionId })`** POSTs `{ messages, tools, systemPrompt, model }` to an endpoint that already speaks the normalized format (the dev/desktop server's `/__studio/ai/chat`). `chatUrl` may also be a function returning the URL (or a promise of it): it is called once, inside the first stream, so a caller can arm its turn's signal before waiting on a lookup, and a stream stopped by the time it answers ends `cancelled` without sending anything; a lookup that rejects rejects the stream. It handles no credentials of its own; `apiKey` and `baseUrl` are passed through as the `X-Api-Key` and `X-Api-Base-URL` headers only when truthy, and the proxy owns the provider key. `sessionId`, likewise only when truthy, is sent as `X-Jx-Ai-Session` (`AI_SESSION_HEADER`) for the backend to hash into a prompt-cache key; the raw id is meant for the backend alone. It is a pass-through, **not** a translator: it re-yields whatever parses out of each `data:` line and returns at the first `done` or `error`, so frames after `done` are dropped.
 
 Three behaviours are easy to get wrong when consuming the stream:
 
 1. **Cancellation is a `done`, not an `error`.** An `AbortError` from `fetch` or mid-read becomes `{ type: "done", stopReason: "cancelled" }`. Treating abort as failure paints a user's Stop button as a crash. The OpenAI client's other stop reasons are `"stop"`, `"tool_calls"` and `"length"`.
-2. **The OpenAI client drains pending tool calls on every termination path.** `[DONE]`, a `finish_reason` of `stop`, `length` or `tool_calls` (any other value is ignored and the stream continues), and a body that simply closes all emit `tool_call_end` for anything still open, then exactly one `done`. The proxy client tracks nothing of its own, so a body that closes without a `done` frame ends its generator with no `done` at all.
-3. **`error` frames may carry an RFC 9457 `problem` beside the human `message`,** because the response has already begun with a 200 and the status can no longer change (a committed standards row in `specs/ai.md` §5). The field is optional and no client in this package sets it. `@jxsuite/server`'s proxy is what populates it, and `createProxyStreamingClient` passes it through. Read `problem.type` for machine dispatch when it is present; keep showing `message`.
+2. **A finish is remembered, not acted on.** A `finish_reason` of `stop`, `length` or `tool_calls` emits `tool_call_end` for every call still open and records the reason (any other value is ignored), but the stream keeps reading, because with `include_usage` the count arrives in a chunk of its own after the finish. The stream ends at `[DONE]` or when the body closes, and ends the same way every time: `tool_call_end` for anything still open, one `usage` frame when the provider reported a count, then exactly one `done` carrying the recorded reason (`stop` when there was none). `usage` precedes `done` because a reader stops at `done`. The proxy client tracks nothing of its own, so a body that closes without a `done` frame ends its generator with no `done` at all.
+3. **`error` frames may carry an RFC 9457 `problem` beside the human `message`,** because the response has already begun with a 200 and the status can no longer change (a committed standards row in `specs/ai.md` §5). The field is optional and no client in this package sets it. The gateway (which `@jxsuite/server`'s proxy runs) is what populates it, and `createProxyStreamingClient` passes it through. Read `problem.type` for machine dispatch when it is present; keep showing `message`.
+
+## The gateway
+
+`@jxsuite/ai/gateway` is the other side of the wire: the server half of the `ai/chat` and `ai/models` routes, which `@jxsuite/server` runs and any other backend can. The host supplies its policy and the gateway does the rest:
+
+```ts
+import { createChatHandler, modelsResponse } from "@jxsuite/ai/gateway";
+
+const chat = createChatHandler<{ signal: AbortSignal }>({
+  // Whose key, which base URL: the host's rules. Return a refusal to answer with a problem instead.
+  resolveUpstream: () => ({
+    apiKey: env.OPENAI_API_KEY,
+    baseUrl: "https://api.openai.com/v1",
+    defaultModel: "gpt-4o",
+    family: "openai-compat",
+    managed: false,
+  }),
+  maxBodyBytes: 1_000_000,
+  maxMessages: 200,
+});
+
+const response = await chat(request, { signal: request.signal });
+const models = modelsResponse({ configured: true, models: [{ id: "gpt-4o" }] });
+```
+
+- **The order is fixed.** `admit` (before the body is read), `resolveUpstream`, the body (`413` once it passes `maxBodyBytes`, counted on the bytes actually read, and `400` when it is not JSON), `messages` (`400` when it is present and not an array), `maxMessages` (`413`), `onAccepted`, then the stream. The first step that refuses answers the request as `application/problem+json`, and the provider is never contacted. `problemResponse` writes `error` beside `detail` and a refusal's `code` beside both, so a problem reader and an older `{error, code}` reader both understand it.
+- **The stream is always a `200`.** The upstream request is on the wire before the response is returned, and anything that goes wrong after that (a network failure, a provider's non-2xx, a reset mid-stream) is an `error` frame, while an abort is `done: cancelled`. `normalizeOpenAIStream` is the server's normalizer, moved here unchanged, so its frames are the ones `packages/server/tests/fixtures/ai-upstream/*.server.json` records.
+- **`fetch` is looked up per request** unless you pass one, so a host (or a test) can swap it.
+- **`modelsResponse(body)` is always a `200`,** because the models route is a capability probe. What the catalogue says is the host's business.
+- **A conversation can keep a warm prompt cache.** Studio's proxy client sends its conversation id as `X-Jx-Ai-Session` (`AI_SESSION_HEADER`). The host checks it with `isAiSessionId`, hashes it with `affinityKey(scope, id)` (the first 32 hex characters of SHA-256 over `jx-ai-affinity/v1:<scope>:<id>`) and returns the hash as the upstream's `sessionAffinity`. `upstreamCacheHints(baseUrl, affinity)` then decides what the provider sees: `prompt_cache_key` in the body for `api.openai.com`, the `x-session-affinity` header for Workers AI (`api.cloudflare.com` with `/ai/` in the path) and AI Gateway (`gateway.ai.cloudflare.com`), and nothing for any other host, so a strict server never meets an unknown field and an arbitrary endpoint never sees a stable id. `createOpenAIStreamingClient` takes the same key as `sessionAffinity` and applies the same rule.
+
+`./gateway` is Worker-safe. Studio never imports it (`packages/studio/tests/ai-import-rules.test.ts`).
 
 ## The tool registry
 
@@ -74,12 +107,13 @@ registry.register(
 - **`validate()` is a lightweight structural check, not a JSON Schema validator.** It reads only `type`, `properties` and `required`: nested schemas, `enum`, `pattern`, `minimum` and every other keyword are ignored. Unknown properties are ignored too (the model may send extra); a required key that is present but `null`/`undefined` counts as missing; numeric strings coerce, so `"42"` validates against `number`. `strict: false`, or a schema without `properties`, skips it entirely.
 - **`execute()` never throws.** An unknown tool, a validation failure, and an exception inside the tool all return `{ success: false, error }`, which is what lets an agent loop feed the failure back to the model and self-correct. It calls `this.validate(...)`, so it is `this`-bound: do not destructure it off the registry.
 - Re-registering a name `console.warn`s and overwrites.
+- **Every call carries a `ToolContext`** as `execute`'s second argument: the call's own signal, its id, the actor, the turn's write ledger, the conversation's session facts and a progress sink (specs/ai.md §3.7). `registry.execute(name, args, ctx)` passes it on; without a `ctx` the call gets `createToolContext()`, a detached one whose signal never aborts and whose ledger nothing files. `linkCallSignal(turnSignal)` gives a call a signal that aborts with the turn and is unlinked by `release()` once the call settles. A definition marks `interactive: true` when its call suspends the turn on a person (`ask_user`). A wrapping registry must forward `ctx`. See `docs/extending/embedding/assistant-harness.md`.
 
 To vary the tool set per turn, wrap rather than mutate. Studio's `createGatedToolRegistry` returns a `ToolRegistry` that filters `list()`/`listForLLM()` through per-tool predicates (`packages/studio/src/services/gated-registry.ts`).
 
 ## The reactive chat store
 
-`createChatState(options?: { model?: string })` returns a `@vue/reactivity` `reactive()` store with its mutators `Object.assign`ed onto the same proxy, so the returned value is both the state and the API. The store holds `messages`, `status`, `streamingContent`, `pendingToolCalls`, `error`, `model`, `tokenCount` and `contextWarning`. The model falls back to `"gpt-4o"` via `||`, so an empty string takes the default too. `status` is `"idle" | "streaming" | "error"`; `toMessagesArray()` serializes back to OpenAI wire shape (assistant `tool_calls`, `tool` messages with `tool_call_id`, everything else `{ role, content }`).
+`createChatState(options?: { model?: string })` returns a `@vue/reactivity` `reactive()` store with its mutators `Object.assign`ed onto the same proxy, so the returned value is both the state and the API. The store holds `messages`, `status`, `streamingContent`, `pendingToolCalls`, `error`, `model`, `tokenCount`, `contextWarning` and `usage`. `recordUsage(event, { systemTokens? })` stores a `usage` frame as a `ChatUsage`: the counts, the context it leaves for the next request (input plus output, less reasoning the provider billed but never streamed), and the id and position of the last message it covers, so a budget can reuse it only while that message is still in place. It also sets `tokenCount` to that context figure. `clearUsage()` forgets it, and `clearChat()` does too. The model falls back to `"gpt-4o"` via `||`, so an empty string takes the default too. `status` is `"idle" | "streaming" | "error"`; `toMessagesArray()` serializes back to OpenAI wire shape (assistant `tool_calls`, `tool` messages with `tool_call_id`, everything else `{ role, content }`).
 
 Two ordering rules inside `beginAssistantTurn` are load-bearing, and any reimplementation must keep them: `status` is set to `"streaming"` **before** the placeholder message is pushed (the push itself triggers effects, and a reader seeing `"idle"` renders the empty placeholder as a finished message), and the pushed object is immediately re-read back through the proxy (`store.messages.at(-1)`), because mutating the raw object would notify nothing. A regression test asserts that an `effect()` reading `messages.at(-1).content` re-runs on every `appendDelta`.
 
@@ -93,4 +127,4 @@ Further surprises:
 
 ## Versioning
 
-Published to npm as `@jxsuite/ai`, following the monorepo's release train. Like every published `@jxsuite` library, it ships TypeScript source: every `exports` subpath resolves to `./src/*.ts`, so a consumer must be able to compile TypeScript out of `node_modules`. Within the monorepo, `@jxsuite/studio` depends on it via `workspace:^` and is its only consumer. `@vue/reactivity` is pinned to an **exact** version, matching `@jxsuite/studio`, `@jxsuite/runtime` and `@jxsuite/compiler`: reactive proxies from one copy do not track in effects from another, and the failure is silent.
+Published to npm as `@jxsuite/ai`, following the monorepo's release train. Like every published `@jxsuite` library, it ships TypeScript source: every `exports` subpath resolves to `./src/*.ts`, so a consumer must be able to compile TypeScript out of `node_modules`. Within the monorepo, `@jxsuite/studio` (the client subpaths) and `@jxsuite/server` (`./gateway`) depend on it via `workspace:^`. `@vue/reactivity` is pinned to an **exact** version, matching `@jxsuite/studio`, `@jxsuite/runtime` and `@jxsuite/compiler`: reactive proxies from one copy do not track in effects from another, and the failure is silent.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jxsuite/ai",
-  "version": "0.37.2",
+  "version": "0.39.0",
   "description": "AI infrastructure for Jx Suite — streaming LLM client abstraction, tool registry, and reactive chat state",
   "license": "MIT",
   "repository": {
@@ -17,7 +17,8 @@
     ".": "./src/index.ts",
     "./streaming-client": "./src/streaming-client.ts",
     "./tools": "./src/tools.ts",
-    "./chat-state": "./src/chat-state.ts"
+    "./chat-state": "./src/chat-state.ts",
+    "./gateway": "./src/gateway/index.ts"
   },
   "publishConfig": {
     "provenance": true
@@ -28,10 +29,11 @@
     "upgrade": "bun update --latest && bun install"
   },
   "dependencies": {
-    "@jxsuite/protocol": "^2.3.1",
+    "@jxsuite/protocol": "^2.5.0",
     "@vue/reactivity": "3.5.43"
   },
   "devDependencies": {
+    "@cloudflare/workers-types": "^5.20260928.1",
     "@happy-dom/global-registrator": "^20.14.5"
   }
 }

package/src/cache-hints.ts

@@ -0,0 +1,111 @@
+/**
+ * Cache-hints.ts — Prompt-cache routing hints for OpenAI-compatible upstreams.
+ *
+ * A provider's prompt cache only pays off when consecutive requests sharing a prefix reach the same
+ * machine. The assistant's requests are built for that (a stable system-prompt prefix, then a
+ * history that only grows), but the providers that route by prefix still need to be told which
+ * requests belong together, and each one is told differently:
+ *
+ * - OpenAI reads a `prompt_cache_key` body field.
+ * - Cloudflare Workers AI (directly, or through AI Gateway) reads an `x-session-affinity` header.
+ *
+ * Neither hint is sent anywhere else. An OpenAI-compatible server is free to reject an unknown body
+ * field, and a browser-side BYOK client adding an unknown header to a cross-origin request turns a
+ * simple request into a CORS preflight the endpoint may not answer. Above all, the value is a
+ * stable per-conversation identifier, and an arbitrary endpoint the user pointed a key at has no
+ * use for one.
+ *
+ * What travels is never the client's own session id. A client sends that only to its own backend
+ * (as {@link AI_SESSION_HEADER}); the backend hashes it with {@link affinityKey}, scoped to itself,
+ * and the hash is the only form a provider sees.
+ *
+ * Worker-safe: `crypto.subtle` and `TextEncoder` are globals in every runtime the gateway targets.
+ *
+ * @module @jxsuite/ai/streaming-client
+ * @license MIT
+ * @docs extending/embedding/backend-protocol
+ */
+
+/**
+ * The request header a client puts its conversation id in, for its OWN backend to hash. Never
+ * forwarded to a provider as sent.
+ */
+export const AI_SESSION_HEADER = "X-Jx-Ai-Session";
+
+/** 1–128 characters of `[A-Za-z0-9_.:-]`: wide enough for any id a client mints, nothing else. */
+const SESSION_ID = /^[\w.:-]{1,128}$/;
+
+/**
+ * Whether a value is a session id a backend may hash. A header that fails this is ignored rather
+ * than refused: the hint is an optimisation, and a request without one is still a valid request.
+ *
+ * @param {unknown} value - Typically the raw {@link AI_SESSION_HEADER} value, or null
+ * @returns {boolean}
+ */
+export function isAiSessionId(value: unknown): value is string {
+  return typeof value === "string" && SESSION_ID.test(value);
+}
+
+/**
+ * The affinity key a backend forwards for a session: the first 32 hex characters of
+ * SHA-256(`jx-ai-affinity/v1:<scope>:<sessionId>`).
+ *
+ * `scope` names the backend (`"local"` for the dev/desktop server), so two backends that happen to
+ * see the same id never share a key, and the version prefix lets the derivation change without
+ * colliding with the old one. 128 bits is plenty to keep unrelated conversations apart, and short
+ * enough for every provider's limit on the field.
+ *
+ * @param {string} scope - The backend's own namespace
+ * @param {string} sessionId - The client's session id, already checked with {@link isAiSessionId}
+ * @returns {Promise<string>} 32 lowercase hex characters
+ */
+export async function affinityKey(scope: string, sessionId: string): Promise<string> {
+  const input = new TextEncoder().encode(`jx-ai-affinity/v1:${scope}:${sessionId}`);
+  const digest = new Uint8Array(await crypto.subtle.digest("SHA-256", input));
+  let hex = "";
+  for (const byte of digest.subarray(0, 16)) {
+    hex += byte.toString(16).padStart(2, "0");
+  }
+  return hex;
+}
+
+/** The hints to merge into one upstream request: extra headers, and extra body members. */
+export interface UpstreamCacheHints {
+  headers: Record<string, string>;
+  body: { prompt_cache_key?: string };
+}
+
+/**
+ * The cache hints for a request to `baseUrl`, in the form that upstream understands, or none.
+ *
+ * Classified by host, because the hint's NAME is per provider: `api.openai.com` gets
+ * `prompt_cache_key` in the body; Workers AI (an `api.cloudflare.com` URL with `/ai/` in its path)
+ * and AI Gateway (`gateway.ai.cloudflare.com`) get the `x-session-affinity` header. Every other
+ * host, an unparseable base URL, and a request with no affinity get nothing.
+ *
+ * @param {string} baseUrl - The upstream's base URL, as the request is about to use it
+ * @param {string} [affinity] - The affinity key, from {@link affinityKey} (or a caller's own id)
+ * @returns {UpstreamCacheHints}
+ */
+export function upstreamCacheHints(baseUrl: string, affinity?: string): UpstreamCacheHints {
+  const hints: UpstreamCacheHints = { body: {}, headers: {} };
+  if (!affinity) {
+    return hints;
+  }
+  let url: URL;
+  try {
+    url = new URL(baseUrl);
+  } catch {
+    return hints;
+  }
+  const host = url.hostname.toLowerCase();
+  if (host === "api.openai.com") {
+    hints.body.prompt_cache_key = affinity;
+  } else if (
+    (host === "api.cloudflare.com" && url.pathname.includes("/ai/")) ||
+    host === "gateway.ai.cloudflare.com"
+  ) {
+    hints.headers["x-session-affinity"] = affinity;
+  }
+  return hints;
+}

package/src/chat-state.ts

@@ -11,6 +11,7 @@
 
 import { reactive } from "@vue/reactivity";
 
+import type { StreamUsageEvent } from "./streaming-client.ts";
 import type { ToolResult } from "./tools.ts";
 
 export type ChatState = "idle" | "streaming" | "error";
@@ -40,6 +41,34 @@ export interface ToolCallRecord {
   result?: ToolResult | null;
 }
 
+/**
+ * The provider's own count for the last request, and where in the transcript it was taken.
+ *
+ * `messageCount` and `lastMessageId` are what make the figure reusable: every message up to and
+ * including `lastMessageId` was part of a request the provider counted, so a later budget needs to
+ * estimate only what was added since, rather than re-estimating the whole history at four
+ * characters per token. A transcript whose message at that position is no longer `lastMessageId`
+ * has been trimmed, rewound, repaired or replaced, and the figure no longer describes it — length
+ * alone cannot say so, because a retry re-grows the transcript to the same length.
+ */
+export interface ChatUsage {
+  inputTokens: number;
+  outputTokens: number;
+  cachedInputTokens?: number;
+  reasoningTokens?: number;
+  messageCount: number;
+  /** The id of the last message the count covers (`messages[messageCount - 1]`). */
+  lastMessageId: string;
+  /**
+   * What the counted request leaves in the context for the NEXT request: input plus output, minus
+   * reasoning the provider generated but never streamed. That reasoning (OpenAI's o-series on chat
+   * completions) is billed as output yet is not replayed, so it occupies nothing on the next send.
+   */
+  contextTokens: number;
+  /** The estimate of the system prompt the counted request carried, when the caller supplied it. */
+  systemTokens?: number;
+}
+
 export interface ChatStore {
   messages: Message[];
   status: ChatState;
@@ -49,6 +78,8 @@ export interface ChatStore {
   model: string;
   tokenCount: number;
   contextWarning: boolean;
+  /** The last reported count, or null when none has been reported since the transcript changed. */
+  usage: ChatUsage | null;
 }
 
 // ─── Types ───────────────────────────────────────────────────────────────────
@@ -63,6 +94,11 @@ function uid() {
   return `msg_${Date.now()}_${_idCounter}`;
 }
 
+/** A token count a budget can use: a finite, non-negative number. */
+function isCount(value: unknown): value is number {
+  return typeof value === "number" && Number.isFinite(value) && value >= 0;
+}
+
 // ─── Factory ─────────────────────────────────────────────────────────────────
 
 /**
@@ -87,6 +123,8 @@ function uid() {
  *   retryLast: () => void;
  *   setModel: (model: string) => void;
  *   setTokenCount: (count: number) => void;
+ *   recordUsage: (usage: StreamUsageEvent, opts?: { systemTokens?: number }) => void;
+ *   clearUsage: () => void;
  *   setContextWarning: (warning: boolean) => void;
  *   toMessagesArray: () => object[];
  * }}
@@ -103,6 +141,7 @@ export function createChatState(opts: { model?: string } = {}) {
     model,
     tokenCount: 0,
     contextWarning: false,
+    usage: null,
   });
 
   let _streamingMessage: Message | null = null;
@@ -257,22 +296,28 @@ export function createChatState(opts: { model?: string } = {}) {
   }
 
   /**
-   * Attach a result to a pending tool call.
+   * Attach a result to the tool call it answers, on the assistant message that made the call.
+   *
+   * The record is looked for on that message rather than on the stream: a loop runs its tools after
+   * the round's stream has finished, when there is no streaming message and no pending call left,
+   * and looking only there attached every result to nothing, so a live chip never showed how its
+   * call ended. The request is the assistant message directly before the tool replies at the end of
+   * the transcript, the one position the wire allows a reply in; a provider may reuse call ids
+   * across rounds, so an earlier request carrying the same id is never the one being answered.
    *
    * @param {string} id
    * @param {ToolResult} result
    */
   function appendToolResult(id: string, result: ToolResult) {
-    const tc = store.pendingToolCalls.find((t) => t.id === id);
-    if (tc) {
-      tc.result = result;
+    let index = store.messages.length - 1;
+    while (index >= 0 && store.messages[index]!.role === "tool") {
+      index -= 1;
     }
-    // Also update in the message record
-    if (_streamingMessage?.toolCalls) {
-      const mtc = _streamingMessage.toolCalls.find((t) => t.id === id);
-      if (mtc) {
-        mtc.result = result;
-      }
+    const request = store.messages[index];
+    const record =
+      request?.role === "assistant" ? request.toolCalls?.find((t) => t.id === id) : null;
+    if (record) {
+      record.result = result;
     }
   }
 
@@ -297,6 +342,7 @@ export function createChatState(opts: { model?: string } = {}) {
     store.status = "error";
     store.error = message;
     store.streamingContent = "";
+    store.pendingToolCalls = [];
     // Remove the partial streaming message — it may contain incomplete tool_calls
     // That would poison the conversation history on the next send.
     if (_streamingMessage) {
@@ -333,6 +379,7 @@ export function createChatState(opts: { model?: string } = {}) {
     store.pendingToolCalls = [];
     store.error = null;
     store.contextWarning = false;
+    store.usage = null;
     _streamingMessage = null;
   }
 
@@ -368,6 +415,47 @@ export function createChatState(opts: { model?: string } = {}) {
     store.tokenCount = count;
   }
 
+  /**
+   * Record the provider's count for the request that just finished. The count covers everything the
+   * request carried plus what it generated, which is the context's size at this point — less any
+   * reasoning the provider billed but never streamed, since the next request cannot replay what it
+   * never received. It replaces the estimate in `tokenCount` rather than adding to it.
+   *
+   * @param {StreamUsageEvent} usage
+   * @param {{ systemTokens?: number }} [extra] - The estimate of the system prompt this request
+   *   carried, so a later budget can add what the prompt has grown by since.
+   */
+  function recordUsage(usage: StreamUsageEvent, extra: { systemTokens?: number } = {}) {
+    /* A backend's frame, trusted no further than its shape: a count that is not a finite,
+       non-negative number would turn every later budget into NaN, and a NaN budget trims history
+       it has no reason to. Such a frame is ignored, and the estimate stands. */
+    if (!isCount(usage.inputTokens) || !isCount(usage.outputTokens)) {
+      return;
+    }
+    const { type: _type, ...counts } = usage;
+    const counted = _streamingMessage ?? store.messages.at(-1);
+    /* Replayed means SENT: `toMessagesArray` drops an assistant turn carrying neither text nor
+       tool calls, reasoning and all, so a turn that only thought is not replayed either. */
+    const replayed = Boolean(
+      counted?.reasoningContent && (counted.content || (counted.toolCalls?.length ?? 0) > 0),
+    );
+    const unreplayed = replayed || !isCount(usage.reasoningTokens) ? 0 : usage.reasoningTokens;
+    const contextTokens = Math.max(0, usage.inputTokens + usage.outputTokens - unreplayed);
+    store.usage = {
+      ...counts,
+      contextTokens,
+      lastMessageId: store.messages.at(-1)?.id ?? "",
+      messageCount: store.messages.length,
+      ...(extra.systemTokens === undefined ? {} : { systemTokens: extra.systemTokens }),
+    };
+    store.tokenCount = contextTokens;
+  }
+
+  /** Forget the last count — the transcript it described has changed underneath it. */
+  function clearUsage() {
+    store.usage = null;
+  }
+
   /**
    * Set the context overflow warning flag.
    *
@@ -465,6 +553,8 @@ export function createChatState(opts: { model?: string } = {}) {
     retryLast,
     setModel,
     setTokenCount,
+    recordUsage,
+    clearUsage,
     setContextWarning,
     toMessagesArray,
   });

package/src/core-types.ts

@@ -0,0 +1,97 @@
+/**
+ * Core-types.ts — the shapes a tool call carries, shared by every Jx harness host.
+ *
+ * Type-only. `./tools` re-exports all of it; the constructors live there (`createLedger`,
+ * `createSessionFacts`, `createToolContext`, `linkCallSignal`).
+ *
+ * A tool used to reach what it needed from its host through module-level slots: the turn's signal
+ * and the call's id from Studio's `ai-turn-signal.ts`, the write ledger from `ai-writes.ts`, the
+ * one-import guard from a flag in `ai-import-tools.ts`. That tied every tool to one Studio window
+ * running one turn. A {@link ToolContext} carries those facts with the call instead, so the same
+ * tool can run in a page, a Worker or an MCP server (specs/ai.md §3.7).
+ *
+ * @license MIT
+ * @module @jxsuite/ai/tools
+ */
+
+/** A value that survives a JSON round trip unchanged. */
+export type JsonValue =
+  | null
+  | boolean
+  | number
+  | string
+  | JsonValue[]
+  | { [key: string]: JsonValue };
+
+/** What a tool call answers. */
+export interface ToolResult {
+  success: boolean;
+  data?: unknown;
+  error?: string;
+  summary?: string;
+}
+
+/** One change the assistant made, as recorded by the tool that made it. */
+export interface AiWrite {
+  /** Project-relative path, or the document path for an in-editor mutation. */
+  path: string;
+  /** The tool that did it: the command title the chip renders. */
+  tool: string;
+  /**
+   * True when the change went to disk with no transaction behind it.
+   *
+   * This is the undo caveat, as a fact rather than a sentence: a `disk` write is NOT reachable by
+   * undo, by the tab's history, or by "Restore to here".
+   */
+  disk: boolean;
+  /** False when the tool reported a failure: a listed attempt that changed nothing. */
+  ok: boolean;
+  /** Why it failed, when it did. */
+  error?: string;
+}
+
+/** Where a turn's tools record their changes. One per turn; the host files it when the turn ends. */
+export interface WriteLedger {
+  record: (write: AiWrite) => void;
+  /** Every write recorded so far, oldest first. */
+  readonly writes: readonly AiWrite[];
+}
+
+/** Who a tool call acts for. */
+export interface Actor {
+  readonly kind: "assistant";
+  /** `assistant:<sessionId|local>:<turnId>`: the key history entries and lanes use. */
+  readonly id: string;
+  readonly turnId: string;
+  readonly sessionId: string | null;
+  readonly model?: string;
+}
+
+/**
+ * Per-session JSON facts: what a tool must remember across turns (an import already ran). The host
+ * decides how long they last; Studio keeps them until New Chat.
+ */
+export interface SessionFacts {
+  readonly sessionId: string | null;
+  get: (key: string) => JsonValue | undefined;
+  set: (key: string, value: JsonValue) => void;
+}
+
+/** Per-call facts. Host services a tool needs are bound when the tool is built, not passed here. */
+export interface ToolContext {
+  /**
+   * A per-CALL child of the turn's signal: aborts with the turn, and is unlinked when the call
+   * settles.
+   */
+  readonly signal: AbortSignal;
+  /**
+   * The provider's id for this call. Equals the chip's id, the question's id and the import run's
+   * key.
+   */
+  readonly callId: string;
+  readonly actor: Actor;
+  readonly ledger: WriteLedger;
+  readonly session: SessionFacts;
+  /** Report progress on a long call (an import's log lines). A no-op when nothing observes it. */
+  progress: (event: JsonValue) => void;
+}