@jxsuite/ai 0.37.2 → 0.38.0

package/src/streaming-client.ts

@@ -5,7 +5,7 @@
  * 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
@@ -21,6 +21,7 @@ export type StreamEvent =
   | StreamToolCallStartEvent
   | StreamToolCallDeltaEvent
   | StreamToolCallEndEvent
+  | StreamUsageEvent
   | StreamDoneEvent
   | StreamErrorEvent;
 
@@ -61,6 +62,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;
@@ -131,6 +153,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 +268,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;
@@ -332,8 +402,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 +454,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 +465,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 +517,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,14 +586,19 @@ 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`
  * @returns {StreamingClient}
  */
 export interface ProxyStreamingClientOptions {
-  chatUrl: string;
+  chatUrl: string | (() => string | Promise<string>);
   model?: string;
   apiKey?: string | undefined;
   baseUrl?: string | undefined;
@@ -532,6 +610,9 @@ export function createProxyStreamingClient({
   apiKey,
   baseUrl,
 }: ProxyStreamingClientOptions): StreamingClient {
+  /** The URL, once the first stream has asked for it. */
+  let resolvedUrl: Promise<string> | null = null;
+
   /**
    * @param {object[]} messages
    * @param {object[]} tools
@@ -556,9 +637,27 @@ export function createProxyStreamingClient({
       headers["X-Api-Base-URL"] = baseUrl;
     }
 
+    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 }),

package/src/tools.ts

@@ -5,16 +5,33 @@
  * can call. Tools have JSON Schema parameter definitions and execute functions.
  * The registry validates arguments before execution.
  *
+ * Every call carries a {@link ToolContext}: its own signal, its call id, who it acts for, the turn's
+ * write ledger and the session's facts (specs/ai.md §3.7). A caller that has none gets a detached
+ * one, so a tool run outside a turn (a test, a command) records into nothing and cannot be stopped.
+ *
+ * @docs extending/embedding/assistant-harness
  * @license MIT
  * @module @jxsuite/ai/tools
  */
 
-export interface ToolResult {
-  success: boolean;
-  data?: unknown;
-  error?: string;
-  summary?: string;
-}
+import type {
+  AiWrite,
+  JsonValue,
+  SessionFacts,
+  ToolContext,
+  ToolResult,
+  WriteLedger,
+} from "./core-types.ts";
+
+export type {
+  Actor,
+  AiWrite,
+  JsonValue,
+  SessionFacts,
+  ToolContext,
+  ToolResult,
+  WriteLedger,
+} from "./core-types.ts";
 
 /**
  * JSON Schema fragment describing a tool's parameters. Only the subset the registry's lightweight
@@ -33,7 +50,12 @@ export interface ToolDefinition {
   parameters: JSONSchema;
   strict: boolean;
   llmStrict: boolean;
-  execute: (args: object) => Promise<ToolResult> | ToolResult;
+  /**
+   * The call suspends the turn on a person rather than doing work (`ask_user`), so a round that
+   * made only such calls does not spend the turn's work budget.
+   */
+  interactive?: boolean;
+  execute: (args: object, ctx: ToolContext) => Promise<ToolResult> | ToolResult;
 }
 
 export interface ToolRegistry {
@@ -41,7 +63,11 @@ export interface ToolRegistry {
   list: () => ToolDefinition[];
   listForLLM: () => object[];
   validate: (toolName: string, args: object) => { valid: boolean; errors?: string[] };
-  execute: (toolName: string, args: object) => Promise<ToolResult>;
+  /**
+   * Run a tool with the call's context; a registry that wraps another forwards `ctx` unchanged.
+   * Omitted, the call gets a detached context ({@link createToolContext}).
+   */
+  execute: (toolName: string, args: object, ctx?: ToolContext) => Promise<ToolResult>;
   getDefinition: (toolName: string) => ToolDefinition | undefined;
 }
 
@@ -82,8 +108,9 @@ export function toolError(error: string): ToolResult {
  * @param {boolean} [opts.llmStrict] - Whether to send OpenAI `strict: true` in the function schema
  *   (default false). Only enable if the parameter schema is OpenAI-strict-compliant
  *   (additionalProperties:false everywhere, all properties in `required`, every property typed).
- * @param {(args: object) => Promise<ToolResult> | ToolResult} opts.execute - The tool
- *   implementation
+ * @param {boolean} [opts.interactive] - The call suspends the turn on a person (`ask_user`)
+ * @param {(args: object, ctx: ToolContext) => Promise<ToolResult> | ToolResult} opts.execute - The
+ *   tool implementation, handed the call's context
  * @returns {ToolDefinition}
  */
 export function createToolDefinition({
@@ -92,6 +119,7 @@ export function createToolDefinition({
   parameters,
   strict = true,
   llmStrict = false,
+  interactive,
   execute,
 }: {
   name: string;
@@ -99,9 +127,119 @@ export function createToolDefinition({
   parameters: JSONSchema;
   strict?: boolean;
   llmStrict?: boolean;
-  execute: (args: object) => Promise<ToolResult> | ToolResult;
+  interactive?: boolean;
+  execute: (args: object, ctx: ToolContext) => Promise<ToolResult> | ToolResult;
 }): ToolDefinition {
-  return { name, description, parameters, strict, llmStrict, execute };
+  return {
+    name,
+    description,
+    parameters,
+    strict,
+    llmStrict,
+    ...(interactive === undefined ? {} : { interactive }),
+    execute,
+  };
+}
+
+// ─── ToolContext ─────────────────────────────────────────────────────────────
+
+/**
+ * A write ledger: the writes one turn's tools record, in order.
+ *
+ * @param {(write: AiWrite) => void} [onRecord] - Told of each write as it is recorded
+ * @returns {WriteLedger}
+ */
+export function createLedger(onRecord?: (write: AiWrite) => void): WriteLedger {
+  const writes: AiWrite[] = [];
+  return {
+    record(write: AiWrite) {
+      writes.push(write);
+      onRecord?.(write);
+    },
+    writes,
+  };
+}
+
+/**
+ * Per-session facts: JSON a tool remembers across one conversation's turns.
+ *
+ * @param {string | null} [sessionId]
+ * @param {Readonly<Record<string, JsonValue>>} [initial] - Facts to start from (a restored session)
+ * @returns {SessionFacts & { toJSON: () => Record<string, JsonValue> }}
+ */
+export function createSessionFacts(
+  sessionId: string | null = null,
+  initial: Readonly<Record<string, JsonValue>> = {},
+): SessionFacts & { toJSON: () => Record<string, JsonValue> } {
+  const facts = new Map<string, JsonValue>(Object.entries(initial));
+  return {
+    sessionId,
+    get: (key: string) => facts.get(key),
+    set: (key: string, value: JsonValue) => {
+      facts.set(key, value);
+    },
+    toJSON: () => Object.fromEntries(facts),
+  };
+}
+
+/** A signal that never aborts: the one a detached call gets. */
+const NEVER = new AbortController().signal;
+
+/**
+ * A tool call's context. Every field not given is detached: a signal that never aborts, an empty
+ * call id, a fresh ledger nothing files, empty session facts and a progress sink that drops
+ * everything. That is what a call made outside any turn gets.
+ *
+ * @param {Partial<ToolContext>} [init]
+ * @returns {ToolContext}
+ */
+export function createToolContext(init: Partial<ToolContext> = {}): ToolContext {
+  const session = init.session ?? createSessionFacts();
+  return {
+    actor: init.actor ?? {
+      id: `assistant:${session.sessionId ?? "local"}:`,
+      kind: "assistant",
+      sessionId: session.sessionId,
+      turnId: "",
+    },
+    callId: init.callId ?? "",
+    ledger: init.ledger ?? createLedger(),
+    progress: init.progress ?? (() => {}),
+    session,
+    signal: init.signal ?? NEVER,
+  };
+}
+
+/**
+ * A call's own signal, linked to its turn's: it aborts when the turn does, with the turn's reason,
+ * and {@link release} unlinks it once the call settles, so a Stop later in the turn reaches no call
+ * that has already finished.
+ *
+ * Linked by hand, with one `{ once: true }` listener, rather than with `AbortSignal.any`, so page
+ * webviews and the Worker's ES2023 typing agree.
+ *
+ * @param {AbortSignal} turn
+ * @returns {{ readonly signal: AbortSignal; release(): void }}
+ */
+export function linkCallSignal(turn: AbortSignal): {
+  readonly signal: AbortSignal;
+  release: () => void;
+} {
+  const call = new AbortController();
+  if (turn.aborted) {
+    call.abort(turn.reason);
+    return { release: () => {}, signal: call.signal };
+  }
+  const onAbort = () => {
+    call.abort(turn.reason);
+  };
+  turn.addEventListener("abort", onAbort, { once: true });
+  return {
+    release: () => {
+      turn.removeEventListener("abort", onAbort);
+    },
+    signal: call.signal,
+  };
 }
 
 // ─── ToolRegistry ────────────────────────────────────────────────────────────
@@ -256,9 +394,10 @@ export function createToolRegistry() {
      *
      * @param {string} toolName
      * @param {object} args
+     * @param {ToolContext} [ctx] - The call's context; omitted, a detached one
      * @returns {Promise<ToolResult>}
      */
-    async execute(toolName: string, args: object) {
+    async execute(toolName: string, args: object, ctx: ToolContext = createToolContext()) {
       const tool = _tools.get(toolName);
       if (!tool) {
         return toolError(`Unknown tool: "${toolName}"`);
@@ -271,7 +410,7 @@ export function createToolRegistry() {
       }
 
       try {
-        return await tool.execute(args);
+        return await tool.execute(args, ctx);
       } catch (error) {
         return toolError(`Tool "${toolName}" execution error: ${(error as Error).message}`);
       }