experimental-a2 0.16.1 → 0.18.0

package/src/ai-harness.ts

@@ -0,0 +1,518 @@
+import type {
+  Experimental_SandboxSession,
+  FinishReason,
+  LanguageModelUsage,
+  ToolSet,
+  UIMessage,
+  UIMessageChunk,
+} from 'ai'
+import type {
+  GenerationCompletedPayload,
+  ReasoningLevel,
+  ToolResultPayload,
+} from './ai.ts'
+import { boundedCache, type BoundedCache } from './bounded-cache.ts'
+import type { ContractEvent, EventDefs } from './contract.ts'
+
+declare const harnessBrand: unique symbol
+
+/**
+ * A coding agent's own loop for `agent({ harness })`, from `fx()` or `pi()`.
+ * `Tools` names the tools the harness brings, which `toolRecovery` may name.
+ */
+export type Harness<
+  Type extends 'fx' | 'pi' = 'fx' | 'pi',
+  Tools extends string = string,
+> = {
+  readonly type: Type
+  readonly [harnessBrand]: { readonly tools: Tools }
+}
+
+/** What a harness's sandbox resolver receives for one step or tool execution. */
+export type HarnessSandboxContext = {
+  session: { readonly id: string }
+  signal: AbortSignal
+}
+
+/** Resolves the AI SDK sandbox session for an agent session, the same one every time. */
+export type HarnessSandbox = (
+  context: HarnessSandboxContext,
+) => Experimental_SandboxSession | Promise<Experimental_SandboxSession>
+
+/** What a harness step streams, as the default generator does. */
+export type HarnessGeneration = {
+  stream: ReadableStream<UIMessageChunk>
+  completion: () => { finishReason?: FinishReason; usage?: LanguageModelUsage }
+}
+
+/** One stored snapshot and the log after it, read in one store operation. */
+export type HarnessLogRead = {
+  snapshot: { index: number; state: unknown } | null
+  events: ContractEvent<EventDefs>[]
+}
+
+/** The session log as a harness step reads it, plus the state saved in it. */
+export type HarnessLog = {
+  /** The latest saved state under `name` and the log after it; `fresh` reads it all. */
+  read(options: {
+    name: string
+    throughIndex: number
+    fresh?: boolean
+  }): Promise<HarnessLogRead>
+  write(options: { name: string; index: number; state: unknown }): Promise<void>
+}
+
+/** Everything one durable generation hands its harness. */
+export type HarnessStep = {
+  agentName: string
+  sessionId: string
+  requestIndex: number
+  /** This attempt at the generation. */
+  generationId: string
+  responseMessageId: string
+  /** The log index this step's context was read through. */
+  throughIndex: number
+  /** The transcript through that index. */
+  messages: readonly UIMessage[]
+  model: string
+  reasoning?: ReasoningLevel
+  speed?: 'fast' | 'standard'
+  instructions: string[]
+  tools: ToolSet
+  log: HarnessLog
+  signal: AbortSignal
+}
+
+/** The internal side of a harness value. */
+export type HarnessDriver = {
+  /** The session's sandbox; every local tool execution receives it. */
+  readonly sandbox?: HarnessSandbox
+  /** Tools the harness brings; A2 executes them like the agent's own. */
+  readonly tools: ToolSet
+  step(step: HarnessStep): Promise<HarnessGeneration>
+}
+
+// A global symbol, so harness values work across separately bundled copies.
+const DRIVER = Symbol.for('experimental-a2.ai.harness')
+
+export function defineHarness<
+  Type extends 'fx' | 'pi',
+  Tools extends string = never,
+>(type: Type, driver: HarnessDriver): Harness<Type, Tools> {
+  const harness = { type }
+  Object.defineProperty(harness, DRIVER, { value: driver })
+  return Object.freeze(harness) as unknown as Harness<Type, Tools>
+}
+
+export function harnessDriver(harness: unknown): HarnessDriver {
+  const driver =
+    typeof harness === 'object' && harness !== null
+      ? (harness as { [DRIVER]?: HarnessDriver })[DRIVER]
+      : undefined
+  if (driver === undefined)
+    throw new TypeError(
+      "harness must be { type: 'ai-sdk' }, or come from fx() or pi()",
+    )
+  return driver
+}
+
+// ------------------------------------------------------- saved state
+
+/**
+ * What an agent's steps saved in this process at the end of a response, per
+ * session and harness. A cache: a cold process rebuilds from the whole log.
+ */
+export type HarnessStates = BoundedCache<
+  string,
+  { index: number; state: unknown }
+>
+
+export const harnessStates = (): HarnessStates =>
+  boundedCache({ maxEntries: 1_000, maxBytes: 64 * 1024 * 1024 })
+
+const savedKey = (sessionId: string, name: string): string =>
+  `${sessionId}\u0000${name}`
+
+// Binary state counts its bytes; everything else its JSON text.
+const sizeOf = (value: unknown): number => {
+  let binary = 0
+  const text = JSON.stringify(value, (_key, item: unknown) => {
+    if (!(item instanceof Uint8Array)) return item
+    binary += item.byteLength
+    return null
+  })
+  return binary + (text?.length ?? 0)
+}
+
+/**
+ * The session log as a harness step reads it: the state this process saved
+ * at the end of the last response, when it has one, and the log after it.
+ * Saved states are copied in and out, so a step cannot change one in place.
+ */
+export function harnessLog(options: {
+  states: HarnessStates
+  session: {
+    readonly id: string
+    history(bounds: {
+      gte?: number
+      lte?: number
+    }): Promise<ContractEvent<EventDefs>[]>
+  }
+}): HarnessLog {
+  const { states, session } = options
+  return {
+    read: async ({ name, throughIndex, fresh }) => {
+      const entry = fresh ? undefined : states.get(savedKey(session.id, name))
+      const snapshot =
+        entry === undefined || entry.index > throughIndex
+          ? null
+          : { index: entry.index, state: structuredClone(entry.state) }
+      const events = await session.history(
+        snapshot === null
+          ? { lte: throughIndex }
+          : { gte: snapshot.index + 1, lte: throughIndex },
+      )
+      return { snapshot, events }
+    },
+    write: async ({ name, index, state }) => {
+      const copy = structuredClone(state)
+      states.set(
+        savedKey(session.id, name),
+        { index, state: copy },
+        sizeOf(copy),
+      )
+    },
+  }
+}
+
+const agentStates = new WeakMap<object, HarnessStates>()
+
+/** Registers an agent server's saved harness states, for `savedHarnessState`. */
+export function registerHarnessStates(
+  server: object,
+  states: HarnessStates,
+): void {
+  agentStates.set(server, states)
+}
+
+/** What an agent's process saved for one session's harness; tests read it. */
+export function savedHarnessState(
+  server: object,
+  options: { sessionId: string; name: string },
+): { index: number; state: unknown } | undefined {
+  return agentStates.get(server)?.get(savedKey(options.sessionId, options.name))
+}
+
+/** A harness's state at the end of a response, from the attempt that ran it. */
+type SavedEnvelope = {
+  generationId: string
+  responseMessageId: string
+  state: unknown
+}
+
+export type HarnessRead<S> = {
+  /** The state saved at the end of the last response, when it is usable. */
+  saved: S | undefined
+  turns: HarnessTurn[]
+  results: Map<string, ToolResultPayload>
+}
+
+/**
+ * Reads a harness's saved state and the log since, as turns to replay. A
+ * snapshot counts only when the attempt that wrote it committed: a stopped,
+ * superseded, or still-running attempt can write one the log never kept.
+ * `replayAll` reads the whole log instead, for a state that fails to restore.
+ */
+export async function readHarness<S>({
+  step,
+  name,
+  parse,
+}: {
+  step: HarnessStep
+  name: string
+  parse: (state: unknown) => S | undefined
+}): Promise<HarnessRead<S> & { replayAll: () => Promise<HarnessRead<S>> }> {
+  const replayAll = async (): Promise<HarnessRead<S>> => {
+    const read = await step.log.read({
+      name,
+      throughIndex: step.throughIndex,
+      fresh: true,
+    })
+    return { saved: undefined, ...harnessTurns(step, read.events, undefined) }
+  }
+  const read = await step.log.read({ name, throughIndex: step.throughIndex })
+  if (read.snapshot === null)
+    return {
+      saved: undefined,
+      ...harnessTurns(step, read.events, undefined),
+      replayAll,
+    }
+  const envelope = savedEnvelope(read.snapshot.state)
+  const saved = envelope === undefined ? undefined : parse(envelope.state)
+  const committed =
+    envelope !== undefined &&
+    read.events.some(
+      (event) =>
+        event.type === 'ai.generation.completed' &&
+        (event.payload as { generationId?: unknown }).generationId ===
+          envelope.generationId,
+    )
+  // A snapshot of this very response belongs to an earlier attempt at it.
+  if (
+    saved === undefined ||
+    !committed ||
+    envelope.responseMessageId === step.responseMessageId
+  )
+    return { ...(await replayAll()), replayAll }
+  return {
+    saved,
+    ...harnessTurns(step, read.events, envelope.responseMessageId),
+    replayAll,
+  }
+}
+
+/** Saves a harness's state at the end of a response; losing it costs a longer replay. */
+export async function saveHarness({
+  step,
+  name,
+  state,
+}: {
+  step: HarnessStep
+  name: string
+  state: unknown
+}): Promise<void> {
+  const envelope: SavedEnvelope = {
+    generationId: step.generationId,
+    responseMessageId: step.responseMessageId,
+    state,
+  }
+  try {
+    await step.log.write({ name, index: step.requestIndex, state: envelope })
+  } catch {
+    // The snapshot is a cache.
+  }
+}
+
+function savedEnvelope(value: unknown): SavedEnvelope | undefined {
+  if (typeof value !== 'object' || value === null) return undefined
+  const { generationId, responseMessageId, state } = value as Record<
+    string,
+    unknown
+  >
+  return typeof generationId === 'string' &&
+    typeof responseMessageId === 'string'
+    ? { generationId, responseMessageId, state }
+    : undefined
+}
+
+// ------------------------------------------------- the log as harness turns
+
+// A2's reason on `ai.message.interrupted` when a steer ends a response.
+const STEERED = 'steered after step'
+
+/** A completed generation as the log recorded it. */
+export type RecordedGeneration = {
+  chunks: UIMessageChunk[]
+  completed: GenerationCompletedPayload
+}
+
+/** One harness turn to rebuild: its prompt, completed generations, and end. */
+export type HarnessTurn = {
+  prompt: string
+  generations: RecordedGeneration[]
+  /** Steering that entered after the generation at that index. */
+  steers: Map<number, string[]>
+  end: 'stop' | 'interrupted' | 'live'
+}
+
+/**
+ * Groups the log since a harness snapshot into harness turns. A response
+ * steered after a step continues its turn, with the next response's prompt as
+ * steering. `covered` is the response the snapshot already holds.
+ */
+export function harnessTurns(
+  step: HarnessStep,
+  events: readonly ContractEvent<EventDefs>[],
+  covered: string | undefined,
+): { turns: HarnessTurn[]; results: Map<string, ToolResultPayload> } {
+  type Generation = {
+    id: string
+    chunks: UIMessageChunk[]
+    completed?: GenerationCompletedPayload
+  }
+  type Response = {
+    id: string
+    prompt: string[]
+    generations: Generation[]
+    end?: 'stop' | 'steered' | 'interrupted'
+  }
+  // Queued messages are admitted together, possibly before the snapshot.
+  const userText = new Map<string, string>(
+    step.messages
+      .filter((message) => message.role === 'user')
+      .map((message) => [message.id, textOf(message)]),
+  )
+  let passive: string[] = []
+  const responses = new Map<string, Response>()
+  const byStart = new Map<number, Generation>()
+  const byId = new Map<string, Generation>()
+  const results = new Map<string, ToolResultPayload>()
+
+  for (const event of events) {
+    const payload = event.payload as Record<string, unknown>
+    switch (event.type) {
+      case 'ai.message.created': {
+        const message = payload['message'] as UIMessage
+        if (message.role !== 'user') break
+        userText.set(message.id, textOf(message))
+        if (payload['generate'] === false) passive.push(message.id)
+        break
+      }
+      case 'ai.generation.requested': {
+        const id = payload['responseMessageId'] as string
+        if (id === covered || responses.has(id)) break
+        const messageId = payload['messageId'] as string
+        const ids = [...passive.filter((p) => p !== messageId), messageId]
+        passive = []
+        responses.set(id, {
+          id,
+          prompt: ids.flatMap((p) => userText.get(p) ?? []),
+          generations: [],
+        })
+        break
+      }
+      case 'ai.generation.started': {
+        const response = responses.get(payload['responseMessageId'] as string)
+        if (response === undefined) break
+        const generation: Generation = {
+          id: payload['generationId'] as string,
+          chunks: [],
+        }
+        response.generations.push(generation)
+        byStart.set(event.index, generation)
+        byId.set(generation.id, generation)
+        break
+      }
+      case 'ai.generation.progress':
+        byStart
+          .get(payload['startIndex'] as number)
+          ?.chunks.push(...(payload['chunks'] as UIMessageChunk[]))
+        break
+      case 'ai.generation.completed': {
+        const generation = byId.get(payload['generationId'] as string)
+        if (generation)
+          generation.completed = event.payload as GenerationCompletedPayload
+        break
+      }
+      case 'ai.tool.result':
+        if (payload['preliminary'] !== true)
+          results.set(
+            payload['toolCallId'] as string,
+            event.payload as ToolResultPayload,
+          )
+        break
+      case 'ai.message.completed': {
+        const response = responses.get(payload['messageId'] as string)
+        if (response) response.end = 'stop'
+        break
+      }
+      case 'ai.message.interrupted': {
+        const response = responses.get(payload['messageId'] as string)
+        if (response)
+          response.end =
+            payload['reason'] === STEERED ? 'steered' : 'interrupted'
+        break
+      }
+    }
+  }
+
+  const turns: HarnessTurn[] = []
+  let steered = false
+  for (const response of responses.values()) {
+    const text = response.prompt.join('\n\n')
+    let turn = turns.at(-1)
+    if (steered && turn !== undefined) {
+      const at = Math.max(0, turn.generations.length - 1)
+      turn.steers.set(at, [...(turn.steers.get(at) ?? []), text])
+    } else {
+      turn = { prompt: text, generations: [], steers: new Map(), end: 'live' }
+      turns.push(turn)
+    }
+    for (const generation of response.generations)
+      if (generation.completed !== undefined)
+        turn.generations.push({
+          chunks: generation.chunks,
+          completed: generation.completed,
+        })
+    turn.end =
+      response.id === step.responseMessageId
+        ? 'live'
+        : response.end === 'stop'
+          ? 'stop'
+          : 'interrupted'
+    steered = response.end === 'steered'
+  }
+  return { turns, results }
+}
+
+function textOf(message: UIMessage): string {
+  return message.parts
+    .flatMap((part) => (part.type === 'text' ? [part.text] : []))
+    .join('\n')
+}
+
+// ------------------------------------------------------------- Gateway
+
+export type GatewayCredential = { token: string; method: 'api-key' | 'oidc' }
+
+/** Read per step: on Vercel the OIDC token rides on each function request. */
+export function gatewayCredential(harness: string): GatewayCredential {
+  const key = process.env['AI_GATEWAY_API_KEY']
+  if (key) return { token: key, method: 'api-key' }
+  const oidc = requestOidcToken() ?? process.env['VERCEL_OIDC_TOKEN']
+  if (oidc) return { token: oidc, method: 'oidc' }
+  throw new Error(`${harness}() needs AI_GATEWAY_API_KEY or VERCEL_OIDC_TOKEN`)
+}
+
+// Where Vercel Functions keep the current request, as @vercel/oidc reads it.
+const REQUEST_CONTEXT = Symbol.for('@vercel/request-context')
+
+function requestOidcToken(): string | undefined {
+  const context = (
+    globalThis as {
+      [REQUEST_CONTEXT]?: { get?: () => { headers?: Record<string, unknown> } }
+    }
+  )[REQUEST_CONTEXT]
+  const token = context?.get?.()?.headers?.['x-vercel-oidc-token']
+  return typeof token === 'string' && token !== '' ? token : undefined
+}
+
+// ------------------------------------------------------------- helpers
+
+export class Divergence extends Error {
+  constructor(harness: string, message: string) {
+    super(`${harness} replay diverged from the log: ${message}`)
+  }
+}
+
+export type Deferred = { promise: Promise<void>; resolve: () => void }
+
+export function deferred(): Deferred {
+  let resolve!: () => void
+  const promise = new Promise<void>((done) => {
+    resolve = done
+  })
+  return { promise, resolve }
+}
+
+/** A result that never settles; the harness is stopped instead. */
+export function pending(
+  signal: AbortSignal | null | undefined,
+): Promise<never> {
+  return new Promise((_, reject) => {
+    if (signal?.aborted) reject(signal.reason)
+    signal?.addEventListener('abort', () => reject(signal.reason), {
+      once: true,
+    })
+  })
+}