experimental-a2 0.18.0 → 0.19.0

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,
@@ -59,6 +62,7 @@ import {
   type RecordedGeneration,
   saveHarness,
 } from './ai-harness.ts'
+import { type SchemaInput, schemaInput } from './ai-tool-input.ts'
 import type { CompactionPartData, ToolResultPayload } from './ai.ts'
 import { boundedCache } from './bounded-cache.ts'
 import { ambientToolScopeStorage } from './internal.ts'
@@ -110,11 +114,14 @@ type PiModules = {
   ai: typeof import('@earendil-works/pi-ai')
 }
 
-/** A content block without the text the transcript already holds. */
+/**
+ * A content block without the text the transcript already holds. A tool call
+ * keeps its arguments only where the transcript's input is not the model's.
+ */
 type StoredBlock =
   | (Omit<TextContent, 'text'> & { ref: string })
   | (Omit<ThinkingContent, 'thinking'> & { ref: string })
-  | Omit<ToolCall, 'arguments'>
+  | (Omit<ToolCall, 'arguments'> & Partial<Pick<ToolCall, 'arguments'>>)
 
 /** What a step's transcript parts do not hold of Pi's response. */
 type StoredResponse = Omit<AssistantMessage, 'content'> & {
@@ -126,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 = {
@@ -228,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
     }
@@ -248,10 +259,13 @@ 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) {
-      live?.(event)
+      await live?.(event)
       if (event.type === 'done' || event.type === 'error') {
         const response = event.type === 'done' ? event.message : event.error
         if (kind === 'compaction') state.compactions.push(response)
@@ -267,11 +281,39 @@ async function runStep(
     }
   }
 
+  /**
+   * A live call's input as Pi's loop, then the AI SDK step, read it. Pi
+   * answers a call to no tool, or whose arguments it cannot prepare or
+   * validate, itself. The agent's tool reads the arguments Pi accepts through
+   * its own schema; A2 runs Pi's own tools from the model's arguments, as
+   * Pi's loop does.
+   */
+  const callInput = async (block: ToolCall): Promise<SchemaInput> => {
+    const definition = tools.find((tool) => tool.name === block.name)
+    if (definition === undefined)
+      return { success: false, errorText: `Tool ${block.name} not found` }
+    let args: unknown
+    try {
+      const prepared = definition.prepareArguments?.(block.arguments)
+      args = modules.ai.validateToolArguments(
+        definition,
+        prepared === undefined
+          ? block
+          : { ...block, arguments: prepared as JsonObject },
+      )
+    } catch (error) {
+      return { success: false, errorText: asError(error).message }
+    }
+    if (builtinOf(step.tools, block.name) !== undefined)
+      return { success: true, value: block.arguments }
+    return schemaInput({ tools: step.tools, toolName: block.name, value: args })
+  }
+
   /** Streams one live step response as UI chunks and records it. */
   const liveStep = () => {
     const ref = (index: number) => `pi:${step.requestIndex}:${index}`
     const started = new Set<string>()
-    return (event: AssistantMessageEvent): void => {
+    return async (event: AssistantMessageEvent): Promise<void> => {
       switch (event.type) {
         case 'start':
           emit({ type: 'start-step' })
@@ -325,6 +367,8 @@ async function runStep(
         case 'done':
         case 'error': {
           const response = event.type === 'done' ? event.message : event.error
+          // Calls whose transcript input is not the model's arguments.
+          const changed = new Set<string>()
           if (event.type === 'done') {
             for (const block of response.content) {
               if (block.type !== 'toolCall') continue
@@ -335,26 +379,41 @@ async function runStep(
                   toolName: block.name,
                 })
               // Pi does not run the calls of a cut-off response.
-              emit(
+              const input: SchemaInput =
                 response.stopReason === 'length'
                   ? {
-                      type: 'tool-input-error',
-                      toolCallId: block.id,
-                      toolName: block.name,
-                      input: block.arguments,
+                      success: false,
                       errorText:
                         'The response hit the output token limit before this call was complete.',
                     }
-                  : {
+                  : await callInput(block)
+              if (
+                input.success &&
+                JSON.stringify(input.value) !== JSON.stringify(block.arguments)
+              )
+                changed.add(block.id)
+              emit(
+                input.success
+                  ? {
                       type: 'tool-input-available',
                       toolCallId: block.id,
                       toolName: block.name,
+                      input: input.value,
+                    }
+                  : {
+                      type: 'tool-input-error',
+                      toolCallId: block.id,
+                      toolName: block.name,
                       input: block.arguments,
+                      errorText: input.errorText,
                     },
               )
             }
           }
-          record({ call: 'step', response: storedResponse(response, ref) })
+          record({
+            call: 'step',
+            response: storedResponse(response, ref, changed),
+          })
           completion = {
             finishReason: finishReason(response.stopReason),
             usage: modelUsage(response.usage),
@@ -498,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,
@@ -508,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) => {
@@ -573,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>()
@@ -620,6 +707,7 @@ function recordedCalls(generation: RecordedGeneration): RecordedCall[] {
 function storedResponse(
   response: AssistantMessage,
   ref: (index: number) => string,
+  changed: ReadonlySet<string>,
 ): StoredResponse {
   return {
     ...response,
@@ -634,6 +722,7 @@ function storedResponse(
           return { ...rest, ref: ref(index) }
         }
         case 'toolCall': {
+          if (changed.has(block.id)) return block
           const { arguments: _, ...rest } = block
           return rest
         }
@@ -653,7 +742,8 @@ function restoredResponse(
       if (block.type === 'toolCall')
         return {
           ...block,
-          arguments: (inputs.get(block.id) ?? {}) as JsonObject,
+          arguments:
+            block.arguments ?? ((inputs.get(block.id) ?? {}) as JsonObject),
         }
       const { ref, ...rest } = block
       const text = texts.get(ref) ?? ''
@@ -826,6 +916,15 @@ async function recordedResult(
   return { ...kept, details: { ...details, truncation } }
 }
 
+/** The Pi tool a name in the agent's tool set runs, unless the agent replaced it. */
+const builtinOf = (tools: ToolSet, name: string): Builtin | undefined =>
+  BUILTINS.find(
+    (candidate) =>
+      candidate === name &&
+      Object.hasOwn(tools, name) &&
+      tools[name] === builtinTools[name],
+  )
+
 // Each definition is typed by its parameters; a session takes any of them.
 const untyped = (definition: object) => definition as ToolDefinition
 
@@ -876,32 +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 = BUILTINS.find(
-          (candidate) => candidate === name && builtinTools[name] === tool,
-        )
-        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'),
+      }
+    }),
   )
 }
 
@@ -1168,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)

package/src/ai-tool-input.ts

@@ -0,0 +1,45 @@
+import {
+  asSchema,
+  InvalidToolInputError,
+  type ToolSet,
+  TypeValidationError,
+} from 'ai'
+
+export type SchemaInput =
+  { success: true; value: unknown } | { success: false; errorText: string }
+
+/**
+ * A call's input through its tool's schema, as the AI SDK step reads it: the
+ * schema's output with its defaults, or the AI SDK's error text. A name the
+ * tool set does not own has no schema.
+ */
+export async function schemaInput({
+  tools,
+  toolName,
+  value,
+}: {
+  tools: ToolSet
+  toolName: string
+  value: unknown
+}): Promise<SchemaInput> {
+  const schema = Object.hasOwn(tools, toolName)
+    ? asSchema(tools[toolName]!.inputSchema)
+    : undefined
+  if (schema?.validate === undefined) return { success: true, value }
+  let cause: unknown
+  try {
+    const result = await schema.validate(value)
+    if (result.success) return { success: true, value: result.value }
+    cause = result.error
+  } catch (error) {
+    cause = error
+  }
+  return {
+    success: false,
+    errorText: new InvalidToolInputError({
+      toolName,
+      toolInput: JSON.stringify(value),
+      cause: TypeValidationError.wrap({ value, cause }),
+    }).message,
+  }
+}