experimental-a2 0.18.1 → 0.19.0

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

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)

package/src/ai-pi.ts

@@ -14,6 +14,7 @@ import type {
   AssistantMessage,
   AssistantMessageEvent,
   AssistantMessageEventStream,
+  ImageContent,
   JsonObject,
   TextContent,
   ThinkingContent,
@@ -50,6 +51,8 @@ import {
   gatewayCredential,
   type Harness,
   type HarnessGeneration,
+  harnessImages,
+  type HarnessInput,
   type HarnessRead,
   type HarnessSandbox,
   type HarnessStep,
@@ -130,11 +133,14 @@ type PiRecord =
   | { call: 'step'; response: StoredResponse }
   | { call: 'compaction'; response: AssistantMessage }
 
+/** A user message as Pi takes it. */
+type PiInput = { text: string; images: ImageContent[] }
+
 type RecordedCall = {
   kind: 'step' | 'compaction'
   response: AssistantMessage
   /** Steering to queue while serving this call. */
-  steers: string[]
+  steers: PiInput[]
 }
 
 type Saved = {
@@ -232,7 +238,8 @@ async function runStep(
         fail(state, diverged(`call ${state.next} is not a ${kind}`))
       state.next++
       // Steering queued while a response streams enters after its tools.
-      for (const text of recorded.steers) await session!.steer(text)
+      for (const { text, images } of recorded.steers)
+        await session!.steer(text, images.length === 0 ? undefined : images)
       replay(out, recorded.response)
       return
     }
@@ -252,6 +259,9 @@ async function runStep(
     const upstream = realStream(model, context, {
       ...options,
       headers: { ...options?.headers, ...gatewayHeaders(credential) },
+      ...(step.speed === 'fast'
+        ? { onPayload: fastTier(options?.onPayload) }
+        : {}),
     })
     const live = kind === 'step' ? liveStep() : undefined
     for await (const event of upstream) {
@@ -547,9 +557,10 @@ async function runStep(
     }
     for (const plan of read.turns) {
       if (step.signal.aborted) return
+      const prompt = await piInput({ input: plan.prompt, signal: step.signal })
       const state: TurnState = {
         plan,
-        calls: recordedTurn(plan),
+        calls: await recordedTurn({ turn: plan, signal: step.signal }),
         next: 0,
         liveStep: false,
         over: false,
@@ -557,8 +568,9 @@ async function runStep(
         compactions: [],
       }
       turn = state
-      const running = session.prompt(plan.prompt, {
+      const running = session.prompt(prompt.text, {
         expandPromptTemplates: false,
+        ...(prompt.images.length === 0 ? {} : { images: prompt.images }),
       })
       if (plan.end !== 'live') {
         await running.catch((error: unknown) => {
@@ -622,17 +634,43 @@ async function runStep(
 // ------------------------------------------------------------ the plan
 
 /** A turn's recorded model calls, in the order Pi made them. */
-function recordedTurn(turn: HarnessTurn): RecordedCall[] {
+async function recordedTurn({
+  turn,
+  signal,
+}: {
+  turn: HarnessTurn
+  signal: AbortSignal
+}): Promise<RecordedCall[]> {
   const calls: RecordedCall[] = []
-  turn.generations.forEach((generation, index) => {
+  for (const [index, generation] of turn.generations.entries()) {
     calls.push(...recordedCalls(generation))
     const steers = turn.steers.get(index)
     const last = calls.findLast((call) => call.kind === 'step')
-    if (steers !== undefined && last !== undefined) last.steers.push(...steers)
-  })
+    if (steers !== undefined && last !== undefined)
+      for (const input of steers)
+        last.steers.push(await piInput({ input, signal }))
+  }
   return calls
 }
 
+async function piInput({
+  input,
+  signal,
+}: {
+  input: HarnessInput
+  signal: AbortSignal
+}): Promise<PiInput> {
+  const images = await harnessImages({ harness: 'pi()', input, signal })
+  return {
+    text: input.text,
+    images: images.map((image) => ({
+      type: 'image',
+      data: image.data,
+      mimeType: image.mediaType,
+    })),
+  }
+}
+
 function recordedCalls(generation: RecordedGeneration): RecordedCall[] {
   const texts = new Map<string, string>()
   const inputs = new Map<string, unknown>()
@@ -937,30 +975,28 @@ async function sessionTools(
   ) => Promise<AgentToolResult<unknown>>,
 ): Promise<ToolDefinition[]> {
   return Promise.all(
-    Object.entries(tools)
-      .filter(([, tool]) => tool.type !== 'provider')
-      .map(async ([name, tool]): Promise<ToolDefinition> => {
-        const builtin = builtinOf(tools, name)
-        if (builtin !== undefined)
-          // Pi's own definition, for its prompt; A2 already ran the call.
-          return Object.assign(builtinDefinition(modules, builtin, cwd), {
-            execute: (
-              toolCallId: string,
-              _params: unknown,
-              signal?: AbortSignal,
-            ) => result(toolCallId, signal, 'pi'),
-          })
-        return {
-          name,
-          label: name,
-          description:
-            typeof tool.description === 'string' ? tool.description : '',
-          parameters: (await asSchema(tool.inputSchema)
-            .jsonSchema) as ToolDefinition['parameters'],
-          execute: (toolCallId, _params, signal) =>
-            result(toolCallId, signal, 'tool'),
-        }
-      }),
+    Object.entries(tools).map(async ([name, tool]): Promise<ToolDefinition> => {
+      const builtin = builtinOf(tools, name)
+      if (builtin !== undefined)
+        // Pi's own definition, for its prompt; A2 already ran the call.
+        return Object.assign(builtinDefinition(modules, builtin, cwd), {
+          execute: (
+            toolCallId: string,
+            _params: unknown,
+            signal?: AbortSignal,
+          ) => result(toolCallId, signal, 'pi'),
+        })
+      return {
+        name,
+        label: name,
+        description:
+          typeof tool.description === 'string' ? tool.description : '',
+        parameters: (await asSchema(tool.inputSchema)
+          .jsonSchema) as ToolDefinition['parameters'],
+        execute: (toolCallId, _params, signal) =>
+          result(toolCallId, signal, 'tool'),
+      }
+    }),
   )
 }
 
@@ -1227,6 +1263,26 @@ function thinkingLevel(step: HarnessStep): {
   return { thinkingLevel: step.reasoning === 'none' ? 'off' : step.reasoning }
 }
 
+type PayloadHook = NonNullable<
+  NonNullable<Parameters<ModelRuntime['streamSimple']>[2]>['onPayload']
+>
+
+/** Asks AI Gateway for the model's fast tier, on the payload Pi's own hook returns. */
+function fastTier(inner: PayloadHook | undefined): PayloadHook {
+  return async (payload, model) => {
+    const body = ((await inner?.(payload, model)) ?? payload) as {
+      providerOptions?: { gateway?: object }
+    }
+    return {
+      ...body,
+      providerOptions: {
+        ...body.providerOptions,
+        gateway: { ...body.providerOptions?.gateway, speed: 'fast' },
+      },
+    }
+  }
+}
+
 function gatewayHeaders(credential: {
   token: string
   method: string

package/src/ai-stored-state.ts

@@ -211,9 +211,7 @@ const storeState = <S extends AIStoredView>(state: S): AIStoredState => {
     if (
       value === undefined ||
       (typeof value === 'string' && value.length > 80) ||
-      (field !== undefined &&
-        payloadFields.has(field) &&
-        !(field === 'input' && isRecord(value) && 'message' in value)) ||
+      (field !== undefined && payloadFields.has(field)) ||
       (isRecord(value) && Object.hasOwn(value, '$a2Value'))
     )
       return intern(value)