experimental-a2 0.3.0 → 0.5.0

package/src/ai.ts

@@ -0,0 +1,2155 @@
+/**
+ * experimental-a2/ai — an isomorphic, extensible agent definition.
+ *
+ * `agent()` assembles a normal A2 contract, a reducer over the built-in
+ * AI protocol, and pure builders for the events applications append.
+ * Model calls and tool implementations stay in experimental-a2/ai/server.
+ */
+
+import type {
+  FinishReason,
+  JSONValue,
+  LanguageModelUsage,
+  UIMessage,
+  UIMessageChunk,
+} from 'ai'
+import { contract as createContract } from './contract.ts'
+import type { AppendInput, Contract, EventDefs } from './contract.ts'
+import { interruptUIMessage, projectUIMessage } from './ai-projector.ts'
+import { ambientToolScopeStorage } from './internal.ts'
+import type { Reducer } from './reducer.ts'
+import type { HandlerContext } from './server.ts'
+import type { StandardSchemaV1 } from './standard-schema.ts'
+import { assertSyncSchema, validateSync } from './validate.ts'
+
+export type AgentStatus =
+  'idle' | 'generating' | 'waiting' | 'failed' | 'closed'
+
+export type GenerationReason = 'message' | 'tool' | 'input' | 'retry'
+
+export type SessionCreatedPayload = { metadata?: JSONValue }
+export type SessionClosedPayload = { reason?: string }
+export type MessageCreatedPayload<M extends UIMessage> = {
+  message: M
+  /** For user messages, false records passive context. Omission or true schedules. */
+  generate?: boolean
+}
+export type MessageCompletedPayload = { messageId: string }
+type MessageInterruption = {
+  messageId: string
+  reason?: string
+  /** Last confirmed event-log index visible to the interrupter. */
+  lastSeenIndex?: number
+}
+export type MessageInterruptedPayload = MessageInterruption &
+  (
+    | { generationId: string; requestId?: never }
+    | { requestId: string; generationId?: never }
+  )
+export type GenerationRequestedPayload = {
+  messageId: string
+  reason: GenerationReason
+  responseMessageId?: string
+}
+export type RetryGenerationOptions = {
+  messageId: string
+  responseMessageId: string
+  retryId: string
+}
+export type RetryRequestedPayload = RetryGenerationOptions
+export type GenerationStartedPayload = {
+  requestId: string
+  messageId: string
+  generationId: string
+  responseMessageId: string
+  attempt: number
+  model: string
+}
+export type GenerationProgressPayload = {
+  requestId: string
+  messageId: string
+  generationId: string
+  responseMessageId: string
+  sequence: number
+  chunks: UIMessageChunk[]
+}
+export type GenerationCompletedPayload = {
+  requestId: string
+  messageId: string
+  generationId: string
+  responseMessageId: string
+  finishReason?: FinishReason
+  usage?: LanguageModelUsage
+}
+export type GenerationFailedPayload = {
+  requestId: string
+  messageId: string
+  generationId: string
+  responseMessageId: string
+  error: string
+  /** True when a later attempt replaced an incomplete attempt. */
+  superseded?: boolean
+  stepLimit?: boolean
+}
+export type ToolCalledPayload = {
+  requestId: string
+  messageId: string
+  generationId: string
+  toolCallId: string
+  toolName: string
+  input: unknown
+  dynamic?: boolean
+  providerExecuted?: boolean
+  supportsDeferredResults?: boolean
+  providerMetadata?: unknown
+  toolMetadata?: unknown
+  title?: string
+}
+export type ToolResultPayload = {
+  requestId: string
+  messageId: string
+  generationId: string
+  toolCallId: string
+  toolName?: string
+  input?: unknown
+  rawInput?: unknown
+  output?: unknown
+  error?: string
+  denied?: boolean
+  preliminary?: boolean
+  phase?: 'input' | 'execution'
+  dynamic?: boolean
+  providerExecuted?: boolean
+  providerMetadata?: unknown
+  toolMetadata?: unknown
+}
+export type ApprovalRequestedPayload = {
+  messageId: string
+  generationId: string
+  approvalId: string
+  toolCallId: string
+  isAutomatic?: boolean
+  signature?: string
+}
+export type ApprovalRespondedPayload = {
+  messageId: string
+  generationId: string
+  approvalId: string
+  approved: boolean
+  reason?: string
+}
+export type InputRequestedPayload = {
+  messageId: string
+  generationId: string
+  inputId: string
+  name: string
+  question?: JSONValue
+}
+export type InputRespondedPayload = {
+  messageId: string
+  generationId: string
+  inputId: string
+  name: string
+  value: JSONValue
+}
+export type CompactionRequestedPayload = {
+  generationId: string
+  throughMessageId: string
+}
+export type CompactionCompletedPayload<M extends UIMessage> = {
+  generationId: string
+  throughMessageId: string
+  messages: M[]
+  retainedMessageIds?: string[]
+}
+
+export type AIEventDefs<M extends UIMessage = UIMessage> = {
+  'ai.session.created': StandardSchemaV1<SessionCreatedPayload>
+  'ai.session.closed': StandardSchemaV1<SessionClosedPayload>
+  'ai.message.created': StandardSchemaV1<MessageCreatedPayload<M>>
+  'ai.message.completed': StandardSchemaV1<MessageCompletedPayload>
+  'ai.message.interrupted': StandardSchemaV1<MessageInterruptedPayload>
+  'ai.generation.requested': StandardSchemaV1<GenerationRequestedPayload>
+  'ai.retry.requested': StandardSchemaV1<RetryRequestedPayload>
+  'ai.generation.started': StandardSchemaV1<GenerationStartedPayload>
+  'ai.generation.progress': StandardSchemaV1<GenerationProgressPayload>
+  'ai.generation.completed': StandardSchemaV1<GenerationCompletedPayload>
+  'ai.generation.failed': StandardSchemaV1<GenerationFailedPayload>
+  'ai.tool.called': StandardSchemaV1<ToolCalledPayload>
+  'ai.tool.result': StandardSchemaV1<ToolResultPayload>
+  'ai.approval.requested': StandardSchemaV1<ApprovalRequestedPayload>
+  'ai.approval.responded': StandardSchemaV1<ApprovalRespondedPayload>
+  'ai.input.requested': StandardSchemaV1<InputRequestedPayload>
+  'ai.input.responded': StandardSchemaV1<InputRespondedPayload>
+  'ai.compaction.requested': StandardSchemaV1<CompactionRequestedPayload>
+  'ai.compaction.completed': StandardSchemaV1<CompactionCompletedPayload<M>>
+}
+
+export type ActiveGeneration = GenerationStartedPayload
+
+export type ActiveGenerationProjection<M extends UIMessage = UIMessage> = {
+  generationId: string
+  responseMessageId: string
+  /** Log index of the completed model step while its response remains interruptible. */
+  completionIndex?: number
+  baseMessage?: M
+  baseCompaction?: CompactionState<M> | null
+  batches: Array<{ index: number; chunks: UIMessageChunk[] }>
+  baseTools: ToolActivity[]
+  toolEvents: Array<
+    | { index: number; type: 'ai.tool.called'; payload: ToolCalledPayload }
+    | { index: number; type: 'ai.tool.result'; payload: ToolResultPayload }
+  >
+  approvalEvents: Array<
+    | {
+        index: number
+        type: 'ai.approval.requested'
+        payload: ApprovalRequestedPayload
+      }
+    | {
+        index: number
+        type: 'ai.approval.responded'
+        payload: ApprovalRespondedPayload
+      }
+  >
+}
+
+export type PendingApproval = ApprovalRequestedPayload
+export type PendingInput = InputRequestedPayload
+
+export type ToolActivity = {
+  requestId: string
+  messageId: string
+  generationId: string
+  toolCallId: string
+  toolName?: string
+  input?: unknown
+  rawInput?: unknown
+  output?: unknown
+  error?: string
+  preliminary?: boolean
+  phase?: 'input' | 'execution'
+  dynamic?: boolean
+  providerExecuted?: boolean
+  callProviderMetadata?: unknown
+  resultProviderMetadata?: unknown
+  toolMetadata?: unknown
+  title?: string
+  status: 'running' | 'completed' | 'failed' | 'denied'
+}
+
+export type CompactionState<M extends UIMessage> = {
+  status: 'running' | 'completed'
+  generationId: string
+  throughMessageId: string
+  messages?: M[]
+  retainedMessageIds?: string[]
+}
+
+export type GenerationUsage = {
+  generationId: string
+  usage: LanguageModelUsage
+}
+
+export type AIState<M extends UIMessage = UIMessage> = {
+  session: { created: boolean; metadata?: JSONValue }
+  messages: M[]
+  status: AgentStatus
+  activeGeneration: ActiveGeneration | null
+  activeRequestId: string | null
+  activeResponseMessageId: string | null
+  activeProjection: ActiveGenerationProjection<M> | null
+  responseGenerationIds: Record<string, string>
+  terminalRequestIds?: Record<string, true>
+  terminalGenerations?: Record<
+    string,
+    'completed' | 'failed' | 'interrupted' | 'superseded'
+  >
+  pendingApprovals: PendingApproval[]
+  pendingInputs: PendingInput[]
+  tools: ToolActivity[]
+  compaction: CompactionState<M> | null
+  usage: GenerationUsage[]
+  error: string | null
+}
+
+export type AgentDefinition<
+  M extends UIMessage = UIMessage,
+  D extends EventDefs = AIEventDefs<M>,
+> = {
+  readonly contract: Contract<D>
+  readonly reducer: Reducer<D, AIState<M>>
+}
+
+export type AgentToolEventDefs = EventDefs &
+  Pick<AIEventDefs, 'ai.tool.called' | 'ai.approval.responded'>
+
+/** The durable handler context supplied when A2 executes a local AI tool. */
+export type AgentToolContext<D extends AgentToolEventDefs> =
+  | HandlerContext<D, 'ai.tool.called'>
+  | HandlerContext<D, 'ai.approval.responded'>
+
+export type AgentOptions<M extends UIMessage, E extends EventDefs> = {
+  name: string
+  events?: E
+  /** Optional stricter schema for custom UIMessage metadata/data/tool types. */
+  messageSchema?: StandardSchemaV1<M>
+  /** Snapshot identity. Bump when application-level projection semantics change. */
+  reducerName?: string
+}
+
+type AIMessageFrom<D extends EventDefs> =
+  D extends Record<'ai.message.created', infer S extends StandardSchemaV1>
+    ? StandardSchemaV1.InferOutput<S> extends MessageCreatedPayload<
+        infer M extends UIMessage
+      >
+      ? M
+      : UIMessage
+    : UIMessage
+
+type SchemaResult<T> = StandardSchemaV1.Result<T>
+
+const issue = <T>(message: string): SchemaResult<T> => ({
+  issues: [{ message }],
+})
+
+const schema = <T>(
+  label: string,
+  parse: (value: unknown) => SchemaResult<T>,
+): StandardSchemaV1<T> => ({
+  '~standard': {
+    version: 1,
+    vendor: 'a2',
+    validate(value) {
+      try {
+        return parse(value)
+      } catch (error) {
+        return issue(
+          `${label}: ${error instanceof Error ? error.message : String(error)}`,
+        )
+      }
+    },
+  },
+})
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null && !Array.isArray(value)
+
+const isJSONCompatible = (
+  value: unknown,
+  seen: Set<object> = new Set(),
+): boolean => {
+  if (
+    value === null ||
+    typeof value === 'string' ||
+    typeof value === 'boolean'
+  ) {
+    return true
+  }
+  if (typeof value === 'number') return Number.isFinite(value)
+  if (typeof value === 'undefined') return true
+  if (typeof value !== 'object') return false
+  if (seen.has(value)) return false
+  seen.add(value)
+  const valid = Array.isArray(value)
+    ? value.every((item) => isJSONCompatible(item, seen))
+    : Object.getPrototypeOf(value) === Object.prototype &&
+      Object.values(value).every((item) => isJSONCompatible(item, seen))
+  seen.delete(value)
+  return valid
+}
+
+const jsonObject = <T extends object>(
+  label: string,
+  check: (value: Record<string, unknown>) => boolean,
+): StandardSchemaV1<T> =>
+  schema(label, (value) => {
+    if (!isRecord(value) || !isJSONCompatible(value) || !check(value)) {
+      return issue(`invalid ${label}`)
+    }
+    return { value: value as T }
+  })
+
+const stringField = (value: Record<string, unknown>, key: string): boolean =>
+  typeof value[key] === 'string' && value[key].length > 0
+
+const optionalStringField = (
+  value: Record<string, unknown>,
+  key: string,
+): boolean => value[key] === undefined || typeof value[key] === 'string'
+
+const messageValue = (value: unknown): value is UIMessage =>
+  isRecord(value) &&
+  typeof value['id'] === 'string' &&
+  (value['role'] === 'user' ||
+    value['role'] === 'assistant' ||
+    value['role'] === 'system') &&
+  Array.isArray(value['parts']) &&
+  value['parts'].every(
+    (part) => isRecord(part) && typeof part['type'] === 'string',
+  ) &&
+  isJSONCompatible(value)
+
+const validateMessage = <M extends UIMessage>(
+  value: unknown,
+  messageSchema?: StandardSchemaV1<M>,
+): SchemaResult<M> => {
+  if (!messageValue(value)) return issue('invalid AI SDK UIMessage')
+  if (!messageSchema) return { value: value as M }
+  return validateSync(messageSchema, value, 'the AI message')
+}
+
+const messagePayloadSchema = <M extends UIMessage>(
+  messageSchema?: StandardSchemaV1<M>,
+): StandardSchemaV1<MessageCreatedPayload<M>> =>
+  schema('ai.message.created', (value) => {
+    if (!isRecord(value)) return issue('invalid ai.message.created payload')
+    const message = validateMessage(value['message'], messageSchema)
+    if (message.issues) return { issues: message.issues }
+    const generate = value['generate']
+    if (generate !== undefined && typeof generate !== 'boolean') {
+      return issue('invalid ai.message.created generate flag')
+    }
+    return {
+      value: {
+        message: message.value,
+        ...(generate === false ? { generate: false } : {}),
+      },
+    }
+  })
+
+const progressPayloadSchema: StandardSchemaV1<GenerationProgressPayload> =
+  schema('ai.generation.progress', (value) => {
+    if (
+      !isRecord(value) ||
+      !isJSONCompatible(value) ||
+      !stringField(value, 'requestId') ||
+      !stringField(value, 'messageId') ||
+      !stringField(value, 'generationId') ||
+      !stringField(value, 'responseMessageId') ||
+      typeof value['sequence'] !== 'number' ||
+      !Number.isInteger(value['sequence']) ||
+      value['sequence'] < 0 ||
+      !Array.isArray(value['chunks']) ||
+      !value['chunks'].every(
+        (chunk) => isRecord(chunk) && typeof chunk['type'] === 'string',
+      )
+    ) {
+      return issue('invalid ai.generation.progress payload')
+    }
+    return { value: value as GenerationProgressPayload }
+  })
+
+const completedPayloadSchema: StandardSchemaV1<GenerationCompletedPayload> =
+  schema('ai.generation.completed', (value) => {
+    if (
+      !isRecord(value) ||
+      !isJSONCompatible(value) ||
+      !stringField(value, 'requestId') ||
+      !stringField(value, 'messageId') ||
+      !stringField(value, 'generationId') ||
+      !stringField(value, 'responseMessageId') ||
+      !optionalStringField(value, 'finishReason')
+    ) {
+      return issue('invalid ai.generation.completed payload')
+    }
+    return { value: value as GenerationCompletedPayload }
+  })
+
+const compactionPayloadSchema = <M extends UIMessage>(
+  messageSchema?: StandardSchemaV1<M>,
+): StandardSchemaV1<CompactionCompletedPayload<M>> =>
+  schema('ai.compaction.completed', (value) => {
+    if (
+      !isRecord(value) ||
+      !isJSONCompatible(value) ||
+      !stringField(value, 'generationId') ||
+      !stringField(value, 'throughMessageId') ||
+      !Array.isArray(value['messages']) ||
+      (value['retainedMessageIds'] !== undefined &&
+        (!Array.isArray(value['retainedMessageIds']) ||
+          !value['retainedMessageIds'].every(
+            (candidate) => typeof candidate === 'string',
+          )))
+    ) {
+      return issue('invalid ai.compaction.completed payload')
+    }
+    const messages: M[] = []
+    for (const candidate of value['messages']) {
+      const result = validateMessage(candidate, messageSchema)
+      if (result.issues) return { issues: result.issues }
+      messages.push(result.value)
+    }
+    return {
+      value: {
+        generationId: value['generationId'] as string,
+        throughMessageId: value['throughMessageId'] as string,
+        messages,
+        ...(value['retainedMessageIds'] === undefined
+          ? {}
+          : { retainedMessageIds: value['retainedMessageIds'] as string[] }),
+      },
+    }
+  })
+
+export function createEvents<M extends UIMessage = UIMessage>(options?: {
+  messageSchema?: StandardSchemaV1<M>
+}): AIEventDefs<M> {
+  const messageSchema = options?.messageSchema
+  if (messageSchema) assertSyncSchema(messageSchema, 'the AI message schema')
+  const definition: AIEventDefs<M> = {
+    'ai.session.created': jsonObject('ai.session.created', () => true),
+    'ai.session.closed': jsonObject('ai.session.closed', (value) =>
+      optionalStringField(value, 'reason'),
+    ),
+    'ai.message.created': messagePayloadSchema(messageSchema),
+    'ai.message.completed': jsonObject('ai.message.completed', (value) =>
+      stringField(value, 'messageId'),
+    ),
+    'ai.message.interrupted': jsonObject(
+      'ai.message.interrupted',
+      (value) =>
+        stringField(value, 'messageId') &&
+        ((stringField(value, 'generationId') &&
+          value['requestId'] === undefined) ||
+          (stringField(value, 'requestId') &&
+            value['generationId'] === undefined)) &&
+        optionalStringField(value, 'reason') &&
+        (value['lastSeenIndex'] === undefined ||
+          (typeof value['lastSeenIndex'] === 'number' &&
+            Number.isInteger(value['lastSeenIndex']) &&
+            value['lastSeenIndex'] >= 0)),
+    ),
+    'ai.generation.requested': jsonObject(
+      'ai.generation.requested',
+      (value) =>
+        stringField(value, 'messageId') &&
+        optionalStringField(value, 'responseMessageId') &&
+        (value['reason'] === 'message' ||
+          value['reason'] === 'tool' ||
+          value['reason'] === 'input' ||
+          value['reason'] === 'retry'),
+    ),
+    'ai.retry.requested': jsonObject(
+      'ai.retry.requested',
+      (value) =>
+        stringField(value, 'messageId') &&
+        stringField(value, 'responseMessageId') &&
+        stringField(value, 'retryId'),
+    ),
+    'ai.generation.started': jsonObject(
+      'ai.generation.started',
+      (value) =>
+        stringField(value, 'requestId') &&
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'responseMessageId') &&
+        typeof value['attempt'] === 'number' &&
+        Number.isInteger(value['attempt']) &&
+        value['attempt'] > 0 &&
+        stringField(value, 'model'),
+    ),
+    'ai.generation.progress': progressPayloadSchema,
+    'ai.generation.completed': completedPayloadSchema,
+    'ai.generation.failed': jsonObject(
+      'ai.generation.failed',
+      (value) =>
+        stringField(value, 'requestId') &&
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'responseMessageId') &&
+        stringField(value, 'error') &&
+        (value['superseded'] === undefined ||
+          typeof value['superseded'] === 'boolean') &&
+        (value['stepLimit'] === undefined ||
+          typeof value['stepLimit'] === 'boolean'),
+    ),
+    'ai.tool.called': jsonObject(
+      'ai.tool.called',
+      (value) =>
+        stringField(value, 'requestId') &&
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'toolCallId') &&
+        stringField(value, 'toolName') &&
+        (value['dynamic'] === undefined ||
+          typeof value['dynamic'] === 'boolean') &&
+        (value['providerExecuted'] === undefined ||
+          typeof value['providerExecuted'] === 'boolean') &&
+        (value['supportsDeferredResults'] === undefined ||
+          typeof value['supportsDeferredResults'] === 'boolean') &&
+        optionalStringField(value, 'title'),
+    ),
+    'ai.tool.result': jsonObject(
+      'ai.tool.result',
+      (value) =>
+        stringField(value, 'requestId') &&
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'toolCallId') &&
+        optionalStringField(value, 'toolName') &&
+        optionalStringField(value, 'error') &&
+        (value['denied'] === undefined ||
+          typeof value['denied'] === 'boolean') &&
+        (value['preliminary'] === undefined ||
+          typeof value['preliminary'] === 'boolean') &&
+        (value['phase'] === undefined ||
+          value['phase'] === 'input' ||
+          value['phase'] === 'execution') &&
+        (value['dynamic'] === undefined ||
+          typeof value['dynamic'] === 'boolean') &&
+        (value['providerExecuted'] === undefined ||
+          typeof value['providerExecuted'] === 'boolean'),
+    ),
+    'ai.approval.requested': jsonObject(
+      'ai.approval.requested',
+      (value) =>
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'approvalId') &&
+        stringField(value, 'toolCallId') &&
+        (value['isAutomatic'] === undefined ||
+          typeof value['isAutomatic'] === 'boolean') &&
+        optionalStringField(value, 'signature'),
+    ),
+    'ai.approval.responded': jsonObject(
+      'ai.approval.responded',
+      (value) =>
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'approvalId') &&
+        typeof value['approved'] === 'boolean' &&
+        optionalStringField(value, 'reason'),
+    ),
+    'ai.input.requested': jsonObject(
+      'ai.input.requested',
+      (value) =>
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'inputId') &&
+        stringField(value, 'name'),
+    ),
+    'ai.input.responded': jsonObject(
+      'ai.input.responded',
+      (value) =>
+        stringField(value, 'messageId') &&
+        stringField(value, 'generationId') &&
+        stringField(value, 'inputId') &&
+        stringField(value, 'name'),
+    ),
+    'ai.compaction.requested': jsonObject(
+      'ai.compaction.requested',
+      (value) =>
+        stringField(value, 'generationId') &&
+        stringField(value, 'throughMessageId'),
+    ),
+    'ai.compaction.completed': compactionPayloadSchema(messageSchema),
+  }
+  return definition
+}
+
+export const events: AIEventDefs<UIMessage> = createEvents()
+
+const initialState = <M extends UIMessage>(): AIState<M> => ({
+  session: { created: false },
+  messages: [],
+  status: 'idle',
+  activeGeneration: null,
+  activeRequestId: null,
+  activeResponseMessageId: null,
+  activeProjection: null,
+  responseGenerationIds: {},
+  pendingApprovals: [],
+  pendingInputs: [],
+  tools: [],
+  compaction: null,
+  usage: [],
+  error: null,
+})
+
+const upsertMessage = <M extends UIMessage>(messages: M[], message: M): M[] => {
+  const index = messages.findIndex((candidate) => candidate.id === message.id)
+  if (index === -1) return [...messages, message]
+  const next = [...messages]
+  next[index] = message
+  return next
+}
+
+const answerApproval = <M extends UIMessage>(
+  messages: M[],
+  response: ApprovalRespondedPayload,
+): M[] =>
+  messages.map((message) => {
+    if (message.id !== response.messageId) return message
+    let changed = false
+    const parts = message.parts.map((part) => {
+      const candidate = part as unknown
+      if (
+        !isRecord(candidate) ||
+        candidate['state'] !== 'approval-requested' ||
+        !isRecord(candidate['approval']) ||
+        candidate['approval']['id'] !== response.approvalId
+      ) {
+        return part
+      }
+      changed = true
+      return {
+        ...part,
+        state: 'approval-responded',
+        approval: {
+          ...candidate['approval'],
+          approved: response.approved,
+          ...(response.reason === undefined ? {} : { reason: response.reason }),
+        },
+      } as typeof part
+    })
+    return changed ? ({ ...message, parts } as M) : message
+  })
+
+const projectMessageChunks = <M extends UIMessage>(
+  messages: M[],
+  messageId: string,
+  chunks: UIMessageChunk[],
+): M[] =>
+  messages.map((message) =>
+    message.id === messageId
+      ? projectUIMessage({ messageId, base: message, chunks })
+      : message,
+  )
+
+const approvalRequestedChunk = (
+  payload: ApprovalRequestedPayload,
+): UIMessageChunk => ({
+  type: 'tool-approval-request',
+  approvalId: payload.approvalId,
+  toolCallId: payload.toolCallId,
+  ...(payload.isAutomatic === undefined
+    ? {}
+    : { isAutomatic: payload.isAutomatic }),
+  ...(payload.signature === undefined ? {} : { signature: payload.signature }),
+})
+
+const approvalRespondedChunk = (
+  payload: ApprovalRespondedPayload,
+): UIMessageChunk => ({
+  type: 'tool-approval-response',
+  approvalId: payload.approvalId,
+  approved: payload.approved,
+  ...(payload.reason === undefined ? {} : { reason: payload.reason }),
+})
+
+const toolResultChunk = (payload: ToolResultPayload): UIMessageChunk => {
+  if (payload.phase === 'input') {
+    return {
+      type: 'tool-input-error',
+      toolCallId: payload.toolCallId,
+      toolName: payload.toolName ?? 'unknown',
+      input: payload.rawInput ?? payload.input,
+      errorText: payload.error ?? 'Invalid tool input',
+      ...(payload.dynamic === undefined ? {} : { dynamic: payload.dynamic }),
+      ...(payload.providerExecuted === undefined
+        ? {}
+        : { providerExecuted: payload.providerExecuted }),
+      ...(payload.providerMetadata === undefined
+        ? {}
+        : { providerMetadata: payload.providerMetadata }),
+      ...(payload.toolMetadata === undefined
+        ? {}
+        : { toolMetadata: payload.toolMetadata }),
+    } as UIMessageChunk
+  }
+  if (payload.denied === true) {
+    return {
+      type: 'tool-output-denied',
+      toolCallId: payload.toolCallId,
+    } as UIMessageChunk
+  }
+  if (payload.error !== undefined) {
+    return {
+      type: 'tool-output-error',
+      toolCallId: payload.toolCallId,
+      errorText: payload.error,
+      ...(payload.dynamic === undefined ? {} : { dynamic: payload.dynamic }),
+      ...(payload.providerExecuted === undefined
+        ? {}
+        : { providerExecuted: payload.providerExecuted }),
+      ...(payload.providerMetadata === undefined
+        ? {}
+        : { providerMetadata: payload.providerMetadata }),
+      ...(payload.toolMetadata === undefined
+        ? {}
+        : { toolMetadata: payload.toolMetadata }),
+    } as UIMessageChunk
+  }
+  return {
+    type: 'tool-output-available',
+    toolCallId: payload.toolCallId,
+    output: payload.output,
+    ...(payload.preliminary === undefined
+      ? {}
+      : { preliminary: payload.preliminary }),
+    ...(payload.dynamic === undefined ? {} : { dynamic: payload.dynamic }),
+    ...(payload.providerExecuted === undefined
+      ? {}
+      : { providerExecuted: payload.providerExecuted }),
+    ...(payload.providerMetadata === undefined
+      ? {}
+      : { providerMetadata: payload.providerMetadata }),
+    ...(payload.toolMetadata === undefined
+      ? {}
+      : { toolMetadata: payload.toolMetadata }),
+  } as UIMessageChunk
+}
+
+const waitingStatus = <M extends UIMessage>(state: AIState<M>): AgentStatus => {
+  if (state.pendingApprovals.length > 0 || state.pendingInputs.length > 0) {
+    return 'waiting'
+  }
+  return state.activeGeneration !== null ||
+    state.tools.some((activity) => activity.status === 'running')
+    ? 'generating'
+    : 'idle'
+}
+
+const lifecycleStatus = <M extends UIMessage>(
+  previous: AIState<M>,
+  next: AIState<M>,
+): AgentStatus =>
+  previous.status === 'failed' ? 'failed' : waitingStatus(next)
+
+const isStaleGeneration = <M extends UIMessage>(
+  state: AIState<M>,
+  messageId: string,
+  generationId: string,
+): boolean => {
+  const current = state.responseGenerationIds[messageId]
+  return current !== undefined && current !== generationId
+}
+
+const projectedMessage = <M extends UIMessage>(
+  projection: ActiveGenerationProjection<M>,
+): M => {
+  const projectionEvents: Array<{ index: number; chunks: UIMessageChunk[] }> = [
+    ...projection.batches,
+  ]
+  for (const event of projection.toolEvents) {
+    if (event.type === 'ai.tool.result') {
+      projectionEvents.push({
+        index: event.index,
+        chunks: [toolResultChunk(event.payload)],
+      })
+    }
+  }
+  for (const event of projection.approvalEvents) {
+    projectionEvents.push({
+      index: event.index,
+      chunks: [
+        event.type === 'ai.approval.requested'
+          ? approvalRequestedChunk(event.payload)
+          : approvalRespondedChunk(event.payload),
+      ],
+    })
+  }
+  return projectUIMessage({
+    messageId: projection.responseMessageId,
+    ...(projection.baseMessage === undefined
+      ? {}
+      : { base: projection.baseMessage }),
+    chunks: projectionEvents
+      .toSorted((left, right) => left.index - right.index)
+      .flatMap((event) => event.chunks),
+  })
+}
+
+const reduceToolActivity = (
+  tools: ToolActivity[],
+  event:
+    | { type: 'ai.tool.called'; payload: ToolCalledPayload }
+    | { type: 'ai.tool.result'; payload: ToolResultPayload },
+): ToolActivity[] => {
+  if (event.type === 'ai.tool.called') {
+    const payload = event.payload
+    const { providerMetadata, ...called } = payload
+    const activity: ToolActivity = {
+      ...called,
+      ...(providerMetadata === undefined
+        ? {}
+        : { callProviderMetadata: providerMetadata }),
+      status: 'running',
+    }
+    return [
+      ...tools.filter(
+        (candidate) =>
+          candidate.generationId !== payload.generationId ||
+          candidate.toolCallId !== payload.toolCallId,
+      ),
+      activity,
+    ]
+  }
+
+  const payload = event.payload
+  const existing = tools.find(
+    (candidate) =>
+      candidate.generationId === payload.generationId &&
+      candidate.toolCallId === payload.toolCallId,
+  )
+  const toolName = payload.toolName ?? existing?.toolName
+  const input = payload.input === undefined ? existing?.input : payload.input
+  const dynamic = payload.dynamic ?? existing?.dynamic
+  const providerExecuted =
+    payload.providerExecuted ?? existing?.providerExecuted
+  const toolMetadata = payload.toolMetadata ?? existing?.toolMetadata
+  const resultProviderMetadata =
+    payload.providerMetadata ?? existing?.resultProviderMetadata
+  const activity: ToolActivity = {
+    requestId: payload.requestId,
+    messageId: payload.messageId,
+    generationId: payload.generationId,
+    toolCallId: payload.toolCallId,
+    ...(toolName === undefined ? {} : { toolName }),
+    ...(input === undefined ? {} : { input }),
+    ...(payload.rawInput === undefined ? {} : { rawInput: payload.rawInput }),
+    ...(payload.output === undefined ? {} : { output: payload.output }),
+    ...(payload.error === undefined ? {} : { error: payload.error }),
+    ...(payload.preliminary === undefined
+      ? {}
+      : { preliminary: payload.preliminary }),
+    ...(payload.phase === undefined ? {} : { phase: payload.phase }),
+    ...(dynamic === undefined ? {} : { dynamic }),
+    ...(providerExecuted === undefined ? {} : { providerExecuted }),
+    ...(existing?.callProviderMetadata === undefined
+      ? {}
+      : { callProviderMetadata: existing.callProviderMetadata }),
+    ...(resultProviderMetadata === undefined ? {} : { resultProviderMetadata }),
+    ...(toolMetadata === undefined ? {} : { toolMetadata }),
+    ...(existing?.title === undefined ? {} : { title: existing.title }),
+    status: payload.denied
+      ? 'denied'
+      : payload.error
+        ? 'failed'
+        : payload.preliminary
+          ? 'running'
+          : 'completed',
+  }
+  return [
+    ...tools.filter(
+      (candidate) =>
+        candidate.generationId !== payload.generationId ||
+        candidate.toolCallId !== payload.toolCallId,
+    ),
+    activity,
+  ]
+}
+
+const projectedTools = <M extends UIMessage>(
+  projection: ActiveGenerationProjection<M>,
+): ToolActivity[] =>
+  projection.toolEvents.reduce(
+    (tools, event) => reduceToolActivity(tools, event),
+    projection.baseTools,
+  )
+
+const restoreGenerationBase = <M extends UIMessage>(
+  state: AIState<M>,
+  generationId: string,
+  responseMessageId: string,
+  reason: 'interrupted' | 'superseded',
+): AIState<M> => {
+  const projection =
+    state.activeProjection?.generationId === generationId
+      ? state.activeProjection
+      : null
+  return {
+    ...state,
+    messages:
+      projection?.baseMessage === undefined
+        ? state.messages.filter((message) => message.id !== responseMessageId)
+        : upsertMessage(state.messages, projection.baseMessage),
+    tools:
+      projection === null
+        ? state.tools.filter(
+            (activity) => activity.generationId !== generationId,
+          )
+        : projection.baseTools,
+    pendingApprovals: state.pendingApprovals.filter(
+      (approval) => approval.generationId !== generationId,
+    ),
+    pendingInputs: state.pendingInputs.filter(
+      (input) => input.generationId !== generationId,
+    ),
+    activeGeneration: null,
+    activeProjection: null,
+    terminalGenerations: {
+      ...state.terminalGenerations,
+      [generationId]: reason,
+    },
+    compaction:
+      projection === null
+        ? state.compaction?.generationId === generationId
+          ? null
+          : state.compaction
+        : (projection.baseCompaction ?? null),
+  }
+}
+
+const stopGenerationTools = (
+  tools: ToolActivity[],
+  generationId: string | undefined,
+  responseMessageId: string,
+  error: string,
+): ToolActivity[] =>
+  tools.map((activity) =>
+    activity.status === 'running' &&
+    (generationId === undefined
+      ? activity.messageId === responseMessageId
+      : activity.generationId === generationId)
+      ? { ...activity, status: 'failed', error, preliminary: false }
+      : activity,
+  )
+
+const interruptGenerationProjection = <M extends UIMessage>(
+  state: AIState<M>,
+  generationId: string,
+  responseMessageId: string,
+  cutoff: number,
+): AIState<M> | null => {
+  const projection =
+    state.activeProjection?.generationId === generationId
+      ? state.activeProjection
+      : null
+  if (projection === null) return null
+  const retained: ActiveGenerationProjection<M> = {
+    ...projection,
+    batches: projection.batches.filter((batch) => batch.index <= cutoff),
+    toolEvents: projection.toolEvents.filter(
+      (toolEvent) => toolEvent.index <= cutoff,
+    ),
+    approvalEvents: projection.approvalEvents.filter(
+      (approvalEvent) => approvalEvent.index <= cutoff,
+    ),
+  }
+  const shouldProject =
+    retained.baseMessage !== undefined || retained.batches.length > 0
+  const restored = restoreGenerationBase(
+    state,
+    generationId,
+    responseMessageId,
+    'interrupted',
+  )
+  return {
+    ...restored,
+    messages: shouldProject
+      ? upsertMessage(
+          restored.messages,
+          interruptUIMessage(projectedMessage(retained)),
+        )
+      : restored.messages,
+    tools: stopGenerationTools(
+      projectedTools(retained),
+      generationId,
+      responseMessageId,
+      'Tool execution was interrupted.',
+    ),
+    usage:
+      projection.completionIndex === undefined ||
+      projection.completionIndex <= cutoff
+        ? restored.usage
+        : restored.usage.filter((entry) => entry.generationId !== generationId),
+  }
+}
+
+const generationRequestId = (generationId: string): string | undefined => {
+  const marker = ':generation:'
+  const markerIndex = generationId.lastIndexOf(marker)
+  return markerIndex === -1 ? undefined : generationId.slice(0, markerIndex)
+}
+
+const terminalGenerationReason = <M extends UIMessage>(
+  state: AIState<M>,
+  generationId: string,
+): 'completed' | 'failed' | 'interrupted' | 'superseded' | undefined => {
+  const terminal = state.terminalGenerations?.[generationId]
+  if (terminal !== undefined) return terminal
+  const requestId = generationRequestId(generationId)
+  return requestId !== undefined &&
+    state.terminalRequestIds?.[requestId] === true
+    ? 'interrupted'
+    : undefined
+}
+
+/** Pure projection of the built-in AI protocol. Unknown extension events are ignored. */
+export function reduceAIState<M extends UIMessage>(
+  state: AIState<M>,
+  event: { type: string; payload: unknown; index: number; id?: string },
+): AIState<M> {
+  if (state.status === 'closed') return state
+  switch (event.type) {
+    case 'ai.session.created': {
+      const payload = event.payload as SessionCreatedPayload
+      return {
+        ...state,
+        session: {
+          created: true,
+          ...(payload.metadata === undefined
+            ? state.session.metadata === undefined
+              ? {}
+              : { metadata: state.session.metadata }
+            : { metadata: payload.metadata }),
+        },
+      }
+    }
+    case 'ai.session.closed':
+      return {
+        ...state,
+        status: 'closed',
+        activeGeneration: null,
+        activeRequestId: null,
+        activeResponseMessageId: null,
+        activeProjection: null,
+      }
+    case 'ai.message.created': {
+      const { message } = event.payload as MessageCreatedPayload<M>
+      return { ...state, messages: upsertMessage(state.messages, message) }
+    }
+    case 'ai.message.completed': {
+      const { messageId } = event.payload as MessageCompletedPayload
+      const generationId = state.responseGenerationIds[messageId]
+      if (
+        generationId === undefined ||
+        terminalGenerationReason(state, generationId) !== undefined
+      ) {
+        return state
+      }
+      const requestId = generationRequestId(generationId)
+      return {
+        ...state,
+        activeRequestId: null,
+        activeResponseMessageId: null,
+        activeProjection:
+          state.activeProjection?.generationId === generationId
+            ? null
+            : state.activeProjection,
+        terminalGenerations: {
+          ...state.terminalGenerations,
+          [generationId]: 'completed',
+        },
+        ...(requestId === undefined
+          ? {}
+          : {
+              terminalRequestIds: {
+                ...state.terminalRequestIds,
+                [requestId]: true,
+              },
+            }),
+      }
+    }
+    case 'ai.generation.requested': {
+      const payload = event.payload as GenerationRequestedPayload
+      const schedulesToolContinuation = payload.reason === 'tool'
+      const responseMessageId =
+        payload.responseMessageId ??
+        (payload.reason === 'message'
+          ? `${payload.messageId}:assistant`
+          : payload.messageId)
+      const ownerGenerationId =
+        state.responseGenerationIds[responseMessageId] ??
+        state.activeGeneration?.generationId
+      const ownerRequestId =
+        ownerGenerationId === undefined
+          ? undefined
+          : generationRequestId(ownerGenerationId)
+      const requestId = (event as { id?: string }).id
+      if (
+        requestId !== undefined &&
+        state.terminalRequestIds?.[requestId] === true
+      ) {
+        return state
+      }
+      const acceptsRequest =
+        payload.reason === 'message' ||
+        (payload.reason === 'tool'
+          ? ownerGenerationId !== undefined &&
+            requestId === `ai.generate:tools:${ownerGenerationId}` &&
+            state.activeRequestId === ownerRequestId &&
+            state.status !== 'failed'
+          : payload.reason === 'retry'
+            ? state.status === 'failed'
+            : ownerGenerationId !== undefined &&
+              requestId?.startsWith(
+                `ai.generate:input:${encodeURIComponent(responseMessageId)}:${encodeURIComponent(ownerGenerationId)}:`,
+              ) === true &&
+              state.activeRequestId === ownerRequestId &&
+              state.status !== 'failed')
+      if (!acceptsRequest) return state
+      return {
+        ...state,
+        messages:
+          payload.reason === 'retry' && payload.responseMessageId !== undefined
+            ? state.messages.filter(
+                (message) => message.id !== payload.responseMessageId,
+              )
+            : state.messages,
+        activeProjection:
+          payload.reason === 'retry' ? null : state.activeProjection,
+        status: schedulesToolContinuation ? state.status : 'generating',
+        activeRequestId: requestId ?? null,
+        activeResponseMessageId: responseMessageId,
+        error: schedulesToolContinuation ? state.error : null,
+      }
+    }
+    case 'ai.generation.started': {
+      const payload = event.payload as GenerationStartedPayload
+      if (
+        state.terminalRequestIds?.[payload.requestId] === true ||
+        terminalGenerationReason(state, payload.generationId) !== undefined
+      ) {
+        return state
+      }
+      if (
+        state.activeRequestId !== null &&
+        state.activeRequestId !== payload.requestId
+      ) {
+        return state
+      }
+      const currentOwner =
+        state.responseGenerationIds[payload.responseMessageId]
+      const currentTerminal =
+        currentOwner === undefined
+          ? undefined
+          : terminalGenerationReason(state, currentOwner)
+      if (
+        currentTerminal === 'completed' ||
+        currentTerminal === 'interrupted'
+      ) {
+        return state
+      }
+      const marker = ':generation:'
+      const currentAttempt = currentOwner?.startsWith(
+        `${payload.requestId}${marker}`,
+      )
+        ? Number(currentOwner.slice(`${payload.requestId}${marker}`.length))
+        : undefined
+      if (
+        currentAttempt !== undefined &&
+        Number.isFinite(currentAttempt) &&
+        currentAttempt >= payload.attempt
+      ) {
+        return state
+      }
+      const activeGeneration = state.activeGeneration
+      const replacesActiveAttempt =
+        activeGeneration !== null &&
+        activeGeneration.requestId === payload.requestId &&
+        activeGeneration.responseMessageId === payload.responseMessageId &&
+        activeGeneration.attempt < payload.attempt &&
+        state.activeProjection?.generationId === activeGeneration.generationId
+      const baseState = replacesActiveAttempt
+        ? restoreGenerationBase(
+            state,
+            activeGeneration.generationId,
+            activeGeneration.responseMessageId,
+            'superseded',
+          )
+        : state
+      const baseMessage = baseState.messages.find(
+        (message) => message.id === payload.responseMessageId,
+      )
+      return {
+        ...baseState,
+        status: 'generating',
+        activeGeneration: payload,
+        activeResponseMessageId: null,
+        responseGenerationIds: {
+          ...baseState.responseGenerationIds,
+          [payload.responseMessageId]: payload.generationId,
+        },
+        activeProjection: {
+          generationId: payload.generationId,
+          responseMessageId: payload.responseMessageId,
+          ...(baseMessage === undefined ? {} : { baseMessage }),
+          baseCompaction: baseState.compaction,
+          batches: [],
+          baseTools: baseState.tools,
+          toolEvents: [],
+          approvalEvents: [],
+        },
+        error: null,
+      }
+    }
+    case 'ai.generation.progress': {
+      const payload = event.payload as GenerationProgressPayload
+      if (terminalGenerationReason(state, payload.generationId) !== undefined) {
+        return state
+      }
+      const activeProjection = state.activeProjection
+      if (
+        state.activeGeneration?.generationId !== payload.generationId ||
+        activeProjection?.generationId !== payload.generationId
+      ) {
+        return state
+      }
+      const projection: ActiveGenerationProjection<M> = {
+        ...activeProjection,
+        batches: [
+          ...activeProjection.batches.filter(
+            (batch) => batch.index !== event.index,
+          ),
+          { index: event.index, chunks: payload.chunks },
+        ].toSorted((left, right) => left.index - right.index),
+      }
+      return {
+        ...state,
+        messages: upsertMessage(state.messages, projectedMessage(projection)),
+        activeProjection: projection,
+        status: 'generating',
+      }
+    }
+    case 'ai.generation.completed': {
+      const payload = event.payload as GenerationCompletedPayload
+      if (terminalGenerationReason(state, payload.generationId) !== undefined) {
+        return state
+      }
+      if (state.activeGeneration?.generationId !== payload.generationId) {
+        return state
+      }
+      const usage = payload.usage
+        ? [
+            ...state.usage.filter(
+              (entry) => entry.generationId !== payload.generationId,
+            ),
+            { generationId: payload.generationId, usage: payload.usage },
+          ]
+        : state.usage
+      const next = {
+        ...state,
+        activeGeneration: null,
+        activeProjection:
+          state.activeProjection?.generationId === payload.generationId
+            ? { ...state.activeProjection, completionIndex: event.index }
+            : state.activeProjection,
+        usage,
+      }
+      return { ...next, status: waitingStatus(next) }
+    }
+    case 'ai.generation.failed': {
+      const payload = event.payload as GenerationFailedPayload
+      if (terminalGenerationReason(state, payload.generationId) !== undefined) {
+        return state
+      }
+      if (
+        state.activeGeneration?.generationId !== payload.generationId &&
+        !(
+          payload.stepLimit === true &&
+          state.responseGenerationIds[payload.responseMessageId] ===
+            payload.generationId &&
+          state.activeRequestId ===
+            `ai.generate:tools:${payload.generationId}` &&
+          state.activeResponseMessageId === payload.responseMessageId
+        )
+      ) {
+        return state
+      }
+      if (payload.superseded === true) {
+        return restoreGenerationBase(
+          state,
+          payload.generationId,
+          payload.responseMessageId,
+          'superseded',
+        )
+      }
+      const projection =
+        state.activeProjection?.generationId === payload.generationId
+          ? state.activeProjection
+          : null
+      const shouldProject =
+        projection !== null &&
+        (projection.baseMessage !== undefined || projection.batches.length > 0)
+      return {
+        ...state,
+        messages: shouldProject
+          ? upsertMessage(
+              state.messages,
+              interruptUIMessage(
+                projectedMessage(projection),
+                `Generation failed before tool completion: ${payload.error}`,
+              ),
+            )
+          : state.messages,
+        status: 'failed',
+        activeGeneration: null,
+        activeResponseMessageId: null,
+        activeProjection: projection,
+        terminalGenerations: {
+          ...state.terminalGenerations,
+          [payload.generationId]: 'failed',
+        },
+        terminalRequestIds: {
+          ...state.terminalRequestIds,
+          [payload.requestId]: true,
+          ...(state.activeRequestId === null
+            ? {}
+            : { [state.activeRequestId]: true as const }),
+        },
+        compaction:
+          projection?.generationId === payload.generationId
+            ? (projection.baseCompaction ?? null)
+            : state.compaction?.generationId === payload.generationId
+              ? null
+              : state.compaction,
+        pendingApprovals: state.pendingApprovals.filter(
+          (approval) => approval.generationId !== payload.generationId,
+        ),
+        pendingInputs: state.pendingInputs.filter(
+          (input) => input.generationId !== payload.generationId,
+        ),
+        tools: stopGenerationTools(
+          state.tools,
+          payload.generationId,
+          payload.responseMessageId,
+          payload.error,
+        ),
+        error: payload.error,
+      }
+    }
+    case 'ai.message.interrupted': {
+      const payload = event.payload as MessageInterruptedPayload
+      const currentGenerationId = state.responseGenerationIds[payload.messageId]
+      const targetRequestId =
+        payload.requestId ?? generationRequestId(payload.generationId)
+      const ownsRequestPhase =
+        payload.requestId !== undefined &&
+        state.activeGeneration === null &&
+        state.activeRequestId === payload.requestId &&
+        state.activeResponseMessageId === payload.messageId
+      if (
+        currentGenerationId !== undefined &&
+        !ownsRequestPhase &&
+        (payload.generationId !== undefined
+          ? currentGenerationId !== payload.generationId
+          : generationRequestId(currentGenerationId) !== payload.requestId)
+      ) {
+        return state
+      }
+      if (
+        state.activeGeneration === null &&
+        state.activeRequestId !== null &&
+        state.activeResponseMessageId === payload.messageId &&
+        targetRequestId !== state.activeRequestId
+      ) {
+        return state
+      }
+      const currentTerminal =
+        currentGenerationId === undefined
+          ? undefined
+          : terminalGenerationReason(state, currentGenerationId)
+      if (
+        currentTerminal === 'completed' ||
+        currentTerminal === 'interrupted' ||
+        currentTerminal === 'superseded'
+      ) {
+        return state
+      }
+      const active = state.activeGeneration
+      if (active === null) {
+        if (
+          payload.requestId !== undefined &&
+          state.activeRequestId !== null &&
+          state.activeRequestId === payload.requestId &&
+          state.activeResponseMessageId === payload.messageId
+        ) {
+          const projectedState =
+            currentGenerationId === undefined
+              ? null
+              : interruptGenerationProjection(
+                  state,
+                  currentGenerationId,
+                  payload.messageId,
+                  payload.lastSeenIndex ?? Number.POSITIVE_INFINITY,
+                )
+          const baseState = projectedState ?? state
+          return {
+            ...baseState,
+            status: 'idle',
+            activeRequestId: null,
+            activeResponseMessageId: null,
+            activeProjection: null,
+            terminalRequestIds: {
+              ...baseState.terminalRequestIds,
+              [payload.requestId]: true,
+            },
+            ...(currentGenerationId === undefined
+              ? {}
+              : {
+                  terminalGenerations: {
+                    ...baseState.terminalGenerations,
+                    [currentGenerationId]: 'interrupted' as const,
+                  },
+                }),
+            error: null,
+          }
+        }
+        const targetGenerationId =
+          payload.generationId ??
+          (currentGenerationId !== undefined &&
+          generationRequestId(currentGenerationId) === payload.requestId
+            ? currentGenerationId
+            : undefined)
+        if (targetGenerationId === undefined) return state
+        if (currentGenerationId !== targetGenerationId) return state
+        const matchesGeneration = (generationId: string): boolean =>
+          targetGenerationId === generationId
+        const hasMessageWork =
+          state.pendingApprovals.some(
+            (approval) =>
+              approval.messageId === payload.messageId &&
+              matchesGeneration(approval.generationId),
+          ) ||
+          state.pendingInputs.some(
+            (input) =>
+              input.messageId === payload.messageId &&
+              matchesGeneration(input.generationId),
+          ) ||
+          state.tools.some(
+            (activity) =>
+              activity.messageId === payload.messageId &&
+              activity.status === 'running' &&
+              matchesGeneration(activity.generationId),
+          )
+        const resolvesFailedGeneration =
+          state.status === 'failed' &&
+          currentGenerationId === targetGenerationId &&
+          (state.activeRequestId === null ||
+            state.activeRequestId === targetRequestId ||
+            state.activeRequestId === `ai.generate:tools:${targetGenerationId}`)
+        const resolvesCurrentGeneration =
+          currentGenerationId === targetGenerationId &&
+          currentTerminal === undefined
+        if (
+          !hasMessageWork &&
+          !resolvesFailedGeneration &&
+          !resolvesCurrentGeneration
+        ) {
+          return state
+        }
+        const pendingRequestId = state.activeRequestId
+        const projectedState = interruptGenerationProjection(
+          state,
+          targetGenerationId,
+          payload.messageId,
+          payload.lastSeenIndex ?? Number.POSITIVE_INFINITY,
+        )
+        const baseState = projectedState ?? state
+        const next = {
+          ...baseState,
+          activeRequestId: null,
+          activeResponseMessageId: null,
+          activeProjection: null,
+          messages:
+            projectedState === null
+              ? state.messages.map((message) =>
+                  message.id === payload.messageId
+                    ? interruptUIMessage(message)
+                    : message,
+                )
+              : baseState.messages,
+          pendingApprovals: baseState.pendingApprovals.filter(
+            (approval) =>
+              approval.messageId !== payload.messageId ||
+              !matchesGeneration(approval.generationId),
+          ),
+          pendingInputs: baseState.pendingInputs.filter(
+            (input) =>
+              input.messageId !== payload.messageId ||
+              !matchesGeneration(input.generationId),
+          ),
+          tools:
+            projectedState === null
+              ? stopGenerationTools(
+                  state.tools,
+                  targetGenerationId,
+                  payload.messageId,
+                  'Tool execution was interrupted.',
+                )
+              : baseState.tools,
+          terminalGenerations: {
+            ...baseState.terminalGenerations,
+            [targetGenerationId]: 'interrupted' as const,
+          },
+          ...(pendingRequestId === null
+            ? {}
+            : {
+                terminalRequestIds: {
+                  ...baseState.terminalRequestIds,
+                  [pendingRequestId]: true as const,
+                },
+              }),
+          error: resolvesFailedGeneration ? null : state.error,
+        }
+        return {
+          ...next,
+          status: waitingStatus(next),
+        }
+      }
+      if (
+        payload.messageId !== active.responseMessageId ||
+        (payload.generationId !== active.generationId &&
+          payload.requestId !== active.requestId)
+      ) {
+        return state
+      }
+      const cutoff = payload.lastSeenIndex ?? Number.POSITIVE_INFINITY
+      const projectedState = interruptGenerationProjection(
+        state,
+        active.generationId,
+        active.responseMessageId,
+        cutoff,
+      )
+      const baseState =
+        projectedState ??
+        restoreGenerationBase(
+          state,
+          active.generationId,
+          active.responseMessageId,
+          'interrupted',
+        )
+      const next = {
+        ...baseState,
+        activeRequestId: null,
+        activeResponseMessageId: null,
+        terminalRequestIds: {
+          ...baseState.terminalRequestIds,
+          [active.requestId]: true as const,
+        },
+        tools:
+          projectedState === null
+            ? stopGenerationTools(
+                baseState.tools,
+                active.generationId,
+                active.responseMessageId,
+                'Tool execution was interrupted.',
+              )
+            : baseState.tools,
+      }
+      return { ...next, status: lifecycleStatus(state, next) }
+    }
+    case 'ai.tool.called': {
+      const payload = event.payload as ToolCalledPayload
+      const projection = state.activeProjection
+      if (
+        isStaleGeneration(state, payload.messageId, payload.generationId) ||
+        terminalGenerationReason(state, payload.generationId) !== undefined
+      ) {
+        return state
+      }
+      if (
+        state.tools.some(
+          (activity) =>
+            activity.generationId === payload.generationId &&
+            activity.toolCallId === payload.toolCallId &&
+            activity.status === 'failed' &&
+            state.activeGeneration?.generationId !== payload.generationId,
+        )
+      ) {
+        return state
+      }
+      if (projection?.generationId !== payload.generationId) {
+        const next = {
+          ...state,
+          tools: reduceToolActivity(state.tools, {
+            type: 'ai.tool.called',
+            payload,
+          }),
+        }
+        return { ...next, status: lifecycleStatus(state, next) }
+      }
+      const activeProjection: ActiveGenerationProjection<M> = {
+        ...projection,
+        toolEvents: [
+          ...projection.toolEvents,
+          { index: event.index, type: 'ai.tool.called', payload },
+        ],
+      }
+      const next = {
+        ...state,
+        activeProjection,
+        tools: projectedTools(activeProjection),
+      }
+      return { ...next, status: lifecycleStatus(state, next) }
+    }
+    case 'ai.tool.result': {
+      const payload = event.payload as ToolResultPayload
+      const projection = state.activeProjection
+      const terminalReason = terminalGenerationReason(
+        state,
+        payload.generationId,
+      )
+      const completesFailedTool =
+        terminalReason === 'failed' &&
+        state.status === 'failed' &&
+        payload.preliminary !== true &&
+        projection?.generationId === payload.generationId &&
+        projectedTools(projection).some(
+          (activity) =>
+            activity.generationId === payload.generationId &&
+            activity.toolCallId === payload.toolCallId &&
+            activity.status === 'running',
+        )
+      if (
+        isStaleGeneration(state, payload.messageId, payload.generationId) ||
+        (terminalReason !== undefined && !completesFailedTool)
+      ) {
+        return state
+      }
+      if (
+        state.tools.some(
+          (activity) =>
+            activity.generationId === payload.generationId &&
+            activity.toolCallId === payload.toolCallId &&
+            activity.status === 'failed' &&
+            state.activeGeneration?.generationId !== payload.generationId &&
+            !completesFailedTool,
+        )
+      ) {
+        return state
+      }
+      if (projection?.generationId !== payload.generationId) {
+        const next = {
+          ...state,
+          messages: projectMessageChunks(state.messages, payload.messageId, [
+            toolResultChunk(payload),
+          ]),
+          tools: reduceToolActivity(state.tools, {
+            type: 'ai.tool.result',
+            payload,
+          }),
+        }
+        return { ...next, status: lifecycleStatus(state, next) }
+      }
+      const activeProjection: ActiveGenerationProjection<M> = {
+        ...projection,
+        toolEvents: [
+          ...projection.toolEvents,
+          { index: event.index, type: 'ai.tool.result', payload },
+        ],
+      }
+      const projected = projectedMessage(activeProjection)
+      const tools = projectedTools(activeProjection)
+      const failureError = state.error ?? 'Generation failed.'
+      const next = {
+        ...state,
+        messages: upsertMessage(
+          state.messages,
+          terminalReason === 'failed'
+            ? interruptUIMessage(
+                projected,
+                `Generation failed before tool completion: ${failureError}`,
+              )
+            : projected,
+        ),
+        activeProjection,
+        tools:
+          terminalReason === 'failed'
+            ? stopGenerationTools(
+                tools,
+                payload.generationId,
+                payload.messageId,
+                failureError,
+              )
+            : tools,
+      }
+      return { ...next, status: lifecycleStatus(state, next) }
+    }
+    case 'ai.approval.requested': {
+      const payload = event.payload as ApprovalRequestedPayload
+      const projection = state.activeProjection
+      if (
+        isStaleGeneration(state, payload.messageId, payload.generationId) ||
+        terminalGenerationReason(state, payload.generationId) !== undefined
+      ) {
+        return state
+      }
+      if (
+        state.tools.some(
+          (activity) =>
+            activity.generationId === payload.generationId &&
+            activity.toolCallId === payload.toolCallId &&
+            activity.status === 'failed' &&
+            state.activeGeneration?.generationId !== payload.generationId,
+        )
+      ) {
+        return state
+      }
+      const activeProjection =
+        projection?.generationId === payload.generationId
+          ? {
+              ...projection,
+              approvalEvents: [
+                ...projection.approvalEvents,
+                {
+                  index: event.index,
+                  type: 'ai.approval.requested' as const,
+                  payload,
+                },
+              ],
+            }
+          : projection
+      const next = {
+        ...state,
+        messages:
+          activeProjection === null || activeProjection === projection
+            ? projectMessageChunks(state.messages, payload.messageId, [
+                approvalRequestedChunk(payload),
+              ])
+            : upsertMessage(state.messages, projectedMessage(activeProjection)),
+        activeProjection,
+        pendingApprovals:
+          payload.isAutomatic === true
+            ? state.pendingApprovals
+            : [
+                ...state.pendingApprovals.filter(
+                  (approval) =>
+                    approval.approvalId !== payload.approvalId ||
+                    approval.messageId !== payload.messageId,
+                ),
+                payload,
+              ],
+      }
+      return {
+        ...next,
+        status: lifecycleStatus(state, next),
+      }
+    }
+    case 'ai.approval.responded': {
+      const payload = event.payload as ApprovalRespondedPayload
+      if (
+        isStaleGeneration(state, payload.messageId, payload.generationId) ||
+        terminalGenerationReason(state, payload.generationId) !== undefined
+      ) {
+        return state
+      }
+      const projection = state.activeProjection
+      const pendingRequest = state.pendingApprovals.find(
+        (candidate) =>
+          candidate.approvalId === payload.approvalId &&
+          candidate.messageId === payload.messageId &&
+          candidate.generationId === payload.generationId,
+      )
+      const activeRequest = projection?.approvalEvents.findLast(
+        (candidate) =>
+          candidate.type === 'ai.approval.requested' &&
+          candidate.payload.approvalId === payload.approvalId &&
+          candidate.payload.messageId === payload.messageId &&
+          candidate.payload.generationId === payload.generationId,
+      )
+      if (
+        state.status === 'failed' &&
+        pendingRequest === undefined &&
+        activeRequest === undefined
+      ) {
+        return state
+      }
+      const activeProjection =
+        projection !== null && activeRequest !== undefined
+          ? {
+              ...projection,
+              approvalEvents: [
+                ...projection.approvalEvents,
+                {
+                  index: event.index,
+                  type: 'ai.approval.responded' as const,
+                  payload,
+                },
+              ],
+            }
+          : projection
+      const next = {
+        ...state,
+        messages:
+          activeProjection === null || activeProjection === projection
+            ? answerApproval(state.messages, payload)
+            : upsertMessage(state.messages, projectedMessage(activeProjection)),
+        activeProjection,
+        pendingApprovals: state.pendingApprovals.filter(
+          (approval) =>
+            approval.approvalId !== payload.approvalId ||
+            approval.messageId !== payload.messageId ||
+            approval.generationId !== payload.generationId,
+        ),
+      }
+      return {
+        ...next,
+        status: lifecycleStatus(state, next),
+      }
+    }
+    case 'ai.input.requested': {
+      const payload = event.payload as InputRequestedPayload
+      if (
+        isStaleGeneration(state, payload.messageId, payload.generationId) ||
+        terminalGenerationReason(state, payload.generationId) !== undefined
+      ) {
+        return state
+      }
+      const next = {
+        ...state,
+        pendingInputs: [
+          ...state.pendingInputs.filter(
+            (input) =>
+              input.messageId !== payload.messageId ||
+              input.generationId !== payload.generationId ||
+              input.inputId !== payload.inputId,
+          ),
+          payload,
+        ],
+      }
+      return { ...next, status: lifecycleStatus(state, next) }
+    }
+    case 'ai.input.responded': {
+      const payload = event.payload as InputRespondedPayload
+      if (
+        isStaleGeneration(state, payload.messageId, payload.generationId) ||
+        terminalGenerationReason(state, payload.generationId) !== undefined
+      ) {
+        return state
+      }
+      const next = {
+        ...state,
+        pendingInputs: state.pendingInputs.filter(
+          (input) =>
+            input.inputId !== payload.inputId ||
+            input.messageId !== payload.messageId ||
+            input.generationId !== payload.generationId,
+        ),
+      }
+      return { ...next, status: lifecycleStatus(state, next) }
+    }
+    case 'ai.compaction.requested': {
+      const payload = event.payload as CompactionRequestedPayload
+      if (state.activeGeneration?.generationId !== payload.generationId) {
+        return state
+      }
+      return {
+        ...state,
+        compaction: { status: 'running', ...payload },
+      }
+    }
+    case 'ai.compaction.completed': {
+      const payload = event.payload as CompactionCompletedPayload<M>
+      if (
+        state.activeGeneration?.generationId !== payload.generationId ||
+        state.compaction?.status !== 'running' ||
+        state.compaction.generationId !== payload.generationId
+      ) {
+        return state
+      }
+      return {
+        ...state,
+        compaction: { status: 'completed', ...payload },
+      }
+    }
+    default:
+      return state
+  }
+}
+
+/** Derive AI SDK UI messages directly from an A2 event history. */
+export function deriveUIMessages<M extends UIMessage = UIMessage>(
+  history: readonly { type: string; payload: unknown; index: number }[],
+): M[] {
+  let state = initialState<M>()
+  for (const event of history) state = reduceAIState(state, event)
+  return state.messages
+}
+
+/** Build the standard AI state projection for any contract containing the protocol. */
+export function createReducer<D extends EventDefs>(options: {
+  contract: Contract<D>
+  name?: string
+}): Reducer<D, AIState<AIMessageFrom<D>>> {
+  type M = AIMessageFrom<D>
+  return options.contract
+    .reducer({
+      name: options.name ?? 'a2.ai.state.v8',
+      initialState: initialState<M>(),
+    })
+    .fold((state, event) => reduceAIState(state, event))
+}
+
+type ProtocolInput<
+  M extends UIMessage,
+  K extends keyof AIEventDefs<M> & string,
+> = AppendInput<Pick<AIEventDefs<M>, K>>
+
+export type AIInputs = {
+  /** Record a message. User messages schedule unless generate is false. */
+  message<M extends UIMessage>(
+    message: M,
+    options?: { generate?: boolean },
+  ): ProtocolInput<M, 'ai.message.created'>[]
+  /** Record a message without requesting a generation. */
+  seed<M extends UIMessage>(
+    message: M,
+  ): ProtocolInput<M, 'ai.message.created'>[]
+  /** Answer an AI SDK tool approval. The tool barrier owns continuation. */
+  approval(
+    response: ApprovalRespondedPayload,
+  ): ProtocolInput<UIMessage, 'ai.approval.responded'>[]
+  /** Answer an application-defined input request. */
+  input(
+    response: InputRespondedPayload,
+  ): ProtocolInput<UIMessage, 'ai.input.responded'>[]
+  /** Record an application-defined request for input. */
+  requestInput(
+    request: InputRequestedPayload,
+  ): ProtocolInput<UIMessage, 'ai.input.requested'>[]
+  /** Request a fresh attempt for a failed response. */
+  retry(
+    options: RetryGenerationOptions,
+  ): ProtocolInput<UIMessage, 'ai.retry.requested'>[]
+  /** Interrupt the active response for a message. */
+  interrupt(
+    interruption: MessageInterruptedPayload,
+  ): ProtocolInput<UIMessage, 'ai.message.interrupted'>[]
+}
+
+const messageInput = <M extends UIMessage>(
+  message: M,
+  generate?: boolean,
+): ProtocolInput<M, 'ai.message.created'>[] => [
+  {
+    type: 'ai.message.created',
+    id: `ai.message:${message.id}`,
+    payload: {
+      message,
+      ...(generate === false ? { generate: false } : {}),
+    },
+  },
+]
+
+/** Pure append inputs for the built-in AI protocol. */
+export const inputs: AIInputs = {
+  message<M extends UIMessage>(message: M, options?: { generate?: boolean }) {
+    return messageInput(message, options?.generate)
+  },
+  seed<M extends UIMessage>(message: M) {
+    return messageInput(message, false)
+  },
+  approval(response) {
+    return [
+      {
+        type: 'ai.approval.responded',
+        id: `ai.approval:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.approvalId)}:response`,
+        payload: response,
+      },
+    ]
+  },
+  input(response) {
+    return [
+      {
+        type: 'ai.input.responded',
+        id: `ai.input:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.inputId)}:response`,
+        payload: response,
+      },
+    ]
+  },
+  requestInput(request) {
+    return [
+      {
+        type: 'ai.input.requested',
+        id: `ai.input:${encodeURIComponent(request.messageId)}:${encodeURIComponent(request.generationId)}:${encodeURIComponent(request.inputId)}:request`,
+        payload: request,
+      },
+    ]
+  },
+  retry(retryOptions) {
+    return [
+      {
+        type: 'ai.retry.requested',
+        id: `ai.retry:${retryOptions.retryId}`,
+        payload: retryOptions,
+      },
+    ]
+  },
+  interrupt(interruption) {
+    return [
+      {
+        type: 'ai.message.interrupted',
+        id: `ai.interrupt:${interruption.generationId ?? interruption.requestId}`,
+        payload: interruption,
+      },
+    ]
+  },
+}
+
+/**
+ * Read the durable handler context ambient to the current AI tool
+ * execution. A2 runs each local tool's `execute` inside this scope, so
+ * plain AI SDK `tool()` definitions reach their `event`, `attempt`,
+ * `session`, and `signal` without an A2-specific wrapper. The agent (or
+ * contract) argument carries the event types and is verified against
+ * the executing server's contract.
+ */
+export function handlerContext<D extends AgentToolEventDefs>(
+  agentOrContract: Contract<D> | { readonly contract: Contract<D> },
+): AgentToolContext<D> {
+  const contract =
+    'contract' in agentOrContract ? agentOrContract.contract : agentOrContract
+  const scope = ambientToolScopeStorage()?.getStore()
+  if (scope === undefined) {
+    throw new TypeError(
+      `handlerContext('${contract.name}') requires an ambient A2 tool execution: call it inside a local tool's execute, or wrap direct tool tests with runWithHandlerContext`,
+    )
+  }
+  if (scope.contract !== undefined && scope.contract.name !== contract.name) {
+    throw new TypeError(
+      `handlerContext('${contract.name}') was called during a tool execution for agent '${scope.contract.name}'`,
+    )
+  }
+  return scope.context as AgentToolContext<D>
+}
+
+/**
+ * Define an agent: built-in AI protocol + application events and one state
+ * reducer. The result is isomorphic.
+ */
+export function agent<
+  M extends UIMessage = UIMessage,
+  E extends EventDefs = Record<never, never>,
+>(options: AgentOptions<M, E>): AgentDefinition<M, AIEventDefs<M> & E> {
+  const builtIns = createEvents(
+    options.messageSchema === undefined
+      ? undefined
+      : { messageSchema: options.messageSchema },
+  )
+  const extensions = options.events ?? ({} as E)
+  for (const key of Object.keys(extensions)) {
+    if (Object.hasOwn(builtIns, key)) {
+      throw new TypeError(
+        `agent event '${key}' conflicts with a built-in event`,
+      )
+    }
+  }
+  const agentContract = createContract({
+    name: options.name,
+    events: { ...builtIns, ...extensions },
+  })
+  const reducer = agentContract
+    .reducer({
+      name: options.reducerName ?? 'a2.ai.state.v8',
+      initialState: initialState<M>(),
+    })
+    .fold((state, event) => reduceAIState(state, event))
+
+  const definition = {
+    contract: agentContract,
+    reducer,
+  }
+  return definition as unknown as AgentDefinition<M, AIEventDefs<M> & E>
+}
+
+export type {
+  FinishReason,
+  JSONValue,
+  LanguageModelUsage,
+  UIMessage,
+  UIMessageChunk,
+}