experimental-a2 0.18.0 → 0.19.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,
@@ -22,6 +22,8 @@ import {
   type GatewayCredential,
   type Harness,
   type HarnessGeneration,
+  harnessImages,
+  type HarnessInput,
   type HarnessRead,
   type HarnessStep,
   type HarnessTurn,
@@ -30,6 +32,7 @@ import {
   saveHarness,
 } from './ai-harness.ts'
 import type { CompactionPartData, GenerationCompletedPayload } from './ai.ts'
+import { schemaInput } from './ai-tool-input.ts'
 import { boundedCache } from './bounded-cache.ts'
 
 export type { Harness } from './ai-harness.ts'
@@ -47,8 +50,11 @@ 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. */
 type ToolCall = { toolCallId: string; toolName: string; input: unknown }
 
 /** A model response A2 recorded, in the Gateway's stream format. */
@@ -60,7 +66,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[]>
@@ -108,6 +114,10 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
 
   let completion: Completion = {}
   let stopped = false
+  // Live responses belong to the step: the kernel aborts its fetch once it
+  // has read a response's finish, while A2 may still be reading its copy.
+  const released = new AbortController()
+  const network = AbortSignal.any([step.signal, released.signal])
   let agent: FxAgent | undefined
   let turn: TurnState | undefined
 
@@ -182,7 +192,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({
@@ -200,8 +210,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)
@@ -245,9 +257,18 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     const response = await fetch(url, {
       ...init,
       headers: withAuth(init.headers, credential),
+      signal: network,
     })
+    // Abandoned before its headers came: the kernel has moved on without it.
+    if (init.signal?.aborted) {
+      void response.body?.cancel()
+      throw init.signal.reason
+    }
     if (!response.ok || response.body === null) return response
-    const [forKernel, forLog] = response.body.tee()
+    const [fromNetwork, forLog] = response.body.tee()
+    const forKernel = init.signal
+      ? fromNetwork.pipeThrough(new TransformStream(), { signal: init.signal })
+      : fromNetwork
     // Live responses reach the log in request order: a compaction's part
     // precedes the step that follows it.
     const previous = state.parsed
@@ -268,8 +289,13 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
       state.parsed = (async () => {
         await previous
         emit({ type: 'start-step' })
-        await streamStep(forLog, emit, (finished) => {
-          completion = finished
+        await streamStep({
+          body: forLog,
+          tools: step.tools,
+          emit,
+          finished: (finished) => {
+            completion = finished
+          },
         })
       })()
     }
@@ -308,7 +334,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
@@ -344,11 +371,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) {
@@ -357,6 +383,7 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
             step.signal.aborted ? step.signal.reason : error,
           )
         } finally {
+          released.abort()
           step.signal.removeEventListener('abort', abort)
           await agent?.close().catch(() => {})
         }
@@ -365,6 +392,7 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
     cancel() {
       closed = true
       stopped = true
+      released.abort()
       turn?.fxTurn.cancel()
     },
   })
@@ -374,11 +402,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) => {
@@ -386,7 +417,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. */
@@ -448,36 +510,35 @@ function recordedCalls(
         })
         break
       case 'tool-input-error':
-        // A call whose input did not parse: FX answered it itself.
+      case 'tool-input-available': {
+        // A call recorded without its text: input that failed to parse is
+        // that text.
+        const text =
+          inputText.get(chunk.toolCallId) ||
+          (chunk.type === 'tool-input-error' && typeof chunk.input === 'string'
+            ? chunk.input
+            : JSON.stringify(chunk.input))
         if (inputText.has(chunk.toolCallId))
           parts.push({ type: 'tool-input-end', id: chunk.toolCallId })
+        // The kernel calls back with this text's JSON, not the input the
+        // tool's schema made of it. A call it answers itself, such as one
+        // whose text is not JSON, is never asked for.
+        const input = parseJson(text)
+        if (input !== undefined)
+          toolCalls.push({
+            toolCallId: chunk.toolCallId,
+            toolName: chunk.toolName,
+            input,
+          })
         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),
+          input: text,
           ...metadata,
         })
         break
+      }
     }
   }
   parts.push({
@@ -505,11 +566,19 @@ function compactionSse(data: CompactionPartData): string {
 // ------------------------------------------------------- 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> {
+async function streamStep({
+  body,
+  tools,
+  emit,
+  finished,
+}: {
+  body: ReadableStream<Uint8Array>
+  tools: ToolSet
+  emit: (chunk: UIMessageChunk) => void
+  finished: (completion: Completion) => void
+}): Promise<void> {
+  // The tool calls whose input started streaming, and whether text arrived.
+  const streamed = new Map<string, boolean>()
   for await (const part of sseParts(body)) {
     const metadata =
       part['providerMetadata'] === undefined
@@ -532,6 +601,7 @@ async function streamStep(
         })
         break
       case 'tool-input-start':
+        streamed.set(part['id'] as string, false)
         emit({
           type: 'tool-input-start',
           toolCallId: part['id'] as string,
@@ -540,6 +610,7 @@ async function streamStep(
         })
         break
       case 'tool-input-delta':
+        if (part['delta'] !== '') streamed.set(part['id'] as string, true)
         emit({
           type: 'tool-input-delta',
           toolCallId: part['id'] as string,
@@ -547,27 +618,39 @@ async function streamStep(
         })
         break
       case 'tool-call': {
+        const toolCallId = part['toolCallId'] as string
+        const toolName = part['toolName'] as string
         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
+        // Replay hands the kernel the text the log holds, so the log keeps
+        // the provider's text when it streamed none.
+        if (streamed.get(toolCallId) !== true && text !== '') {
+          if (!streamed.has(toolCallId))
+            emit({ type: 'tool-input-start', toolCallId, toolName })
+          emit({ type: 'tool-input-delta', toolCallId, inputTextDelta: text })
         }
-        emit({
-          type: 'tool-input-available',
-          toolCallId: part['toolCallId'] as string,
-          toolName: part['toolName'] as string,
-          input,
-          ...metadata,
+        const input = await parseToolInput({
+          tools,
+          toolName,
+          text,
         })
+        emit(
+          input.success
+            ? {
+                type: 'tool-input-available',
+                toolCallId,
+                toolName,
+                input: input.value,
+                ...metadata,
+              }
+            : {
+                type: 'tool-input-error',
+                toolCallId,
+                toolName,
+                input: input.input,
+                errorText: input.errorText,
+                ...metadata,
+              },
+        )
         break
       }
       case 'finish': {
@@ -587,6 +670,34 @@ async function streamStep(
   }
 }
 
+type ToolInput =
+  | { success: true; value: unknown }
+  | { success: false; input: unknown; errorText: string }
+
+/**
+ * A live call's input as the AI SDK step reads it: JSON, then the tool's
+ * schema with its defaults. Input that fails either is an input error, which
+ * A2 answers without running the tool.
+ */
+async function parseToolInput({
+  tools,
+  toolName,
+  text,
+}: {
+  tools: ToolSet
+  toolName: string
+  text: string
+}): Promise<ToolInput> {
+  let value: unknown
+  try {
+    value = JSON.parse(text)
+  } catch (error) {
+    return { success: false, input: text, errorText: errorText(error) }
+  }
+  const input = await schemaInput({ tools, toolName, value })
+  return input.success ? input : { ...input, input: value }
+}
+
 async function readParts(
   body: ReadableStream<Uint8Array>,
 ): Promise<StreamPart[]> {
@@ -693,15 +804,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),
+    })),
   )
 }
 
@@ -760,6 +869,14 @@ function decodeBody(body: BodyInit | null | undefined): string {
   throw new TypeError('FX sent an unexpected request body')
 }
 
+function parseJson(text: string): unknown {
+  try {
+    return JSON.parse(text)
+  } catch {
+    return undefined
+  }
+}
+
 function sameJson(a: unknown, b: unknown): boolean {
   return JSON.stringify(a) === JSON.stringify(b)
 }

package/src/ai-harness.ts

@@ -1,16 +1,20 @@
-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 { STEERED_REASON } from './ai-control.ts'
 import { boundedCache, type BoundedCache } from './bounded-cache.ts'
 import type { ContractEvent, EventDefs } from './contract.ts'
 
@@ -307,21 +311,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 +346,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 +369,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 +382,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 +426,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 +435,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 +464,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