experimental-a2 0.15.0 → 0.16.1

package/src/ai-server.ts

@@ -9,10 +9,20 @@
 // 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 {
+  createContextReducer,
+  contextForAttempt,
+  generationBase,
+  type AIContextState,
+  type AIModelContext,
+} from './ai-context.ts'
+import type { Reducer } from './reducer.ts'
+import { aiStateReducer } from './ai-stored-state.ts'
+import { aiEventId, AI_EVENT_ID_PREFIX } from './ai-id.ts'
+import type { ControlCommit, ControlCommand } from './ai-control.ts'
 import { AsyncLocalStorage } from 'node:async_hooks'
 import { createHash } from 'node:crypto'
-import { asSchema, convertToModelMessages } from 'ai'
+import { convertToModelMessages, isToolUIPart } from 'ai'
 import type {
   FinishReason,
   Instructions,
@@ -26,14 +36,15 @@ import type {
 } from 'ai'
 import { generateAISDKStep, type AISDKStepSettings } from './ai-sdk-step.ts'
 import { gatewayModelId, readModelLimits } from './ai-model-metadata.ts'
+import { estimateInputTokens, measuredInputTokens } from './ai-input-tokens.ts'
 import type {
   AIEventDefs,
-  AIState,
   AgentDefinition,
   AgentToolContext,
   AgentToolEventDefs,
   ApprovalRequestedPayload,
   CompactionCompletedPayload,
+  CompactionRequestedPayload,
   GenerationCompletedPayload,
   GenerationFailedPayload,
   GenerationProgressPayload,
@@ -42,10 +53,7 @@ import type {
   ToolCalledPayload,
   ToolResultPayload,
 } from './ai.ts'
-import {
-  aiCoordinatorReducer,
-  type AICoordinatorState,
-} from './ai-coordinator.ts'
+import type { AICoordinatorState } from './ai-coordinator.ts'
 import type { AppendInput, ContractEvent, EventDefs } from './contract.ts'
 import { A2Error } from './errors.ts'
 import {
@@ -116,7 +124,6 @@ export type AgentGenerateContext<
   responseMessageId: string
   messages: M[]
   modelMessages: ModelMessage[]
-  state: AIState<M>
   session: Pick<HandlerContext<D>['session'], 'state'>
   signal: AbortSignal
   model: LanguageModel
@@ -149,7 +156,7 @@ export type AgentResolverContext<
   D extends AIEventDefs<M> & EventDefs,
 > = {
   event: ContractEvent<D, 'ai.generation.requested'>
-  state: AIState<M>
+  messages: M[]
   session: Pick<HandlerContext<D>['session'], 'state'>
   signal: AbortSignal
 }
@@ -472,13 +479,12 @@ const foldAIEvents = <
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
 >(options: {
-  agent: AgentDefinition<M, D>
-  state: AIState<M>
+  reducer: Reducer<D, AIContextState<M>>
+  state: AIContextState<M>
   events: ContractEvent<D>[]
-}): AIState<M> => {
+}): AIContextState<M> => {
   let state = options.state
-  for (const event of options.events)
-    state = options.agent.reducer.fold(state, event)
+  for (const event of options.events) state = options.reducer.fold(state, event)
   return state
 }
 
@@ -539,7 +545,10 @@ const summarize = async <
         ...options.context.modelMessages,
         { role: 'user', content: instruction },
       ],
-      responseMessageId: `${options.context.generationId}:summary`,
+      responseMessageId: await aiEventId(
+        'a2.ai.summary',
+        options.context.generationId,
+      ),
     },
     messageMetadata: undefined,
   })
@@ -573,98 +582,26 @@ const summarize = async <
 }
 
 const activeCompaction = <M extends UIMessage>(
-  state: AIState<M>,
-): AIState<M>['compaction'] => {
+  state: AIModelContext<M>,
+): AIModelContext<M>['compaction'] => {
   const compaction = state.compaction
-  return compaction?.status === 'completed' &&
-    compaction.messages !== undefined &&
-    state.messages.some((message) => message.id === compaction.throughMessageId)
-    ? compaction
-    : null
+  return compaction?.status === 'completed' ? compaction : null
 }
 
-const contextMessages = async <
-  M extends UIMessage,
-  D extends AIEventDefs<M> & EventDefs,
->(options: {
-  agent: AgentDefinition<M, D>
-  session: HandlerContext<D>['session']
-  snapshot: { state: AIState<M>; index: number }
+const contextMessages = <M extends UIMessage>({
+  state,
+  coordinator,
+}: {
+  state: AIModelContext<M>
   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)
-  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 throughIndex = compaction.throughIndex
-  if (throughIndex !== undefined) {
-    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 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 visible([
-    ...compaction.messages!,
-    ...state.messages
-      .slice(0, boundary + 1)
-      .filter((message) => retained.has(message.id)),
-    ...state.messages.slice(boundary + 1),
-  ])
+}): M[] => {
+  const queued = new Set(coordinator.queued.map((item) => item.messageId))
+  return state.messages.filter((message) => !queued.has(message.id))
 }
 
 const modelContext = async <M extends UIMessage>(options: {
   messages: M[]
-  state: AIState<M>
+  state: AIModelContext<M>
   tools: ToolSet
 }): Promise<ModelMessage[]> => {
   const messages = await convertToModelMessages(options.messages, {
@@ -676,45 +613,6 @@ const modelContext = async <M extends UIMessage>(options: {
     : [{ role: 'user', content: summary }, ...messages]
 }
 
-const estimateInputTokens = async (options: {
-  messages: ModelMessage[]
-  instructions: Instructions | undefined
-  tools: ToolSet
-}): Promise<number> => {
-  const tools = await Promise.all(
-    Object.entries(options.tools).map(async ([name, tool]) => ({
-      name,
-      description: tool.description,
-      inputSchema: await asSchema(tool.inputSchema).jsonSchema,
-    })),
-  )
-  return Math.ceil(
-    new TextEncoder().encode(
-      JSON.stringify({
-        messages: options.messages,
-        instructions: options.instructions,
-        tools,
-      }),
-    ).byteLength / 4,
-  )
-}
-
-const measuredInputTokens = (options: {
-  estimate: number
-  model: string
-  calibration: AICoordinatorState['calibration']
-}): number => {
-  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'
 
 type PendingToolCall = {
@@ -722,12 +620,16 @@ type PendingToolCall = {
   classification: ToolCallClassification
 }
 
-const toolCalledEvent = <D extends EventDefs>(
+const toolCalledEvent = async <D extends EventDefs>(
   payload: ToolCalledPayload,
-): AppendInput<D> =>
+): Promise<AppendInput<D>> =>
   ({
     type: 'ai.tool.called',
-    id: `${payload.generationId}:tool:${payload.toolCallId}:called`,
+    id: await aiEventId(
+      'a2.ai.tool.called',
+      payload.generationId,
+      payload.toolCallId,
+    ),
     payload,
   }) as AppendInput<D>
 
@@ -757,7 +659,10 @@ const approvalClassification = <T extends ToolSet>(
   return tool?.needsApproval === true ? 'approval' : 'automatic'
 }
 
-const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
+const lifecycleEvents = async <
+  D extends EventDefs,
+  T extends ToolSet,
+>(options: {
   requestId: string
   messageId: string
   generationId: string
@@ -767,7 +672,7 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
   tools: T
   generation: AgentGenerationSettings<T>
   custom: boolean
-}): { events: AppendInput<D>[]; pending: PendingToolCall[] } => {
+}): Promise<{ events: AppendInput<D>[]; pending: PendingToolCall[] }> => {
   const result: AppendInput<D>[] = []
   const pending = [...options.pending]
   const take = (toolCallId: string): PendingToolCall | undefined => {
@@ -777,12 +682,18 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
     if (index === -1) return undefined
     return pending.splice(index, 1)[0]
   }
-  const flush = (toolCallId: string): void => {
+  const flush = async (toolCallId: string): Promise<void> => {
     const pendingCall = take(toolCallId)
-    if (pendingCall) result.push(toolCalledEvent<D>(pendingCall.call))
+    if (pendingCall) result.push(await toolCalledEvent<D>(pendingCall.call))
   }
   for (const [chunkIndex, chunk] of options.chunks.entries()) {
-    const resultId = `${options.generationId}:tool-result:${options.sequence}:${chunkIndex}`
+    const resultId = (): Promise<string> =>
+      aiEventId(
+        'a2.ai.tool.stream-result',
+        options.generationId,
+        String(options.sequence),
+        String(chunkIndex),
+      )
     if (chunk.type === 'tool-input-available') {
       const payload: ToolCalledPayload = {
         requestId: options.requestId,
@@ -816,7 +727,7 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
         chunk.toolName,
       )
       if (classification === 'automatic') {
-        result.push(toolCalledEvent<D>(payload))
+        result.push(await toolCalledEvent<D>(payload))
       } else {
         pending.push({ call: payload, classification })
       }
@@ -850,13 +761,13 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
         input: chunk.input,
       }
       take(chunk.toolCallId)
-      result.push(toolCalledEvent<D>(called), {
+      result.push(await toolCalledEvent<D>(called), {
         type: 'ai.tool.result',
-        id: resultId,
+        id: await resultId(),
         payload,
       } as AppendInput<D>)
     } else if (chunk.type === 'tool-output-available') {
-      flush(chunk.toolCallId)
+      await flush(chunk.toolCallId)
       const payload: ToolResultPayload = {
         requestId: options.requestId,
         messageId: options.messageId,
@@ -880,11 +791,11 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
       }
       result.push({
         type: 'ai.tool.result',
-        id: resultId,
+        id: await resultId(),
         payload,
       } as AppendInput<D>)
     } else if (chunk.type === 'tool-output-error') {
-      flush(chunk.toolCallId)
+      await flush(chunk.toolCallId)
       const payload: ToolResultPayload = {
         requestId: options.requestId,
         messageId: options.messageId,
@@ -905,11 +816,11 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
       }
       result.push({
         type: 'ai.tool.result',
-        id: resultId,
+        id: await resultId(),
         payload,
       } as AppendInput<D>)
     } else if (chunk.type === 'tool-output-denied') {
-      flush(chunk.toolCallId)
+      await flush(chunk.toolCallId)
       const payload: ToolResultPayload = {
         requestId: options.requestId,
         messageId: options.messageId,
@@ -920,12 +831,13 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
       }
       result.push({
         type: 'ai.tool.result',
-        id: resultId,
+        id: await resultId(),
         payload,
       } as AppendInput<D>)
     } else if (chunk.type === 'tool-approval-request') {
-      flush(chunk.toolCallId)
+      await flush(chunk.toolCallId)
       const payload: ApprovalRequestedPayload = {
+        requestId: options.requestId,
         messageId: options.messageId,
         generationId: options.generationId,
         approvalId: chunk.approvalId,
@@ -939,14 +851,23 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
       }
       result.push({
         type: 'ai.approval.requested',
-        id: `${options.generationId}:approval:${chunk.approvalId}`,
+        id: await aiEventId(
+          'a2.ai.approval.requested',
+          options.generationId,
+          chunk.approvalId,
+        ),
         payload,
       } as AppendInput<D>)
     } else if (chunk.type === 'tool-approval-response') {
       result.push({
         type: 'ai.approval.responded',
-        id: `${options.generationId}:approval:${chunk.approvalId}:automatic-response`,
+        id: await aiEventId(
+          'a2.ai.approval.automatic-response',
+          options.generationId,
+          chunk.approvalId,
+        ),
         payload: {
+          requestId: options.requestId,
           messageId: options.messageId,
           generationId: options.generationId,
           approvalId: chunk.approvalId,
@@ -959,73 +880,6 @@ const lifecycleEvents = <D extends EventDefs, T extends ToolSet>(options: {
   return { events: result, pending }
 }
 
-function withoutGenerationLifecycle<D extends EventDefs>(
-  history: ContractEvent<D>[],
-  generationIds: ReadonlySet<string>,
-): ContractEvent<D>[] {
-  if (generationIds.size === 0) return history
-  const approvalKeys = new Set<string>()
-  const inputKeys = new Set<string>()
-  for (const event of history) {
-    const payload = event.payload as {
-      generationId?: unknown
-      messageId?: unknown
-      approvalId?: unknown
-      inputId?: unknown
-    }
-    if (
-      typeof payload.generationId !== 'string' ||
-      !generationIds.has(payload.generationId) ||
-      typeof payload.messageId !== 'string'
-    ) {
-      continue
-    }
-    if (
-      event.type === 'ai.approval.requested' &&
-      typeof payload.approvalId === 'string'
-    ) {
-      approvalKeys.add(`${payload.messageId}\u001f${payload.approvalId}`)
-    }
-    if (
-      event.type === 'ai.input.requested' &&
-      typeof payload.inputId === 'string'
-    ) {
-      inputKeys.add(`${payload.messageId}\u001f${payload.inputId}`)
-    }
-  }
-  return history.filter((event) => {
-    const payload = event.payload as {
-      generationId?: unknown
-      messageId?: unknown
-      approvalId?: unknown
-      inputId?: unknown
-    }
-    if (
-      typeof payload.generationId === 'string' &&
-      generationIds.has(payload.generationId)
-    ) {
-      return false
-    }
-    if (
-      event.type === 'ai.approval.responded' &&
-      typeof payload.messageId === 'string' &&
-      typeof payload.approvalId === 'string'
-    ) {
-      return !approvalKeys.has(
-        `${payload.messageId}\u001f${payload.approvalId}`,
-      )
-    }
-    if (
-      event.type === 'ai.input.responded' &&
-      typeof payload.messageId === 'string' &&
-      typeof payload.inputId === 'string'
-    ) {
-      return !inputKeys.has(`${payload.messageId}\u001f${payload.inputId}`)
-    }
-    return true
-  })
-}
-
 const validateCompaction = <
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
@@ -1135,6 +989,14 @@ export function createHandlers<
   const maxSteps = options.maxSteps ?? Number.POSITIVE_INFINITY
   const control = createControlRuntime({ agent: options.agent })
   const coordinator = control.coordinator
+  const contextReducer = createContextReducer({ agent: options.agent })
+  const readContext = async (
+    session: HandlerContext<D>['session'],
+    through: number | 'latest',
+  ) => {
+    const snapshot = await session.state(contextReducer, { through })
+    return { state: contextReducer.view(snapshot.state), index: snapshot.index }
+  }
   const promptCache = new Map<string, Promise<ModelMessage[]>>()
 
   type RuntimeTool = {
@@ -1150,7 +1012,7 @@ export function createHandlers<
 
   type ToolHandlerContext = HandlerContext<D, 'ai.tool.execution.requested'>
 
-  const resultEvent = (
+  const resultEvent = async (
     call: ToolCalledPayload,
     suffix: string,
     result: {
@@ -1159,10 +1021,15 @@ export function createHandlers<
       denied?: boolean
       preliminary?: boolean
     },
-  ): AppendInput<D> =>
+  ): Promise<AppendInput<D>> =>
     ({
       type: 'ai.tool.result',
-      id: `${call.generationId}:tool:${call.toolCallId}:${suffix}`,
+      id: await aiEventId(
+        'a2.ai.tool.result',
+        call.generationId,
+        call.toolCallId,
+        suffix,
+      ),
       payload: {
         requestId: call.requestId,
         messageId: call.messageId,
@@ -1200,32 +1067,20 @@ export function createHandlers<
       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.type === 'ai.generation.started' ||
-            event.type === 'ai.compaction.requested' ||
-            event.type === 'ai.compaction.completed') &&
-          (event.payload as { generationId: string }).generationId ===
-            call.generationId,
+      const snapshot = await readContext(ctx.session, ctx.event.index)
+      if (
+        snapshot.state.generation?.owner.generationId !== owner.generationId ||
+        snapshot.state.generation.owner.requestId !== owner.requestId ||
+        snapshot.state.activeProjection?.generationId !== owner.generationId
       )
-      const state = foldAIEvents({
-        agent: options.agent,
-        events: appended,
-        state: snapshot.state,
-      })
+        throw new TypeError(
+          'Tool prompt checkpoint does not own its generation',
+        )
+      const state = generationBase(snapshot.state)
       return modelContext({
-        messages: await contextMessages({
-          agent: options.agent,
-          session: ctx.session,
-          snapshot,
-          coordinator: atPrompt.state,
-          appended,
-        }),
+        messages: state.messages.filter(
+          (message) => state.messageIndexes[message.id]! <= frontier,
+        ),
         state,
         tools,
       })
@@ -1237,11 +1092,11 @@ export function createHandlers<
     return computation
   }
 
-  const toolExecutionFailure = (
+  const toolExecutionFailure = async (
     ctx: ToolHandlerContext,
     call: ToolCalledPayload,
     error: unknown,
-  ): AppendInput<D> | void => {
+  ): Promise<AppendInput<D> | void> => {
     const schedulerFailure = consumeSchedulerSendFailure(error)
     if (checkAbort(ctx.signal)) return
     if (schedulerFailure === 'retryable') throw error
@@ -1297,7 +1152,7 @@ export function createHandlers<
           ctx,
           name: `tool:${call.toolCallId}:preliminary:${sequence}`,
           events: [
-            resultEvent(
+            await resultEvent(
               call,
               `${ctx.event.id}:execution:${ctx.attempt}:${sequence}:preliminary`,
               {
@@ -1364,6 +1219,7 @@ export function createHandlers<
       )
       if (
         state.active?.turnId !== request.turnId ||
+        Object.values(state.pendingStops).includes(request.turnId) ||
         state.active.suspended ||
         state.active.version !== request.version ||
         state.coordinator.response?.failure !== undefined ||
@@ -1393,17 +1249,23 @@ export function createHandlers<
     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',
-    })
+    const snapshot = await readContext(ctx.session, 'latest')
     if (
-      snapshot.state.active?.phase === 'paused' ||
-      snapshot.state.active?.phase === 'pausing'
+      snapshot.state.closed ||
+      snapshot.state.authority?.requestId !== requestId ||
+      snapshot.state.authority?.phase === 'paused' ||
+      snapshot.state.authority?.phase === 'pausing'
     )
       return
-    const coordinatorState = (
-      await ctx.session.state(coordinator, { through: snapshot.index })
+    const controlState = (
+      await ctx.session.state(control.reducer, { through: snapshot.index })
     ).state
+    if (
+      request.control !== undefined &&
+      Object.values(controlState.pendingStops).includes(request.control.turnId)
+    )
+      return
+    const coordinatorState = controlState.coordinator
     if (
       coordinatorState.closed ||
       coordinatorState.response?.activeRequestId !== requestId
@@ -1423,54 +1285,29 @@ export function createHandlers<
       return
     if (current !== undefined && current.attempt >= ctx.attempt) return
     const attempt = ctx.attempt
-    const generationId = `${requestId}:generation:${attempt}`
-    const responseMessageId =
-      request.responseMessageId ??
-      (request.reason === 'message'
-        ? `${request.messageId}:assistant`
-        : request.messageId)
+    const generationId = await aiEventId(
+      'a2.ai.generation',
+      requestId,
+      String(attempt),
+    )
+    const responseMessageId = request.responseMessageId
     const responseStepCount = response.stepCount
     if (
       request.reason === 'tool' &&
-      (requestId !==
-        `ai.generate:tools:${response.source?.generation.generationId}` ||
+      (request.sourceGenerationId !==
+        response.source?.generation.generationId ||
         response.responseMessageId !== request.messageId ||
         response.responseMessageId !== request.responseMessageId)
     )
       return
     const previous = current
-    const replaced =
+    const state =
       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 replacedGenerationIds = new Set(
-      replaced.map(({ generation: owner }) => owner.generationId),
-    )
-    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,
-      })
-    }
+        ? snapshot.state
+        : contextForAttempt(snapshot.state)
     const resolverContext: AgentResolverContext<M, D> = {
       event: ctx.event,
-      state,
+      messages: contextMessages({ state, coordinator: coordinatorState }),
       session: {
         state: (reducer, readOptions) =>
           ctx.session.state(reducer, {
@@ -1509,7 +1346,7 @@ export function createHandlers<
         if (source === undefined) return
         return {
           type: 'ai.generation.failed',
-          id: `${requestId}:step-limit`,
+          id: await aiEventId('a2.ai.step-limit', requestId),
           payload: {
             requestId: source.requestId,
             messageId: source.messageId,
@@ -1548,7 +1385,12 @@ export function createHandlers<
       ) {
         startEvents.push({
           type: 'ai.model.metadata.requested',
-          id: `ai.model.metadata:${encodeURIComponent(options.agent.contract.name)}:${encodeURIComponent(ctx.event.sessionId)}:${encodeURIComponent(catalogModelId)}`,
+          id: await aiEventId(
+            'a2.ai.model.metadata',
+            options.agent.contract.name,
+            ctx.event.sessionId,
+            catalogModelId,
+          ),
           payload: { modelId: catalogModelId },
         } as AppendInput<D>)
       }
@@ -1564,7 +1406,7 @@ export function createHandlers<
         }
         startEvents.push({
           type: 'ai.generation.failed',
-          id: `${payload.generationId}:superseded`,
+          id: await aiEventId('a2.ai.superseded', payload.generationId),
           payload: superseded,
         } as AppendInput<D>)
       }
@@ -1583,12 +1425,9 @@ export function createHandlers<
       if (checkAbort(ctx.signal)) return
 
       const promptCoordinatorState = coordinatorState
-      const messages = await contextMessages({
-        agent: options.agent,
-        session: ctx.session,
-        snapshot: { state, index: snapshot.index },
+      const messages = contextMessages({
+        state,
         coordinator: promptCoordinatorState,
-        excluded: replacedGenerationIds,
       })
       let generationMessages = messages
       let modelMessages = await modelContext({ messages, state, tools })
@@ -1599,7 +1438,6 @@ export function createHandlers<
         responseMessageId,
         messages,
         modelMessages,
-        state,
         session: resolverContext.session,
         signal: ctx.signal,
         model: resolvedModel,
@@ -1624,8 +1462,24 @@ export function createHandlers<
           : undefined
       let inputTokenEstimate: number | undefined
       let compacted = false
-      const canCompact = response.source?.canCompact ?? true
-      if (policy && canCompact) {
+      let compactionEstimate:
+        | Pick<
+            CompactionRequestedPayload,
+            'inputTokenEstimate' | 'inputTokens' | 'thresholdTokens'
+          >
+        | undefined
+      const hasPendingTools = messages.some((message) =>
+        message.parts.some(
+          (part) =>
+            isToolUIPart(part) &&
+            !(
+              part.state === 'output-error' ||
+              part.state === 'output-denied' ||
+              (part.state === 'output-available' && part.preliminary !== true)
+            ),
+        ),
+      )
+      if (policy && !hasPendingTools) {
         const compactionContext = {
           ...resolverContext,
           messages,
@@ -1665,10 +1519,24 @@ export function createHandlers<
             )
           }
           shouldCompact = threshold !== undefined && inputTokens >= threshold
+          if (threshold !== undefined && shouldCompact) {
+            compactionEstimate = {
+              inputTokenEstimate,
+              inputTokens,
+              thresholdTokens: threshold,
+            }
+          }
         } else {
           shouldCompact = false
         }
         if (checkAbort(ctx.signal)) return
+        if (shouldCompact) {
+          const latest = (
+            await ctx.session.state(coordinator, { through: 'latest' })
+          ).state
+          shouldCompact = Object.keys(latest.unsettledTools ?? {}).length === 0
+        }
+        if (checkAbort(ctx.signal)) return
         if (shouldCompact) {
           const throughMessageId =
             state.messages.findLast(
@@ -1685,8 +1553,16 @@ export function createHandlers<
               events: [
                 {
                   type: 'ai.compaction.requested',
-                  id: `${generationId}:compaction:requested`,
-                  payload: { generationId, throughMessageId, throughIndex },
+                  id: await aiEventId(
+                    'a2.ai.compaction.requested',
+                    generationId,
+                  ),
+                  payload: {
+                    generationId,
+                    throughMessageId,
+                    throughIndex,
+                    ...compactionEstimate,
+                  },
                 } as AppendInput<D>,
               ],
             })),
@@ -1722,7 +1598,10 @@ export function createHandlers<
               events: [
                 {
                   type: 'ai.compaction.completed',
-                  id: `${generationId}:compaction:completed`,
+                  id: await aiEventId(
+                    'a2.ai.compaction.completed',
+                    generationId,
+                  ),
                   payload: completed,
                 } as AppendInput<D>,
               ],
@@ -1739,17 +1618,14 @@ export function createHandlers<
       }
 
       const currentState = foldAIEvents({
-        agent: options.agent,
+        reducer: contextReducer[aiStateReducer],
         events: generationEvents,
-        state: snapshot.state,
+        state,
       })
       if (request.reason === 'retry' || previous !== undefined) {
-        generationMessages = await contextMessages({
-          agent: options.agent,
-          session: ctx.session,
-          snapshot,
+        generationMessages = contextMessages({
+          state: currentState,
           coordinator: promptCoordinatorState,
-          appended: generationEvents,
         })
       }
       if (
@@ -1785,7 +1661,6 @@ export function createHandlers<
         responseMessageId,
         messages: generationMessages,
         modelMessages,
-        state: currentState,
         session: resolverContext.session,
         signal: ctx.signal,
         model: resolvedModel,
@@ -1827,7 +1702,7 @@ export function createHandlers<
           sequence,
           chunks: update.chunks,
         }
-        const lifecycle = lifecycleEvents<D, T>({
+        const lifecycle = await lifecycleEvents<D, T>({
           requestId,
           messageId: responseMessageId,
           generationId,
@@ -1845,7 +1720,11 @@ export function createHandlers<
           events: [
             {
               type: 'ai.generation.progress',
-              id: `${generationId}:progress:${sequence}`,
+              id: await aiEventId(
+                'a2.ai.progress',
+                generationId,
+                String(sequence),
+              ),
               payload: progress,
             } as AppendInput<D>,
             ...lifecycle.events,
@@ -1868,7 +1747,7 @@ export function createHandlers<
       }
       const completionEvent = {
         type: 'ai.generation.completed',
-        id: `${generationId}:completed`,
+        id: await aiEventId('a2.ai.completed', generationId),
         payload: completed,
       } as AppendInput<D>
       const unresolved = options.generate
@@ -1881,14 +1760,16 @@ export function createHandlers<
           `generation ended with unresolved tool authorization for '${unresolved[0]!.call.toolName}'`,
         )
       }
-      const pendingEvents = pendingToolCalls
-        .filter(
-          (pending) =>
-            pending.classification === 'automatic' ||
-            (options.generate === undefined &&
-              pending.classification === 'unknown'),
-        )
-        .map((pending) => toolCalledEvent<D>(pending.call))
+      const pendingEvents = await Promise.all(
+        pendingToolCalls
+          .filter(
+            (pending) =>
+              pending.classification === 'automatic' ||
+              (options.generate === undefined &&
+                pending.classification === 'unknown'),
+          )
+          .map((pending) => toolCalledEvent<D>(pending.call)),
+      )
       if (finish.finishReason === 'tool-calls') {
         return [...pendingEvents, completionEvent]
       }
@@ -1897,7 +1778,7 @@ export function createHandlers<
         completionEvent,
         {
           type: 'ai.message.completed',
-          id: `${generationId}:message:completed`,
+          id: await aiEventId('a2.ai.message.completed', generationId),
           payload: {
             messageId: responseMessageId,
           } satisfies MessageCompletedPayload,
@@ -1920,7 +1801,7 @@ export function createHandlers<
       }
       return {
         type: 'ai.generation.failed',
-        id: `${generationId}:failed`,
+        id: await aiEventId('a2.ai.failed', generationId),
         payload: failed,
       } as AppendInput<D>
     }
@@ -1943,12 +1824,36 @@ export function createHandlers<
         ctx.signal.throwIfAborted()
         return {
           type: 'ai.model.metadata.resolved',
-          id: `${ctx.event.id}: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': { lane: 'a2.ai.control', handler: control.handler },
+    'ai.control.requested': {
+      lane: ({ event }) => {
+        const command = event.payload as ControlCommand<M>
+        if (event.id === '')
+          throw new A2Error(
+            'INVALID_PAYLOAD',
+            'AI command IDs must be nonempty',
+          )
+        if (event.id?.startsWith(AI_EVENT_ID_PREFIX))
+          throw new A2Error(
+            'INVALID_PAYLOAD',
+            `AI command IDs cannot use the reserved '${AI_EVENT_ID_PREFIX}' prefix`,
+          )
+        if (
+          (command.action === 'send' || command.action === 'steer') &&
+          event.id !== command.message.id
+        )
+          throw new A2Error(
+            'INVALID_PAYLOAD',
+            'AI input event ID must equal message.id',
+          )
+        return 'a2.ai.control'
+      },
+      handler: control.handler,
+    },
     'ai.work.reported': { lane: 'a2.ai.control', handler: control.handler },
     'ai.message.completed': {
       handler: async (ctx) => {
@@ -1978,6 +1883,9 @@ export function createHandlers<
     'ai.generation.requested': {
       lane: 'a2.ai.model',
       abortOn: {
+        'ai.control.requested': (event, trigger) =>
+          event.payload.action === 'stop' &&
+          event.payload.turnId === trigger.payload.control?.turnId,
         'ai.control.committed': (event, trigger) => {
           const active = (event.payload as ControlCommit).view
           return (
@@ -1999,6 +1907,9 @@ export function createHandlers<
     },
     'ai.tool.execution.requested': {
       abortOn: {
+        'ai.control.requested': (event, trigger) =>
+          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) => {