experimental-a2 0.18.1 → 0.20.0

package/src/ai-fx.ts

@@ -12,7 +12,7 @@ import {
   type ToolSet,
   type UIMessageChunk,
 } from 'ai'
-import type { FxAgent, FxHostTool, FxTurn } from 'libfx'
+import type { FxAgent, FxHostTool, FxPromptBlock, FxTurn } from 'libfx'
 import {
   type Deferred,
   defineHarness,
@@ -20,8 +20,11 @@ import {
   Divergence,
   gatewayCredential,
   type GatewayCredential,
+  gatewaySessionAffinity,
   type Harness,
   type HarnessGeneration,
+  harnessImages,
+  type HarnessInput,
   type HarnessRead,
   type HarnessStep,
   type HarnessTurn,
@@ -48,6 +51,8 @@ export function fx(): Harness<'fx', never> {
 
 const SNAPSHOT_VERSION = 1
 const CATALOG_TTL_MS = 60 * 60 * 1000
+/** The Gateway model catalog libfx reads for each new kernel. */
+const CATALOG_PATH = '/coding-agent/v1/models'
 
 type StreamPart = { type: string } & Record<string, unknown>
 /** A recorded tool call, with the input the kernel reads from its text. */
@@ -62,7 +67,7 @@ type RecordedCall = {
 
 /** One FX turn to replay: its prompt, recorded calls, and how it ended. */
 type TurnPlan = {
-  prompt: string
+  prompt: string | FxPromptBlock[]
   calls: RecordedCall[]
   /** Steering to queue while serving the call at that index. */
   steers: Map<number, string[]>
@@ -188,7 +193,7 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     try {
       checkpoint = await agent!.checkpoint()
     } catch {
-      // Over libfx's bound: later steps replay from the previous snapshot.
+      // Later steps replay from the previous snapshot.
       return
     }
     await saveHarness({
@@ -206,8 +211,10 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     init: RequestInit = {},
   ): Promise<Response> => {
     const url = input instanceof Request ? input.url : String(input)
-    if ((init.method ?? 'GET').toUpperCase() === 'GET')
+    const method = (init.method ?? 'GET').toUpperCase()
+    if (method === 'GET' && new URL(url).pathname === CATALOG_PATH)
       return catalog(url, credential, withAuth(init.headers, credential))
+    if (method === 'GET') return fetch(input, init)
     const state = turn!
     if (state.failure) throw state.failure
     const body = decodeBody(init.body)
@@ -238,19 +245,27 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
       return pending(init.signal)
     }
     state.live = true
-    return live(state, url, init, compaction)
+    return live(state, url, { ...init, body }, compaction)
   }
 
+  const affinity = gatewaySessionAffinity({
+    agentName: step.agentName,
+    sessionId: step.sessionId,
+  })
   let compactions = 0
   const live = async (
     state: TurnState,
     url: string,
-    init: RequestInit,
+    init: RequestInit & { body: string },
     compaction: boolean,
   ): Promise<Response> => {
+    const headers = withAuth(init.headers, credential)
+    // The kernel's own affinity is new with every rebuild; A2's lasts the session.
+    headers.set('x-session-affinity', affinity)
     const response = await fetch(url, {
       ...init,
-      headers: withAuth(init.headers, credential),
+      body: withGatewayCaching(init.body),
+      headers,
       signal: network,
     })
     // Abandoned before its headers came: the kernel has moved on without it.
@@ -328,7 +343,8 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     return fail(state, diverged(`no recorded result for ${toolName}`))
   })
 
-  let turns = read.turns.map(turnPlan)
+  const plans = (turns: HarnessTurn[]): Promise<TurnPlan[]> =>
+    Promise.all(turns.map((recorded) => turnPlan(recorded, step.signal)))
   let results = read.results
   let controller: ReadableStreamDefaultController<UIMessageChunk> | undefined
   let closed = false
@@ -364,11 +380,10 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
             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)
+          await run(streamController, await plans(read.turns))
           closed = true
           streamController.close()
         } catch (error) {
@@ -396,11 +411,14 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
 
 // ------------------------------------------------------------ the plan
 
-function turnPlan(turn: HarnessTurn): TurnPlan {
+async function turnPlan(
+  turn: HarnessTurn,
+  signal: AbortSignal,
+): Promise<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)
+  const steer = (at: number, inputs: HarnessInput[] | undefined) => {
+    if (inputs !== undefined) steers.set(Math.max(0, at), inputs.map(steerText))
   }
   steer(0, turn.generations.length === 0 ? turn.steers.get(0) : undefined)
   turn.generations.forEach((generation, index) => {
@@ -408,7 +426,38 @@ function turnPlan(turn: HarnessTurn): TurnPlan {
     // 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 }
+  const images = await harnessImages({
+    harness: 'fx()',
+    input: turn.prompt,
+    signal,
+  })
+  return {
+    prompt:
+      images.length === 0
+        ? turn.prompt.text
+        : [
+            ...images.map((image): FxPromptBlock => ({
+              type: 'image',
+              data: image.data,
+              mimeType: image.mediaType,
+            })),
+            ...(turn.prompt.text === ''
+              ? []
+              : [{ type: 'text' as const, text: turn.prompt.text }]),
+          ],
+    calls,
+    steers,
+    end: turn.end,
+  }
+}
+
+/** libfx steers with text alone, so a steer's files fail the step. */
+function steerText(input: HarnessInput): string {
+  if (input.files.length > 0)
+    throw new Error(
+      'fx() steers with text only: send files in a message of their own',
+    )
+  return input.text
 }
 
 /** A completed generation's model calls: its compactions, then its step. */
@@ -734,6 +783,31 @@ async function loadLibfx(): Promise<typeof import('libfx')> {
   }
 }
 
+/**
+ * A model request with Gateway's automatic caching, as the AI SDK harness
+ * asks for it; options the request already has take precedence. Every other
+ * byte stays as the kernel wrote it, a tool input's spacing included.
+ */
+function withGatewayCaching(body: string): string {
+  const request = JSON.parse(body) as {
+    providerOptions?: Record<string, Record<string, unknown> | undefined>
+  }
+  const options = request.providerOptions
+  if (options === undefined)
+    return body.replace(
+      /^\{/,
+      '{"providerOptions":{"gateway":{"caching":"auto"}},',
+    )
+  // The kernel writes its own options last; any other shape goes unchanged.
+  const own = `"providerOptions":${JSON.stringify(options)}}`
+  if (!body.endsWith(own)) return body
+  const merged = {
+    ...options,
+    gateway: { caching: 'auto', ...options['gateway'] },
+  }
+  return `${body.slice(0, -own.length)}"providerOptions":${JSON.stringify(merged)}}`
+}
+
 function withAuth(
   headers: HeadersInit | undefined,
   credential: GatewayCredential,
@@ -764,15 +838,13 @@ async function kernelTools(
   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),
-      })),
+    Object.entries(tools).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),
+    })),
   )
 }
 

package/src/ai-harness.ts

@@ -1,16 +1,21 @@
-import type {
-  Experimental_SandboxSession,
-  FinishReason,
-  LanguageModelUsage,
-  ToolSet,
-  UIMessage,
-  UIMessageChunk,
+import {
+  createDownload,
+  type Experimental_SandboxSession,
+  type FileUIPart,
+  type FinishReason,
+  type LanguageModelUsage,
+  type ToolSet,
+  type UIMessage,
+  type UIMessageChunk,
 } from 'ai'
 import type {
   GenerationCompletedPayload,
+  MessageCreatedPayload,
   ReasoningLevel,
   ToolResultPayload,
 } from './ai.ts'
+import { createHash } from 'node:crypto'
+import { STEERED_REASON } from './ai-control.ts'
 import { boundedCache, type BoundedCache } from './bounded-cache.ts'
 import type { ContractEvent, EventDefs } from './contract.ts'
 
@@ -307,21 +312,21 @@ function savedEnvelope(value: unknown): SavedEnvelope | 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
 }
 
+/** What user messages hand a harness: their text, and the files they attach. */
+export type HarnessInput = { text: string; files: FileUIPart[] }
+
 /** One harness turn to rebuild: its prompt, completed generations, and end. */
 export type HarnessTurn = {
-  prompt: string
+  prompt: HarnessInput
   generations: RecordedGeneration[]
   /** Steering that entered after the generation at that index. */
-  steers: Map<number, string[]>
+  steers: Map<number, HarnessInput[]>
   end: 'stop' | 'interrupted' | 'live'
 }
 
@@ -342,15 +347,18 @@ export function harnessTurns(
   }
   type Response = {
     id: string
-    prompt: string[]
+    prompt: HarnessInput[]
     generations: Generation[]
     end?: 'stop' | 'steered' | 'interrupted'
   }
-  // Queued messages are admitted together, possibly before the snapshot.
-  const userText = new Map<string, string>(
+  // Prompts take their content from the transcript: a harness agent never
+  // runs A2's compaction, so the step's messages hold every input admitted by
+  // its read, before the snapshot or after it. A session another loop
+  // compacted can't move to a harness.
+  const userInput = new Map<string, HarnessInput>(
     step.messages
       .filter((message) => message.role === 'user')
-      .map((message) => [message.id, textOf(message)]),
+      .map((message) => [message.id, inputOf(message)]),
   )
   let passive: string[] = []
   const responses = new Map<string, Response>()
@@ -362,10 +370,9 @@ export function harnessTurns(
     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)
+        const admitted = event.payload as MessageCreatedPayload
+        if (admitted.role === 'user' && admitted.generate === false)
+          passive.push(admitted.messageId)
         break
       }
       case 'ai.generation.requested': {
@@ -376,7 +383,7 @@ export function harnessTurns(
         passive = []
         responses.set(id, {
           id,
-          prompt: ids.flatMap((p) => userText.get(p) ?? []),
+          prompt: ids.flatMap((p) => userInput.get(p) ?? []),
           generations: [],
         })
         break
@@ -420,7 +427,7 @@ export function harnessTurns(
         const response = responses.get(payload['messageId'] as string)
         if (response)
           response.end =
-            payload['reason'] === STEERED ? 'steered' : 'interrupted'
+            payload['reason'] === STEERED_REASON ? 'steered' : 'interrupted'
         break
       }
     }
@@ -429,13 +436,16 @@ export function harnessTurns(
   const turns: HarnessTurn[] = []
   let steered = false
   for (const response of responses.values()) {
-    const text = response.prompt.join('\n\n')
+    const input: HarnessInput = {
+      text: response.prompt.map((prompt) => prompt.text).join('\n\n'),
+      files: response.prompt.flatMap((prompt) => prompt.files),
+    }
     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])
+      turn.steers.set(at, [...(turn.steers.get(at) ?? []), input])
     } else {
-      turn = { prompt: text, generations: [], steers: new Map(), end: 'live' }
+      turn = { prompt: input, generations: [], steers: new Map(), end: 'live' }
       turns.push(turn)
     }
     for (const generation of response.generations)
@@ -455,10 +465,125 @@ export function harnessTurns(
   return { turns, results }
 }
 
-function textOf(message: UIMessage): string {
-  return message.parts
-    .flatMap((part) => (part.type === 'text' ? [part.text] : []))
-    .join('\n')
+function inputOf(message: UIMessage): HarnessInput {
+  return {
+    text: message.parts
+      .flatMap((part) => (part.type === 'text' ? [part.text] : []))
+      .join('\n'),
+    files: message.parts.filter(
+      (part): part is FileUIPart => part.type === 'file',
+    ),
+  }
+}
+
+// ------------------------------------------------------------- attachments
+
+/** An attached image as an engine takes it: its bytes in base64. */
+export type HarnessImage = { mediaType: string; data: string }
+
+// What one message's files hold in a step's memory, at most.
+const MAX_ATTACHMENT_BYTES = 32 * 1024 * 1024
+
+type Attachment = { data: string; bytes: number }
+
+// Every step of a response replays its turn, and a URL's file doesn't change,
+// so the process keeps what it fetched.
+const attachments = boundedCache<string, Attachment>({
+  maxEntries: 256,
+  maxBytes: 64 * 1024 * 1024,
+})
+
+/**
+ * The images an input attaches, as its engine takes them: engines take
+ * bytes, not URLs. A file it can't attach fails the step, naming the file,
+ * and is never left out: the failure may pass, and the step is retried like
+ * any other.
+ */
+export async function harnessImages({
+  harness,
+  input,
+  signal,
+}: {
+  /** As the error names it: `fx()` or `pi()`. */
+  harness: string
+  input: HarnessInput
+  signal: AbortSignal
+}): Promise<HarnessImage[]> {
+  const images: HarnessImage[] = []
+  let budget = MAX_ATTACHMENT_BYTES
+  // One at a time, so the budget bounds what the step holds.
+  for (const file of input.files) {
+    try {
+      if (!file.mediaType.startsWith('image/'))
+        throw new Error(`it takes images only, not ${file.mediaType}`)
+      const attached = await attachment({ url: file.url, budget, signal })
+      budget -= attached.bytes
+      images.push({
+        mediaType: file.mediaType,
+        data: attached.data,
+      })
+    } catch (error) {
+      if (signal.aborted) throw error
+      throw new Error(
+        `${harness} couldn't attach ${file.filename ?? 'a file'}: ${
+          error instanceof Error ? error.message : String(error)
+        }`,
+        { cause: error },
+      )
+    }
+  }
+  return images
+}
+
+async function attachment({
+  url,
+  budget,
+  signal,
+}: {
+  url: string
+  budget: number
+  signal: AbortSignal
+}): Promise<Attachment> {
+  const tooLarge = () =>
+    new Error(
+      `a message's files may hold ${MAX_ATTACHMENT_BYTES / 1024 / 1024} MiB in a step`,
+    )
+  const cached = attachments.get(url)
+  if (cached !== undefined) {
+    if (cached.bytes > budget) throw tooLarge()
+    return cached
+  }
+  const { protocol } = new URL(url)
+  if (protocol !== 'https:' && protocol !== 'data:')
+    throw new Error(`attachments must be https: or data: URLs, not ${protocol}`)
+  // AI SDK's download checks an https: URL's address and its redirects,
+  // which a message's sender chooses; a data: URL holds its bytes.
+  const bytes =
+    protocol === 'data:'
+      ? dataBytes(url)
+      : (
+          await createDownload({ maxBytes: budget })({
+            url: new URL(url),
+            abortSignal: signal,
+          })
+        ).data
+  if (bytes.byteLength > budget) throw tooLarge()
+  const attached: Attachment = {
+    data: Buffer.from(bytes).toString('base64'),
+    bytes: bytes.byteLength,
+  }
+  if (protocol === 'https:')
+    attachments.set(url, attached, attached.data.length)
+  return attached
+}
+
+function dataBytes(url: string): Uint8Array {
+  const comma = url.indexOf(',')
+  if (comma === -1) throw new Error('its data: URL has no data')
+  const body = url.slice(comma + 1)
+  return url.slice(0, comma).endsWith(';base64')
+    ? Buffer.from(body, 'base64')
+    : Buffer.from(decodeURIComponent(body))
 }
 
 // ------------------------------------------------------------- Gateway
@@ -474,6 +599,23 @@ export function gatewayCredential(harness: string): GatewayCredential {
   throw new Error(`${harness}() needs AI_GATEWAY_API_KEY or VERCEL_OIDC_TOKEN`)
 }
 
+/**
+ * The `x-session-affinity` an agent session's Gateway calls carry: stable for
+ * the session, so supporting providers can reuse their cache, and free of its
+ * raw ID.
+ */
+export function gatewaySessionAffinity({
+  agentName,
+  sessionId,
+}: {
+  agentName: string
+  sessionId: string
+}): string {
+  return createHash('sha256')
+    .update(JSON.stringify([agentName, sessionId]))
+    .digest('hex')
+}
+
 // Where Vercel Functions keep the current request, as @vercel/oidc reads it.
 const REQUEST_CONTEXT = Symbol.for('@vercel/request-context')
 

package/src/ai-message-shape.ts

@@ -0,0 +1,41 @@
+import type { UIMessage } from 'ai'
+
+export const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null && !Array.isArray(value)
+
+export const isJSONCompatible = (
+  value: unknown,
+  seen: Set<object> = new Set(),
+): boolean => {
+  if (
+    value === null ||
+    typeof value === 'string' ||
+    typeof value === 'boolean'
+  ) {
+    return true
+  }
+  if (typeof value === 'number') return Number.isFinite(value)
+  if (typeof value === 'undefined') return true
+  if (typeof value !== 'object') return false
+  if (seen.has(value)) return false
+  seen.add(value)
+  const valid = Array.isArray(value)
+    ? value.every((item) => isJSONCompatible(item, seen))
+    : Object.getPrototypeOf(value) === Object.prototype &&
+      Object.values(value).every((item) => isJSONCompatible(item, seen))
+  seen.delete(value)
+  return valid
+}
+
+/** Whether a value has an AI SDK `UIMessage`'s shape. */
+export const isUIMessage = (value: unknown): value is UIMessage =>
+  isRecord(value) &&
+  typeof value['id'] === 'string' &&
+  (value['role'] === 'user' ||
+    value['role'] === 'assistant' ||
+    value['role'] === 'system') &&
+  Array.isArray(value['parts']) &&
+  value['parts'].every(
+    (part) => isRecord(part) && typeof part['type'] === 'string',
+  ) &&
+  isJSONCompatible(value)