experimental-a2 0.17.0 → 0.18.0

package/src/ai-fx.ts

@@ -0,0 +1,833 @@
+/**
+ * experimental-a2/ai/fx — `fx()`: the FX harness for `agent({ harness })`.
+ * Server-only; FX's kernel comes from the optional `libfx` peer dependency.
+ */
+
+// oxlint-disable no-await-in-loop -- turns replay in order, each on the kernel the last one left.
+
+import {
+  asSchema,
+  type FinishReason,
+  type LanguageModelUsage,
+  type ToolSet,
+  type UIMessageChunk,
+} from 'ai'
+import type { FxAgent, FxHostTool, FxTurn } from 'libfx'
+import {
+  type Deferred,
+  defineHarness,
+  deferred,
+  Divergence,
+  gatewayCredential,
+  type GatewayCredential,
+  type Harness,
+  type HarnessGeneration,
+  type HarnessRead,
+  type HarnessStep,
+  type HarnessTurn,
+  pending,
+  readHarness,
+  saveHarness,
+} from './ai-harness.ts'
+import type { CompactionPartData, GenerationCompletedPayload } from './ai.ts'
+import { boundedCache } from './bounded-cache.ts'
+
+export type { Harness } from './ai-harness.ts'
+
+/**
+ * The FX harness. Every model call FX makes is one durable step; A2 rebuilds
+ * the kernel for each step from its latest checkpoint and the log since.
+ */
+export function fx(): Harness<'fx', never> {
+  return defineHarness('fx', {
+    tools: {},
+    step: runStep,
+  })
+}
+
+const SNAPSHOT_VERSION = 1
+const CATALOG_TTL_MS = 60 * 60 * 1000
+
+type StreamPart = { type: string } & Record<string, unknown>
+type ToolCall = { toolCallId: string; toolName: string; input: unknown }
+
+/** A model response A2 recorded, in the Gateway's stream format. */
+type RecordedCall = {
+  kind: 'step' | 'compaction'
+  sse: string
+  toolCalls: ToolCall[]
+}
+
+/** One FX turn to replay: its prompt, recorded calls, and how it ended. */
+type TurnPlan = {
+  prompt: string
+  calls: RecordedCall[]
+  /** Steering to queue while serving the call at that index. */
+  steers: Map<number, string[]>
+  end: 'stop' | 'interrupted' | 'live'
+}
+
+type Completion = { finishReason?: FinishReason; usage?: LanguageModelUsage }
+
+/** The turn being replayed or run live, as the fetch and tool hooks see it. */
+type TurnState = {
+  plan: TurnPlan
+  fxTurn: FxTurn
+  next: number
+  /** Tool calls of the last served step and the ones already answered. */
+  served: ToolCall[]
+  answered: Set<string>
+  live: boolean
+  liveStep: boolean
+  failure?: Error
+  /** Resolves when the live step is over: FX asked for its tools. */
+  toolsRequested: Deferred
+  parsed?: Promise<void>
+}
+
+function fail(state: TurnState, error: Error): never {
+  state.failure ??= error
+  state.fxTurn.cancel()
+  throw state.failure
+}
+
+const diverged = (message: string): Divergence => new Divergence('FX', message)
+
+// ---------------------------------------------------------------- step
+
+async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
+  const { createFxAgent } = await loadLibfx()
+  const credential = gatewayCredential('fx')
+  const name = `${step.agentName}.a2.harness.fx`
+  const { replayAll, ...first } = await readHarness({
+    step,
+    name,
+    parse: savedCheckpoint,
+  })
+  let read: HarnessRead<Uint8Array> = first
+
+  let completion: Completion = {}
+  let stopped = false
+  let agent: FxAgent | undefined
+  let turn: TurnState | undefined
+
+  const run = async (
+    controller: ReadableStreamDefaultController<UIMessageChunk>,
+    turns: TurnPlan[],
+  ): Promise<void> => {
+    controller.enqueue({ type: 'start', messageId: step.responseMessageId })
+    for (const plan of turns) {
+      if (stopped) return
+      const fxTurn = agent!.prompt(plan.prompt, { signal: step.signal })
+      const state: TurnState = {
+        plan,
+        fxTurn,
+        next: 0,
+        served: [],
+        answered: new Set(),
+        live: false,
+        liveStep: false,
+        toolsRequested: deferred(),
+      }
+      turn = state
+      const drained = (async () => {
+        for await (const _ of fxTurn) {
+          // Live output reaches A2 from the Gateway response itself.
+        }
+      })()
+      // A failed turn surfaces through its result.
+      drained.catch(() => {})
+      if (plan.end !== 'live') {
+        const result = await fxTurn.result.catch((error: unknown) => {
+          throw state.failure ?? error
+        })
+        await drained
+        if (state.failure) throw state.failure
+        const expected = plan.end === 'stop' ? 'end_turn' : 'cancelled'
+        if (result.stopReason !== expected)
+          throw diverged(
+            `a recorded turn ended with ${result.stopReason}, not ${expected}`,
+          )
+        continue
+      }
+      const outcome = await Promise.race([
+        fxTurn.result.then((result) => ({ ended: true, result })),
+        state.toolsRequested.promise.then(() => ({ ended: false as const })),
+      ]).catch((error: unknown) => {
+        throw state.failure ?? error
+      })
+      if (state.failure) throw state.failure
+      await state.parsed
+      if (outcome.ended) {
+        await drained
+        // Only a turn this attempt finished describes the log.
+        if (outcome.result.stopReason === 'end_turn' && !step.signal.aborted)
+          await saveCheckpoint()
+      } else {
+        // FX now waits on tools that A2's tool handlers run.
+        fxTurn.cancel()
+      }
+      controller.enqueue({ type: 'finish-step' })
+      controller.enqueue({
+        type: 'finish',
+        ...(completion.finishReason === undefined
+          ? {}
+          : { finishReason: completion.finishReason }),
+      })
+    }
+  }
+
+  const saveCheckpoint = async (): Promise<void> => {
+    let checkpoint: Uint8Array
+    try {
+      checkpoint = await agent!.checkpoint()
+    } catch {
+      // Over libfx's bound: later steps replay from the previous snapshot.
+      return
+    }
+    await saveHarness({
+      step,
+      name,
+      state: {
+        version: SNAPSHOT_VERSION,
+        checkpoint: Buffer.from(checkpoint).toString('base64'),
+      },
+    })
+  }
+
+  const hostFetch = async (
+    input: string | URL | Request,
+    init: RequestInit = {},
+  ): Promise<Response> => {
+    const url = input instanceof Request ? input.url : String(input)
+    if ((init.method ?? 'GET').toUpperCase() === 'GET')
+      return catalog(url, credential, withAuth(init.headers, credential))
+    const state = turn!
+    if (state.failure) throw state.failure
+    const body = decodeBody(init.body)
+    const compaction = isCompaction(body)
+    const recorded = state.plan.calls[state.next]
+    if (recorded !== undefined) {
+      if ((recorded.kind === 'compaction') !== compaction)
+        fail(state, diverged(`request ${state.next} is not a ${recorded.kind}`))
+      const index = state.next++
+      if (recorded.kind === 'step') {
+        state.served = recorded.toolCalls
+        state.answered = new Set()
+      }
+      const steers = state.plan.steers.get(index) ?? []
+      return relay(recorded.sse, async () => {
+        for (const text of steers) await state.fxTurn.steer(text)
+      })
+    }
+    if (state.plan.end === 'stop')
+      fail(state, diverged('FX continued a turn the log ended'))
+    if (state.plan.end === 'interrupted') {
+      state.fxTurn.cancel()
+      throw new DOMException('The recorded turn was interrupted', 'AbortError')
+    }
+    if (state.liveStep && !compaction) {
+      // One generation is one step; the next step is A2's next generation.
+      state.toolsRequested.resolve()
+      return pending(init.signal)
+    }
+    state.live = true
+    return live(state, url, init, compaction)
+  }
+
+  let compactions = 0
+  const live = async (
+    state: TurnState,
+    url: string,
+    init: RequestInit,
+    compaction: boolean,
+  ): Promise<Response> => {
+    const response = await fetch(url, {
+      ...init,
+      headers: withAuth(init.headers, credential),
+    })
+    if (!response.ok || response.body === null) return response
+    const [forKernel, forLog] = response.body.tee()
+    // Live responses reach the log in request order: a compaction's part
+    // precedes the step that follows it.
+    const previous = state.parsed
+    if (compaction) {
+      // Unique across the response's generations: the transcript keys data parts by ID.
+      const id = `fx-compaction:${step.requestIndex}:${compactions++}`
+      state.parsed = (async () => {
+        const all = await readParts(forLog)
+        await previous
+        emit({
+          type: 'data-a2-compaction',
+          id,
+          data: compactionData(all),
+        })
+      })()
+    } else {
+      state.liveStep = true
+      state.parsed = (async () => {
+        await previous
+        emit({ type: 'start-step' })
+        await streamStep(forLog, emit, (finished) => {
+          completion = finished
+        })
+      })()
+    }
+    // A stopped step abandons its parse; its failure surfaces through the turn.
+    state.parsed.catch(() => {})
+    return new Response(forKernel, {
+      status: response.status,
+      headers: response.headers,
+    })
+  }
+
+  const hostTools = await kernelTools(step.tools, (toolName, toolInput) => {
+    const state = turn!
+    if (state.failure) throw state.failure
+    const call = state.served.find(
+      (candidate) =>
+        candidate.toolName === toolName &&
+        !state.answered.has(candidate.toolCallId) &&
+        sameJson(candidate.input, toolInput),
+    )
+    const result = call === undefined ? undefined : results.get(call.toolCallId)
+    if (call !== undefined && result !== undefined) {
+      state.answered.add(call.toolCallId)
+      if (result.error !== undefined) throw new Error(result.error)
+      if (result.denied === true) throw new Error('The tool call was denied.')
+      return result.output ?? null
+    }
+    if (state.live) {
+      state.toolsRequested.resolve()
+      return pending(undefined)
+    }
+    if (state.plan.end === 'interrupted') {
+      state.fxTurn.cancel()
+      return pending(undefined)
+    }
+    return fail(state, diverged(`no recorded result for ${toolName}`))
+  })
+
+  let turns = read.turns.map(turnPlan)
+  let results = read.results
+  let controller: ReadableStreamDefaultController<UIMessageChunk> | undefined
+  let closed = false
+  const emit = (chunk: UIMessageChunk): void => {
+    if (!closed) controller!.enqueue(chunk)
+  }
+
+  const create = (checkpoint: Uint8Array | undefined) =>
+    createFxAgent({
+      apiKey: credential.token,
+      model: kernelModel(step),
+      ...(step.instructions.length === 0
+        ? {}
+        : { instructions: step.instructions }),
+      tools: hostTools,
+      ...(checkpoint === undefined ? {} : { checkpoint }),
+      fetch: hostFetch,
+    })
+
+  const stream = new ReadableStream<UIMessageChunk>({
+    start(streamController) {
+      controller = streamController
+      const abort = () => {
+        stopped = true
+        turn?.fxTurn.cancel()
+      }
+      step.signal.addEventListener('abort', abort, { once: true })
+      void (async () => {
+        try {
+          try {
+            agent = await create(read.saved)
+          } catch (error) {
+            if (read.saved === undefined) throw error
+            // An unusable checkpoint: replay the whole log instead.
+            read = await replayAll()
+            turns = read.turns.map(turnPlan)
+            results = read.results
+            agent = await create(undefined)
+          }
+          await run(streamController, turns)
+          closed = true
+          streamController.close()
+        } catch (error) {
+          closed = true
+          streamController.error(
+            step.signal.aborted ? step.signal.reason : error,
+          )
+        } finally {
+          step.signal.removeEventListener('abort', abort)
+          await agent?.close().catch(() => {})
+        }
+      })()
+    },
+    cancel() {
+      closed = true
+      stopped = true
+      turn?.fxTurn.cancel()
+    },
+  })
+
+  return { stream, completion: () => completion }
+}
+
+// ------------------------------------------------------------ the plan
+
+function turnPlan(turn: HarnessTurn): TurnPlan {
+  const calls: RecordedCall[] = []
+  const steers = new Map<number, string[]>()
+  const steer = (at: number, texts: string[] | undefined) => {
+    if (texts !== undefined) steers.set(Math.max(0, at), texts)
+  }
+  steer(0, turn.generations.length === 0 ? turn.steers.get(0) : undefined)
+  turn.generations.forEach((generation, index) => {
+    calls.push(...recordedCalls(generation.chunks, generation.completed))
+    // A generation's step is its last call.
+    steer(calls.length - 1, turn.steers.get(index))
+  })
+  return { prompt: turn.prompt, calls, steers, end: turn.end }
+}
+
+/** A completed generation's model calls: its compactions, then its step. */
+function recordedCalls(
+  chunks: readonly UIMessageChunk[],
+  completed: GenerationCompletedPayload,
+): RecordedCall[] {
+  const calls: RecordedCall[] = []
+  const parts: StreamPart[] = []
+  const toolCalls: ToolCall[] = []
+  // A streamed call's input is the provider's exact text; the kernel copies it.
+  const inputText = new Map<string, string>()
+  for (const chunk of chunks) {
+    const metadata =
+      'providerMetadata' in chunk && chunk.providerMetadata !== undefined
+        ? { providerMetadata: chunk.providerMetadata }
+        : {}
+    switch (chunk.type) {
+      case 'data-a2-compaction':
+        calls.push({
+          kind: 'compaction',
+          sse: compactionSse(chunk.data as CompactionPartData),
+          toolCalls: [],
+        })
+        break
+      case 'text-start':
+      case 'text-end':
+      case 'reasoning-start':
+      case 'reasoning-end':
+        parts.push({ type: chunk.type, id: chunk.id, ...metadata })
+        break
+      case 'text-delta':
+      case 'reasoning-delta':
+        parts.push({
+          type: chunk.type,
+          id: chunk.id,
+          delta: chunk.delta,
+          ...metadata,
+        })
+        break
+      case 'tool-input-start':
+        inputText.set(chunk.toolCallId, '')
+        parts.push({
+          type: 'tool-input-start',
+          id: chunk.toolCallId,
+          toolName: chunk.toolName,
+          ...metadata,
+        })
+        break
+      case 'tool-input-delta':
+        inputText.set(
+          chunk.toolCallId,
+          (inputText.get(chunk.toolCallId) ?? '') + chunk.inputTextDelta,
+        )
+        parts.push({
+          type: 'tool-input-delta',
+          id: chunk.toolCallId,
+          delta: chunk.inputTextDelta,
+        })
+        break
+      case 'tool-input-error':
+        // A call whose input did not parse: FX answered it itself.
+        if (inputText.has(chunk.toolCallId))
+          parts.push({ type: 'tool-input-end', id: chunk.toolCallId })
+        parts.push({
+          type: 'tool-call',
+          toolCallId: chunk.toolCallId,
+          toolName: chunk.toolName,
+          input:
+            inputText.get(chunk.toolCallId) ||
+            (typeof chunk.input === 'string'
+              ? chunk.input
+              : JSON.stringify(chunk.input)),
+        })
+        break
+      case 'tool-input-available':
+        if (inputText.has(chunk.toolCallId))
+          parts.push({ type: 'tool-input-end', id: chunk.toolCallId })
+        toolCalls.push({
+          toolCallId: chunk.toolCallId,
+          toolName: chunk.toolName,
+          input: chunk.input,
+        })
+        parts.push({
+          type: 'tool-call',
+          toolCallId: chunk.toolCallId,
+          toolName: chunk.toolName,
+          input: inputText.get(chunk.toolCallId) || JSON.stringify(chunk.input),
+          ...metadata,
+        })
+        break
+    }
+  }
+  parts.push({
+    type: 'finish',
+    finishReason: { unified: completed.finishReason ?? 'other' },
+    usage: streamUsage(completed.usage),
+  })
+  calls.push({ kind: 'step', sse: sse(parts), toolCalls })
+  return calls
+}
+
+function compactionSse(data: CompactionPartData): string {
+  return sse([
+    { type: 'text-start', id: '0' },
+    { type: 'text-delta', id: '0', delta: data.summary ?? '' },
+    { type: 'text-end', id: '0' },
+    {
+      type: 'finish',
+      finishReason: { unified: 'stop' },
+      usage: streamUsage(data.usage),
+    },
+  ])
+}
+
+// ------------------------------------------------------- live responses
+
+/** Streams one live step response as UI chunks and reports its finish. */
+async function streamStep(
+  body: ReadableStream<Uint8Array>,
+  emit: (chunk: UIMessageChunk) => void,
+  finished: (completion: Completion) => void,
+): Promise<void> {
+  for await (const part of sseParts(body)) {
+    const metadata =
+      part['providerMetadata'] === undefined
+        ? {}
+        : { providerMetadata: part['providerMetadata'] as never }
+    switch (part.type) {
+      case 'text-start':
+      case 'text-end':
+      case 'reasoning-start':
+      case 'reasoning-end':
+        emit({ type: part.type, id: part['id'] as string, ...metadata })
+        break
+      case 'text-delta':
+      case 'reasoning-delta':
+        emit({
+          type: part.type,
+          id: part['id'] as string,
+          delta: part['delta'] as string,
+          ...metadata,
+        })
+        break
+      case 'tool-input-start':
+        emit({
+          type: 'tool-input-start',
+          toolCallId: part['id'] as string,
+          toolName: part['toolName'] as string,
+          ...metadata,
+        })
+        break
+      case 'tool-input-delta':
+        emit({
+          type: 'tool-input-delta',
+          toolCallId: part['id'] as string,
+          inputTextDelta: part['delta'] as string,
+        })
+        break
+      case 'tool-call': {
+        const text = part['input'] as string
+        let input: unknown
+        try {
+          input = JSON.parse(text)
+        } catch (error) {
+          emit({
+            type: 'tool-input-error',
+            toolCallId: part['toolCallId'] as string,
+            toolName: part['toolName'] as string,
+            input: text,
+            errorText: error instanceof Error ? error.message : String(error),
+          })
+          break
+        }
+        emit({
+          type: 'tool-input-available',
+          toolCallId: part['toolCallId'] as string,
+          toolName: part['toolName'] as string,
+          input,
+          ...metadata,
+        })
+        break
+      }
+      case 'finish': {
+        const reason = part['finishReason'] as { unified?: FinishReason }
+        finished({
+          ...(reason?.unified === undefined
+            ? {}
+            : { finishReason: reason.unified }),
+          usage: modelUsage(part['usage']),
+        })
+        break
+      }
+      case 'error':
+        emit({ type: 'error', errorText: errorText(part['error']) })
+        break
+    }
+  }
+}
+
+async function readParts(
+  body: ReadableStream<Uint8Array>,
+): Promise<StreamPart[]> {
+  const all: StreamPart[] = []
+  for await (const part of sseParts(body)) all.push(part)
+  return all
+}
+
+function compactionData(all: readonly StreamPart[]): CompactionPartData {
+  const finish = all.find((part) => part.type === 'finish')
+  return {
+    status: 'completed',
+    summary: all
+      .filter((part) => part.type === 'text-delta')
+      .map((part) => part['delta'] as string)
+      .join(''),
+    ...(finish === undefined ? {} : { usage: modelUsage(finish['usage']) }),
+  }
+}
+
+async function* sseParts(
+  body: ReadableStream<Uint8Array>,
+): AsyncGenerator<StreamPart> {
+  const decoder = new TextDecoder()
+  let pendingText = ''
+  for await (const chunk of body) {
+    pendingText += decoder.decode(chunk, { stream: true })
+    let end: number
+    while ((end = pendingText.indexOf('\n\n')) !== -1) {
+      const frame = pendingText.slice(0, end)
+      pendingText = pendingText.slice(end + 2)
+      if (frame.startsWith('data: {')) yield JSON.parse(frame.slice(6))
+    }
+  }
+}
+
+/**
+ * Hands a recorded response to the kernel frame by frame. Steering due after
+ * it is queued before its finish part, while the kernel still waits on it.
+ */
+function relay(text: string, beforeFinish: () => Promise<void>): Response {
+  const encoder = new TextEncoder()
+  const frames = text.split(/(?<=\n\n)/)
+  let index = 0
+  let steered = false
+  return new Response(
+    new ReadableStream({
+      async pull(controller) {
+        const frame = frames[index++]
+        if (frame === undefined) {
+          controller.close()
+          return
+        }
+        if (!steered && frame.includes('"type":"finish"')) {
+          steered = true
+          await beforeFinish()
+        }
+        controller.enqueue(encoder.encode(frame))
+      },
+    }),
+    { headers: { 'content-type': 'text/event-stream' } },
+  )
+}
+
+// ----------------------------------------------------------- the kernel
+
+async function loadLibfx(): Promise<typeof import('libfx')> {
+  try {
+    return await import('libfx')
+  } catch (error) {
+    throw new Error('fx() needs the libfx package: install libfx', {
+      cause: error,
+    })
+  }
+}
+
+function withAuth(
+  headers: HeadersInit | undefined,
+  credential: GatewayCredential,
+): Headers {
+  const result = new Headers(headers)
+  result.set('ai-gateway-auth-method', credential.method)
+  return result
+}
+
+function kernelModel(
+  step: HarnessStep,
+): string | { id: string; effort?: string; fast?: boolean } {
+  const effort =
+    step.reasoning === undefined || step.reasoning === 'provider-default'
+      ? undefined
+      : step.reasoning
+  const fast = step.speed === 'fast'
+  if (effort === undefined && !fast) return step.model
+  return {
+    id: step.model,
+    ...(effort === undefined ? {} : { effort }),
+    ...(fast ? { fast } : {}),
+  }
+}
+
+async function kernelTools(
+  tools: ToolSet,
+  execute: (toolName: string, input: unknown) => unknown,
+): Promise<FxHostTool[]> {
+  return Promise.all(
+    Object.entries(tools)
+      .filter(([, tool]) => tool.type !== 'provider')
+      .map(async ([toolName, tool]) => ({
+        name: toolName,
+        description:
+          typeof tool.description === 'string' ? tool.description : toolName,
+        inputSchema: (await asSchema(tool.inputSchema).jsonSchema) as object,
+        execute: (input: unknown) => execute(toolName, input),
+      })),
+  )
+}
+
+// Keyed by credential, and Vercel's OIDC tokens rotate, so it is bounded.
+const catalogCache = boundedCache<
+  string,
+  { expires: number; body: Promise<string> }
+>({ maxEntries: 16, maxBytes: 16 })
+
+/**
+ * Model catalog reads, shared for an hour by every step in the process that
+ * uses the same credential; credentials can see different catalogs.
+ */
+async function catalog(
+  url: string,
+  credential: GatewayCredential,
+  headers: Headers,
+): Promise<Response> {
+  const key = `${url}\n${credential.token}`
+  const now = Date.now()
+  let entry = catalogCache.get(key)
+  if (entry === undefined || entry.expires <= now) {
+    const body = fetch(url, { headers }).then(async (response) => {
+      if (!response.ok)
+        throw new Error(`catalog request failed: ${response.status}`)
+      return response.text()
+    })
+    entry = { expires: now + CATALOG_TTL_MS, body }
+    catalogCache.set(key, entry, 1)
+    body.catch(() => catalogCache.delete(key))
+  }
+  return new Response(await entry.body, {
+    headers: { 'content-type': 'application/json' },
+  })
+}
+
+// ------------------------------------------------------------- helpers
+
+function savedCheckpoint(state: unknown): Uint8Array | undefined {
+  if (typeof state !== 'object' || state === null) return undefined
+  const { version, checkpoint } = state as Record<string, unknown>
+  if (version !== SNAPSHOT_VERSION || typeof checkpoint !== 'string')
+    return undefined
+  return new Uint8Array(Buffer.from(checkpoint, 'base64'))
+}
+
+function isCompaction(body: string): boolean {
+  const request = JSON.parse(body) as { toolChoice?: { type?: string } }
+  return request.toolChoice?.type === 'none'
+}
+
+function decodeBody(body: BodyInit | null | undefined): string {
+  if (typeof body === 'string') return body
+  if (body instanceof Uint8Array || body instanceof ArrayBuffer)
+    return new TextDecoder().decode(body)
+  throw new TypeError('FX sent an unexpected request body')
+}
+
+function sameJson(a: unknown, b: unknown): boolean {
+  return JSON.stringify(a) === JSON.stringify(b)
+}
+
+function sse(all: readonly object[]): string {
+  return (
+    all.map((part) => `data: ${JSON.stringify(part)}\n\n`).join('') +
+    'data: [DONE]\n\n'
+  )
+}
+
+function errorText(error: unknown): string {
+  if (typeof error === 'string') return error
+  if (error instanceof Error) return error.message
+  return JSON.stringify(error)
+}
+
+type StreamUsage = {
+  inputTokens?: {
+    total?: number | undefined
+    noCache?: number | undefined
+    cacheRead?: number | undefined
+    cacheWrite?: number | undefined
+  }
+  outputTokens?: {
+    total?: number | undefined
+    text?: number | undefined
+    reasoning?: number | undefined
+  }
+}
+
+/** A Gateway stream usage, as AI SDK normalizes it. */
+function modelUsage(value: unknown): LanguageModelUsage {
+  const usage = (value ?? {}) as StreamUsage
+  const input = usage.inputTokens ?? {}
+  const output = usage.outputTokens ?? {}
+  return {
+    inputTokens: input.total,
+    inputTokenDetails: {
+      noCacheTokens: input.noCache,
+      cacheReadTokens: input.cacheRead,
+      cacheWriteTokens: input.cacheWrite,
+    },
+    outputTokens: output.total,
+    outputTokenDetails: {
+      textTokens: output.text,
+      reasoningTokens: output.reasoning,
+    },
+    totalTokens:
+      input.total === undefined && output.total === undefined
+        ? undefined
+        : (input.total ?? 0) + (output.total ?? 0),
+  }
+}
+
+/** The inverse of `modelUsage`, for replay. */
+function streamUsage(usage: LanguageModelUsage | undefined): StreamUsage {
+  return {
+    inputTokens: {
+      total: usage?.inputTokens,
+      noCache: usage?.inputTokenDetails?.noCacheTokens,
+      cacheRead: usage?.inputTokenDetails?.cacheReadTokens,
+      cacheWrite: usage?.inputTokenDetails?.cacheWriteTokens,
+    },
+    outputTokens: {
+      total: usage?.outputTokens,
+      text: usage?.outputTokenDetails?.textTokens,
+      reasoning: usage?.outputTokenDetails?.reasoningTokens,
+    },
+  }
+}