@mlx-node/server 0.0.12 → 0.0.15

package/src/stop-sequence-buffer.ts

@@ -0,0 +1,161 @@
+/**
+ * Buffers streaming text to detect configured stop sequences. Text that
+ * cannot be part of a partial stop sequence is released immediately; a
+ * trailing suffix that could be the start of a stop sequence is held back
+ * until a later push resolves it or the stream is flushed. Once a full stop
+ * sequence is seen, everything after it is suppressed.
+ */
+export class StopSequenceBuffer {
+  private readonly stopSequences: string[];
+  private readonly maxLength: number;
+  private pending_ = '';
+  private _matched: string | null = null;
+
+  constructor(stopSequences: string[]) {
+    // Drop empty AND whitespace-only entries: a whitespace-only stop would
+    // truncate normal output at the first space/newline, and the real
+    // Anthropic API rejects such stops outright. Mirrors the same trim filter
+    // in the request mapper so a whitespace-only configuration is a no-op.
+    this.stopSequences = stopSequences.filter((s) => s.trim().length > 0);
+    this.maxLength = this.stopSequences.reduce((max, s) => Math.max(max, s.length), 0);
+  }
+
+  /**
+   * Earliest index wins; on a tie at the same index the longest wins. Returns
+   * `{ idx, seq }` for the winning stop, or `{ idx: -1, seq: null }` when none
+   * is present in `pending`.
+   */
+  private findMatch(): { idx: number; seq: string | null } {
+    let matchIdx = -1;
+    let matchSeq: string | null = null;
+    for (const seq of this.stopSequences) {
+      const idx = this.pending_.indexOf(seq);
+      if (idx < 0) {
+        continue;
+      }
+      if (matchIdx < 0 || idx < matchIdx || (idx === matchIdx && seq.length > (matchSeq?.length ?? 0))) {
+        matchIdx = idx;
+        matchSeq = seq;
+      }
+    }
+    return { idx: matchIdx, seq: matchSeq };
+  }
+
+  /** The stop sequence that has matched so far, or `null` if none has. */
+  get matched(): string | null {
+    return this._matched;
+  }
+
+  /**
+   * The text currently held back (received but neither emitted nor matched).
+   * The streaming done-path reads this so it can scan the terminal/recovered
+   * text on the SAME buffer with the held partial still in place, and so it
+   * can reconstruct the full received-but-unemitted prefix for overlap math.
+   */
+  get pending(): string {
+    return this.pending_;
+  }
+
+  /**
+   * The earliest start index `j` in `[0, limit]` such that `pending.slice(j)`
+   * is a non-empty STRICT prefix of some configured stop — i.e. a partial that
+   * a later push could still grow into that stop. Returns -1 when no pending
+   * suffix at or before `limit` is viable. Scanning from the front yields the
+   * earliest start index, which is the one whose completed stop would win the
+   * earliest-index tiebreak. A suffix longer than `maxLength - 1` can never be
+   * a strict prefix of any stop, so the search starts no earlier than that.
+   */
+  private earliestViablePrefixIndex(limit: number): number {
+    if (this.pending_.length === 0) {
+      return -1;
+    }
+    const lowerBound = Math.max(0, this.pending_.length - (this.maxLength - 1));
+    const upper = Math.min(limit, this.pending_.length - 1);
+    for (let j = lowerBound; j <= upper; j++) {
+      const suffix = this.pending_.slice(j);
+      if (this.stopSequences.some((seq) => suffix.length < seq.length && seq.startsWith(suffix))) {
+        return j;
+      }
+    }
+    return -1;
+  }
+
+  /**
+   * Feed text in. Returns `safeText` (emit as delta) and `matched` (the stop
+   * sequence that has been matched, or `null`). After a match every push
+   * returns empty `safeText` and keeps reporting the matched sequence.
+   */
+  push(text: string): { safeText: string; matched: string | null } {
+    if (this._matched !== null) {
+      return { safeText: '', matched: this._matched };
+    }
+
+    // Transparent pass-through when there is nothing to detect.
+    if (this.stopSequences.length === 0) {
+      return { safeText: text, matched: null };
+    }
+
+    this.pending_ += text;
+
+    const { idx: matchIdx, seq: matchSeq } = this.findMatch();
+
+    // Earliest start index of a still-growable stop prefix. When a full match
+    // exists we only look at or before it (`limit = matchIdx`): a viable prefix
+    // beginning AFTER the match would complete at a later index and lose the
+    // earliest-index tiebreak, and the bytes from the match onward are
+    // suppressed anyway. A viable prefix at or before the match could still
+    // complete into a stop that WINS (earlier index, or longer at the same
+    // index), so the match must be held. With no full match we consider the
+    // whole pending text.
+    const limit = matchIdx >= 0 ? matchIdx : this.pending_.length - 1;
+    const holdIdx = this.earliestViablePrefixIndex(limit);
+
+    if (matchIdx >= 0 && matchSeq !== null) {
+      if (holdIdx >= 0) {
+        // A longer/earlier stop could still complete from `holdIdx`; emit only
+        // the bytes before it and keep the rest pending for a later push or
+        // `flush()` to resolve.
+        const safeText = this.pending_.slice(0, holdIdx);
+        this.pending_ = this.pending_.slice(holdIdx);
+        return { safeText, matched: null };
+      }
+      const safeText = this.pending_.slice(0, matchIdx);
+      this._matched = matchSeq;
+      this.pending_ = '';
+      return { safeText, matched: matchSeq };
+    }
+
+    // No full match: release everything before the earliest viable prefix and
+    // hold that suffix back, since a later push could complete it.
+    const safeLen = holdIdx >= 0 ? holdIdx : this.pending_.length;
+    const safeText = this.pending_.slice(0, safeLen);
+    this.pending_ = this.pending_.slice(safeLen);
+    return { safeText, matched: null };
+  }
+
+  /**
+   * Release any held-back text at stream end. If a stop sequence already
+   * matched, nothing more is emitted; otherwise the residue could not
+   * complete any sequence and is released.
+   */
+  flush(): { safeText: string; matched: string | null } {
+    if (this._matched !== null) {
+      return { safeText: '', matched: this._matched };
+    }
+    // The stream has ended, so any match `push()` held back for a possible
+    // longer same-index stop can no longer be extended — resolve it now.
+    // Re-scan `pending` for the earliest match (longest on tie) and commit it
+    // if present; otherwise the residue could not complete any sequence and is
+    // released verbatim.
+    const { idx: matchIdx, seq: matchSeq } = this.findMatch();
+    if (matchIdx >= 0 && matchSeq !== null) {
+      const safeText = this.pending_.slice(0, matchIdx);
+      this._matched = matchSeq;
+      this.pending_ = '';
+      return { safeText, matched: matchSeq };
+    }
+    const safeText = this.pending_;
+    this.pending_ = '';
+    return { safeText, matched: null };
+  }
+}

package/src/streaming.ts

@@ -0,0 +1,205 @@
+/** SSE writer utilities. */
+
+import type { IncomingMessage, ServerResponse } from 'node:http';
+
+/** Default time a connected SSE peer may remain continuously backpressured. */
+export const DEFAULT_SSE_DRAIN_TIMEOUT_MS = 30_000;
+
+/**
+ * Responses that have committed to SSE (`beginSSE`) but have not yet been
+ * ended (`endSSE`) or had their connection torn down.
+ *
+ * Purely an accounting aid for {@link activeSSEStreamCount}, which graceful
+ * shutdown reads to report how many streams a forced close cut short. Nothing
+ * in the request path branches on membership, so a miscount can only skew a
+ * diagnostic number — never the behaviour of a live stream.
+ *
+ * Module-scoped because `@mlx-node/server` is loaded once per process and
+ * `beginSSE`/`endSSE` are free functions called from both endpoints.
+ */
+const activeSSEResponses = new Set<ServerResponse>();
+
+/** Disconnect state whose listeners stay armed for one complete SSE handler. */
+export interface SSEClientAbortTracker {
+  readonly aborted: boolean;
+  markAborted(): void;
+  dispose(): void;
+}
+
+/**
+ * Track request/response/socket disconnects until the caller's outermost
+ * `finally`. Keeping this lifetime outside the decode loop matters: a final
+ * native item can expand into backpressured residual protocol frames after the
+ * iterator has already closed, and a disconnect during that drain must still
+ * prevent a success terminal and session adoption.
+ */
+export function trackSSEClientAbort(res: ServerResponse, httpReq: IncomingMessage | undefined): SSEClientAbortTracker {
+  let aborted = false;
+  let disposed = false;
+  const onClose = (): void => {
+    aborted = true;
+  };
+  const onError = (_err: unknown): void => {
+    aborted = true;
+  };
+  const socket = res.socket;
+
+  if (httpReq != null) {
+    httpReq.once('close', onClose);
+    httpReq.on('error', onError);
+  }
+  res.once('close', onClose);
+  res.on('error', onError);
+  if (socket != null) {
+    socket.once('close', onClose);
+    socket.on('error', onError);
+  }
+
+  return {
+    get aborted(): boolean {
+      return aborted;
+    },
+    markAborted(): void {
+      aborted = true;
+    },
+    dispose(): void {
+      if (disposed) return;
+      disposed = true;
+      if (httpReq != null) {
+        httpReq.removeListener('close', onClose);
+        httpReq.removeListener('error', onError);
+      }
+      res.removeListener('close', onClose);
+      res.removeListener('error', onError);
+      if (socket != null) {
+        socket.removeListener('close', onClose);
+        socket.removeListener('error', onError);
+      }
+    },
+  };
+}
+
+interface SSEDrainWaitOptions {
+  /** Override used by focused tests; production reads the env/default. */
+  timeoutMs?: number;
+  /** Mark the owning handler's sticky abort state before the transport closes. */
+  onTimeout?: () => void;
+}
+
+function resolveSSEDrainTimeoutMs(): number {
+  const raw = process.env.MLX_SSE_DRAIN_TIMEOUT_MS;
+  if (raw == null || raw.trim() === '') return DEFAULT_SSE_DRAIN_TIMEOUT_MS;
+  const parsed = Number(raw);
+  return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : DEFAULT_SSE_DRAIN_TIMEOUT_MS;
+}
+
+export function beginSSE(res: ServerResponse): void {
+  activeSSEResponses.add(res);
+  // Belt-and-braces cleanup: a stream torn down by a client disconnect (or by
+  // `server.closeAllConnections()`) may unwind through an error path that
+  // never reaches `endSSE`. Without this the entry would leak for the life of
+  // the process and inflate the count forever.
+  res.once('close', () => {
+    activeSSEResponses.delete(res);
+  });
+  res.writeHead(200, {
+    'Content-Type': 'text/event-stream',
+    'Cache-Control': 'no-cache',
+    Connection: 'keep-alive',
+  });
+}
+
+/** Write one SSE event. Injects `type: eventType` into the payload (data's own `type` wins) for OpenAI SDK compatibility. */
+export function writeSSEEvent(res: ServerResponse, eventType: string, data: object): boolean {
+  const payload = { type: eventType, ...data };
+  return res.write(`event: ${eventType}\ndata: ${JSON.stringify(payload)}\n\n`);
+}
+
+/**
+ * Wait until a backpressured response can accept more data, or until its
+ * transport closes. Close and error resolve rather than reject: the endpoint's
+ * outer abort tracker owns the sticky state, and the next loop check exits
+ * before another native item is written.
+ *
+ * Call this synchronously after `writeSSEEvent` returns false. In particular,
+ * do not defer listener installation until the next iterator turn: `drain`
+ * could fire while that turn is being fetched and leave the handler parked on
+ * an event that already happened.
+ */
+export function awaitDrainOrClose(res: ServerResponse, options: SSEDrainWaitOptions = {}): Promise<void> {
+  return new Promise<void>((resolve) => {
+    let settled = false;
+    const socket = res.socket;
+    const timeoutMs = options.timeoutMs ?? resolveSSEDrainTimeoutMs();
+    const timer = setTimeout(() => {
+      // Set the handler-visible state synchronously. `destroy()` emits close on
+      // a later turn, which is too late for the success gate immediately after
+      // this promise resolves.
+      options.onTimeout?.();
+      if (!res.destroyed) res.destroy();
+      settle();
+    }, timeoutMs);
+    timer.unref();
+    const settle = (): void => {
+      if (settled) return;
+      settled = true;
+      clearTimeout(timer);
+      res.removeListener('drain', onDrain);
+      res.removeListener('close', onClose);
+      res.removeListener('error', onError);
+      if (socket != null) socket.removeListener('close', onClose);
+      resolve();
+    };
+    const onDrain = (): void => {
+      settle();
+    };
+    const onClose = (): void => {
+      settle();
+    };
+    const onError = (_err: unknown): void => {
+      settle();
+    };
+
+    // A destroyed peer cannot emit a future useful drain. Do not include
+    // `writableEnded` here: write-after-end returns false and reports
+    // ERR_STREAM_WRITE_AFTER_END asynchronously through the error listener.
+    if (res.destroyed || (socket != null && socket.destroyed)) {
+      settle();
+      return;
+    }
+
+    res.once('drain', onDrain);
+    res.once('close', onClose);
+    res.once('error', onError);
+    if (socket != null) socket.once('close', onClose);
+  });
+}
+
+export function endSSE(res: ServerResponse): void {
+  activeSSEResponses.delete(res);
+  res.end();
+}
+
+/**
+ * Number of SSE streams currently open process-wide. Diagnostics and
+ * shutdown accounting only.
+ */
+export function activeSSEStreamCount(): number {
+  return activeSSEResponses.size;
+}
+
+/**
+ * Number of active SSE streams among a caller-owned collection of responses.
+ *
+ * `createServer()` uses this to intersect the process-wide SSE registry with
+ * the responses accepted by one `node:http` Server. Keep the no-argument
+ * {@link activeSSEStreamCount} above for process-wide diagnostics and
+ * standalone `createHandler()` consumers.
+ */
+export function activeSSEStreamCountForResponses(responses: WeakSet<ServerResponse>): number {
+  let count = 0;
+  for (const response of activeSSEResponses) {
+    if (responses.has(response)) count += 1;
+  }
+  return count;
+}

package/src/text-recovery.ts

@@ -0,0 +1,41 @@
+/**
+ * Streaming text-recovery helpers shared between the `/v1/messages` and
+ * `/v1/responses` endpoints.
+ *
+ * Both endpoints have a tool-call streaming recovery branch that has to
+ * compute the unsent suffix of `finalText` given that some prefix of the
+ * model's output may already have been streamed to the wire, but native-side
+ * string normalization makes `finalText` diverge from the streamed-prefix
+ * verbatim. Concrete divergences seen in practice:
+ *
+ *   * The native side trims leading whitespace after `</think>` via
+ *     `split_at_think_end`, so the streamed text can end in `"\n\n"` while
+ *     `finalText` starts at `"<tool_call>"` (no overlap — emit `finalText`
+ *     whole).
+ *   * The native side `.trim()`s tool-tag-bracketed content boundaries, so
+ *     the streamed text can have a trailing space that `finalText` lacks
+ *     (also no overlap — emit `finalText` whole).
+ *
+ * Internal-only — not exported from `packages/server/src/index.ts`.
+ */
+
+/**
+ * Find the largest k such that `streamed.endsWith(final.slice(0, k))`.
+ *
+ * Returns 0 when there is no overlap (caller emits `final` whole).
+ * Returns `final.length` when `final` is fully contained as a suffix of
+ * `streamed` (caller emits nothing).
+ *
+ * Used by the `/v1/messages` and `/v1/responses` streaming tool-call
+ * recovery branches to decide how much of `finalText` is already on the
+ * wire when native-side normalization (e.g. `.trim()`, post-`</think>`
+ * whitespace stripping) makes the streamed prefix diverge from the
+ * `finalText` prefix verbatim.
+ */
+export function longestSuffixPrefixOverlap(streamed: string, final: string): number {
+  const max = Math.min(streamed.length, final.length);
+  for (let k = max; k > 0; k--) {
+    if (streamed.endsWith(final.slice(0, k))) return k;
+  }
+  return 0;
+}

package/src/timing.ts

@@ -0,0 +1,236 @@
+/** Wire-safe server timing extensions derived from native performance metrics. */
+
+export interface PerformanceMetricsForUsage {
+  ttftMs?: number;
+  prefillTokensPerSecond?: number;
+  decodeTokensPerSecond?: number;
+}
+
+interface TimingUsageExtensions {
+  /** Server-extension: native time-to-first-token in milliseconds. */
+  time_to_first_token_ms?: number;
+  /** Server-extension: prompt-token throughput for the tokens actually prefetched this turn. */
+  prefill_tokens_per_second?: number;
+  /** Server-extension: generated-token throughput during decode. */
+  decode_tokens_per_second?: number;
+  /** Server-extension: native/server inference elapsed, excluding HTTP transport and logging overhead. */
+  server_inference_elapsed_ms?: number;
+  /** Server-extension alias for disambiguating native TTFT from request/HTTP elapsed time. */
+  server_time_to_first_token_ms?: number;
+  /** Server-extension: handler-start to first native token, including model resolve/load and queue wait. */
+  server_total_time_to_first_token_ms?: number;
+  /** Server-extension alias for native prefill throughput. */
+  server_prefill_tokens_per_second?: number;
+  /** Server-extension alias for native decode throughput. */
+  server_decode_tokens_per_second?: number;
+  /** Server-extension: prompt tokens actually prefetched this turn after cached-prefix reuse. */
+  prefill_input_tokens?: number;
+  /** Server-extension: prompt tokens skipped because a cached prefix was reused. */
+  cached_prefix_tokens?: number;
+  /**
+   * Server-extension: time spent resolving/loading/aliasing the requested model
+   * before registry lookup. Includes both the synchronous lookup AND any time
+   * spent driving the load. Excludes time spent blocked behind a peer
+   * request's in-flight load — that wait is reported separately via
+   * `server_load_wait_ms` so a fast follower request is not mis-attributed
+   * a long resolve when it merely inherited a cold-load wait.
+   */
+  server_model_resolve_ms?: number;
+  /**
+   * Server-extension: wall-clock time this request spent blocked on the
+   * process-wide model-load writer lock. Set when the load was already
+   * in flight when this request arrived (a peer drove the load and we
+   * waited for it to finish) AND when this request itself drove the
+   * load. Inspect `server_load_owner` to disambiguate. When the writer
+   * lock was free and the resolve was a no-op (model already loaded),
+   * this field is elided.
+   */
+  server_load_wait_ms?: number;
+  /**
+   * Server-extension: `true` when this request acquired the model-load
+   * writer lock with no contention (i.e. either the model was already
+   * loaded and the call was a no-op, or this request itself drove the
+   * load). `false` when the request was parked behind a peer's in-flight
+   * load. Elided when no load coordinator is wired (single-process
+   * tests, embedded callers).
+   */
+  server_load_owner?: boolean;
+  /** Server-extension: time spent waiting behind the per-model execution mutex. */
+  server_queue_ms?: number;
+  /** Server-extension: handler time before native inference begins, including resolve and queue wait. */
+  server_pre_inference_ms?: number;
+  /** Server-extension: effective process-level paged-prefill chunk size. */
+  server_paged_prefill_chunk_size?: number;
+  /** Server-extension: effective process-level paged-prefill eval/clear cadence. */
+  server_paged_prefill_eval_interval?: number;
+  /** Server-extension: effective process-level paged-decode cache-clear cadence. */
+  server_paged_decode_cache_clear_interval?: number;
+}
+
+function finitePositive(value: number | undefined): number | undefined {
+  return value != null && Number.isFinite(value) && value > 0 ? value : undefined;
+}
+
+function finiteNonNegativeInteger(value: number | undefined): number | undefined {
+  return value != null && Number.isFinite(value) && value >= 0 ? Math.max(0, Math.floor(value)) : undefined;
+}
+
+function finiteNonNegative(value: number | undefined): number | undefined {
+  return value != null && Number.isFinite(value) && value >= 0 ? value : undefined;
+}
+
+export interface ServerTimingForUsage {
+  server_model_resolve_ms?: number;
+  server_load_wait_ms?: number;
+  server_load_owner?: boolean;
+  server_queue_ms?: number;
+  server_pre_inference_ms?: number;
+  server_paged_prefill_chunk_size?: number;
+  server_paged_prefill_eval_interval?: number;
+  server_paged_decode_cache_clear_interval?: number;
+}
+
+const I32_MAX = 0x7fff_ffff;
+
+function parseI32(value: string | undefined): number | undefined {
+  if (value == null) return undefined;
+  const trimmed = value.trim();
+  if (!/^[+-]?\d+$/.test(trimmed)) return undefined;
+  const parsed = Number.parseInt(trimmed, 10);
+  return Number.isSafeInteger(parsed) && parsed >= -0x8000_0000 && parsed <= I32_MAX ? parsed : undefined;
+}
+
+function parseNonNegativeI32(value: string | undefined, fallback: number): number {
+  const parsed = parseI32(value);
+  return parsed != null && parsed >= 0 ? parsed : fallback;
+}
+
+function parsePositiveI32(value: string | undefined, fallback: number): number {
+  const parsed = parseI32(value);
+  return parsed != null && parsed > 0 ? parsed : fallback;
+}
+
+export function resolveServerTuningForUsage(
+  env: Record<string, string | undefined> = process.env,
+): Pick<
+  ServerTimingForUsage,
+  'server_paged_prefill_chunk_size' | 'server_paged_prefill_eval_interval' | 'server_paged_decode_cache_clear_interval'
+> {
+  return {
+    server_paged_prefill_chunk_size: parseNonNegativeI32(env.MLX_PAGED_PREFILL_CHUNK_SIZE, 0),
+    server_paged_prefill_eval_interval: parsePositiveI32(env.MLX_PAGED_PREFILL_EVAL_INTERVAL, 8),
+    server_paged_decode_cache_clear_interval: parsePositiveI32(env.MLX_PAGED_DECODE_CACHE_CLEAR_INTERVAL, 1024),
+  };
+}
+
+function computeServerInferenceElapsedMs(
+  ttftMs: number | undefined,
+  decodeTokensPerSecond: number | undefined,
+  outputTokens: number | undefined,
+): number | undefined {
+  if (ttftMs == null) return undefined;
+
+  const generatedTokens = finiteNonNegativeInteger(outputTokens);
+  if (generatedTokens == null) {
+    return ttftMs;
+  }
+  if (generatedTokens <= 1) {
+    return ttftMs;
+  }
+  if (decodeTokensPerSecond == null) {
+    return undefined;
+  }
+  return ttftMs + ((generatedTokens - 1) / decodeTokensPerSecond) * 1000;
+}
+
+function buildTimingUsageExtensions(
+  performance: PerformanceMetricsForUsage | undefined,
+  promptTokens: number | undefined,
+  outputTokens: number | undefined,
+  cachedTokens: number | undefined,
+  serverTiming?: ServerTimingForUsage,
+): TimingUsageExtensions {
+  const ttftMs = finitePositive(performance?.ttftMs);
+  const prefillTokensPerSecond = finitePositive(performance?.prefillTokensPerSecond);
+  const decodeTokensPerSecond = finitePositive(performance?.decodeTokensPerSecond);
+  const serverInferenceElapsedMs = computeServerInferenceElapsedMs(ttftMs, decodeTokensPerSecond, outputTokens);
+  const preInferenceMs = finiteNonNegative(serverTiming?.server_pre_inference_ms);
+
+  const extensions: TimingUsageExtensions = {};
+  if (ttftMs != null) {
+    extensions.time_to_first_token_ms = ttftMs;
+    extensions.server_time_to_first_token_ms = ttftMs;
+    if (preInferenceMs != null) {
+      extensions.server_total_time_to_first_token_ms = preInferenceMs + ttftMs;
+    }
+  }
+  if (prefillTokensPerSecond != null) {
+    extensions.prefill_tokens_per_second = prefillTokensPerSecond;
+    extensions.server_prefill_tokens_per_second = prefillTokensPerSecond;
+  }
+  if (decodeTokensPerSecond != null) {
+    extensions.decode_tokens_per_second = decodeTokensPerSecond;
+    extensions.server_decode_tokens_per_second = decodeTokensPerSecond;
+  }
+  if (serverInferenceElapsedMs != null && Number.isFinite(serverInferenceElapsedMs) && serverInferenceElapsedMs > 0) {
+    extensions.server_inference_elapsed_ms = serverInferenceElapsedMs;
+  }
+
+  if (performance != null) {
+    const prompt = finiteNonNegativeInteger(promptTokens);
+    const cached = finiteNonNegativeInteger(cachedTokens);
+    if (prompt != null) {
+      const cachedPrefix = cached == null ? 0 : Math.min(cached, prompt);
+      extensions.prefill_input_tokens = prompt - cachedPrefix;
+      if (cachedPrefix > 0) {
+        extensions.cached_prefix_tokens = cachedPrefix;
+      }
+    }
+  }
+
+  const modelResolveMs = finiteNonNegative(serverTiming?.server_model_resolve_ms);
+  if (modelResolveMs != null) {
+    extensions.server_model_resolve_ms = modelResolveMs;
+  }
+  const loadWaitMs = finiteNonNegative(serverTiming?.server_load_wait_ms);
+  if (loadWaitMs != null) {
+    extensions.server_load_wait_ms = loadWaitMs;
+  }
+  if (typeof serverTiming?.server_load_owner === 'boolean') {
+    extensions.server_load_owner = serverTiming.server_load_owner;
+  }
+  const queueMs = finiteNonNegative(serverTiming?.server_queue_ms);
+  if (queueMs != null) {
+    extensions.server_queue_ms = queueMs;
+  }
+  if (preInferenceMs != null) {
+    extensions.server_pre_inference_ms = preInferenceMs;
+  }
+  const pagedPrefillChunkSize = finiteNonNegativeInteger(serverTiming?.server_paged_prefill_chunk_size);
+  if (pagedPrefillChunkSize != null) {
+    extensions.server_paged_prefill_chunk_size = pagedPrefillChunkSize;
+  }
+  const pagedPrefillEvalInterval = finiteNonNegativeInteger(serverTiming?.server_paged_prefill_eval_interval);
+  if (pagedPrefillEvalInterval != null) {
+    extensions.server_paged_prefill_eval_interval = pagedPrefillEvalInterval;
+  }
+  const pagedDecodeCacheClearInterval = finiteNonNegativeInteger(
+    serverTiming?.server_paged_decode_cache_clear_interval,
+  );
+  if (pagedDecodeCacheClearInterval != null) {
+    extensions.server_paged_decode_cache_clear_interval = pagedDecodeCacheClearInterval;
+  }
+
+  return extensions;
+}
+
+export function mergeTimingUsageExtensions<T extends TimingUsageExtensions>(
+  usage: T,
+  performance: PerformanceMetricsForUsage | undefined,
+  promptTokens: number | undefined,
+  outputTokens: number | undefined,
+  cachedTokens: number | undefined,
+  serverTiming?: ServerTimingForUsage,
+): void {
+  Object.assign(usage, buildTimingUsageExtensions(performance, promptTokens, outputTokens, cachedTokens, serverTiming));
+}