experimental-a2 0.12.0 → 0.14.0

package/src/ai.ts

@@ -13,19 +13,54 @@ import type {
   UIMessage,
   UIMessageChunk,
 } from 'ai'
+import {
+  isQueueCommand,
+  projectAIClientState,
+  type PendingQueueCommand,
+} from './ai-client-state.ts'
+export type { PendingQueueCommand } from './ai-client-state.ts'
+import { applyControlChanges } from './ai-control-state.ts'
+import { initialControlState } from './ai-control.ts'
+import type {
+  ControlCommit,
+  ActiveTurn,
+  ControlCommand,
+  ControlReceipt,
+  InboxItem,
+  ToolExecutionRequest,
+  WorkReport,
+} from './ai-control.ts'
+export type {
+  ActiveTurn,
+  ControlCommand,
+  ControlReceipt,
+  InboxItem,
+} from './ai-control.ts'
 import { contract as createContract } from './contract.ts'
-import type { AppendInput, Contract, EventDefs } from './contract.ts'
-import { interruptUIMessage, projectUIMessage } from './ai-projector.ts'
+import type { Contract, EventDefs } from './contract.ts'
+import {
+  interruptUIMessage,
+  pauseUIMessage,
+  projectUIMessage,
+  projectUIMessageStream,
+  type UIMessageCursor,
+} from './ai-projector.ts'
 import { ambientToolScopeStorage } from './internal.ts'
-import type { Reducer } from './reducer.ts'
+import { clientStateProjection, type Reducer } from './reducer.ts'
+import {
+  appendProgressBatch,
+  flattenProgressBatches,
+  progressBatches,
+  type ProgressBatches,
+} from './ai-progress-batches.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'
+  'idle' | 'generating' | 'waiting' | 'failed' | 'closed' | 'paused'
 
-export type GenerationReason = 'message' | 'tool' | 'input' | 'retry'
+export type GenerationReason = 'message' | 'tool' | 'input' | 'retry' | 'resume'
 
 export type SessionCreatedPayload = { metadata?: JSONValue }
 export type SessionClosedPayload = { reason?: string }
@@ -47,6 +82,7 @@ export type MessageInterruptedPayload = MessageInterruption &
     | { requestId: string; generationId?: never }
   )
 export type GenerationRequestedPayload = {
+  control?: { turnId: string; version: number }
   messageId: string
   reason: GenerationReason
   responseMessageId?: string
@@ -64,6 +100,7 @@ export type GenerationStartedPayload = {
   responseMessageId: string
   attempt: number
   model: string
+  promptThroughIndex?: number
 }
 export type GenerationProgressPayload = {
   requestId: string
@@ -80,6 +117,7 @@ export type GenerationCompletedPayload = {
   responseMessageId: string
   finishReason?: FinishReason
   usage?: LanguageModelUsage
+  inputTokenEstimate?: number
 }
 export type GenerationFailedPayload = {
   requestId: string
@@ -152,18 +190,45 @@ export type InputRespondedPayload = {
   name: string
   value: JSONValue
 }
+export type ModelLimits = {
+  contextWindow: number
+  maxOutputTokens: number
+}
+export type ModelMetadataRequestedPayload = { modelId: string }
+export type ModelMetadataResolvedPayload = {
+  modelId: string
+  limits: ModelLimits | null
+}
+export type ModelMetadataState =
+  | { status: 'pending' }
+  | { status: 'resolved'; limits: ModelLimits }
+  | { status: 'unavailable' }
+
 export type CompactionRequestedPayload = {
   generationId: string
   throughMessageId: string
+  throughIndex?: number
 }
 export type CompactionCompletedPayload<M extends UIMessage> = {
   generationId: string
   throughMessageId: string
+  throughIndex?: number
   messages: M[]
   retainedMessageIds?: string[]
+  summary?: string
+  usage?: LanguageModelUsage
 }
 
 export type AIEventDefs<M extends UIMessage = UIMessage> = {
+  'ai.control.requested': StandardSchemaV1<ControlCommand<M>>
+  'ai.control.committed': StandardSchemaV1<ControlCommit>
+  'ai.control.decided': StandardSchemaV1<ControlReceipt>
+  'ai.work.reported': StandardSchemaV1<WorkReport>
+  'ai.tool.execution.requested': StandardSchemaV1<ToolExecutionRequest>
+  'ai.generation.paused': StandardSchemaV1<{
+    generationId: string
+    messageId: string
+  }>
   'ai.session.created': StandardSchemaV1<SessionCreatedPayload>
   'ai.session.closed': StandardSchemaV1<SessionClosedPayload>
   'ai.message.created': StandardSchemaV1<MessageCreatedPayload<M>>
@@ -183,6 +248,8 @@ export type AIEventDefs<M extends UIMessage = UIMessage> = {
   'ai.input.responded': StandardSchemaV1<InputRespondedPayload>
   'ai.compaction.requested': StandardSchemaV1<CompactionRequestedPayload>
   'ai.compaction.completed': StandardSchemaV1<CompactionCompletedPayload<M>>
+  'ai.model.metadata.requested': StandardSchemaV1<ModelMetadataRequestedPayload>
+  'ai.model.metadata.resolved': StandardSchemaV1<ModelMetadataResolvedPayload>
 }
 
 export type ActiveGeneration = GenerationStartedPayload
@@ -193,8 +260,10 @@ export type ActiveGenerationProjection<M extends UIMessage = UIMessage> = {
   /** Log index of the completed model step while its response remains interruptible. */
   completionIndex?: number
   baseMessage?: M
+  promptThroughIndex?: number
+  cursor?: { throughIndex: number; stream: UIMessageCursor } | null
   baseCompaction?: CompactionState<M> | null
-  batches: Array<{ index: number; chunks: UIMessageChunk[] }>
+  batches: ProgressBatches
   baseTools: ToolActivity[]
   toolEvents: Array<
     | { index: number; type: 'ai.tool.called'; payload: ToolCalledPayload }
@@ -242,8 +311,11 @@ export type CompactionState<M extends UIMessage> = {
   status: 'running' | 'completed'
   generationId: string
   throughMessageId: string
+  throughIndex?: number
   messages?: M[]
   retainedMessageIds?: string[]
+  summary?: string
+  usage?: LanguageModelUsage
 }
 
 export type GenerationUsage = {
@@ -252,6 +324,12 @@ export type GenerationUsage = {
 }
 
 export type AIState<M extends UIMessage = UIMessage> = {
+  inbox: { paused: boolean; items: InboxItem<M>[] }
+  pendingQueueCommands?: PendingQueueCommand<M>[]
+  active: ActiveTurn<M> | null
+  starting: boolean
+  receipt: ControlReceipt | null
+  pausedGenerationIds?: Record<string, true>
   session: { created: boolean; metadata?: JSONValue }
   messages: M[]
   status: AgentStatus
@@ -269,6 +347,7 @@ export type AIState<M extends UIMessage = UIMessage> = {
   pendingInputs: PendingInput[]
   tools: ToolActivity[]
   compaction: CompactionState<M> | null
+  modelMetadata: Record<string, ModelMetadataState>
   usage: GenerationUsage[]
   error: string | null
 }
@@ -282,12 +361,13 @@ export type AgentDefinition<
 }
 
 export type AgentToolEventDefs = EventDefs &
-  Pick<AIEventDefs, 'ai.tool.called' | 'ai.approval.responded'>
+  Pick<AIEventDefs, 'ai.tool.execution.requested'>
 
 /** 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 AgentToolContext<D extends AgentToolEventDefs> = HandlerContext<
+  D,
+  'ai.tool.execution.requested'
+>
 
 export type AgentOptions<M extends UIMessage, E extends EventDefs> = {
   name: string
@@ -373,6 +453,15 @@ const jsonObject = <T extends object>(
 const stringField = (value: Record<string, unknown>, key: string): boolean =>
   typeof value[key] === 'string' && value[key].length > 0
 
+const optionalIndexField = (
+  value: Record<string, unknown>,
+  key: string,
+): boolean =>
+  value[key] === undefined ||
+  (typeof value[key] === 'number' &&
+    Number.isSafeInteger(value[key]) &&
+    value[key] >= 0)
+
 const optionalStringField = (
   value: Record<string, unknown>,
   key: string,
@@ -449,6 +538,7 @@ const completedPayloadSchema: StandardSchemaV1<GenerationCompletedPayload> =
       !stringField(value, 'messageId') ||
       !stringField(value, 'generationId') ||
       !stringField(value, 'responseMessageId') ||
+      !optionalIndexField(value, 'inputTokenEstimate') ||
       !optionalStringField(value, 'finishReason')
     ) {
       return issue('invalid ai.generation.completed payload')
@@ -465,6 +555,8 @@ const compactionPayloadSchema = <M extends UIMessage>(
       !isJSONCompatible(value) ||
       !stringField(value, 'generationId') ||
       !stringField(value, 'throughMessageId') ||
+      !optionalIndexField(value, 'throughIndex') ||
+      !optionalStringField(value, 'summary') ||
       !Array.isArray(value['messages']) ||
       (value['retainedMessageIds'] !== undefined &&
         (!Array.isArray(value['retainedMessageIds']) ||
@@ -484,7 +576,16 @@ const compactionPayloadSchema = <M extends UIMessage>(
       value: {
         generationId: value['generationId'] as string,
         throughMessageId: value['throughMessageId'] as string,
+        ...(value['throughIndex'] === undefined
+          ? {}
+          : { throughIndex: value['throughIndex'] as number }),
         messages,
+        ...(value['summary'] === undefined
+          ? {}
+          : { summary: value['summary'] as string }),
+        ...(value['usage'] === undefined
+          ? {}
+          : { usage: value['usage'] as LanguageModelUsage }),
         ...(value['retainedMessageIds'] === undefined
           ? {}
           : { retainedMessageIds: value['retainedMessageIds'] as string[] }),
@@ -498,6 +599,143 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
   const messageSchema = options?.messageSchema
   if (messageSchema) assertSyncSchema(messageSchema, 'the AI message schema')
   const definition: AIEventDefs<M> = {
+    'ai.control.requested': schema('ai.control.requested', (value) => {
+      if (!isRecord(value) || !isJSONCompatible(value))
+        return issue('invalid AI command')
+      const action = value['action']
+      const forwarded = {
+        approval: ['ai.approval.responded', 'response'],
+        'tool-result': ['ai.tool.result', 'result'],
+        input: ['ai.input.responded', 'response'],
+        'request-input': ['ai.input.requested', 'request'],
+        retry: ['ai.retry.requested', 'request'],
+        interrupt: ['ai.message.interrupted', 'request'],
+      } as const
+      if (typeof action === 'string' && Object.hasOwn(forwarded, action)) {
+        const [type, key] = forwarded[action as keyof typeof forwarded]
+        const payload = validateSync(
+          definition[type] as StandardSchemaV1<unknown>,
+          value[key],
+          'AI command',
+        )
+        if (payload.issues) return { issues: payload.issues }
+        return { value: { action, [key]: payload.value } as ControlCommand<M> }
+      }
+      if (action === 'send' || action === 'edit' || action === 'steer') {
+        const parsed = validateMessage(value['message'], messageSchema)
+        if (parsed.issues) return parsed
+        if (
+          (action === 'edit' &&
+            (!Number.isSafeInteger(value['expectedRevision']) ||
+              Number(value['expectedRevision']) < 0)) ||
+          (action === 'steer' &&
+            (!stringField(value, 'turnId') ||
+              !optionalIndexField(value, 'lastSeenIndex'))) ||
+          (action === 'send' &&
+            value['generate'] !== undefined &&
+            typeof value['generate'] !== 'boolean')
+        )
+          return issue('invalid AI command')
+        return {
+          value: { ...value, message: parsed.value } as ControlCommand<M>,
+        }
+      }
+      if (
+        (action === 'remove' && stringField(value, 'inputId')) ||
+        (action === 'send-now' &&
+          stringField(value, 'inputId') &&
+          (value['turnId'] === null || stringField(value, 'turnId')) &&
+          optionalIndexField(value, 'lastSeenIndex')) ||
+        (action === 'move' &&
+          stringField(value, 'inputId') &&
+          (value['beforeId'] === null || stringField(value, 'beforeId'))) ||
+        (action === 'stop' &&
+          stringField(value, 'turnId') &&
+          optionalIndexField(value, 'lastSeenIndex')) ||
+        (action === 'pause' &&
+          (value['when'] === 'now' || value['when'] === 'after-turn')) ||
+        action === 'resume'
+      )
+        return { value: value as ControlCommand<M> }
+      return issue('invalid AI command')
+    }),
+    'ai.control.committed': jsonObject(
+      'ai.control.committed',
+      (value) => Array.isArray(value['changes']) && isRecord(value['view']),
+    ),
+    'ai.control.decided': jsonObject(
+      'ai.control.decided',
+      (value) =>
+        stringField(value, 'commandId') &&
+        (value['outcome'] === 'applied' || value['outcome'] === 'rejected'),
+    ),
+    'ai.work.reported': schema('ai.work.reported', (value) => {
+      if (
+        !isRecord(value) ||
+        !isJSONCompatible(value) ||
+        !stringField(value, 'workId') ||
+        !stringField(value, 'turnId') ||
+        typeof value['version'] !== 'number' ||
+        !optionalIndexField(value, 'version') ||
+        typeof value['attempt'] !== 'number' ||
+        !Number.isSafeInteger(value['attempt']) ||
+        value['attempt'] < 1 ||
+        typeof value['settled'] !== 'boolean' ||
+        (value['kind'] !== 'model' && value['kind'] !== 'tool') ||
+        !Array.isArray(value['events'])
+      )
+        return issue('invalid AI work report')
+      const allowed = new Set([
+        'ai.generation.started',
+        'ai.generation.completed',
+        'ai.generation.failed',
+        'ai.tool.called',
+        'ai.tool.result',
+        'ai.approval.requested',
+        'ai.approval.responded',
+        'ai.compaction.requested',
+        'ai.compaction.completed',
+        'ai.model.metadata.requested',
+      ])
+      const events: WorkReport['events'] = []
+      for (const event of value['events']) {
+        if (
+          !isRecord(event) ||
+          !stringField(event, 'id') ||
+          typeof event['type'] !== 'string' ||
+          !allowed.has(event['type']) ||
+          (value['kind'] === 'tool' && event['type'] !== 'ai.tool.result')
+        )
+          return issue('invalid AI work fact')
+        const parsed = validateSync(
+          definition[
+            event['type'] as keyof AIEventDefs<M>
+          ] as StandardSchemaV1<unknown>,
+          event['payload'],
+          'AI work fact',
+        )
+        if (parsed.issues) return { issues: parsed.issues }
+        events.push({
+          id: event['id'] as string,
+          type: event['type'],
+          payload: parsed.value,
+        })
+      }
+      return { value: { ...value, events } as WorkReport }
+    }),
+    'ai.tool.execution.requested': jsonObject(
+      'ai.tool.execution.requested',
+      (value) =>
+        stringField(value, 'turnId') &&
+        optionalIndexField(value, 'version') &&
+        isRecord(value['call']) &&
+        isRecord(value['generation']),
+    ),
+    'ai.generation.paused': jsonObject(
+      'ai.generation.paused',
+      (value) =>
+        stringField(value, 'generationId') && stringField(value, 'messageId'),
+    ),
     'ai.session.created': jsonObject('ai.session.created', () => true),
     'ai.session.closed': jsonObject('ai.session.closed', (value) =>
       optionalStringField(value, 'reason'),
@@ -528,7 +766,8 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
         (value['reason'] === 'message' ||
           value['reason'] === 'tool' ||
           value['reason'] === 'input' ||
-          value['reason'] === 'retry'),
+          value['reason'] === 'retry' ||
+          value['reason'] === 'resume'),
     ),
     'ai.retry.requested': jsonObject(
       'ai.retry.requested',
@@ -547,7 +786,8 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
         typeof value['attempt'] === 'number' &&
         Number.isInteger(value['attempt']) &&
         value['attempt'] > 0 &&
-        stringField(value, 'model'),
+        stringField(value, 'model') &&
+        optionalIndexField(value, 'promptThroughIndex'),
     ),
     'ai.generation.progress': progressPayloadSchema,
     'ai.generation.completed': completedPayloadSchema,
@@ -641,9 +881,27 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
       'ai.compaction.requested',
       (value) =>
         stringField(value, 'generationId') &&
-        stringField(value, 'throughMessageId'),
+        stringField(value, 'throughMessageId') &&
+        optionalIndexField(value, 'throughIndex'),
     ),
     'ai.compaction.completed': compactionPayloadSchema(messageSchema),
+    'ai.model.metadata.requested': jsonObject(
+      'ai.model.metadata.requested',
+      (value) => stringField(value, 'modelId'),
+    ),
+    'ai.model.metadata.resolved': jsonObject(
+      'ai.model.metadata.resolved',
+      (value) =>
+        stringField(value, 'modelId') &&
+        (value['limits'] === null ||
+          (isRecord(value['limits']) &&
+            typeof value['limits']['contextWindow'] === 'number' &&
+            Number.isSafeInteger(value['limits']['contextWindow']) &&
+            value['limits']['contextWindow'] > 0 &&
+            typeof value['limits']['maxOutputTokens'] === 'number' &&
+            Number.isSafeInteger(value['limits']['maxOutputTokens']) &&
+            value['limits']['maxOutputTokens'] > 0)),
+    ),
   }
   return definition
 }
@@ -651,6 +909,11 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
 export const events: AIEventDefs<UIMessage> = createEvents()
 
 const initialState = <M extends UIMessage>(): AIState<M> => ({
+  inbox: initialControlState<M>().inbox,
+  pendingQueueCommands: [],
+  active: null,
+  starting: false,
+  receipt: null,
   session: { created: false },
   messages: [],
   status: 'idle',
@@ -663,6 +926,7 @@ const initialState = <M extends UIMessage>(): AIState<M> => ({
   pendingInputs: [],
   tools: [],
   compaction: null,
+  modelMetadata: {},
   usage: [],
   error: null,
 })
@@ -675,6 +939,23 @@ const upsertMessage = <M extends UIMessage>(messages: M[], message: M): M[] => {
   return next
 }
 
+const upsertResponse = <M extends UIMessage>(options: {
+  messages: M[]
+  message: M
+  inputId: string
+}): M[] => {
+  const { messages, message, inputId } = options
+  if (messages.some((item) => item.id === message.id))
+    return upsertMessage(messages, message)
+  const inputIndex = messages.findIndex((item) => item.id === inputId)
+  if (inputIndex === -1) return [...messages, message]
+  return [
+    ...messages.slice(0, inputIndex + 1),
+    message,
+    ...messages.slice(inputIndex + 1),
+  ]
+}
+
 const answerApproval = <M extends UIMessage>(
   messages: M[],
   response: ApprovalRespondedPayload,
@@ -826,12 +1107,10 @@ const isStaleGeneration = <M extends UIMessage>(
   return current !== undefined && current !== generationId
 }
 
-const projectedMessage = <M extends UIMessage>(
+const projectedMessageStream = <M extends UIMessage>(
   projection: ActiveGenerationProjection<M>,
-): M => {
-  const projectionEvents: Array<{ index: number; chunks: UIMessageChunk[] }> = [
-    ...projection.batches,
-  ]
+): { message: M; cursor: UIMessageCursor } => {
+  const projectionEvents = flattenProgressBatches(projection.batches)
   for (const event of projection.toolEvents) {
     if (event.type === 'ai.tool.result') {
       projectionEvents.push({
@@ -850,7 +1129,7 @@ const projectedMessage = <M extends UIMessage>(
       ],
     })
   }
-  return projectUIMessage({
+  return projectUIMessageStream({
     messageId: projection.responseMessageId,
     ...(projection.baseMessage === undefined
       ? {}
@@ -861,6 +1140,25 @@ const projectedMessage = <M extends UIMessage>(
   })
 }
 
+const projectedMessage = <M extends UIMessage>(
+  projection: ActiveGenerationProjection<M>,
+): M => projectedMessageStream(projection).message
+
+const projectedResponse = <M extends UIMessage>(options: {
+  state: AIState<M>
+  projection: ActiveGenerationProjection<M>
+}): M => {
+  const message = projectedMessage(options.projection)
+  return options.state.pausedGenerationIds?.[options.projection.generationId]
+    ? pauseUIMessage({
+        message,
+        toolCallIds: new Set(
+          options.state.tools.map((tool) => tool.toolCallId),
+        ),
+      })
+    : message
+}
+
 const reduceToolActivity = (
   tools: ToolActivity[],
   event:
@@ -1020,7 +1318,11 @@ const interruptGenerationProjection = <M extends UIMessage>(
   if (projection === null) return null
   const retained: ActiveGenerationProjection<M> = {
     ...projection,
-    batches: projection.batches.filter((batch) => batch.index <= cutoff),
+    batches: progressBatches(
+      flattenProgressBatches(projection.batches).filter(
+        (batch) => batch.index <= cutoff,
+      ),
+    ),
     toolEvents: projection.toolEvents.filter(
       (toolEvent) => toolEvent.index <= cutoff,
     ),
@@ -1040,7 +1342,7 @@ const interruptGenerationProjection = <M extends UIMessage>(
     ...restored,
     messages: shouldProject
       ? upsertMessage(
-          restored.messages,
+          state.messages,
           interruptUIMessage(projectedMessage(retained)),
         )
       : restored.messages,
@@ -1084,6 +1386,141 @@ export function reduceAIState<M extends UIMessage>(
 ): AIState<M> {
   if (state.status === 'closed') return state
   switch (event.type) {
+    case 'ai.control.requested': {
+      const command = event.payload as ControlCommand<M>
+      if (!event.id || !isQueueCommand(command)) return state
+      const pending = state.pendingQueueCommands ?? []
+      if (pending.some((item) => item.commandId === event.id)) return state
+      if (
+        (command.action === 'send' || command.action === 'steer') &&
+        (state.messages.some((message) => message.id === command.message.id) ||
+          state.inbox.items.some((item) => item.id === command.message.id))
+      )
+        return state
+      return {
+        ...state,
+        pendingQueueCommands: [...pending, { commandId: event.id, command }],
+      }
+    }
+    case 'ai.control.committed': {
+      const commit = event.payload as ControlCommit
+      const visible = applyControlChanges({
+        state: { inbox: state.inbox, active: state.active },
+        changes: commit.changes.filter(
+          (change) => change.path[0] === 'inbox' || change.path[0] === 'active',
+        ),
+      })
+      const { phase } = commit.view
+      const status: AgentStatus = commit.view.closed
+        ? 'closed'
+        : phase === 'paused' ||
+            phase === 'pausing' ||
+            (commit.view.paused && phase === null)
+          ? 'paused'
+          : phase === 'failed'
+            ? 'failed'
+            : phase === 'starting' || phase === 'generating'
+              ? 'generating'
+              : phase
+                ? 'waiting'
+                : 'idle'
+      return {
+        ...state,
+        ...visible,
+        status,
+        activeRequestId: commit.view.requestId,
+        activeResponseMessageId: commit.view.responseMessageId,
+      }
+    }
+    case 'ai.control.decided': {
+      const receipt = event.payload as ControlReceipt
+      if (receipt.work) return state
+      const pending = state.pendingQueueCommands ?? []
+      const settled = pending.find(
+        (item) => item.commandId === receipt.commandId,
+      )
+      return {
+        ...state,
+        receipt,
+        ...(settled === undefined
+          ? {}
+          : {
+              pendingQueueCommands: pending.filter(
+                (item) => item.commandId !== receipt.commandId,
+              ),
+            }),
+      }
+    }
+    case 'ai.generation.paused': {
+      const payload = event.payload as {
+        generationId: string
+        messageId: string
+      }
+      const toolCallIds = new Set(
+        state.tools
+          .filter((tool) => tool.messageId === payload.messageId)
+          .map((tool) => tool.toolCallId),
+      )
+      const projection = state.activeProjection
+      return {
+        ...state,
+        activeGeneration: null,
+        activeProjection:
+          projection?.generationId === payload.generationId
+            ? {
+                ...projection,
+                cursor: null,
+                batches: progressBatches(
+                  flattenProgressBatches(projection.batches).map((batch) => ({
+                    index: batch.index,
+                    chunks: batch.chunks.filter(
+                      (chunk) =>
+                        !('toolCallId' in chunk) ||
+                        toolCallIds.has(chunk.toolCallId),
+                    ),
+                  })),
+                ),
+              }
+            : projection,
+        compaction:
+          state.compaction?.status === 'running'
+            ? (projection?.baseCompaction ?? null)
+            : state.compaction,
+        messages: state.messages.map((message) =>
+          message.id === payload.messageId
+            ? pauseUIMessage({ message, toolCallIds })
+            : message,
+        ),
+        pausedGenerationIds: {
+          ...state.pausedGenerationIds,
+          [payload.generationId]: true,
+        },
+      }
+    }
+    case 'ai.model.metadata.requested': {
+      const { modelId } = event.payload as ModelMetadataRequestedPayload
+      if (Object.hasOwn(state.modelMetadata, modelId)) return state
+      return {
+        ...state,
+        modelMetadata: {
+          ...state.modelMetadata,
+          [modelId]: { status: 'pending' },
+        },
+      }
+    }
+    case 'ai.model.metadata.resolved': {
+      const { modelId, limits } = event.payload as ModelMetadataResolvedPayload
+      return {
+        ...state,
+        modelMetadata: {
+          ...state.modelMetadata,
+          [modelId]:
+            limits === null
+              ? { status: 'unavailable' }
+              : { status: 'resolved', limits },
+        },
+      }
+    }
     case 'ai.session.created': {
       const payload = event.payload as SessionCreatedPayload
       return {
@@ -1101,7 +1538,14 @@ export function reduceAIState<M extends UIMessage>(
     case 'ai.session.closed':
       return {
         ...state,
+        active: null,
+        inbox: { paused: true, items: [] },
+        pendingQueueCommands: [],
         status: 'closed',
+        compaction:
+          state.compaction?.status === 'running'
+            ? (state.activeProjection?.baseCompaction ?? null)
+            : state.compaction,
         activeGeneration: null,
         activeRequestId: null,
         activeResponseMessageId: null,
@@ -1109,7 +1553,14 @@ export function reduceAIState<M extends UIMessage>(
       }
     case 'ai.message.created': {
       const { message } = event.payload as MessageCreatedPayload<M>
-      return { ...state, messages: upsertMessage(state.messages, message) }
+      return {
+        ...state,
+        messages: upsertMessage(state.messages, message),
+        activeProjection:
+          state.activeProjection?.responseMessageId === message.id
+            ? { ...state.activeProjection, cursor: null }
+            : state.activeProjection,
+      }
     }
     case 'ai.message.completed': {
       const { messageId } = event.payload as MessageCompletedPayload
@@ -1166,6 +1617,7 @@ export function reduceAIState<M extends UIMessage>(
         return state
       }
       const acceptsRequest =
+        (payload.control !== undefined && state.activeRequestId === event.id) ||
         payload.reason === 'message' ||
         (payload.reason === 'tool'
           ? ownerGenerationId !== undefined &&
@@ -1266,9 +1718,10 @@ export function reduceAIState<M extends UIMessage>(
         activeProjection: {
           generationId: payload.generationId,
           responseMessageId: payload.responseMessageId,
+          promptThroughIndex: payload.promptThroughIndex ?? event.index - 1,
           ...(baseMessage === undefined ? {} : { baseMessage }),
           baseCompaction: baseState.compaction,
-          batches: [],
+          batches: progressBatches(),
           baseTools: baseState.tools,
           toolEvents: [],
           approvalEvents: [],
@@ -1278,6 +1731,7 @@ export function reduceAIState<M extends UIMessage>(
     }
     case 'ai.generation.progress': {
       const payload = event.payload as GenerationProgressPayload
+      if (state.pausedGenerationIds?.[payload.generationId]) return state
       if (terminalGenerationReason(state, payload.generationId) !== undefined) {
         return state
       }
@@ -1288,18 +1742,63 @@ export function reduceAIState<M extends UIMessage>(
       ) {
         return state
       }
+      const ordered =
+        event.index > (activeProjection.batches.tail.at(-1)?.index ?? -1)
       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),
+        batches: ordered
+          ? appendProgressBatch({
+              batches: activeProjection.batches,
+              batch: { index: event.index, chunks: payload.chunks },
+            })
+          : progressBatches(
+              [
+                ...flattenProgressBatches(activeProjection.batches).filter(
+                  (batch) => batch.index !== event.index,
+                ),
+                { index: event.index, chunks: payload.chunks },
+              ].toSorted((left, right) => left.index - right.index),
+            ),
+      }
+      const base = state.messages.find(
+        (message) => message.id === projection.responseMessageId,
+      )
+      const incremental =
+        base !== undefined &&
+        activeProjection.cursor != null &&
+        event.index > activeProjection.cursor.throughIndex
+      const projected = incremental
+        ? projectUIMessageStream({
+            messageId: projection.responseMessageId,
+            base,
+            cursor: activeProjection.cursor!.stream,
+            chunks: payload.chunks,
+          })
+        : projectedMessageStream(projection)
+      projection.cursor = {
+        throughIndex: incremental
+          ? event.index
+          : Math.max(
+              event.index,
+              projection.batches.tail.at(-1)?.index ?? 0,
+              projection.toolEvents.reduce(
+                (index, entry) => Math.max(index, entry.index),
+                0,
+              ),
+              projection.approvalEvents.reduce(
+                (index, entry) => Math.max(index, entry.index),
+                0,
+              ),
+            ),
+        stream: projected.cursor,
       }
       return {
         ...state,
-        messages: upsertMessage(state.messages, projectedMessage(projection)),
+        messages: upsertResponse({
+          messages: state.messages,
+          message: projected.message,
+          inputId: state.activeGeneration.messageId,
+        }),
         activeProjection: projection,
         status: 'generating',
       }
@@ -1727,12 +2226,16 @@ export function reduceAIState<M extends UIMessage>(
       }
       const activeProjection: ActiveGenerationProjection<M> = {
         ...projection,
+        cursor: null,
         toolEvents: [
           ...projection.toolEvents,
           { index: event.index, type: 'ai.tool.result', payload },
         ],
       }
-      const projected = projectedMessage(activeProjection)
+      const projected = projectedResponse({
+        state,
+        projection: activeProjection,
+      })
       const tools = projectedTools(activeProjection)
       const failureError = state.error ?? 'Generation failed.'
       const next = {
@@ -1783,6 +2286,7 @@ export function reduceAIState<M extends UIMessage>(
         projection?.generationId === payload.generationId
           ? {
               ...projection,
+              cursor: null,
               approvalEvents: [
                 ...projection.approvalEvents,
                 {
@@ -1800,7 +2304,10 @@ export function reduceAIState<M extends UIMessage>(
             ? projectMessageChunks(state.messages, payload.messageId, [
                 approvalRequestedChunk(payload),
               ])
-            : upsertMessage(state.messages, projectedMessage(activeProjection)),
+            : upsertMessage(
+                state.messages,
+                projectedResponse({ state, projection: activeProjection }),
+              ),
         activeProjection,
         pendingApprovals:
           payload.isAutomatic === true
@@ -1852,6 +2359,7 @@ export function reduceAIState<M extends UIMessage>(
         projection !== null && activeRequest !== undefined
           ? {
               ...projection,
+              cursor: null,
               approvalEvents: [
                 ...projection.approvalEvents,
                 {
@@ -1867,8 +2375,14 @@ export function reduceAIState<M extends UIMessage>(
         messages:
           activeProjection === null || activeProjection === projection
             ? answerApproval(state.messages, payload)
-            : upsertMessage(state.messages, projectedMessage(activeProjection)),
-        activeProjection,
+            : upsertMessage(
+                state.messages,
+                projectedResponse({ state, projection: activeProjection }),
+              ),
+        activeProjection:
+          activeProjection === null
+            ? null
+            : { ...activeProjection, cursor: null },
         pendingApprovals: state.pendingApprovals.filter(
           (approval) =>
             approval.approvalId !== payload.approvalId ||
@@ -1943,7 +2457,12 @@ export function reduceAIState<M extends UIMessage>(
       }
       return {
         ...state,
-        compaction: { status: 'completed', ...payload },
+        compaction: {
+          status: 'completed',
+          ...payload,
+          throughIndex:
+            payload.throughIndex ?? state.activeProjection!.promptThroughIndex!,
+        },
       }
     }
     default:
@@ -1966,59 +2485,84 @@ export function createReducer<D extends EventDefs>(options: {
   name?: string
 }): Reducer<D, AIState<AIMessageFrom<D>>> {
   type M = AIMessageFrom<D>
-  return options.contract
+  const reducer = options.contract
     .reducer({
-      name: options.name ?? 'a2.ai.state.v8',
+      name: options.name ?? 'a2.ai.state.v17',
       initialState: initialState<M>(),
     })
     .fold((state, event) => reduceAIState(state, event))
+  return { ...reducer, [clientStateProjection]: projectAIClientState }
 }
 
-type ProtocolInput<
+type CommandInput<C> = { type: 'ai.control.requested'; payload: C; id?: string }
+type ControlInput<
   M extends UIMessage,
-  K extends keyof AIEventDefs<M> & string,
-> = AppendInput<Pick<AIEventDefs<M>, K>>
+  A extends ControlCommand<M>['action'],
+> = CommandInput<Extract<ControlCommand<M>, { action: A }>>
 
 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. */
+  ): ControlInput<M, 'send'>[]
+  seed<M extends UIMessage>(message: M): ControlInput<M, 'send'>[]
+  queue: {
+    remove(options: { inputId: string }): ControlInput<UIMessage, 'remove'>[]
+    sendNow(options: {
+      inputId: string
+      turnId: string | null
+      lastSeenIndex?: number
+    }): ControlInput<UIMessage, 'send-now'>[]
+    edit<M extends UIMessage>(options: {
+      message: M
+      expectedRevision: number
+    }): ControlInput<M, 'edit'>[]
+    move(options: {
+      inputId: string
+      beforeId: string | null
+    }): ControlInput<UIMessage, 'move'>[]
+  }
+  stop(options: {
+    turnId: string
+    lastSeenIndex?: number
+  }): ControlInput<UIMessage, 'stop'>[]
+  steer<M extends UIMessage>(options: {
+    turnId: string
+    message: M
+    lastSeenIndex?: number
+  }): ControlInput<M, 'steer'>[]
+  pause(options: {
+    when: 'now' | 'after-turn'
+  }): ControlInput<UIMessage, 'pause'>[]
+  resume(): ControlInput<UIMessage, 'resume'>[]
   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. */
+  ): ControlInput<UIMessage, 'approval'>[]
+  input(response: InputRespondedPayload): ControlInput<UIMessage, '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. */
+  ): ControlInput<UIMessage, 'request-input'>[]
+  retry(request: RetryGenerationOptions): ControlInput<UIMessage, 'retry'>[]
   interrupt(
-    interruption: MessageInterruptedPayload,
-  ): ProtocolInput<UIMessage, 'ai.message.interrupted'>[]
+    request: MessageInterruptedPayload,
+  ): ControlInput<UIMessage, 'interrupt'>[]
+  toolResult(
+    result: ToolResultPayload,
+  ): ControlInput<UIMessage, 'tool-result'>[]
 }
 
+const commandInput = <C extends ControlCommand<UIMessage>>(
+  payload: C,
+): CommandInput<C>[] => [{ type: 'ai.control.requested', payload }]
 const messageInput = <M extends UIMessage>(
   message: M,
   generate?: boolean,
-): ProtocolInput<M, 'ai.message.created'>[] => [
+): ControlInput<M, 'send'>[] => [
   {
-    type: 'ai.message.created',
+    type: 'ai.control.requested',
     id: `ai.message:${message.id}`,
     payload: {
+      action: 'send',
       message,
       ...(generate === false ? { generate: false } : {}),
     },
@@ -2027,57 +2571,53 @@ const messageInput = <M extends UIMessage>(
 
 /** Pure append inputs for the built-in AI protocol. */
 export const inputs: AIInputs = {
+  toolResult: (result) => commandInput({ action: 'tool-result', result }),
+  queue: {
+    remove: (options) => commandInput({ action: 'remove', ...options }),
+    sendNow: (options) => commandInput({ action: 'send-now', ...options }),
+    edit: (options) => commandInput({ action: 'edit', ...options }),
+    move: (options) => commandInput({ action: 'move', ...options }),
+  },
+  stop: (options) => commandInput({ action: 'stop', ...options }),
+  steer: (options) => commandInput({ action: 'steer', ...options }),
+  pause: (options) => commandInput({ action: 'pause', ...options }),
+  resume: () => commandInput({ action: 'resume' }),
   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,
-      },
-    ]
-  },
+  approval: (response) => [
+    {
+      ...commandInput({ action: 'approval', response })[0]!,
+      id: `ai.approval:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.approvalId)}:response`,
+    },
+  ],
+  input: (response) => [
+    {
+      ...commandInput({ action: 'input', response })[0]!,
+      id: `ai.input:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.inputId)}:response`,
+    },
+  ],
+  requestInput: (request) => [
+    {
+      ...commandInput({ action: 'request-input', request })[0]!,
+      id: `ai.input:${encodeURIComponent(request.messageId)}:${encodeURIComponent(request.generationId)}:${encodeURIComponent(request.inputId)}:request`,
+    },
+  ],
+  retry: (request) => [
+    {
+      ...commandInput({ action: 'retry', request })[0]!,
+      id: `ai.retry:${request.retryId}`,
+    },
+  ],
+  interrupt: (request) => [
+    {
+      ...commandInput({ action: 'interrupt', request })[0]!,
+      id: `ai.interrupt:${request.generationId ?? request.requestId}`,
+    },
+  ],
 }
 
 /**
@@ -2134,14 +2674,14 @@ export function agent<
   })
   const reducer = agentContract
     .reducer({
-      name: options.reducerName ?? 'a2.ai.state.v8',
+      name: options.reducerName ?? 'a2.ai.state.v17',
       initialState: initialState<M>(),
     })
     .fold((state, event) => reduceAIState(state, event))
 
   const definition = {
     contract: agentContract,
-    reducer,
+    reducer: { ...reducer, [clientStateProjection]: projectAIClientState<M> },
   }
   return definition as unknown as AgentDefinition<M, AIEventDefs<M> & E>
 }