experimental-a2 0.20.0 → 0.22.0

package/src/ai-agent.ts

@@ -1,6 +1,6 @@
 /**
- * The server-only implementation of the experimental-a2/ai protocol: the
- * handler table and `agent()`, its assembly as an ordinary A2 server.
+ * The server-only implementation of the experimental-a2/ai protocol: its
+ * consumers and `agent()`, its assembly as an ordinary A2 server.
  */
 
 // oxlint-disable no-await-in-loop -- stream chunks and durable appends are
@@ -75,14 +75,17 @@ import type {
 import { A2Error } from './errors.ts'
 import {
   consumeSchedulerSendFailure,
+  drainSettledHooks,
   failedPreviousAttempt,
   installAmbientToolScopeStorage,
   type AmbientToolScope,
 } from './internal.ts'
+import { runtimeConsumer } from './runtime-consumer.ts'
 import {
   createServer,
   type A2Server,
-  type HandlerContext,
+  type Consumer,
+  type ConsumerContext,
   type ServerOptions,
 } from './server.ts'
 import {
@@ -191,7 +194,7 @@ export type AgentResolverContext<
    * in the checkpoint the generation already reads, so it costs no I/O.
    */
   state: X
-  session: Pick<HandlerContext<D>['session'], 'state'>
+  session: Pick<ConsumerContext<D>['session'], 'state'>
   signal: AbortSignal
 }
 
@@ -239,7 +242,7 @@ export type CompactionOptions<
 
 export type AgentGenerationSettings<T extends ToolSet> = AISDKStepSettings<T>
 
-export type HandlersOptions<
+export type ConsumersOptions<
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
   T extends ToolSet = ToolSet,
@@ -344,27 +347,24 @@ type AgentServerOptions<
   T extends ToolSet,
   X extends object = Record<never, never>,
   H extends string = string,
+  C extends readonly unknown[] = readonly Consumer<D>[],
 > = AgentModelOptions<M, D, T, X, H> &
-  Omit<
-    ServerOptions<D, Record<never, never>, Reducer<D, S>>,
-    'contract' | 'handlers'
-  > & {
-    /** Handlers for application events. */
-    handlers?: ServerOptions<D, Record<never, never>, Reducer<D, S>>['handlers']
-  }
+  Omit<ServerOptions<D, Record<never, never>, Reducer<D, S>, C>, 'contract'>
 
 export type AgentOptions<
   M extends UIMessage,
   E extends EventDefs,
   T extends ToolSet,
   H extends string = string,
+  C extends readonly unknown[] = readonly Consumer<AIEventDefs<M> & E>[],
 > = AgentServerOptions<
   NoInfer<M>,
   AIEventDefs<NoInfer<M>> & E,
   AIState<NoInfer<M>>,
   T,
   Record<never, never>,
-  H
+  H,
+  C
 > & {
   /** Identity: prefixes storage keys and queue messages. */
   name: string
@@ -384,7 +384,8 @@ export type AgentFromContractOptions<
   Name extends string = string,
   X extends object = Record<never, never>,
   H extends string = string,
-> = AgentServerOptions<AIMessageOf<D>, D, S, T, X, H> & {
+  C extends readonly unknown[] = readonly Consumer<D>[],
+> = AgentServerOptions<AIMessageOf<D>, D, S, T, X, H, C> & {
   /**
    * A shared `agentContract()`. Its application state reaches every resolver
    * as `state`; a fully custom contract is core A2, served by `createServer()`.
@@ -480,8 +481,8 @@ const toolScopeStorage = installAmbientToolScopeStorage(
   () => new AsyncLocalStorage<AmbientToolScope>(),
 )
 
-/** Run `fn` with `ctx` as the ambient tool handler context, for direct tool tests. */
-export function runWithHandlerContext<D extends AgentToolEventDefs, R>(
+/** Run `fn` with `ctx` as the ambient tool consumer context, for direct tool tests. */
+export function runWithConsumerContext<D extends AgentToolEventDefs, R>(
   ctx: AgentToolContext<D>,
   fn: () => R,
 ): R {
@@ -683,7 +684,7 @@ async function* consumeGeneration(options: {
   const reader = options.source.stream.getReader()
   // The next read starts before a batch is yielded, and the consumer may
   // take a while to ask again: a stream that fails meanwhile must not leave
-  // that read rejected with no handler, which ends the process.
+  // that read rejected with nothing to catch it, which ends the process.
   const readAhead = (): Promise<StreamRead<UIMessageChunk>> => {
     const read = reader.read()
     read.catch(() => {})
@@ -1232,14 +1233,12 @@ const assertNoLoopOptions = (options: { harness?: { type: string } }): void => {
   }
 }
 
-export function createHandlers<
+export function createConsumers<
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
   T extends ToolSet = ToolSet,
   X extends object = Record<never, never>,
->(
-  options: HandlersOptions<M, D, T, X>,
-): NonNullable<ServerOptions<D>['handlers']> {
+>(options: ConsumersOptions<M, D, T, X>): Consumer<D>[] {
   if (options.model === undefined) {
     throw new TypeError('agent requires model')
   }
@@ -1349,7 +1348,7 @@ export function createHandlers<
     contract: options.contract,
   })
   const readContext = async (
-    session: HandlerContext<D>['session'],
+    session: ConsumerContext<D>['session'],
     through: number | 'latest',
   ) => {
     const snapshot = await session.state(contextReducer, { through })
@@ -1390,7 +1389,7 @@ export function createHandlers<
     ) => unknown
   }
 
-  type ToolHandlerContext = HandlerContext<D, 'ai.tool.execution.requested'>
+  type ToolConsumerContext = ConsumerContext<D, 'ai.tool.execution.requested'>
 
   const resultEvent = async (
     call: ToolCalledPayload,
@@ -1435,7 +1434,7 @@ export function createHandlers<
     }) as AppendInput<D>
 
   const promptMessages = (input: {
-    ctx: ToolHandlerContext
+    ctx: ToolConsumerContext
     call: ToolCalledPayload
     generation: GenerationStartedPayload
   }): Promise<ModelMessage[]> => {
@@ -1485,7 +1484,7 @@ export function createHandlers<
   }
 
   const toolExecutionFailure = async (
-    ctx: ToolHandlerContext,
+    ctx: ToolConsumerContext,
     call: ToolCalledPayload,
     error: unknown,
   ): Promise<AppendInput<D> | void> => {
@@ -1498,7 +1497,7 @@ export function createHandlers<
   }
 
   const runToolExecution = async (
-    ctx: ToolHandlerContext,
+    ctx: ToolConsumerContext,
     call: ToolCalledPayload,
     execute: NonNullable<RuntimeTool['execute']>,
     messages: ModelMessage[],
@@ -1575,7 +1574,7 @@ export function createHandlers<
   }
 
   const executeTool = async (input: {
-    ctx: ToolHandlerContext
+    ctx: ToolConsumerContext
     call: ToolCalledPayload
     generation: GenerationStartedPayload
   }): Promise<AppendInput<D> | void> => {
@@ -1629,7 +1628,7 @@ export function createHandlers<
   }
 
   const handleToolExecution = async (
-    ctx: ToolHandlerContext,
+    ctx: ToolConsumerContext,
   ): Promise<AppendInput<D>> => {
     try {
       const state = (
@@ -1666,7 +1665,7 @@ export function createHandlers<
   }
 
   const generationHandler = async (
-    ctx: HandlerContext<D, 'ai.generation.requested'>,
+    ctx: ConsumerContext<D, 'ai.generation.requested'>,
   ): Promise<void | AppendInput<D> | readonly AppendInput<D>[]> => {
     if (checkAbort(ctx.signal)) return
     const requestId = ctx.event.id
@@ -2150,6 +2149,7 @@ export function createHandlers<
       const source =
         harness !== undefined
           ? await harness.step({
+              scope: states,
               agentName: options.contract.name,
               sessionId: ctx.event.sessionId,
               requestIndex: ctx.event.index,
@@ -2315,23 +2315,41 @@ export function createHandlers<
     }
   }
 
-  return {
-    'ai.model.metadata.requested': {
-      handler: async (ctx) => {
-        const limits = await readModelLimits({
-          modelId: ctx.event.payload.modelId,
-          signal: ctx.signal,
-        })
-        ctx.signal.throwIfAborted()
-        return {
-          type: 'ai.model.metadata.resolved',
-          id: await aiEventId('a2.ai.model.metadata.resolved', ctx.event.id),
-          payload: { modelId: ctx.event.payload.modelId, limits },
-        } as AppendInput<D>
-      },
-    },
-    'ai.control.requested': {
+  const contract = options.contract as unknown as Contract<
+    D,
+    Record<never, never>,
+    undefined
+  >
+  const { toolLane } = options
+  type Of<K extends string> = ContractEvent<D, K & keyof D & string>
+  type Generation = Of<'ai.generation.requested'>
+  type ToolCall = Of<'ai.tool.execution.requested'>
+  return [
+    runtimeConsumer({
+      contract,
+      name: 'a2.ai.metadata',
+      types: ['ai.model.metadata.requested'],
+    }).run(async (ctx) => {
+      const limits = await readModelLimits({
+        modelId: ctx.event.payload.modelId,
+        signal: ctx.signal,
+      })
+      ctx.signal.throwIfAborted()
+      return {
+        type: 'ai.model.metadata.resolved',
+        id: await aiEventId('a2.ai.model.metadata.resolved', ctx.event.id),
+        payload: { modelId: ctx.event.payload.modelId, limits },
+      } as AppendInput<D>
+    }),
+    // One lane for every control event: commands, work reports, and closes
+    // apply one at a time. Its lane function also rejects a bad command at
+    // append, before anything is written.
+    runtimeConsumer({
+      contract,
+      name: 'a2.ai.control',
+      types: ['ai.control.requested', 'ai.work.reported', 'ai.session.closed'],
       lane: ({ event }) => {
+        if (event.type !== 'ai.control.requested') return 'a2.ai.control'
         const command = event.payload as ControlCommand<M>
         if (event.id === '')
           throw new A2Error(
@@ -2353,17 +2371,25 @@ export function createHandlers<
           )
         return 'a2.ai.control'
       },
-      handler: controlHandler,
-    },
-    'ai.work.reported': { lane: 'a2.ai.control', handler: controlHandler },
-    'ai.session.closed': { lane: 'a2.ai.control', handler: controlHandler },
-    'ai.generation.requested': {
-      lane: 'a2.ai.model',
+    }).run(controlHandler as never),
+    runtimeConsumer({
+      contract,
+      name: 'a2.ai.model',
+      types: ['ai.generation.requested'],
+      lane: () => 'a2.ai.model',
+      // The contract is generic here, so the map's keys widen to string and
+      // each predicate names its own event types.
       abortOn: {
-        'ai.control.requested': (event, trigger) =>
+        'ai.control.requested': (
+          event: Of<'ai.control.requested'>,
+          trigger: Generation,
+        ) =>
           event.payload.action === 'stop' &&
           event.payload.turnId === trigger.payload.control?.turnId,
-        'ai.control.committed': (event, trigger) => {
+        'ai.control.committed': (
+          event: Of<'ai.control.committed'>,
+          trigger: Generation,
+        ) => {
           const active = (event.payload as ControlCommit).view
           return (
             active?.turnId !== trigger.payload.control?.turnId ||
@@ -2371,39 +2397,50 @@ export function createHandlers<
           )
         },
         'ai.session.closed': true,
-      },
-      handler: (ctx) =>
-        telemetry.span(
-          'a2.ai.generate',
-          {
-            'a2.contract': options.contract.name,
-            'a2.session_id': ctx.event.sessionId,
-            'a2.ai.generate.request_id': ctx.event.id,
-            'a2.ai.generate.attempt': ctx.attempt,
-          },
-          async () => {
-            try {
-              return control.settled({
-                ctx,
-                events: await generationHandler(ctx),
-              })
-            } catch (error) {
-              if (control.cancelled({ error, signal: ctx.signal }))
-                return control.settled({ ctx, events: undefined })
-              throw error
-            }
-          },
-        ),
-    },
-    'ai.tool.execution.requested': {
-      ...(options.toolLane === undefined ? {} : { lane: options.toolLane }),
+      } as never,
+    }).run((ctx) =>
+      telemetry.span(
+        'a2.ai.generate',
+        {
+          'a2.contract': options.contract.name,
+          'a2.session_id': ctx.event.sessionId,
+          'a2.ai.generate.request_id': ctx.event.id,
+          'a2.ai.generate.attempt': ctx.attempt,
+        },
+        async () => {
+          try {
+            return control.settled({
+              ctx,
+              events: await generationHandler(ctx),
+            })
+          } catch (error) {
+            if (control.cancelled({ error, signal: ctx.signal }))
+              return control.settled({ ctx, events: undefined })
+            throw error
+          }
+        },
+      ),
+    ),
+    runtimeConsumer({
+      contract,
+      name: 'a2.ai.tools',
+      types: ['ai.tool.execution.requested'],
+      ...(toolLane === undefined ? {} : { lane: () => toolLane }),
       abortOn: {
-        'ai.control.requested': (event, trigger) =>
+        'ai.control.requested': (
+          event: Of<'ai.control.requested'>,
+          trigger: ToolCall,
+        ) =>
           event.payload.action === 'stop' &&
           event.payload.turnId === trigger.payload.turnId,
-        'ai.generation.failed': (event, trigger) =>
-          event.payload.generationId === trigger.payload.call.generationId,
-        'ai.control.committed': (event, trigger) => {
+        'ai.generation.failed': (
+          event: Of<'ai.generation.failed'>,
+          trigger: ToolCall,
+        ) => event.payload.generationId === trigger.payload.call.generationId,
+        'ai.control.committed': (
+          event: Of<'ai.control.committed'>,
+          trigger: ToolCall,
+        ) => {
           const active = (event.payload as ControlCommit).view
           return (
             active?.turnId !== trigger.payload.turnId ||
@@ -2411,26 +2448,26 @@ export function createHandlers<
           )
         },
         'ai.session.closed': true,
-      },
-      handler: (ctx) =>
-        telemetry.span(
-          'a2.ai.tool',
-          {
-            'a2.contract': options.contract.name,
-            'a2.session_id': ctx.event.sessionId,
-            'a2.ai.tool.request_id': ctx.event.payload.call.requestId,
-            'a2.ai.tool.call_id': ctx.event.payload.call.toolCallId,
-            'a2.ai.tool.name': ctx.event.payload.call.toolName,
-            'a2.ai.tool.attempt': ctx.attempt,
-          },
-          () => handleToolExecution(ctx),
-        ),
-    },
-  } as NonNullable<ServerOptions<D>['handlers']>
+      } as never,
+    }).run((ctx) =>
+      telemetry.span(
+        'a2.ai.tool',
+        {
+          'a2.contract': options.contract.name,
+          'a2.session_id': ctx.event.sessionId,
+          'a2.ai.tool.request_id': ctx.event.payload.call.requestId,
+          'a2.ai.tool.call_id': ctx.event.payload.call.toolCallId,
+          'a2.ai.tool.name': ctx.event.payload.call.toolName,
+          'a2.ai.tool.attempt': ctx.attempt,
+        },
+        () => handleToolExecution(ctx),
+      ),
+    ),
+  ] as Consumer<D>[]
 }
 
 /**
- * Serve an AI agent: the built-in protocol, its handlers, and A2's store,
+ * Serve an AI agent: the built-in protocol, its consumers, and A2's store,
  * scheduler, and telemetry, as one ordinary A2 server. Pass `name` for the
  * built-in protocol, or a shared `agentContract()` when clients fold
  * application state.
@@ -2440,8 +2477,9 @@ export function agent<
   M extends UIMessage = UIMessage,
   E extends EventDefs = Record<never, never>,
   H extends string = never,
+  const C extends readonly unknown[] = readonly Consumer<AIEventDefs<M> & E>[],
 >(
-  options: AgentOptions<M, E, T, H>,
+  options: AgentOptions<M, E, T, H, C>,
 ): Agent<AIEventDefs<M> & E, AIState<M>, typeof AI_STATE_REDUCER>
 export function agent<
   D extends AIEventDefs<AIMessageOf<D>> & EventDefs,
@@ -2450,7 +2488,8 @@ export function agent<
   Name extends string = string,
   X extends object = Record<never, never>,
   H extends string = never,
->(options: AgentFromContractOptions<D, S, T, Name, X, H>): Agent<D, S, Name>
+  const C extends readonly unknown[] = readonly Consumer<D>[],
+>(options: AgentFromContractOptions<D, S, T, Name, X, H, C>): Agent<D, S, Name>
 export function agent(
   options:
     | AgentOptions<UIMessage, EventDefs, ToolSet>
@@ -2458,6 +2497,12 @@ export function agent(
 ): Agent<AIEventDefs & EventDefs, AIState> {
   type D = AIEventDefs & EventDefs
   assertNoLoopOptions(options)
+  if (
+    options.toolLane !== undefined &&
+    (typeof options.toolLane !== 'string' || options.toolLane.length === 0)
+  ) {
+    throw new TypeError('toolLane must be a non-empty string')
+  }
   const {
     model,
     tools,
@@ -2467,13 +2512,12 @@ export function agent(
     toolRecovery,
     progress,
     toolLane,
-    handlers,
     ...rest
   } = options
   let contract: Contract<D, Record<never, never>, Reducer<D, AIState>>
   let serverOptions: Omit<
     ServerOptions<D, Record<never, never>, Reducer<D, AIState>>,
-    'contract' | 'handlers'
+    'contract'
   >
   if ('contract' in rest) {
     const { contract: shared, ...remaining } = rest
@@ -2493,7 +2537,7 @@ export function agent(
     serverOptions = remaining
   }
   const states = harnessStates()
-  const builtIns = createHandlers<UIMessage, D>({
+  const builtIns = createConsumers<UIMessage, D>({
     contract,
     model,
     harnessStates: states,
@@ -2507,17 +2551,24 @@ export function agent(
     ...(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`)
-  }
+  } as ConsumersOptions<UIMessage, D>)
   const server = createServer({
     ...serverOptions,
     contract,
-    handlers: { ...builtIns, ...handlers },
+    consumers: [...builtIns, ...(serverOptions.consumers ?? [])],
   })
   registerHarnessStates(server, states)
+  const driver =
+    harness === undefined || harness.type === 'ai-sdk'
+      ? undefined
+      : harnessDriver(harness)
+  const release = driver?.release
+  // What a harness kept for a session's next step goes once nothing is left
+  // to run for the session here.
+  if (release !== undefined)
+    drainSettledHooks.set(server, (sessionId) =>
+      release({ scope: states, sessionId }),
+    )
   const applicationEvents = 'contract' in rest
   setServerFetchHooks(server, {
     validateIngress: (context) =>

package/src/ai-control-server.ts

@@ -31,7 +31,7 @@ import type {
   ContractEvent,
   EventDefs,
 } from './contract.ts'
-import type { HandlerContext } from './server.ts'
+import type { ConsumerContext } from './server.ts'
 import type { Reducer } from './reducer.ts'
 import type { AICoordinatorState, CoordinatorCall } from './ai-coordinator.ts'
 import { validateSync } from './validate.ts'
@@ -56,18 +56,18 @@ export type AIControlRuntime<
   D extends AIEventDefs<M> & EventDefs,
 > = {
   handler(
-    ctx: HandlerContext<
+    ctx: ConsumerContext<
       D,
       'ai.control.requested' | 'ai.work.reported' | 'ai.session.closed'
     >,
   ): Promise<AppendInput<D>[]>
   append(options: {
-    ctx: HandlerContext<D>
+    ctx: ConsumerContext<D>
     name: string
     events: AppendInput<D>[]
   }): Promise<ContractEvent<D>[]>
   settled(options: {
-    ctx: HandlerContext<D>
+    ctx: ConsumerContext<D>
     events: void | AppendInput<D> | readonly AppendInput<D>[]
   }): Promise<AppendInput<D>>
   cancelled(options: { error: unknown; signal: AbortSignal }): boolean
@@ -117,7 +117,7 @@ export const createControlRuntime = <
   const coordinator = aiCoordinatorReducer(options.contract)
 
   const handler = async (
-    ctx: HandlerContext<
+    ctx: ConsumerContext<
       D,
       'ai.control.requested' | 'ai.work.reported' | 'ai.session.closed'
     >,
@@ -761,7 +761,7 @@ export const createControlRuntime = <
     settled,
     name,
   }: {
-    ctx: HandlerContext<D>
+    ctx: ConsumerContext<D>
     events: readonly AppendInput<D>[]
     settled: boolean
     name: string
@@ -828,7 +828,7 @@ export const createControlRuntime = <
     name,
     events,
   }: {
-    ctx: HandlerContext<D>
+    ctx: ConsumerContext<D>
     name: string
     events: AppendInput<D>[]
   }): Promise<ContractEvent<D>[]> => {
@@ -896,7 +896,7 @@ export const createControlRuntime = <
     ctx,
     events,
   }: {
-    ctx: HandlerContext<D>
+    ctx: ConsumerContext<D>
     events: void | AppendInput<D> | readonly AppendInput<D>[]
   }): Promise<AppendInput<D>> =>
     report({