experimental-a2 0.17.0 → 0.18.0

package/src/ai-agent.ts

@@ -23,9 +23,9 @@ import { AsyncLocalStorage } from 'node:async_hooks'
 import { createHash } from 'node:crypto'
 import { convertToModelMessages, isToolUIPart } from 'ai'
 import type {
+  Experimental_SandboxSession,
   FinishReason,
   Instructions,
-  LanguageModel,
   LanguageModelUsage,
   ModelMessage,
   TextStreamPart,
@@ -34,6 +34,15 @@ import type {
   UIMessageChunk,
 } from 'ai'
 import { generateAISDKStep, type AISDKStepSettings } from './ai-sdk-step.ts'
+import {
+  harnessDriver,
+  harnessLog,
+  harnessStates,
+  registerHarnessStates,
+  type Harness,
+  type HarnessDriver,
+  type HarnessStates,
+} from './ai-harness.ts'
 import { gatewayModelId, readModelLimits } from './ai-model-metadata.ts'
 import { estimateInputTokens, measuredInputTokens } from './ai-input-tokens.ts'
 import { AI_STATE_REDUCER, agentContract, events as aiEvents } from './ai.ts'
@@ -66,6 +75,7 @@ import type {
 import { A2Error } from './errors.ts'
 import {
   consumeSchedulerSendFailure,
+  failedPreviousAttempt,
   installAmbientToolScopeStorage,
   type AmbientToolScope,
 } from './internal.ts'
@@ -141,7 +151,8 @@ export type AgentGenerateContext<
 > = AgentResolverContext<M, D, X> & {
   generationId: string
   modelMessages: ModelMessage[]
-  model: LanguageModel
+  /** The resolved AI Gateway model ID. */
+  model: string
   tools: T
   instructions?: Instructions
   generation: AgentGenerationSettings<T>
@@ -188,7 +199,8 @@ export type Resolvable<T, C> = T | ((context: C) => T | Promise<T>)
 
 /** A model with its reasoning level and speed, chosen together. */
 export type ModelSettings = {
-  model: LanguageModel
+  /** An AI Gateway model ID, such as `'anthropic/claude-sonnet-5.5'`. */
+  model: string
   /** Overrides `generation.reasoning`. */
   reasoning?: ReasoningLevel
   /** `fast` prefers AI Gateway's faster tier when the model offers one. */
@@ -235,36 +247,29 @@ export type HandlersOptions<
 > = AgentModelOptions<M, D, T, X> & {
   contract: Contract<D, Record<never, never>, unknown>
   telemetry?: A2Telemetry
+  /** Where harness steps keep their saved states; one per agent. */
+  harnessStates?: HarnessStates
 }
 
-export type AgentModelOptions<
-  M extends UIMessage,
-  D extends AIEventDefs<M> & EventDefs,
+/** What a tool call does when a crash or a cancellation may have cut it short. */
+export type ToolRecovery = 'rerun' | 'report'
+
+/**
+ * The AI SDK harness, the default: one AI SDK `streamText()` step per
+ * durable generation, with A2 running tools, approvals, and compaction.
+ */
+export type AISDKHarness<
+  M extends UIMessage = UIMessage,
+  D extends AIEventDefs<M> & EventDefs = AIEventDefs<M> & EventDefs,
   T extends ToolSet = ToolSet,
   X extends object = Record<never, never>,
 > = {
-  /**
-   * AI SDK model string or provider model, or `ModelSettings` to choose its
-   * reasoning level and speed with it, optionally resolved per generation.
-   */
-  model: Resolvable<
-    LanguageModel | ModelSettings,
-    AgentResolverContext<M, D, X>
-  >
-  tools?: T
-  instructions?: Resolvable<
-    Instructions | undefined,
-    AgentResolverContext<M, D, X>
-  >
+  type: 'ai-sdk'
   /** Settings for each one-step AI SDK streamText call, resolved per generation. */
   generation?: Resolvable<
     AgentGenerationSettings<T>,
     AgentResolverContext<M, D, X>
   >
-  /**
-   * Maximum model steps in one assistant response. Defaults to Infinity.
-   */
-  maxSteps?: number
   /** Map default AI SDK stream parts to typed assistant message metadata. */
   messageMetadata?: AgentMessageMetadata<M, D, T, X>
   /**
@@ -276,6 +281,53 @@ export type AgentModelOptions<
     CompactionOptions<M, D, X>,
     AgentResolverContext<M, D, X>
   >
+}
+
+type AgentLoopOptions<
+  M extends UIMessage,
+  D extends AIEventDefs<M> & EventDefs,
+  T extends ToolSet,
+  X extends object,
+  H extends string,
+> =
+  | {
+      /** The agent's loop: the AI SDK's (the default), or `fx()` or `pi()`. */
+      harness?: AISDKHarness<M, D, T, X>
+      /** Per tool name: rerun an interrupted call (the default) or report it to the model. */
+      toolRecovery?: Readonly<
+        Partial<Record<NoInfer<keyof T & string>, ToolRecovery>>
+      >
+    }
+  | {
+      harness: Harness<'fx' | 'pi', H>
+      /** Per tool name, the harness's own tools included: rerun an interrupted call (the default) or report it. */
+      toolRecovery?: Readonly<
+        Partial<Record<NoInfer<(keyof T & string) | H>, ToolRecovery>>
+      >
+    }
+
+export type AgentModelOptions<
+  M extends UIMessage,
+  D extends AIEventDefs<M> & EventDefs,
+  T extends ToolSet = ToolSet,
+  X extends object = Record<never, never>,
+  /** The tool names a harness from `fx()` or `pi()` brings. */
+  H extends string = string,
+> = {
+  /**
+   * An AI Gateway model ID, or `ModelSettings` to choose its reasoning level
+   * and speed with it, optionally resolved per generation.
+   */
+  model: Resolvable<string | ModelSettings, AgentResolverContext<M, D, X>>
+  tools?: T
+  instructions?: Resolvable<
+    Instructions | undefined,
+    AgentResolverContext<M, D, X>
+  >
+  /**
+   * Maximum model steps in one assistant response. Defaults to Infinity.
+   */
+  maxSteps?: number
   /** Progress is durably flushed at either limit, whichever is reached first. */
   progress?: { maxChunks?: number; maxDelayMs?: number }
   /**
@@ -283,7 +335,7 @@ export type AgentModelOptions<
    * independent tools run concurrently.
    */
   toolLane?: string
-}
+} & AgentLoopOptions<M, D, T, X, H>
 
 type AgentServerOptions<
   M extends UIMessage,
@@ -291,7 +343,8 @@ type AgentServerOptions<
   S,
   T extends ToolSet,
   X extends object = Record<never, never>,
-> = AgentModelOptions<M, D, T, X> &
+  H extends string = string,
+> = AgentModelOptions<M, D, T, X, H> &
   Omit<
     ServerOptions<D, Record<never, never>, Reducer<D, S>>,
     'contract' | 'handlers'
@@ -304,11 +357,14 @@ export type AgentOptions<
   M extends UIMessage,
   E extends EventDefs,
   T extends ToolSet,
+  H extends string = string,
 > = AgentServerOptions<
   NoInfer<M>,
   AIEventDefs<NoInfer<M>> & E,
   AIState<NoInfer<M>>,
-  T
+  T,
+  Record<never, never>,
+  H
 > & {
   /** Identity: prefixes storage keys and queue messages. */
   name: string
@@ -327,7 +383,8 @@ export type AgentFromContractOptions<
   T extends ToolSet,
   Name extends string = string,
   X extends object = Record<never, never>,
-> = AgentServerOptions<AIMessageOf<D>, D, S, T, X> & {
+  H extends string = string,
+> = AgentServerOptions<AIMessageOf<D>, D, S, T, X, H> & {
   /**
    * A shared `agentContract()`. Its application state reaches every resolver
    * as `state`; a fully custom contract is core A2, served by `createServer()`.
@@ -407,13 +464,14 @@ const withModelSettings = <T extends ToolSet>(options: {
   }
 }
 
-const modelName = (model: LanguageModel): string => {
-  if (typeof model === 'string') return model
-  if ('modelId' in model && typeof model.modelId === 'string') {
-    return model.modelId
-  }
-  return 'custom'
-}
+const instructionTexts = (instructions: Instructions | undefined): string[] =>
+  instructions === undefined
+    ? []
+    : typeof instructions === 'string'
+      ? [instructions]
+      : (Array.isArray(instructions) ? instructions : [instructions]).map(
+          (message) => message.content,
+        )
 
 const errorMessage = (error: unknown): string =>
   error instanceof Error ? error.message : String(error)
@@ -877,7 +935,8 @@ const lifecycleEvents = async <
   pending: PendingToolCall[]
   tools: T
   generation: AgentGenerationSettings<T>
-  custom: boolean
+  /** A harness's tool calls run without approval. */
+  approvals?: boolean
 }): Promise<{ events: AppendInput<D>[]; pending: PendingToolCall[] }> => {
   const result: AppendInput<D>[] = []
   const pending = [...options.pending]
@@ -927,11 +986,14 @@ const lifecycleEvents = async <
         ...(chunk.title === undefined ? {} : { title: chunk.title }),
       }
       take(payload.toolCallId)
-      const classification = approvalClassification(
-        options.tools,
-        options.generation,
-        chunk.toolName,
-      )
+      const classification =
+        options.approvals === false
+          ? 'automatic'
+          : approvalClassification(
+              options.tools,
+              options.generation,
+              chunk.toolName,
+            )
       if (classification === 'automatic') {
         result.push(await toolCalledEvent<D>(payload))
       } else {
@@ -1142,6 +1204,26 @@ const validateCompaction = <
   return compaction
 }
 
+// Options of the AI SDK harness, which live in it, not on agent().
+const AI_SDK_HARNESS_OPTIONS = [
+  'generation',
+  'generate',
+  'compaction',
+  'messageMetadata',
+] as const
+
+const assertNoLoopOptions = (options: { harness?: { type: string } }): void => {
+  const harness = options.harness?.type ?? 'ai-sdk'
+  for (const key of AI_SDK_HARNESS_OPTIONS) {
+    if (!Object.hasOwn(options, key)) continue
+    throw new TypeError(
+      harness === 'ai-sdk'
+        ? `${key} belongs to the harness: agent({ harness: { type: 'ai-sdk', ${key} } })`
+        : `${key} is an option of the AI SDK harness; ${harness}() runs its own loop`,
+    )
+  }
+}
+
 export function createHandlers<
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
@@ -1175,22 +1257,63 @@ export function createHandlers<
     throw new TypeError('maxSteps must be a positive integer or Infinity')
   }
 
+  assertNoLoopOptions(options)
+  const configured = options.harness
+  const loop: Omit<AISDKHarness<M, D, T, X>, 'type'> = configured?.type ===
+  'ai-sdk'
+    ? configured
+    : {}
+  const harness: HarnessDriver | undefined =
+    configured === undefined || configured.type === 'ai-sdk'
+      ? undefined
+      : harnessDriver(configured)
+  const states = options.harnessStates ?? harnessStates()
+  if (harness !== undefined) {
+    for (const [name, tool] of Object.entries(options.tools ?? {})) {
+      const needsApproval = (tool as { needsApproval?: unknown }).needsApproval
+      if (needsApproval !== undefined && needsApproval !== false)
+        throw new TypeError(
+          `a harness runs tools without approval; '${name}' sets needsApproval`,
+        )
+    }
+  }
+
   const configuredCompaction =
-    options.compaction ?? (options.generate === undefined ? {} : false)
+    loop.compaction ??
+    (loop.generate === undefined && harness === undefined ? {} : false)
   const staticCompaction =
     typeof configuredCompaction === 'function'
       ? undefined
       : validateCompaction({
           compaction: configuredCompaction,
-          explicit: options.compaction !== undefined,
+          explicit: loop.compaction !== undefined,
           generation:
-            typeof options.generation === 'function'
-              ? undefined
-              : options.generation,
+            typeof loop.generation === 'function' ? undefined : loop.generation,
           tools: options.tools,
         })
 
-  const tools = options.tools ?? ({} as T)
+  // The harness's own tools come first, so an agent tool with a name replaces them.
+  const tools = (
+    harness === undefined
+      ? (options.tools ?? {})
+      : { ...harness.tools, ...options.tools }
+  ) as T
+  const toolRecovery: Readonly<Record<string, ToolRecovery | undefined>> =
+    options.toolRecovery ?? {}
+  for (const [name, recovery] of Object.entries(toolRecovery)) {
+    if (!Object.hasOwn(tools, name))
+      throw new TypeError(
+        `toolRecovery names '${name}', which is not one of the agent's tools`,
+      )
+    if (typeof (tools[name] as RuntimeTool).execute !== 'function')
+      throw new TypeError(
+        `toolRecovery names '${name}', which has no execute function`,
+      )
+    if (recovery !== 'rerun' && recovery !== 'report')
+      throw new TypeError(
+        `toolRecovery for '${name}' must be 'rerun' or 'report'`,
+      )
+  }
   const telemetry = options.telemetry ?? CONSOLE_TELEMETRY
   const maxSteps = options.maxSteps ?? Number.POSITIVE_INFINITY
   const control = createControlRuntime({
@@ -1250,6 +1373,7 @@ export function createHandlers<
         toolCallId: string
         messages: ModelMessage[]
         abortSignal?: AbortSignal
+        experimental_sandbox?: Experimental_SandboxSession
       },
     ) => unknown
   }
@@ -1366,6 +1490,7 @@ export function createHandlers<
     call: ToolCalledPayload,
     execute: NonNullable<RuntimeTool['execute']>,
     messages: ModelMessage[],
+    sandbox: Experimental_SandboxSession | undefined,
   ): Promise<AppendInput<D> | void> => {
     let output: unknown
     try {
@@ -1373,6 +1498,7 @@ export function createHandlers<
         toolCallId: call.toolCallId,
         messages,
         abortSignal: ctx.signal,
+        ...(sandbox === undefined ? {} : { experimental_sandbox: sandbox }),
       })
     } catch (error) {
       return toolExecutionFailure(ctx, call, error)
@@ -1449,7 +1575,25 @@ export function createHandlers<
         error: `Tool '${call.toolName}' has no server executor`,
       })
     }
+    const report = toolRecovery[call.toolName] === 'report'
+    const interrupted = () =>
+      resultEvent(call, `${ctx.event.id}:execution:error`, {
+        error: `Tool '${call.toolName}' was interrupted and may have partly run.`,
+      })
+    // A thrown failure never follows `execute` for a reported tool (see
+    // below), so only a lapsed attempt may have reached it.
+    if (report && ctx.attempt > 1 && !failedPreviousAttempt.has(ctx))
+      return interrupted()
     const messages = await promptMessages(input)
+    // A sandbox that fails to resolve fails the attempt, which A2 retries;
+    // it is not the tool's failure to report to the model.
+    const sandbox =
+      harness?.sandbox === undefined
+        ? undefined
+        : await resolve(harness.sandbox, {
+            session: { id: ctx.session.id },
+            signal: ctx.signal,
+          })
     if (checkAbort(ctx.signal)) return
     const scope: AmbientToolScope = {
       contract: options.contract,
@@ -1457,9 +1601,19 @@ export function createHandlers<
     }
     // Async iteration resumes in the caller's context, so the whole
     // execution, not just the execute call, runs inside the scope.
-    return toolScopeStorage.run(scope, () =>
-      runToolExecution(ctx, call, execute, messages),
-    )
+    const run = () =>
+      toolScopeStorage.run(scope, () =>
+        runToolExecution(ctx, call, execute, messages, sandbox),
+      )
+    if (!report) return run()
+    try {
+      // Undefined means a user cancellation after `execute` started.
+      return (await run()) ?? (await interrupted())
+    } catch {
+      // A lapsed claim surrenders; the next attempt reports the call.
+      checkAbort(ctx.signal)
+      return interrupted()
+    }
   }
 
   const handleToolExecution = async (
@@ -1632,19 +1786,17 @@ export function createHandlers<
       } as AppendInput<D>)
       return events
     }
-    const startedPayload = (
-      model?: LanguageModel,
-    ): GenerationStartedPayload => ({
+    const startedPayload = (model?: string): GenerationStartedPayload => ({
       requestId,
       messageId: request.messageId,
       generationId,
       responseMessageId,
       attempt,
-      ...(model === undefined ? {} : { model: modelName(model) }),
+      ...(model === undefined ? {} : { model }),
       promptThroughIndex: snapshot.index,
     })
 
-    let resolvedModel: LanguageModel | undefined
+    let resolvedModel: string | undefined
     let resolving = true
     let generationStarted = false
     try {
@@ -1653,15 +1805,10 @@ export function createHandlers<
         throw new TypeError('the model resolver returned undefined')
       }
       const settings: ModelSettings =
-        typeof choice === 'string' || 'specificationVersion' in choice
-          ? { model: choice }
-          : choice
+        typeof choice === 'string' ? { model: choice } : choice
       const model = settings.model
-      if (
-        typeof model !== 'string' &&
-        (typeof model !== 'object' || model === null)
-      )
-        throw new TypeError('the model resolver must return a model')
+      if (typeof model !== 'string' || model === '')
+        throw new TypeError('the model must be an AI Gateway model ID')
       resolvedModel = model
       if (checkAbort(ctx.signal)) return
       const resolvedInstructions =
@@ -1670,7 +1817,7 @@ export function createHandlers<
           : await resolve(options.instructions, resolverContext)
       if (checkAbort(ctx.signal)) return
       const resolvedGeneration: AgentGenerationSettings<T> = await resolve(
-        options.generation ?? {},
+        loop.generation ?? {},
         resolverContext,
       )
       if (checkAbort(ctx.signal)) return
@@ -1686,13 +1833,13 @@ export function createHandlers<
       })
       const compaction =
         typeof configuredCompaction === 'function' ||
-        typeof options.generation === 'function'
+        typeof loop.generation === 'function'
           ? validateCompaction({
               compaction:
                 typeof configuredCompaction === 'function'
                   ? await resolve(configuredCompaction, resolverContext)
                   : staticCompaction!,
-              explicit: options.compaction !== undefined,
+              explicit: loop.compaction !== undefined,
               generation,
               tools: options.tools,
             })
@@ -1802,7 +1949,7 @@ export function createHandlers<
           })
           const inputTokens = measuredInputTokens({
             estimate: inputTokenEstimate,
-            model: modelName(resolvedModel),
+            model: resolvedModel,
             calibration: coordinatorState.calibration,
           })
           const threshold =
@@ -1988,15 +2135,40 @@ export function createHandlers<
         messages: generationMessages,
         modelMessages,
       }
-      const custom = options.generate !== undefined
       const source =
-        options.generate === undefined
-          ? await generateWithAISDK({
+        harness !== undefined
+          ? await harness.step({
               agentName: options.contract.name,
-              context: generateContext,
-              messageMetadata: options.messageMetadata,
+              sessionId: ctx.event.sessionId,
+              requestIndex: ctx.event.index,
+              generationId,
+              responseMessageId,
+              throughIndex: snapshot.index,
+              messages: generationMessages,
+              model,
+              ...(settings.reasoning === undefined
+                ? {}
+                : { reasoning: settings.reasoning }),
+              ...(settings.speed === undefined
+                ? {}
+                : { speed: settings.speed }),
+              instructions: instructionTexts(resolvedInstructions),
+              tools,
+              log: harnessLog({
+                states,
+                session: ctx.session as Parameters<
+                  typeof harnessLog
+                >[0]['session'],
+              }),
+              signal: ctx.signal,
             })
-          : { stream: await options.generate(generateContext) }
+          : loop.generate === undefined
+            ? await generateWithAISDK({
+                agentName: options.contract.name,
+                context: generateContext,
+                messageMetadata: loop.messageMetadata,
+              })
+            : { stream: await loop.generate(generateContext) }
       const updates = consumeGeneration({
         source,
         ...(options.progress === undefined
@@ -2025,7 +2197,7 @@ export function createHandlers<
           pending: pendingToolCalls,
           tools,
           generation,
-          custom,
+          approvals: harness === undefined,
         })
         pendingToolCalls = lifecycle.pending
         await control.append({
@@ -2064,7 +2236,7 @@ export function createHandlers<
         id: await aiEventId('a2.ai.completed', generationId),
         payload: completed,
       } as AppendInput<D>
-      const unresolved = options.generate
+      const unresolved = loop.generate
         ? pendingToolCalls.filter(
             (pending) => pending.classification !== 'automatic',
           )
@@ -2079,7 +2251,7 @@ export function createHandlers<
           .filter(
             (pending) =>
               pending.classification === 'automatic' ||
-              (options.generate === undefined &&
+              (loop.generate === undefined &&
                 pending.classification === 'unknown'),
           )
           .map((pending) => toolCalledEvent<D>(pending.call)),
@@ -2255,8 +2427,9 @@ export function agent<
   T extends ToolSet = ToolSet,
   M extends UIMessage = UIMessage,
   E extends EventDefs = Record<never, never>,
+  H extends string = never,
 >(
-  options: AgentOptions<M, E, T>,
+  options: AgentOptions<M, E, T, H>,
 ): Agent<AIEventDefs<M> & E, AIState<M>, typeof AI_STATE_REDUCER>
 export function agent<
   D extends AIEventDefs<AIMessageOf<D>> & EventDefs,
@@ -2264,22 +2437,22 @@ export function agent<
   T extends ToolSet = ToolSet,
   Name extends string = string,
   X extends object = Record<never, never>,
->(options: AgentFromContractOptions<D, S, T, Name, X>): Agent<D, S, Name>
+  H extends string = never,
+>(options: AgentFromContractOptions<D, S, T, Name, X, H>): Agent<D, S, Name>
 export function agent(
   options:
     | AgentOptions<UIMessage, EventDefs, ToolSet>
     | AgentFromContractOptions<AIEventDefs & EventDefs, AIState, ToolSet>,
 ): Agent<AIEventDefs & EventDefs, AIState> {
   type D = AIEventDefs & EventDefs
+  assertNoLoopOptions(options)
   const {
     model,
     tools,
     instructions,
-    generation,
     maxSteps,
-    messageMetadata,
-    generate,
-    compaction,
+    harness,
+    toolRecovery,
     progress,
     toolLane,
     handlers,
@@ -2307,22 +2480,22 @@ export function agent(
     })
     serverOptions = remaining
   }
+  const states = harnessStates()
   const builtIns = createHandlers<UIMessage, D>({
     contract,
     model,
+    harnessStates: states,
     ...(serverOptions.telemetry === undefined
       ? {}
       : { telemetry: serverOptions.telemetry }),
     ...(tools === undefined ? {} : { tools }),
     ...(instructions === undefined ? {} : { instructions }),
-    ...(generation === undefined ? {} : { generation }),
     ...(maxSteps === undefined ? {} : { maxSteps }),
-    ...(messageMetadata === undefined ? {} : { messageMetadata }),
-    ...(generate === undefined ? {} : { generate }),
-    ...(compaction === undefined ? {} : { compaction }),
+    ...(harness === undefined ? {} : { harness }),
+    ...(toolRecovery === undefined ? {} : { toolRecovery }),
     ...(progress === undefined ? {} : { progress }),
     ...(toolLane === undefined ? {} : { toolLane }),
-  })
+  } as HandlersOptions<UIMessage, D>)
   for (const type of Object.keys(handlers ?? {})) {
     if (Object.hasOwn(builtIns, type))
       throw new TypeError(`agent handles '${type}' itself`)
@@ -2332,6 +2505,7 @@ export function agent(
     contract,
     handlers: { ...builtIns, ...handlers },
   })
+  registerHarnessStates(server, states)
   const applicationEvents = 'contract' in rest
   setServerFetchHooks(server, {
     validateIngress: (context) =>
@@ -2340,4 +2514,4 @@ export function agent(
   return server
 }
 
-export type { Instructions, LanguageModel, ToolSet }
+export type { Instructions, ToolSet }

package/src/ai-fx.browser.ts

@@ -0,0 +1,5 @@
+/** experimental-a2/ai/fx is server-only; browsers reach an agent through its route. */
+throw new Error(
+  'experimental-a2/ai/fx is server-only — a client bundle imported it. Keep agent({ harness: fx() }) in server code.',
+)
+export {}