experimental-a2 0.13.0 → 0.14.1

package/src/ai-server.ts

@@ -8,7 +8,10 @@
 // oxlint-disable no-await-in-loop -- stream chunks and durable appends are
 // ordered; parallel consumption would corrupt progress sequence and state.
 
+import { createControlRuntime, ControlCancelled } from './ai-control-server.ts'
+import type { ControlCommit } from './ai-control.ts'
 import { AsyncLocalStorage } from 'node:async_hooks'
+import { createHash } from 'node:crypto'
 import { asSchema, convertToModelMessages } from 'ai'
 import type {
   FinishReason,
@@ -34,17 +37,13 @@ import type {
   GenerationCompletedPayload,
   GenerationFailedPayload,
   GenerationProgressPayload,
-  GenerationRequestedPayload,
   GenerationStartedPayload,
   MessageCompletedPayload,
-  MessageCreatedPayload,
-  MessageInterruptedPayload,
   ToolCalledPayload,
   ToolResultPayload,
 } from './ai.ts'
 import {
   aiCoordinatorReducer,
-  continuationReady,
   type AICoordinatorState,
 } from './ai-coordinator.ts'
 import type { AppendInput, ContractEvent, EventDefs } from './contract.ts'
@@ -67,24 +66,23 @@ import {
 
 function validateAgentIngress(context: ServerIngressContext): void {
   const rejected = context.events.find((event) => {
-    if (event.type !== 'ai.message.created') {
-      return !(
-        event.type === 'ai.approval.responded' ||
-        event.type === 'ai.input.responded' ||
-        event.type === 'ai.message.interrupted' ||
-        event.type === 'ai.retry.requested'
+    if (event.type !== 'ai.control.requested') return true
+    const command = event.payload as Record<string, unknown>
+    if (
+      command['action'] === 'request-input' ||
+      command['action'] === 'tool-result'
+    )
+      return true
+    if (
+      command['action'] === 'send' ||
+      command['action'] === 'edit' ||
+      command['action'] === 'steer'
+    ) {
+      return (
+        (command['message'] as { role?: unknown } | undefined)?.role !== 'user'
       )
     }
-    const payload = event.payload
-    return (
-      typeof payload !== 'object' ||
-      payload === null ||
-      !('message' in payload) ||
-      typeof payload.message !== 'object' ||
-      payload.message === null ||
-      !('role' in payload.message) ||
-      payload.message.role !== 'user'
-    )
+    return false
   })
   if (rejected !== undefined) {
     throw new A2Error(
@@ -119,7 +117,7 @@ export type AgentGenerateContext<
   messages: M[]
   modelMessages: ModelMessage[]
   state: AIState<M>
-  history: ContractEvent<D>[]
+  session: Pick<HandlerContext<D>['session'], 'state'>
   signal: AbortSignal
   model: LanguageModel
   tools: T
@@ -152,7 +150,7 @@ export type AgentResolverContext<
 > = {
   event: ContractEvent<D, 'ai.generation.requested'>
   state: AIState<M>
-  history: ContractEvent<D>[]
+  session: Pick<HandlerContext<D>['session'], 'state'>
   signal: AbortSignal
 }
 
@@ -214,9 +212,7 @@ export type CreateHandlersOptions<
    * The custom stream owns its metadata chunks.
    */
   generate?: AgentGenerate<M, D, T>
-  compaction?:
-    | CompactionOptions<M, D>
-    | ((context: AgentResolverContext<M, D>) => CompactionOptions<M, D>)
+  compaction?: Resolvable<CompactionOptions<M, D>, AgentResolverContext<M, D>>
   /** Progress is durably flushed at either limit, whichever is reached first. */
   progress?: { maxChunks?: number; maxDelayMs?: number }
 }
@@ -341,10 +337,38 @@ async function generateWithAISDK<
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
   T extends ToolSet,
->(
-  context: AgentGenerateContext<M, D, T>,
-  messageMetadata: AgentMessageMetadata<M, D, T> | undefined,
-): Promise<GenerationSource> {
+>(options: {
+  agentName: string
+  context: AgentGenerateContext<M, D, T>
+  messageMetadata: AgentMessageMetadata<M, D, T> | undefined
+}): Promise<GenerationSource> {
+  const { agentName, context, messageMetadata } = options
+  const headers = context.generation.headers ?? {}
+  const generation =
+    gatewayModelId(context.model) === undefined
+      ? context.generation
+      : {
+          ...context.generation,
+          headers: Object.keys(headers).some(
+            (key) => key.toLowerCase() === 'x-session-affinity',
+          )
+            ? headers
+            : {
+                ...headers,
+                'x-session-affinity': createHash('sha256')
+                  .update(
+                    JSON.stringify([agentName, context.request.sessionId]),
+                  )
+                  .digest('hex'),
+              },
+          providerOptions: {
+            ...context.generation.providerOptions,
+            gateway: {
+              caching: 'auto',
+              ...context.generation.providerOptions?.['gateway'],
+            },
+          },
+        }
   const step = await generateAISDKStep<M, T>({
     model: context.model,
     tools: context.tools,
@@ -355,11 +379,12 @@ async function generateWithAISDK<
     ...(context.instructions === undefined
       ? {}
       : { instructions: context.instructions }),
-    settings: context.generation,
+    settings: generation,
     ...(messageMetadata === undefined
       ? {}
       : {
-          messageMetadata: ({ part }) => messageMetadata({ ...context, part }),
+          messageMetadata: ({ part }) =>
+            messageMetadata({ ...context, generation, part }),
         }),
   })
   return {
@@ -443,12 +468,17 @@ async function* consumeGeneration(options: {
   }
 }
 
-const replay = <M extends UIMessage, D extends AIEventDefs<M> & EventDefs>(
-  definition: AgentDefinition<M, D>,
-  history: ContractEvent<D>[],
-): AIState<M> => {
-  let state = definition.reducer.initialState
-  for (const event of history) state = definition.reducer.fold(state, event)
+const foldAIEvents = <
+  M extends UIMessage,
+  D extends AIEventDefs<M> & EventDefs,
+>(options: {
+  agent: AgentDefinition<M, D>
+  state: AIState<M>
+  events: ContractEvent<D>[]
+}): AIState<M> => {
+  let state = options.state
+  for (const event of options.events)
+    state = options.agent.reducer.fold(state, event)
   return state
 }
 
@@ -457,6 +487,7 @@ const summarize = async <
   D extends AIEventDefs<M> & EventDefs,
   T extends ToolSet,
 >(options: {
+  agentName: string
   context: AgentGenerateContext<M, D, T>
   instructions: string | undefined
 }): Promise<{ summary: string; usage?: LanguageModelUsage }> => {
@@ -498,8 +529,9 @@ const summarize = async <
       return [name, definition]
     }),
   ) as T
-  const source = await generateWithAISDK(
-    {
+  const source = await generateWithAISDK({
+    agentName: options.agentName,
+    context: {
       ...options.context,
       generation,
       tools,
@@ -509,8 +541,8 @@ const summarize = async <
       ],
       responseMessageId: `${options.context.generationId}:summary`,
     },
-    undefined,
-  )
+    messageMetadata: undefined,
+  })
   let summary = ''
   let calledTool = false
   let finish: AgentGenerationFinish | undefined
@@ -551,71 +583,83 @@ const activeCompaction = <M extends UIMessage>(
     : null
 }
 
-const contextMessages = <
+const contextMessages = async <
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
 >(options: {
   agent: AgentDefinition<M, D>
-  state: AIState<M>
-  history: ContractEvent<D>[]
-}): M[] => {
-  const { agent, state, history } = options
+  session: HandlerContext<D>['session']
+  snapshot: { state: AIState<M>; index: number }
+  coordinator: AICoordinatorState
+  appended?: ContractEvent<D>[]
+  excluded?: ReadonlySet<string>
+}): Promise<M[]> => {
+  const { agent, session, snapshot } = options
+  const appended = options.appended ?? []
+  const state = foldAIEvents({ agent, events: appended, state: snapshot.state })
   const compaction = activeCompaction(state)
-  if (compaction === null) return state.messages
+  const queued = new Set(
+    options.coordinator.queued.map((item) => item.messageId),
+  )
+  const visible = (messages: M[]): M[] =>
+    messages.filter((message) => !queued.has(message.id))
+  if (compaction === null) return visible(state.messages)
   const retained = new Set(compaction.retainedMessageIds ?? [])
-  const startedIndex = history.find(
-    (event) =>
-      event.type === 'ai.generation.started' &&
-      (event.payload as GenerationStartedPayload).generationId ===
-        compaction.generationId,
-  )?.index
-  const throughIndex =
-    compaction.throughIndex ??
-    (startedIndex === undefined ? undefined : startedIndex - 1)
+  const throughIndex = compaction.throughIndex
   if (throughIndex !== undefined) {
-    const prefix = history.filter((event) => event.index <= throughIndex)
-    for (const queued of queuedMessagesAt(prefix))
-      retained.add(queued.messageId)
-    let context = replay(agent, prefix)
-    context = {
-      ...context,
-      messages: [
-        ...compaction.messages!,
-        ...context.messages.filter((message) => retained.has(message.id)),
-      ],
-      activeProjection: null,
-    }
-    for (const event of history) {
-      if (event.index > throughIndex)
-        context = agent.reducer.fold(context, event)
+    const [prefix, priorCoordinator, tail] = await Promise.all([
+      throughIndex === snapshot.index
+        ? Promise.resolve(snapshot)
+        : session.state(agent.reducer, { through: throughIndex }),
+      session.state(aiCoordinatorReducer(agent.contract), {
+        through: throughIndex,
+      }),
+      throughIndex >= snapshot.index
+        ? Promise.resolve([])
+        : session.history({ gte: throughIndex + 1, lte: snapshot.index }),
+    ])
+    for (const item of priorCoordinator.state.queued)
+      retained.add(item.messageId)
+    const events =
+      options.excluded === undefined
+        ? tail
+        : withoutGenerationLifecycle(tail, options.excluded)
+    const context = foldAIEvents({
+      agent,
+      events: [...events, ...appended],
+      state: {
+        ...prefix.state,
+        messages: [
+          ...compaction.messages!,
+          ...prefix.state.messages.filter((message) =>
+            retained.has(message.id),
+          ),
+        ],
+        activeProjection: null,
+      },
+    })
+    const positions = new Map<string, number>()
+    for (const message of [...compaction.messages!, ...state.messages]) {
+      if (!positions.has(message.id)) positions.set(message.id, positions.size)
     }
-    return context.messages
+    return visible(
+      context.messages.toSorted(
+        (left, right) =>
+          (positions.get(left.id) ?? positions.size) -
+          (positions.get(right.id) ?? positions.size),
+      ),
+    )
   }
   const boundary = state.messages.findIndex(
     (message) => message.id === compaction.throughMessageId,
   )
-  return [
+  return visible([
     ...compaction.messages!,
     ...state.messages
       .slice(0, boundary + 1)
       .filter((message) => retained.has(message.id)),
     ...state.messages.slice(boundary + 1),
-  ]
-}
-
-const activeContextMessages = <
-  M extends UIMessage,
-  D extends AIEventDefs<M> & EventDefs,
->(options: {
-  agent: AgentDefinition<M, D>
-  state: AIState<M>
-  history: ContractEvent<D>[]
-  coordinator: AICoordinatorState
-}): M[] => {
-  const queued = new Set(
-    options.coordinator.queued.map((item) => item.messageId),
-  )
-  return contextMessages(options).filter((message) => !queued.has(message.id))
+  ])
 }
 
 const modelContext = async <M extends UIMessage>(options: {
@@ -655,51 +699,20 @@ const estimateInputTokens = async (options: {
   )
 }
 
-const measuredInputTokens = <D extends EventDefs>(options: {
+const measuredInputTokens = (options: {
   estimate: number
   model: string
-  history: ContractEvent<D>[]
+  calibration: AICoordinatorState['calibration']
 }): number => {
-  for (const event of options.history.toReversed()) {
-    if (
-      event.type === 'ai.compaction.completed' ||
-      event.type === 'ai.retry.requested'
-    )
-      break
-    if (event.type !== 'ai.generation.completed') continue
-    const completed = event.payload as GenerationCompletedPayload
-    const started = options.history.find(
-      (candidate) =>
-        candidate.type === 'ai.generation.started' &&
-        (candidate.payload as GenerationStartedPayload).generationId ===
-          completed.generationId,
-    )
-    const input = completed.usage?.inputTokens
-    const estimate = completed.inputTokenEstimate
-    if (
-      (started?.payload as GenerationStartedPayload | undefined)?.model !==
-      options.model
-    )
-      break
-    if (
-      input === undefined ||
-      estimate === undefined ||
-      !Number.isFinite(input) ||
-      input < 0 ||
-      options.estimate < estimate
-    )
-      break
-    return Math.max(
-      options.estimate,
-      Math.ceil(input + options.estimate - estimate),
-    )
-  }
-  return options.estimate
-}
-
-const generationRequestId = (generationId: string): string | undefined => {
-  const markerIndex = generationId.lastIndexOf(':generation:')
-  return markerIndex === -1 ? undefined : generationId.slice(0, markerIndex)
+  const previous = options.calibration
+  return previous !== undefined &&
+    previous.model === options.model &&
+    options.estimate >= previous.estimate
+    ? Math.max(
+        options.estimate,
+        Math.ceil(previous.inputTokens + options.estimate - previous.estimate),
+      )
+    : options.estimate
 }
 
 type ToolCallClassification = 'automatic' | 'approval' | 'unknown'
@@ -946,36 +959,6 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
   return { events: result, pending }
 }
 
-function queuedMessagesAt<D extends EventDefs>(
-  history: ContractEvent<D>[],
-): AICoordinatorState['queued'] {
-  const queued: AICoordinatorState['queued'] = []
-  for (const event of history) {
-    if (event.type === 'ai.message.created') {
-      const payload = event.payload as MessageCreatedPayload<UIMessage>
-      if (payload.message.role !== 'user') continue
-      const duplicate = queued.findIndex(
-        (item) => item.messageId === payload.message.id,
-      )
-      if (duplicate !== -1) queued.splice(duplicate, 1)
-      queued.push({
-        index: event.index,
-        messageId: payload.message.id,
-        generate: payload.generate !== false,
-      })
-      continue
-    }
-    if (event.type !== 'ai.generation.requested') continue
-    const request = event.payload as GenerationRequestedPayload
-    if (request.reason !== 'message') continue
-    const requested = queued.findIndex(
-      (item) => item.messageId === request.messageId,
-    )
-    if (requested !== -1) queued.splice(0, requested + 1)
-  }
-  return queued
-}
-
 function withoutGenerationLifecycle<D extends EventDefs>(
   history: ContractEvent<D>[],
   generationIds: ReadonlySet<string>,
@@ -1062,7 +1045,7 @@ const validateCompaction = <
   }
   if (compaction !== false && 'then' in compaction) {
     void Promise.resolve(compaction).catch(() => {})
-    throw new TypeError('compaction options must resolve synchronously')
+    throw new TypeError('compaction options must be a policy or a resolver')
   }
   if (compaction !== false && !('shouldCompact' in compaction)) {
     if (
@@ -1150,89 +1133,10 @@ export function createHandlers<
   const tools = options.tools ?? ({} as T)
   const generation: AgentGenerationSettings<T> = options.generation ?? {}
   const maxSteps = options.maxSteps ?? Number.POSITIVE_INFINITY
-  const coordinator = aiCoordinatorReducer(options.agent.contract)
+  const control = createControlRuntime({ agent: options.agent })
+  const coordinator = control.coordinator
   const promptCache = new Map<string, Promise<ModelMessage[]>>()
 
-  const coordinatorStateAt = (
-    history: ContractEvent<D>[],
-    frontier = Number.POSITIVE_INFINITY,
-  ): AICoordinatorState => {
-    let state = coordinator.initialState
-    for (const event of history) {
-      if (event.index >= frontier) break
-      state = coordinator.fold(state, event)
-    }
-    return state
-  }
-
-  const requestForNextMessage = (
-    state: AICoordinatorState,
-  ): AppendInput<D> | undefined => {
-    if (state.closed || state.response !== undefined) return undefined
-    const next = state.queued.find((item) => item.generate !== false)
-    if (next === undefined) return undefined
-    return {
-      type: 'ai.generation.requested',
-      id: `ai.generate:message:${next.messageId}`,
-      payload: { messageId: next.messageId, reason: 'message' },
-    } as AppendInput<D>
-  }
-
-  const scheduleNext = async (
-    ctx: Pick<HandlerContext<D>, 'session'>,
-  ): Promise<AppendInput<D> | void> =>
-    requestForNextMessage(
-      (await ctx.session.state(coordinator, { through: 'latest' })).state,
-    )
-
-  const continueIfReady = async (
-    ctx: Pick<HandlerContext<D>, 'session'>,
-    generationId: string,
-  ): Promise<void> => {
-    const state = (await ctx.session.state(coordinator, { through: 'latest' }))
-      .state
-    const response = state.response
-    if (response?.generation?.generationId !== generationId) {
-      return
-    }
-    const input = response.inputResponse
-    const inputReady =
-      response.completion?.generationId === generationId &&
-      response.failure === undefined &&
-      input !== undefined &&
-      response.inputs.length === 0 &&
-      response.calls.every(
-        (call) =>
-          call.terminal ||
-          (call.call.providerExecuted === true &&
-            call.call.supportsDeferredResults !== true &&
-            call.approval !== undefined &&
-            call.response !== undefined),
-      )
-    if (inputReady) {
-      await ctx.session.append('continue-after-input', {
-        type: 'ai.generation.requested',
-        id: `ai.generate:input:${encodeURIComponent(response.responseMessageId)}:${encodeURIComponent(input.generationId)}:${encodeURIComponent(input.inputId)}`,
-        payload: {
-          messageId: response.responseMessageId,
-          responseMessageId: response.responseMessageId,
-          reason: 'input',
-        },
-      } as AppendInput<D>)
-      return
-    }
-    if (!continuationReady(state)) return
-    await ctx.session.append('continue-after-tools', {
-      type: 'ai.generation.requested',
-      id: `ai.generate:tools:${generationId}`,
-      payload: {
-        messageId: response.responseMessageId,
-        responseMessageId: response.responseMessageId,
-        reason: 'tool',
-      },
-    } as AppendInput<D>)
-  }
-
   type RuntimeTool = {
     execute?: (
       input: unknown,
@@ -1244,9 +1148,7 @@ export function createHandlers<
     ) => unknown
   }
 
-  type ToolHandlerContext =
-    | HandlerContext<D, 'ai.tool.called'>
-    | HandlerContext<D, 'ai.approval.responded'>
+  type ToolHandlerContext = HandlerContext<D, 'ai.tool.execution.requested'>
 
   const resultEvent = (
     call: ToolCalledPayload,
@@ -1285,43 +1187,44 @@ export function createHandlers<
       },
     }) as AppendInput<D>
 
-  const promptMessages = (
-    sessionId: string,
-    call: ToolCalledPayload,
-    readHistory: () => Promise<ContractEvent<D>[]>,
-  ): Promise<ModelMessage[]> => {
-    const key = promptCacheKey(sessionId, call.generationId)
+  const promptMessages = (input: {
+    ctx: ToolHandlerContext
+    call: ToolCalledPayload
+    generation: GenerationStartedPayload
+  }): Promise<ModelMessage[]> => {
+    const { ctx, call, generation: owner } = input
+    const key = promptCacheKey(ctx.event.sessionId, call.generationId)
     const cached = promptCache.get(key)
     if (cached) return cached
     const computation = (async () => {
-      const history = await readHistory()
-      const started = history.find(
-        (event) =>
-          event.type === 'ai.generation.started' &&
-          (event.payload as GenerationStartedPayload).generationId ===
-            call.generationId,
-      )
-      const startedPayload = started?.payload as
-        GenerationStartedPayload | undefined
-      const frontier =
-        startedPayload?.promptThroughIndex ??
-        (started?.index ?? Number.POSITIVE_INFINITY) - 1
-      const promptHistory = history.filter(
+      const frontier = owner.promptThroughIndex
+      if (frontier === undefined)
+        throw new TypeError('AI work requires a prompt checkpoint')
+      const [snapshot, atPrompt, tail] = await Promise.all([
+        ctx.session.state(options.agent.reducer, { through: frontier }),
+        ctx.session.state(coordinator, { through: frontier }),
+        ctx.session.history({ gte: frontier + 1, lte: ctx.event.index }),
+      ])
+      const appended = tail.filter(
         (event) =>
-          event.index <= frontier ||
-          ((event.type === 'ai.generation.started' ||
+          (event.type === 'ai.generation.started' ||
             event.type === 'ai.compaction.requested' ||
             event.type === 'ai.compaction.completed') &&
-            (event.payload as { generationId: string }).generationId ===
-              call.generationId),
+          (event.payload as { generationId: string }).generationId ===
+            call.generationId,
       )
-      const state = replay(options.agent, promptHistory)
+      const state = foldAIEvents({
+        agent: options.agent,
+        events: appended,
+        state: snapshot.state,
+      })
       return modelContext({
-        messages: activeContextMessages({
+        messages: await contextMessages({
           agent: options.agent,
-          state,
-          history: promptHistory,
-          coordinator: coordinatorStateAt(promptHistory),
+          session: ctx.session,
+          snapshot,
+          coordinator: atPrompt.state,
+          appended,
         }),
         state,
         tools,
@@ -1342,7 +1245,7 @@ export function createHandlers<
     const schedulerFailure = consumeSchedulerSendFailure(error)
     if (checkAbort(ctx.signal)) return
     if (schedulerFailure === 'retryable') throw error
-    return resultEvent(call, 'execution:error', {
+    return resultEvent(call, `${ctx.event.id}:execution:error`, {
       error: errorMessage(error),
     })
   }
@@ -1370,7 +1273,7 @@ export function createHandlers<
       return toolExecutionFailure(ctx, call, error)
     }
     if (iterator === undefined) {
-      return resultEvent(call, 'execution:0', { output })
+      return resultEvent(call, `${ctx.event.id}:execution:0`, { output })
     }
     let last: unknown
     let sequence = 0
@@ -1383,23 +1286,27 @@ export function createHandlers<
         } catch (error) {
           return toolExecutionFailure(ctx, call, error)
         }
-        if (checkAbort(ctx.signal)) return
         if (result.done) {
+          if (ctx.signal.reason instanceof A2Error) throw ctx.signal.reason
           done = true
           break
         }
+        if (checkAbort(ctx.signal)) return
         last = result.value
-        await ctx.session.append(
-          `tool:${call.toolCallId}:preliminary:${sequence}`,
-          resultEvent(
-            call,
-            `execution:${ctx.attempt}:${sequence}:preliminary`,
-            {
-              output: result.value,
-              preliminary: true,
-            },
-          ),
-        )
+        await control.append({
+          ctx,
+          name: `tool:${call.toolCallId}:preliminary:${sequence}`,
+          events: [
+            resultEvent(
+              call,
+              `${ctx.event.id}:execution:${ctx.attempt}:${sequence}:preliminary`,
+              {
+                output: result.value,
+                preliminary: true,
+              },
+            ),
+          ],
+        })
         sequence += 1
       }
     } finally {
@@ -1413,27 +1320,25 @@ export function createHandlers<
     }
     return resultEvent(
       call,
-      `execution:${sequence}:final`,
+      `${ctx.event.id}:execution:${sequence}:final`,
       sequence === 0 ? {} : { output: last },
     )
   }
 
-  const executeTool = async (
-    ctx: ToolHandlerContext,
-    call: ToolCalledPayload,
-  ): Promise<AppendInput<D> | void> => {
+  const executeTool = async (input: {
+    ctx: ToolHandlerContext
+    call: ToolCalledPayload
+    generation: GenerationStartedPayload
+  }): Promise<AppendInput<D> | void> => {
+    const { ctx, call } = input
     const tool = tools[call.toolName] as RuntimeTool | undefined
     const execute = tool?.execute
     if (execute === undefined) {
-      return resultEvent(call, 'execution:error', {
+      return resultEvent(call, `${ctx.event.id}:execution:error`, {
         error: `Tool '${call.toolName}' has no server executor`,
       })
     }
-    const messages = await promptMessages(
-      ctx.event.sessionId,
-      call,
-      ctx.session.history,
-    )
+    const messages = await promptMessages(input)
     if (checkAbort(ctx.signal)) return
     const scope: AmbientToolScope = {
       contract: options.agent.contract,
@@ -1446,100 +1351,77 @@ export function createHandlers<
     )
   }
 
-  const handleToolCall = async (
-    ctx: HandlerContext<D, 'ai.tool.called'>,
-  ): Promise<AppendInput<D> | void> => {
-    const state = (await ctx.session.state(coordinator, { through: 'latest' }))
-      .state
-    const response = state.response
-    const current = response?.calls.find(
-      (candidate) => candidate.index === ctx.event.index,
-    )
-    if (
-      current === undefined ||
-      response?.generation?.generationId !== ctx.event.payload.generationId ||
-      response.failure !== undefined ||
-      current.terminal ||
-      current.approval !== undefined
-    ) {
-      return
-    }
-    if (current.call.providerExecuted === true) {
-      await continueIfReady(ctx, current.call.generationId)
-      return
-    }
-    return executeTool(ctx, current.call)
-  }
-
-  const handleApproval = async (
-    ctx: HandlerContext<D, 'ai.approval.responded'>,
-  ): Promise<AppendInput<D> | void> => {
-    const state = (await ctx.session.state(coordinator, { through: 'latest' }))
-      .state
-    const response = state.response
-    const current = response?.calls.find(
-      (candidate) =>
-        candidate.approval?.approvalId === ctx.event.payload.approvalId &&
-        candidate.approval.messageId === ctx.event.payload.messageId &&
-        candidate.approval.generationId === ctx.event.payload.generationId &&
-        candidate.responseIndex === ctx.event.index,
-    )
-    if (
-      current === undefined ||
-      response?.generation?.generationId !== current.call.generationId ||
-      response.failure !== undefined ||
-      current.terminal
-    ) {
-      return
-    }
-    if (current.call.providerExecuted === true) {
-      await continueIfReady(ctx, current.call.generationId)
-      return
-    }
-    if (!ctx.event.payload.approved) {
-      return resultEvent(current.call, 'execution:denied', { denied: true })
+  const handleToolExecution = async (
+    ctx: ToolHandlerContext,
+  ): Promise<AppendInput<D>> => {
+    try {
+      const state = (
+        await ctx.session.state(control.reducer, { through: 'latest' })
+      ).state
+      const request = ctx.event.payload
+      const call = state.coordinator.response?.calls.find(
+        (candidate) => candidate.work?.id === ctx.event.id,
+      )
+      if (
+        state.active?.turnId !== request.turnId ||
+        state.active.suspended ||
+        state.active.version !== request.version ||
+        state.coordinator.response?.failure !== undefined ||
+        !call ||
+        call.terminal ||
+        call.work?.settled
+      )
+        return control.settled({ ctx, events: undefined })
+      return control.settled({
+        ctx,
+        events: await executeTool({
+          ctx,
+          call: request.call,
+          generation: request.generation,
+        }),
+      })
+    } catch (error) {
+      if (control.cancelled({ error, signal: ctx.signal }))
+        return control.settled({ ctx, events: undefined })
+      throw error
     }
-    return executeTool(ctx, current.call)
   }
 
   const generationHandler = async (
     ctx: HandlerContext<D, 'ai.generation.requested'>,
   ): Promise<void | AppendInput<D> | readonly AppendInput<D>[]> => {
-    const history = await ctx.session.history()
+    if (checkAbort(ctx.signal)) return
     const requestId = ctx.event.id
     const request = ctx.event.payload
+    const snapshot = await ctx.session.state(options.agent.reducer, {
+      through: 'latest',
+    })
+    if (
+      snapshot.state.active?.phase === 'paused' ||
+      snapshot.state.active?.phase === 'pausing'
+    )
+      return
     const coordinatorState = (
-      await ctx.session.state(coordinator, { through: 'latest' })
+      await ctx.session.state(coordinator, { through: snapshot.index })
     ).state
     if (
       coordinatorState.closed ||
       coordinatorState.response?.activeRequestId !== requestId
-    ) {
-      return
-    }
-    const terminal = history.some((event) => {
-      if (event.type === 'ai.generation.completed') {
-        return (
-          (event.payload as GenerationCompletedPayload).requestId === requestId
-        )
-      }
-      if (event.type !== 'ai.generation.failed') return false
-      const payload = event.payload as GenerationFailedPayload
-      return payload.requestId === requestId && payload.superseded !== true
-    })
-    if (terminal) return
-
-    const previousStarts = history.filter(
-      (event) =>
-        event.type === 'ai.generation.started' &&
-        (event.payload as GenerationStartedPayload).requestId === requestId,
     )
-    const priorAttempts = previousStarts.map(
-      (event) => (event.payload as GenerationStartedPayload).attempt,
+      return
+    const response = coordinatorState.response
+    const current =
+      response.generation?.requestId === requestId
+        ? response.generation
+        : undefined
+    if (
+      current !== undefined &&
+      (response.completion !== undefined ||
+        (response.failure !== undefined &&
+          response.failure.superseded !== true))
     )
-    if (priorAttempts.some((priorAttempt) => priorAttempt >= ctx.attempt)) {
       return
-    }
+    if (current !== undefined && current.attempt >= ctx.attempt) return
     const attempt = ctx.attempt
     const generationId = `${requestId}:generation:${attempt}`
     const responseMessageId =
@@ -1547,65 +1429,55 @@ export function createHandlers<
       (request.reason === 'message'
         ? `${request.messageId}:assistant`
         : request.messageId)
-
-    const responseStepCount = history.filter(
-      (event) =>
-        event.type === 'ai.generation.completed' &&
-        (event.payload as GenerationCompletedPayload).responseMessageId ===
-          responseMessageId,
-    ).length
-
-    const sourceCoordinatorState = coordinatorStateAt(history, ctx.event.index)
-    if (request.reason === 'tool') {
-      const prefix = 'ai.generate:tools:'
-      if (!requestId.startsWith(prefix)) return
-      const sourceGenerationId = requestId.slice(prefix.length)
-      if (
-        !continuationReady(sourceCoordinatorState) ||
-        sourceCoordinatorState.response?.generation?.generationId !==
-          sourceGenerationId ||
-        sourceCoordinatorState.response.responseMessageId !==
-          request.messageId ||
-        sourceCoordinatorState.response.responseMessageId !==
-          request.responseMessageId
-      ) {
-        return
-      }
-    }
-
-    const previous = previousStarts
-      .filter(
-        (event) =>
-          (event.payload as GenerationStartedPayload).attempt < attempt,
-      )
-      .toSorted(
-        (left, right) =>
-          (right.payload as GenerationStartedPayload).attempt -
-          (left.payload as GenerationStartedPayload).attempt,
-      )[0]
-
-    const incompleteId = previous
-      ? (previous.payload as GenerationStartedPayload).generationId
-      : undefined
-    const replacedGenerationIds = new Set<string>()
-    if (incompleteId !== undefined) replacedGenerationIds.add(incompleteId)
+    const responseStepCount = response.stepCount
     if (
-      request.reason === 'retry' &&
-      sourceCoordinatorState.response?.failure?.generationId !== undefined
-    ) {
-      replacedGenerationIds.add(
-        sourceCoordinatorState.response.failure.generationId,
-      )
+      request.reason === 'tool' &&
+      (requestId !==
+        `ai.generate:tools:${response.source?.generation.generationId}` ||
+        response.responseMessageId !== request.messageId ||
+        response.responseMessageId !== request.responseMessageId)
+    )
+      return
+    const previous = current
+    const replaced =
+      previous === undefined
+        ? []
+        : [{ generation: previous, frontier: response.promptThroughIndex! }]
+    if (request.reason === 'retry' && response.source?.failed) {
+      replaced.push({
+        generation: response.source.generation,
+        frontier: response.source.promptThroughIndex,
+      })
     }
-    const promptHistory = withoutGenerationLifecycle(
-      history,
-      replacedGenerationIds,
+    const replacedGenerationIds = new Set(
+      replaced.map(({ generation: owner }) => owner.generationId),
     )
-    const state = replay(options.agent, promptHistory)
+    let state = snapshot.state
+    if (replaced.length > 0) {
+      const through = Math.min(...replaced.map(({ frontier }) => frontier))
+      const baseline = await ctx.session.state(options.agent.reducer, {
+        through,
+      })
+      const tail = await ctx.session.history({
+        gte: through + 1,
+        lte: snapshot.index,
+      })
+      state = foldAIEvents({
+        agent: options.agent,
+        events: withoutGenerationLifecycle(tail, replacedGenerationIds),
+        state: baseline.state,
+      })
+    }
     const resolverContext: AgentResolverContext<M, D> = {
       event: ctx.event,
       state,
-      history,
+      session: {
+        state: (reducer, readOptions) =>
+          ctx.session.state(reducer, {
+            ...readOptions,
+            through: readOptions?.through ?? snapshot.index,
+          }),
+      },
       signal: ctx.signal,
     }
 
@@ -1624,7 +1496,7 @@ export function createHandlers<
       const compaction =
         typeof configuredCompaction === 'function'
           ? validateCompaction({
-              compaction: configuredCompaction(resolverContext),
+              compaction: await resolve(configuredCompaction, resolverContext),
               explicit: true,
               generation: options.generation,
               tools: options.tools,
@@ -1633,7 +1505,7 @@ export function createHandlers<
       if (checkAbort(ctx.signal)) return
 
       if (request.reason === 'tool' && responseStepCount >= maxSteps) {
-        const source = sourceCoordinatorState.response?.generation
+        const source = response.source?.generation
         if (source === undefined) return
         return {
           type: 'ai.generation.failed',
@@ -1656,7 +1528,7 @@ export function createHandlers<
         responseMessageId,
         attempt,
         model: modelName(resolvedModel),
-        promptThroughIndex: history.at(-1)?.index ?? ctx.event.index,
+        promptThroughIndex: snapshot.index,
       }
       const catalogModelId = gatewayModelId(resolvedModel)
       const gatewayOptions = options.generation?.providerOptions?.['gateway']
@@ -1681,7 +1553,7 @@ export function createHandlers<
         } as AppendInput<D>)
       }
       if (previous) {
-        const payload = previous.payload as GenerationStartedPayload
+        const payload = previous
         const superseded: GenerationFailedPayload = {
           requestId,
           messageId: payload.messageId,
@@ -1701,33 +1573,24 @@ export function createHandlers<
         id: generationId,
         payload: started,
       } as AppendInput<D>)
-      const generationHistory = [
-        ...history,
-        ...(await ctx.session.append('generation-start', ...startEvents)),
-      ]
+      const generationEvents = await control.append({
+        ctx,
+        name: 'generation-start',
+        events: startEvents,
+      })
       generationStarted = true
 
-      const startedCoordinatorState = (
-        await ctx.session.state(coordinator, { through: 'latest' })
-      ).state
-      if (
-        startedCoordinatorState.response?.activeRequestId !== requestId ||
-        startedCoordinatorState.response.generation?.generationId !==
-          generationId
-      ) {
-        return
-      }
+      if (checkAbort(ctx.signal)) return
 
-      const promptCoordinatorState = {
-        ...coordinatorStateAt(history),
-        queued: queuedMessagesAt(history),
-      }
-      const messages = activeContextMessages({
+      const promptCoordinatorState = coordinatorState
+      const messages = await contextMessages({
         agent: options.agent,
-        state,
-        history: promptHistory,
+        session: ctx.session,
+        snapshot: { state, index: snapshot.index },
         coordinator: promptCoordinatorState,
+        excluded: replacedGenerationIds,
       })
+      let generationMessages = messages
       let modelMessages = await modelContext({ messages, state, tools })
       const baseContext: AgentGenerateContext<M, D, T> = {
         request: ctx.event,
@@ -1737,7 +1600,7 @@ export function createHandlers<
         messages,
         modelMessages,
         state,
-        history,
+        session: resolverContext.session,
         signal: ctx.signal,
         model: resolvedModel,
         tools,
@@ -1761,9 +1624,7 @@ export function createHandlers<
           : undefined
       let inputTokenEstimate: number | undefined
       let compacted = false
-      const canCompact =
-        sourceCoordinatorState.response?.calls.every((call) => call.terminal) ??
-        true
+      const canCompact = response.source?.canCompact ?? true
       if (policy && canCompact) {
         const compactionContext = {
           ...resolverContext,
@@ -1782,7 +1643,7 @@ export function createHandlers<
           const inputTokens = measuredInputTokens({
             estimate: inputTokenEstimate,
             model: modelName(resolvedModel),
-            history,
+            calibration: coordinatorState.calibration,
           })
           const threshold =
             policy.thresholdTokens ??
@@ -1816,18 +1677,25 @@ export function createHandlers<
                   (queued) => queued.messageId === message.id,
                 ),
             )?.id ?? request.messageId
-          const throughIndex = history.at(-1)?.index ?? ctx.event.index
-          generationHistory.push(
-            ...(await ctx.session.append('compaction-requested', {
-              type: 'ai.compaction.requested',
-              id: `${generationId}:compaction:requested`,
-              payload: { generationId, throughMessageId, throughIndex },
+          const throughIndex = snapshot.index
+          generationEvents.push(
+            ...(await control.append({
+              ctx,
+              name: 'compaction-requested',
+              events: [
+                {
+                  type: 'ai.compaction.requested',
+                  id: `${generationId}:compaction:requested`,
+                  payload: { generationId, throughMessageId, throughIndex },
+                } as AppendInput<D>,
+              ],
             })),
           )
           const result = !('shouldCompact' in policy)
             ? {
                 messages: [] as M[],
                 ...(await summarize({
+                  agentName: options.agent.contract.name,
                   context: baseContext,
                   instructions: policy.instructions,
                 })),
@@ -1847,30 +1715,55 @@ export function createHandlers<
                   ),
                 }),
           }
-          generationHistory.push(
-            ...(await ctx.session.append('compaction-completed', {
-              type: 'ai.compaction.completed',
-              id: `${generationId}:compaction:completed`,
-              payload: completed,
+          generationEvents.push(
+            ...(await control.append({
+              ctx,
+              name: 'compaction-completed',
+              events: [
+                {
+                  type: 'ai.compaction.completed',
+                  id: `${generationId}:compaction:completed`,
+                  payload: completed,
+                } as AppendInput<D>,
+              ],
             })),
           )
+          const queued = new Set(
+            promptCoordinatorState.queued.map((item) => item.messageId),
+          )
+          generationMessages = result.messages.filter(
+            (message) => !queued.has(message.id),
+          )
           compacted = true
         }
       }
 
-      const currentState = replay(options.agent, generationHistory)
-      const currentCoordinatorState = coordinatorStateAt(generationHistory)
-      const generationMessages = activeContextMessages({
+      const currentState = foldAIEvents({
         agent: options.agent,
-        state: currentState,
-        history: generationHistory,
-        coordinator: currentCoordinatorState,
-      })
-      modelMessages = await modelContext({
-        messages: generationMessages,
-        state: currentState,
-        tools,
+        events: generationEvents,
+        state: snapshot.state,
       })
+      if (request.reason === 'retry' || previous !== undefined) {
+        generationMessages = await contextMessages({
+          agent: options.agent,
+          session: ctx.session,
+          snapshot,
+          coordinator: promptCoordinatorState,
+          appended: generationEvents,
+        })
+      }
+      if (
+        (policy && 'shouldCompact' in policy) ||
+        generationMessages !== messages ||
+        activeCompaction(currentState)?.summary !==
+          activeCompaction(state)?.summary
+      ) {
+        modelMessages = await modelContext({
+          messages: generationMessages,
+          state: currentState,
+          tools,
+        })
+      }
       promptCache.set(
         promptCacheKey(ctx.event.sessionId, generationId),
         Promise.resolve(modelMessages),
@@ -1893,7 +1786,7 @@ export function createHandlers<
         messages: generationMessages,
         modelMessages,
         state: currentState,
-        history,
+        session: resolverContext.session,
         signal: ctx.signal,
         model: resolvedModel,
         tools,
@@ -1905,7 +1798,11 @@ export function createHandlers<
       const custom = options.generate !== undefined
       const source =
         options.generate === undefined
-          ? await generateWithAISDK(generateContext, options.messageMetadata)
+          ? await generateWithAISDK({
+              agentName: options.agent.contract.name,
+              context: generateContext,
+              messageMetadata: options.messageMetadata,
+            })
           : { stream: await options.generate(generateContext) }
       const updates = consumeGeneration({
         source,
@@ -1942,15 +1839,18 @@ export function createHandlers<
           custom,
         })
         pendingToolCalls = lifecycle.pending
-        await ctx.session.append(
-          `generation-progress:${sequence}`,
-          {
-            type: 'ai.generation.progress',
-            id: `${generationId}:progress:${sequence}`,
-            payload: progress,
-          },
-          ...lifecycle.events,
-        )
+        await control.append({
+          ctx,
+          name: `generation-progress:${sequence}`,
+          events: [
+            {
+              type: 'ai.generation.progress',
+              id: `${generationId}:progress:${sequence}`,
+              payload: progress,
+            } as AppendInput<D>,
+            ...lifecycle.events,
+          ],
+        })
         sequence += 1
       }
 
@@ -2004,22 +1904,12 @@ export function createHandlers<
         } as AppendInput<D>,
       ]
     } catch (error) {
+      if (error instanceof ControlCancelled) throw error
       if (error instanceof A2Error) {
         checkAbort(ctx.signal)
         throw error
       }
-      if (checkAbort(ctx.signal)) {
-        const interrupted: MessageInterruptedPayload = {
-          messageId: responseMessageId,
-          generationId,
-          reason: 'aborted',
-        }
-        return {
-          type: 'ai.message.interrupted',
-          id: `${generationId}:interrupted`,
-          payload: interrupted,
-        } as AppendInput<D>
-      }
+      if (checkAbort(ctx.signal)) return
       if (!generationStarted) throw error
       const failed: GenerationFailedPayload = {
         requestId,
@@ -2036,59 +1926,6 @@ export function createHandlers<
     }
   }
 
-  const handleMessageCreated = async (
-    ctx: HandlerContext<D, 'ai.message.created'>,
-  ): Promise<AppendInput<D> | void> => {
-    if (
-      ctx.event.payload.message.role !== 'user' ||
-      ctx.event.payload.generate === false
-    ) {
-      return
-    }
-    return scheduleNext(ctx)
-  }
-
-  const handleRetry = async (
-    ctx: HandlerContext<D, 'ai.retry.requested'>,
-  ): Promise<AppendInput<D> | void> => {
-    const state = (await ctx.session.state(coordinator, { through: 'latest' }))
-      .state
-    const response = state.response
-    if (
-      response?.status !== 'failed' ||
-      response.rootMessageId !== ctx.event.payload.messageId ||
-      response.responseMessageId !== ctx.event.payload.responseMessageId
-    ) {
-      return
-    }
-    return {
-      type: 'ai.generation.requested',
-      id: `ai.generate:retry:${ctx.event.payload.retryId}`,
-      payload: {
-        messageId: response.rootMessageId,
-        responseMessageId: response.responseMessageId,
-        reason: 'retry',
-      },
-    } as AppendInput<D>
-  }
-
-  const handleInputResponse = async (
-    ctx: HandlerContext<D, 'ai.input.responded'>,
-  ): Promise<void> => {
-    const state = (await ctx.session.state(coordinator, { through: 'latest' }))
-      .state
-    const response = state.response
-    if (
-      response?.responseMessageId !== ctx.event.payload.messageId ||
-      response.inputResponse?.index !== ctx.event.index ||
-      response.inputResponse.generationId !== ctx.event.payload.generationId ||
-      response.inputResponse.inputId !== ctx.event.payload.inputId
-    ) {
-      return
-    }
-    await continueIfReady(ctx, ctx.event.payload.generationId)
-  }
-
   const clearPromptCache = (sessionId: string): void => {
     const prefix = `${sessionId}\u001f`
     for (const key of promptCache.keys()) {
@@ -2096,13 +1933,6 @@ export function createHandlers<
     }
   }
 
-  const handleResponseEnded = async (
-    ctx: HandlerContext<D, 'ai.message.completed' | 'ai.message.interrupted'>,
-  ): Promise<AppendInput<D> | void> => {
-    clearPromptCache(ctx.event.sessionId)
-    return scheduleNext(ctx)
-  }
-
   return {
     'ai.model.metadata.requested': {
       handler: async (ctx) => {
@@ -2118,24 +1948,23 @@ export function createHandlers<
         } as AppendInput<D>
       },
     },
-    'ai.message.created': { lane: 'a2.ai.turn', handler: handleMessageCreated },
-    'ai.retry.requested': { lane: 'a2.ai.turn', handler: handleRetry },
-    'ai.input.responded': {
-      lane: 'a2.ai.turn',
-      handler: handleInputResponse,
-    },
+    'ai.control.requested': { lane: 'a2.ai.control', handler: control.handler },
+    'ai.work.reported': { lane: 'a2.ai.control', handler: control.handler },
     'ai.message.completed': {
-      lane: 'a2.ai.turn',
-      handler: handleResponseEnded,
+      handler: async (ctx) => {
+        clearPromptCache(ctx.event.sessionId)
+      },
     },
     'ai.message.interrupted': {
-      lane: 'a2.ai.turn',
-      handler: handleResponseEnded,
+      handler: async (ctx) => {
+        clearPromptCache(ctx.event.sessionId)
+      },
     },
     'ai.session.closed': {
-      handler: (ctx) => {
+      lane: 'a2.ai.control',
+      handler: async (ctx) => {
         clearPromptCache(ctx.event.sessionId)
-        return Promise.resolve()
+        return control.handler(ctx)
       },
     },
     'ai.generation.failed': {
@@ -2147,60 +1976,41 @@ export function createHandlers<
       },
     },
     'ai.generation.requested': {
-      lane: 'a2.ai.turn',
+      lane: 'a2.ai.model',
       abortOn: {
-        'ai.message.interrupted': (event, trigger, context) => {
-          const responseMessageId =
-            trigger.payload.responseMessageId ??
-            (trigger.payload.reason === 'message'
-              ? `${trigger.payload.messageId}:assistant`
-              : trigger.payload.messageId)
+        'ai.control.committed': (event, trigger) => {
+          const active = (event.payload as ControlCommit).view
           return (
-            event.payload.messageId === responseMessageId &&
-            (event.payload.requestId === trigger.id ||
-              event.payload.generationId ===
-                `${trigger.id}:generation:${context.attempt}`)
+            active?.turnId !== trigger.payload.control?.turnId ||
+            active?.version !== trigger.payload.control?.version
           )
         },
         'ai.session.closed': true,
       },
-      handler: generationHandler,
-    },
-    'ai.generation.completed': {
-      handler: async (ctx) =>
-        continueIfReady(ctx, ctx.event.payload.generationId),
-    },
-    'ai.tool.called': {
-      abortOn: {
-        'ai.generation.failed': (event, trigger) =>
-          event.payload.generationId === trigger.payload.generationId,
-        'ai.message.interrupted': (event, trigger) =>
-          event.payload.messageId === trigger.payload.messageId &&
-          (event.payload.requestId === trigger.payload.requestId ||
-            event.payload.generationId === trigger.payload.generationId),
-        'ai.session.closed': true,
+      handler: async (ctx) => {
+        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
+        }
       },
-      handler: handleToolCall,
     },
-    'ai.approval.responded': {
+    'ai.tool.execution.requested': {
       abortOn: {
         'ai.generation.failed': (event, trigger) =>
-          event.payload.responseMessageId === trigger.payload.messageId &&
-          event.payload.generationId === trigger.payload.generationId,
-        'ai.message.interrupted': (event, trigger) =>
-          event.payload.messageId === trigger.payload.messageId &&
-          (event.payload.requestId ===
-            generationRequestId(trigger.payload.generationId) ||
-            event.payload.generationId === trigger.payload.generationId),
+          event.payload.generationId === trigger.payload.call.generationId,
+        'ai.control.committed': (event, trigger) => {
+          const active = (event.payload as ControlCommit).view
+          return (
+            active?.turnId !== trigger.payload.turnId ||
+            active?.version !== trigger.payload.version
+          )
+        },
         'ai.session.closed': true,
       },
-      handler: handleApproval,
-    },
-    'ai.tool.result': {
-      handler: async (ctx) => {
-        if (ctx.event.payload.preliminary === true) return
-        await continueIfReady(ctx, ctx.event.payload.generationId)
-      },
+      handler: handleToolExecution,
     },
   } as NonNullable<ServerOptions<D>['handlers']>
 }