experimental-a2 0.12.0 → 0.13.0

package/src/ai-server.ts

@@ -9,7 +9,7 @@
 // ordered; parallel consumption would corrupt progress sequence and state.
 
 import { AsyncLocalStorage } from 'node:async_hooks'
-import { convertToModelMessages } from 'ai'
+import { asSchema, convertToModelMessages } from 'ai'
 import type {
   FinishReason,
   Instructions,
@@ -22,6 +22,7 @@ import type {
   UIMessageChunk,
 } from 'ai'
 import { generateAISDKStep, type AISDKStepSettings } from './ai-sdk-step.ts'
+import { gatewayModelId, readModelLimits } from './ai-model-metadata.ts'
 import type {
   AIEventDefs,
   AIState,
@@ -116,6 +117,7 @@ export type AgentGenerateContext<
   generationId: string
   responseMessageId: string
   messages: M[]
+  modelMessages: ModelMessage[]
   state: AIState<M>
   history: ContractEvent<D>[]
   signal: AbortSignal
@@ -161,13 +163,29 @@ export type CompactionPolicy<
   D extends AIEventDefs<M> & EventDefs,
 > = {
   shouldCompact(
-    context: AgentResolverContext<M, D> & { messages: M[] },
+    context: AgentResolverContext<M, D> & {
+      messages: M[]
+      modelMessages: ModelMessage[]
+    },
   ): boolean | Promise<boolean>
   compact(
-    context: AgentResolverContext<M, D> & { messages: M[] },
+    context: AgentResolverContext<M, D> & {
+      messages: M[]
+      modelMessages: ModelMessage[]
+    },
   ): M[] | Promise<M[]>
 }
 
+export type AutomaticCompaction = {
+  thresholdTokens?: number
+  instructions?: string
+}
+
+export type CompactionOptions<
+  M extends UIMessage,
+  D extends AIEventDefs<M> & EventDefs,
+> = false | AutomaticCompaction | CompactionPolicy<M, D>
+
 export type AgentGenerationSettings<T extends ToolSet> = AISDKStepSettings<T>
 
 export type CreateHandlersOptions<
@@ -196,7 +214,9 @@ export type CreateHandlersOptions<
    * The custom stream owns its metadata chunks.
    */
   generate?: AgentGenerate<M, D, T>
-  compaction?: CompactionPolicy<M, D>
+  compaction?:
+    | CompactionOptions<M, D>
+    | ((context: AgentResolverContext<M, D>) => CompactionOptions<M, D>)
   /** Progress is durably flushed at either limit, whichever is reached first. */
   progress?: { maxChunks?: number; maxDelayMs?: number }
 }
@@ -210,13 +230,35 @@ export type CreateAgentServerOptions<
     handlers?: ServerOptions<D>['handlers']
   }
 
-const resolve = async <T, C>(
+const checkAbort = (signal: AbortSignal): boolean => {
+  if (!signal.aborted) return false
+  const reason: unknown = signal.reason
+  if (
+    reason instanceof A2Error &&
+    (reason.code === 'CLAIM_EXPIRED' || reason.code === 'SUPERSEDED_ATTEMPT')
+  ) {
+    throw reason
+  }
+  return true
+}
+
+const resolve = async <T, C extends { signal: AbortSignal }>(
   value: Resolvable<T, C>,
   context: C,
-): Promise<T> =>
-  typeof value === 'function'
-    ? await (value as (context: C) => T | Promise<T>)(context)
-    : value
+): Promise<T> => {
+  let result: T
+  try {
+    result =
+      typeof value === 'function'
+        ? await (value as (context: C) => T | Promise<T>)(context)
+        : value
+  } catch (error) {
+    checkAbort(context.signal)
+    throw error
+  }
+  checkAbort(context.signal)
+  return result
+}
 
 const modelName = (model: LanguageModel): string => {
   if (typeof model === 'string') return model
@@ -265,10 +307,12 @@ const isBoundaryChunk = (chunk: UIMessageChunk | undefined): boolean =>
     chunk.type === 'abort' ||
     chunk.type === 'error')
 
+type StreamRead<T> = Awaited<ReturnType<ReadableStreamDefaultReader<T>['read']>>
+
 const nextOrTimer = async <T>(
-  next: Promise<IteratorResult<T>>,
+  next: Promise<StreamRead<T>>,
   delayMs: number,
-): Promise<{ type: 'next'; result: IteratorResult<T> } | { type: 'timer' }> => {
+): Promise<{ type: 'next'; result: StreamRead<T> } | { type: 'timer' }> => {
   let timer: ReturnType<typeof setTimeout> | undefined
   try {
     return await Promise.race([
@@ -305,6 +349,7 @@ async function generateWithAISDK<
     model: context.model,
     tools: context.tools,
     messages: context.messages,
+    modelMessages: context.modelMessages,
     responseMessageId: context.responseMessageId,
     abortSignal: context.signal,
     ...(context.instructions === undefined
@@ -333,11 +378,11 @@ async function* consumeGeneration(options: {
   const maxChunks = options.progress?.maxChunks ?? 16
   const maxDelayMs = options.progress?.maxDelayMs ?? 30
   let lastFlush = Date.now()
-  const chunks = options.source.stream[Symbol.asyncIterator]()
-  let pendingChunk = chunks.next()
+  const reader = options.source.stream.getReader()
+  let pendingChunk = reader.read()
   try {
     for (;;) {
-      let result: IteratorResult<UIMessageChunk>
+      let result: StreamRead<UIMessageChunk>
       if (pending.length > 0) {
         const remaining = Math.max(0, maxDelayMs - (Date.now() - lastFlush))
         const outcome = await nextOrTimer(pendingChunk, remaining)
@@ -355,7 +400,7 @@ async function* consumeGeneration(options: {
       pending.push(chunk)
       if (chunk.type === 'finish') streamedFinishReason = chunk.finishReason
       if (chunk.type === 'error') streamedError = chunk.errorText
-      pendingChunk = chunks.next()
+      pendingChunk = reader.read()
 
       if (pending.length >= maxChunks || isBoundaryChunk(pending.at(-1))) {
         yield { type: 'progress', chunks: pending.splice(0) }
@@ -366,6 +411,15 @@ async function* consumeGeneration(options: {
     if (pending.length > 0)
       yield { type: 'progress', chunks: pending.splice(0) }
     throw error
+  } finally {
+    void pendingChunk.catch(() => {})
+    try {
+      await reader.cancel()
+    } catch {
+      // Cleanup must not replace the error that ended generation.
+    } finally {
+      reader.releaseLock()
+    }
   }
 
   let completion: GenerationCompletion | undefined
@@ -398,32 +452,249 @@ const replay = <M extends UIMessage, D extends AIEventDefs<M> & EventDefs>(
   return state
 }
 
-const contextMessages = <M extends UIMessage>(state: AIState<M>): M[] => {
+const summarize = async <
+  M extends UIMessage,
+  D extends AIEventDefs<M> & EventDefs,
+  T extends ToolSet,
+>(options: {
+  context: AgentGenerateContext<M, D, T>
+  instructions: string | undefined
+}): Promise<{ summary: string; usage?: LanguageModelUsage }> => {
+  const instruction = [
+    'Summarize this conversation so you can continue the task after the earlier context is replaced by your summary.',
+    'Preserve the objective, constraints, decisions, exact identifiers, completed actions and their results, unresolved questions, and next steps. Distinguish observations from assumptions.',
+    'Respond only with the summary as text. Do not call any tools or continue working on the task.',
+    options.instructions,
+  ]
+    .filter(Boolean)
+    .join('\n')
+  const generation = { ...options.context.generation }
+  for (const key of Object.keys(generation)) {
+    if (
+      key.startsWith('on') ||
+      key.startsWith('experimental_on') ||
+      [
+        'repairToolCall',
+        'experimental_repairToolCall',
+        'experimental_refineToolInput',
+        'experimental_transform',
+        'toolApproval',
+      ].includes(key)
+    ) {
+      Reflect.deleteProperty(generation, key)
+    }
+  }
+  const tools = Object.fromEntries(
+    Object.entries(options.context.tools).map(([name, tool]) => {
+      const definition = { ...tool, needsApproval: false }
+      for (const key of [
+        'execute',
+        'onInputStart',
+        'onInputDelta',
+        'onInputAvailable',
+      ]) {
+        Reflect.deleteProperty(definition, key)
+      }
+      return [name, definition]
+    }),
+  ) as T
+  const source = await generateWithAISDK(
+    {
+      ...options.context,
+      generation,
+      tools,
+      modelMessages: [
+        ...options.context.modelMessages,
+        { role: 'user', content: instruction },
+      ],
+      responseMessageId: `${options.context.generationId}:summary`,
+    },
+    undefined,
+  )
+  let summary = ''
+  let calledTool = false
+  let finish: AgentGenerationFinish | undefined
+  for await (const update of consumeGeneration({ source })) {
+    options.context.signal.throwIfAborted()
+    if (update.type === 'finish') {
+      finish = update
+      continue
+    }
+    for (const chunk of update.chunks) {
+      if (chunk.type === 'text-delta') summary += chunk.delta
+      if (chunk.type.startsWith('tool-')) calledTool = true
+    }
+  }
+  if (
+    calledTool ||
+    summary.trim().length === 0 ||
+    finish?.finishReason !== 'stop'
+  ) {
+    throw new Error(
+      'compaction must produce a complete text summary without calling tools',
+    )
+  }
+  return {
+    summary,
+    ...(finish.usage === undefined ? {} : { usage: finish.usage }),
+  }
+}
+
+const activeCompaction = <M extends UIMessage>(
+  state: AIState<M>,
+): AIState<M>['compaction'] => {
   const compaction = state.compaction
-  if (compaction?.status !== 'completed' || !compaction.messages) {
-    return state.messages
+  return compaction?.status === 'completed' &&
+    compaction.messages !== undefined &&
+    state.messages.some((message) => message.id === compaction.throughMessageId)
+    ? compaction
+    : null
+}
+
+const contextMessages = <
+  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
+  const compaction = activeCompaction(state)
+  if (compaction === null) return 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)
+  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)
+    }
+    return context.messages
   }
   const boundary = state.messages.findIndex(
     (message) => message.id === compaction.throughMessageId,
   )
-  const retained = new Set(compaction.retainedMessageIds ?? [])
-  return boundary === -1
-    ? state.messages
-    : [
-        ...compaction.messages,
-        ...state.messages
-          .slice(0, boundary + 1)
-          .filter((message) => retained.has(message.id)),
-        ...state.messages.slice(boundary + 1),
-      ]
+  return [
+    ...compaction.messages!,
+    ...state.messages
+      .slice(0, boundary + 1)
+      .filter((message) => retained.has(message.id)),
+    ...state.messages.slice(boundary + 1),
+  ]
 }
 
-const activeContextMessages = <M extends UIMessage>(
-  state: AIState<M>,
-  coordinator: AICoordinatorState,
-): M[] => {
-  const queued = new Set(coordinator.queued.map((item) => item.messageId))
-  return contextMessages(state).filter((message) => !queued.has(message.id))
+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: {
+  messages: M[]
+  state: AIState<M>
+  tools: ToolSet
+}): Promise<ModelMessage[]> => {
+  const messages = await convertToModelMessages(options.messages, {
+    tools: options.tools,
+  })
+  const summary = activeCompaction(options.state)?.summary
+  return summary === undefined
+    ? messages
+    : [{ 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 = <D extends EventDefs>(options: {
+  estimate: number
+  model: string
+  history: ContractEvent<D>[]
+}): 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 => {
@@ -772,6 +1043,62 @@ function withoutGenerationLifecycle<D extends EventDefs>(
   })
 }
 
+const validateCompaction = <
+  M extends UIMessage,
+  D extends AIEventDefs<M> & EventDefs,
+  T extends ToolSet,
+>(options: {
+  compaction: CompactionOptions<M, D>
+  explicit: boolean
+  generation: AgentGenerationSettings<T> | undefined
+  tools: T | undefined
+}): CompactionOptions<M, D> => {
+  let compaction = options.compaction
+  if (
+    compaction !== false &&
+    (typeof compaction !== 'object' || compaction === null)
+  ) {
+    throw new TypeError('compaction must resolve to false or a policy object')
+  }
+  if (compaction !== false && 'then' in compaction) {
+    void Promise.resolve(compaction).catch(() => {})
+    throw new TypeError('compaction options must resolve synchronously')
+  }
+  if (compaction !== false && !('shouldCompact' in compaction)) {
+    if (
+      compaction.thresholdTokens !== undefined &&
+      (!Number.isSafeInteger(compaction.thresholdTokens) ||
+        compaction.thresholdTokens < 1)
+    ) {
+      throw new TypeError(
+        'compaction.thresholdTokens must be a positive safe integer',
+      )
+    }
+    const choice = options.generation?.toolChoice
+    const unsupported =
+      options.generation?.output !== undefined
+        ? 'structured output'
+        : choice !== undefined && choice !== 'auto' && choice !== 'none'
+          ? 'forced tool choice'
+          : Object.values(options.tools ?? {}).some(
+                (tool) =>
+                  tool.type === 'provider' && tool.isProviderExecuted === true,
+              )
+            ? 'provider-executed tools'
+            : undefined
+    if (unsupported !== undefined) {
+      if (options.explicit) {
+        throw new TypeError(
+          `automatic compaction does not support ${unsupported}; use a custom policy or compaction: false`,
+        )
+      }
+      compaction = false
+    }
+  }
+
+  return compaction
+}
+
 /**
  * Build the ordinary A2 handler table for the built-in agent protocol.
  * Application handlers can be spread beside this table.
@@ -808,8 +1135,20 @@ export function createHandlers<
     throw new TypeError('maxSteps must be a positive integer or Infinity')
   }
 
+  const configuredCompaction =
+    options.compaction ?? (options.generate === undefined ? {} : false)
+  const staticCompaction =
+    typeof configuredCompaction === 'function'
+      ? undefined
+      : validateCompaction({
+          compaction: configuredCompaction,
+          explicit: options.compaction !== undefined,
+          generation: options.generation,
+          tools: options.tools,
+        })
+
   const tools = options.tools ?? ({} as T)
-  const generation = options.generation ?? {}
+  const generation: AgentGenerationSettings<T> = options.generation ?? {}
   const maxSteps = options.maxSteps ?? Number.POSITIVE_INFINITY
   const coordinator = aiCoordinatorReducer(options.agent.contract)
   const promptCache = new Map<string, Promise<ModelMessage[]>>()
@@ -956,39 +1295,37 @@ export function createHandlers<
     if (cached) return cached
     const computation = (async () => {
       const history = await readHistory()
-      const coordinatorState = coordinatorStateAt(history)
-      const compaction = history.find(
-        (event) =>
-          event.type === 'ai.compaction.completed' &&
-          (event.payload as CompactionCompletedPayload<M>).generationId ===
-            call.generationId,
-      )
-      if (compaction !== undefined) {
-        const messages = (
-          compaction.payload as CompactionCompletedPayload<M>
-        ).messages.filter(
-          (message) =>
-            !coordinatorState.queued.some(
-              (queued) => queued.messageId === message.id,
-            ),
-        )
-        return convertToModelMessages(messages, { tools })
-      }
       const started = history.find(
         (event) =>
           event.type === 'ai.generation.started' &&
           (event.payload as GenerationStartedPayload).generationId ===
             call.generationId,
       )
-      const frontier = started?.index ?? Number.POSITIVE_INFINITY
-      const state = replay(
-        options.agent,
-        history.filter((event) => event.index < frontier),
-      )
-      return convertToModelMessages(
-        activeContextMessages(state, coordinatorState),
-        { tools },
+      const startedPayload = started?.payload as
+        GenerationStartedPayload | undefined
+      const frontier =
+        startedPayload?.promptThroughIndex ??
+        (started?.index ?? Number.POSITIVE_INFINITY) - 1
+      const promptHistory = history.filter(
+        (event) =>
+          event.index <= frontier ||
+          ((event.type === 'ai.generation.started' ||
+            event.type === 'ai.compaction.requested' ||
+            event.type === 'ai.compaction.completed') &&
+            (event.payload as { generationId: string }).generationId ===
+              call.generationId),
       )
+      const state = replay(options.agent, promptHistory)
+      return modelContext({
+        messages: activeContextMessages({
+          agent: options.agent,
+          state,
+          history: promptHistory,
+          coordinator: coordinatorStateAt(promptHistory),
+        }),
+        state,
+        tools,
+      })
     })()
     promptCache.set(key, computation)
     void computation.catch(() => {
@@ -1003,7 +1340,7 @@ export function createHandlers<
     error: unknown,
   ): AppendInput<D> | void => {
     const schedulerFailure = consumeSchedulerSendFailure(error)
-    if (ctx.signal.aborted) return
+    if (checkAbort(ctx.signal)) return
     if (schedulerFailure === 'retryable') throw error
     return resultEvent(call, 'execution:error', {
       error: errorMessage(error),
@@ -1037,40 +1374,42 @@ export function createHandlers<
     }
     let last: unknown
     let sequence = 0
-    for (;;) {
-      let result: IteratorResult<unknown>
-      try {
-        result = await iterator.next()
-      } catch (error) {
-        return toolExecutionFailure(ctx, call, error)
-      }
-      if (result.done) break
-      if (ctx.signal.aborted) {
+    let done = false
+    try {
+      for (;;) {
+        let result: IteratorResult<unknown>
         try {
-          await iterator.return?.()
-        } catch {
-          // The aborted handler cannot record an iterator cleanup failure.
+          result = await iterator.next()
+        } catch (error) {
+          return toolExecutionFailure(ctx, call, error)
         }
-        return
-      }
-      last = result.value
-      try {
+        if (checkAbort(ctx.signal)) return
+        if (result.done) {
+          done = true
+          break
+        }
+        last = result.value
         await ctx.session.append(
           `tool:${call.toolCallId}:preliminary:${sequence}`,
-          resultEvent(call, `execution:${sequence}:preliminary`, {
-            output: result.value,
-            preliminary: true,
-          }),
+          resultEvent(
+            call,
+            `execution:${ctx.attempt}:${sequence}:preliminary`,
+            {
+              output: result.value,
+              preliminary: true,
+            },
+          ),
         )
-      } catch (error) {
+        sequence += 1
+      }
+    } finally {
+      if (!done) {
         try {
           await iterator.return?.()
         } catch {
-          // Preserve the append failure that makes the durable handler retry.
+          // Preserve the failure or cancellation that ended execution.
         }
-        throw error
       }
-      sequence += 1
     }
     return resultEvent(
       call,
@@ -1095,7 +1434,7 @@ export function createHandlers<
       call,
       ctx.session.history,
     )
-    if (ctx.signal.aborted) return
+    if (checkAbort(ctx.signal)) return
     const scope: AmbientToolScope = {
       contract: options.agent.contract,
       context: ctx,
@@ -1262,7 +1601,7 @@ export function createHandlers<
       history,
       replacedGenerationIds,
     )
-    let state = replay(options.agent, promptHistory)
+    const state = replay(options.agent, promptHistory)
     const resolverContext: AgentResolverContext<M, D> = {
       event: ctx.event,
       state,
@@ -1270,138 +1609,289 @@ export function createHandlers<
       signal: ctx.signal,
     }
 
-    const resolvedModel = await resolve(options.model, resolverContext)
-    if (resolvedModel === undefined) {
-      throw new TypeError('the model resolver returned undefined')
-    }
-    if (ctx.signal.aborted) return
-
-    if (request.reason === 'tool' && responseStepCount >= maxSteps) {
-      const source = sourceCoordinatorState.response?.generation
-      if (source === undefined) return
-      return {
-        type: 'ai.generation.failed',
-        id: `${requestId}:step-limit`,
-        payload: {
-          requestId: source.requestId,
-          messageId: source.messageId,
-          generationId: source.generationId,
-          responseMessageId,
-          error: `agent exceeded the ${maxSteps}-step limit`,
-          stepLimit: true,
-        },
-      } as AppendInput<D>
-    }
+    let generationStarted = false
+    try {
+      const resolvedModel = await resolve(options.model, resolverContext)
+      if (resolvedModel === undefined) {
+        throw new TypeError('the model resolver returned undefined')
+      }
+      if (checkAbort(ctx.signal)) return
+      const resolvedInstructions =
+        options.instructions === undefined
+          ? undefined
+          : await resolve(options.instructions, resolverContext)
+      if (checkAbort(ctx.signal)) return
+      const compaction =
+        typeof configuredCompaction === 'function'
+          ? validateCompaction({
+              compaction: configuredCompaction(resolverContext),
+              explicit: true,
+              generation: options.generation,
+              tools: options.tools,
+            })
+          : staticCompaction!
+      if (checkAbort(ctx.signal)) return
+
+      if (request.reason === 'tool' && responseStepCount >= maxSteps) {
+        const source = sourceCoordinatorState.response?.generation
+        if (source === undefined) return
+        return {
+          type: 'ai.generation.failed',
+          id: `${requestId}:step-limit`,
+          payload: {
+            requestId: source.requestId,
+            messageId: source.messageId,
+            generationId: source.generationId,
+            responseMessageId,
+            error: `agent exceeded the ${maxSteps}-step limit`,
+            stepLimit: true,
+          },
+        } as AppendInput<D>
+      }
 
-    const started: GenerationStartedPayload = {
-      requestId,
-      messageId: request.messageId,
-      generationId,
-      responseMessageId,
-      attempt,
-      model: modelName(resolvedModel),
-    }
-    const startEvents: AppendInput<D>[] = []
-    if (previous) {
-      const payload = previous.payload as GenerationStartedPayload
-      const superseded: GenerationFailedPayload = {
+      const started: GenerationStartedPayload = {
         requestId,
-        messageId: payload.messageId,
-        generationId: payload.generationId,
-        responseMessageId: payload.responseMessageId,
-        error: 'generation attempt was superseded after an incomplete run',
-        superseded: true,
+        messageId: request.messageId,
+        generationId,
+        responseMessageId,
+        attempt,
+        model: modelName(resolvedModel),
+        promptThroughIndex: history.at(-1)?.index ?? ctx.event.index,
+      }
+      const catalogModelId = gatewayModelId(resolvedModel)
+      const gatewayOptions = options.generation?.providerOptions?.['gateway']
+      const usesFallbackModels =
+        Array.isArray(gatewayOptions?.['models']) &&
+        gatewayOptions['models'].length > 0
+      const discoversMetadata =
+        compaction !== false &&
+        !('shouldCompact' in compaction) &&
+        compaction.thresholdTokens === undefined &&
+        !usesFallbackModels &&
+        catalogModelId !== undefined
+      const startEvents: AppendInput<D>[] = []
+      if (
+        discoversMetadata &&
+        !Object.hasOwn(state.modelMetadata, catalogModelId)
+      ) {
+        startEvents.push({
+          type: 'ai.model.metadata.requested',
+          id: `ai.model.metadata:${encodeURIComponent(options.agent.contract.name)}:${encodeURIComponent(ctx.event.sessionId)}:${encodeURIComponent(catalogModelId)}`,
+          payload: { modelId: catalogModelId },
+        } as AppendInput<D>)
+      }
+      if (previous) {
+        const payload = previous.payload as GenerationStartedPayload
+        const superseded: GenerationFailedPayload = {
+          requestId,
+          messageId: payload.messageId,
+          generationId: payload.generationId,
+          responseMessageId: payload.responseMessageId,
+          error: 'generation attempt was superseded after an incomplete run',
+          superseded: true,
+        }
+        startEvents.push({
+          type: 'ai.generation.failed',
+          id: `${payload.generationId}:superseded`,
+          payload: superseded,
+        } as AppendInput<D>)
       }
       startEvents.push({
-        type: 'ai.generation.failed',
-        id: `${payload.generationId}:superseded`,
-        payload: superseded,
+        type: 'ai.generation.started',
+        id: generationId,
+        payload: started,
       } as AppendInput<D>)
-    }
-    startEvents.push({
-      type: 'ai.generation.started',
-      id: generationId,
-      payload: started,
-    } as AppendInput<D>)
-    await ctx.session.append('generation-start', ...startEvents)
+      const generationHistory = [
+        ...history,
+        ...(await ctx.session.append('generation-start', ...startEvents)),
+      ]
+      generationStarted = true
 
-    const startedCoordinatorState = (
-      await ctx.session.state(coordinator, { through: 'latest' })
-    ).state
-    if (
-      startedCoordinatorState.response?.activeRequestId !== requestId ||
-      startedCoordinatorState.response.generation?.generationId !== generationId
-    ) {
-      return
-    }
+      const startedCoordinatorState = (
+        await ctx.session.state(coordinator, { through: 'latest' })
+      ).state
+      if (
+        startedCoordinatorState.response?.activeRequestId !== requestId ||
+        startedCoordinatorState.response.generation?.generationId !==
+          generationId
+      ) {
+        return
+      }
 
-    const promptCoordinatorState = {
-      ...coordinatorStateAt(history),
-      queued: queuedMessagesAt(history),
-    }
-    let messages = activeContextMessages(state, promptCoordinatorState)
-    if (options.compaction) {
-      const compactionContext = { ...resolverContext, messages }
-      if (await options.compaction.shouldCompact(compactionContext)) {
-        const throughMessageId = messages.at(-1)?.id ?? request.messageId
-        await ctx.session.append('compaction-requested', {
-          type: 'ai.compaction.requested',
-          id: `${generationId}:compaction:requested`,
-          payload: { generationId, throughMessageId },
-        })
-        messages = await options.compaction.compact(compactionContext)
-        const completed: CompactionCompletedPayload<M> = {
-          generationId,
-          throughMessageId,
+      const promptCoordinatorState = {
+        ...coordinatorStateAt(history),
+        queued: queuedMessagesAt(history),
+      }
+      const messages = activeContextMessages({
+        agent: options.agent,
+        state,
+        history: promptHistory,
+        coordinator: promptCoordinatorState,
+      })
+      let modelMessages = await modelContext({ messages, state, tools })
+      const baseContext: AgentGenerateContext<M, D, T> = {
+        request: ctx.event,
+        requestId,
+        generationId,
+        responseMessageId,
+        messages,
+        modelMessages,
+        state,
+        history,
+        signal: ctx.signal,
+        model: resolvedModel,
+        tools,
+        ...(resolvedInstructions === undefined
+          ? {}
+          : { instructions: resolvedInstructions }),
+        generation,
+      }
+      const policy = compaction
+      const tracksInput =
+        policy !== false &&
+        !('shouldCompact' in policy) &&
+        (policy.thresholdTokens !== undefined || discoversMetadata)
+      const metadata =
+        catalogModelId === undefined
+          ? undefined
+          : state.modelMetadata[catalogModelId]
+      const limits =
+        metadata?.status === 'resolved' && !usesFallbackModels
+          ? metadata.limits
+          : undefined
+      let inputTokenEstimate: number | undefined
+      let compacted = false
+      const canCompact =
+        sourceCoordinatorState.response?.calls.every((call) => call.terminal) ??
+        true
+      if (policy && canCompact) {
+        const compactionContext = {
+          ...resolverContext,
           messages,
-          ...(promptCoordinatorState.queued.length === 0
-            ? {}
-            : {
-                retainedMessageIds: promptCoordinatorState.queued.map(
-                  (item) => item.messageId,
-                ),
-              }),
+          modelMessages,
         }
-        await ctx.session.append('compaction-completed', {
-          type: 'ai.compaction.completed',
-          id: `${generationId}:compaction:completed`,
-          payload: completed,
-        })
-        state = {
-          ...state,
-          compaction: { status: 'completed', ...completed },
+        let shouldCompact: boolean
+        if ('shouldCompact' in policy) {
+          shouldCompact = await policy.shouldCompact(compactionContext)
+        } else if (tracksInput) {
+          inputTokenEstimate = await estimateInputTokens({
+            messages: modelMessages,
+            instructions: resolvedInstructions,
+            tools,
+          })
+          const inputTokens = measuredInputTokens({
+            estimate: inputTokenEstimate,
+            model: modelName(resolvedModel),
+            history,
+          })
+          const threshold =
+            policy.thresholdTokens ??
+            (limits === undefined
+              ? undefined
+              : Math.floor(
+                  Math.min(
+                    limits.contextWindow * 0.75,
+                    limits.contextWindow - (generation.maxOutputTokens ?? 0),
+                  ),
+                ))
+          if (
+            limits !== undefined &&
+            threshold !== undefined &&
+            threshold <= 0
+          ) {
+            throw new Error(
+              `output allowance exhausts the ${limits.contextWindow}-token context window`,
+            )
+          }
+          shouldCompact = threshold !== undefined && inputTokens >= threshold
+        } else {
+          shouldCompact = false
+        }
+        if (checkAbort(ctx.signal)) return
+        if (shouldCompact) {
+          const throughMessageId =
+            state.messages.findLast(
+              (message) =>
+                !promptCoordinatorState.queued.some(
+                  (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 result = !('shouldCompact' in policy)
+            ? {
+                messages: [] as M[],
+                ...(await summarize({
+                  context: baseContext,
+                  instructions: policy.instructions,
+                })),
+              }
+            : { messages: await policy.compact(compactionContext) }
+          if (checkAbort(ctx.signal)) return
+          const completed: CompactionCompletedPayload<M> = {
+            generationId,
+            throughMessageId,
+            throughIndex,
+            ...result,
+            ...(promptCoordinatorState.queued.length === 0
+              ? {}
+              : {
+                  retainedMessageIds: promptCoordinatorState.queued.map(
+                    (item) => item.messageId,
+                  ),
+                }),
+          }
+          generationHistory.push(
+            ...(await ctx.session.append('compaction-completed', {
+              type: 'ai.compaction.completed',
+              id: `${generationId}:compaction:completed`,
+              payload: completed,
+            })),
+          )
+          compacted = true
         }
       }
-    }
 
-    const currentState = (
-      await ctx.session.state(options.agent.reducer, { through: 'latest' })
-    ).state
-    const currentCoordinatorState = (
-      await ctx.session.state(coordinator, { through: 'latest' })
-    ).state
-    const generationMessages = activeContextMessages(
-      currentState,
-      currentCoordinatorState,
-    )
-    promptCache.set(
-      promptCacheKey(ctx.event.sessionId, generationId),
-      convertToModelMessages(generationMessages, { tools }),
-    )
+      const currentState = replay(options.agent, generationHistory)
+      const currentCoordinatorState = coordinatorStateAt(generationHistory)
+      const generationMessages = activeContextMessages({
+        agent: options.agent,
+        state: currentState,
+        history: generationHistory,
+        coordinator: currentCoordinatorState,
+      })
+      modelMessages = await modelContext({
+        messages: generationMessages,
+        state: currentState,
+        tools,
+      })
+      promptCache.set(
+        promptCacheKey(ctx.event.sessionId, generationId),
+        Promise.resolve(modelMessages),
+      )
+      if (tracksInput && (inputTokenEstimate === undefined || compacted)) {
+        inputTokenEstimate = await estimateInputTokens({
+          messages: modelMessages,
+          instructions: resolvedInstructions,
+          tools,
+        })
+      }
 
-    let pendingToolCalls: PendingToolCall[] = []
-    try {
-      const resolvedInstructions =
-        options.instructions === undefined
-          ? undefined
-          : await resolve(options.instructions, resolverContext)
-      if (ctx.signal.aborted) return
+      let pendingToolCalls: PendingToolCall[] = []
+      if (checkAbort(ctx.signal)) return
       const generateContext: AgentGenerateContext<M, D, T> = {
         request: ctx.event,
         requestId,
         generationId,
         responseMessageId,
         messages: generationMessages,
+        modelMessages,
         state: currentState,
         history,
         signal: ctx.signal,
@@ -1470,6 +1960,7 @@ export function createHandlers<
         messageId: request.messageId,
         generationId,
         responseMessageId,
+        ...(inputTokenEstimate === undefined ? {} : { inputTokenEstimate }),
         ...(finish.finishReason === undefined
           ? {}
           : { finishReason: finish.finishReason }),
@@ -1513,8 +2004,11 @@ export function createHandlers<
         } as AppendInput<D>,
       ]
     } catch (error) {
-      if (error instanceof A2Error) throw error
-      if (ctx.signal.aborted) {
+      if (error instanceof A2Error) {
+        checkAbort(ctx.signal)
+        throw error
+      }
+      if (checkAbort(ctx.signal)) {
         const interrupted: MessageInterruptedPayload = {
           messageId: responseMessageId,
           generationId,
@@ -1526,6 +2020,7 @@ export function createHandlers<
           payload: interrupted,
         } as AppendInput<D>
       }
+      if (!generationStarted) throw error
       const failed: GenerationFailedPayload = {
         requestId,
         messageId: request.messageId,
@@ -1609,6 +2104,20 @@ 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: `${ctx.event.id}:resolved`,
+          payload: { modelId: ctx.event.payload.modelId, limits },
+        } as AppendInput<D>
+      },
+    },
     'ai.message.created': { lane: 'a2.ai.turn', handler: handleMessageCreated },
     'ai.retry.requested': { lane: 'a2.ai.turn', handler: handleRetry },
     'ai.input.responded': {