experimental-a2 0.15.0 → 0.16.1

package/src/ai.ts

@@ -21,6 +21,7 @@ import {
 export type { PendingQueueCommand } from './ai-client-state.ts'
 import { applyControlChanges } from './ai-control-state.ts'
 import { initialControlState } from './ai-control.ts'
+import { AI_EVENT_ID_PREFIX } from './ai-id.ts'
 import type {
   ControlCommit,
   ActiveTurn,
@@ -42,13 +43,27 @@ import {
   interruptUIMessage,
   pauseUIMessage,
   projectUIMessage,
-  projectUIMessageStream,
   type UIMessageCursor,
 } from './ai-projector.ts'
 import { ambientToolScopeStorage } from './internal.ts'
+import {
+  upsertMessage,
+  upsertResponse,
+  answerApproval,
+  approvalRequestedChunk,
+  toolResultChunk,
+  advanceProjection,
+  projectedTools,
+  restoreGenerationBase,
+  stopGenerationTools,
+  interruptGenerationProjection,
+  projectedMessage,
+  reduceToolActivity,
+} from './ai-message-projection.ts'
 import { clientStateProjection, type Reducer } from './reducer.ts'
+import { storedAIReducer, type AIStoredReducer } from './ai-stored-state.ts'
+export type { AIStoredState, AIStoredReducer } from './ai-stored-state.ts'
 import {
-  appendProgressBatch,
   flattenProgressBatches,
   progressBatches,
   type ProgressBatches,
@@ -85,7 +100,8 @@ export type GenerationRequestedPayload = {
   control?: { turnId: string; version: number }
   messageId: string
   reason: GenerationReason
-  responseMessageId?: string
+  responseMessageId: string
+  sourceGenerationId?: string
 }
 export type RetryGenerationOptions = {
   messageId: string
@@ -162,6 +178,7 @@ export type ToolResultPayload = {
   toolMetadata?: unknown
 }
 export type ApprovalRequestedPayload = {
+  requestId: string
   messageId: string
   generationId: string
   approvalId: string
@@ -170,6 +187,7 @@ export type ApprovalRequestedPayload = {
   signature?: string
 }
 export type ApprovalRespondedPayload = {
+  requestId: string
   messageId: string
   generationId: string
   approvalId: string
@@ -177,6 +195,7 @@ export type ApprovalRespondedPayload = {
   reason?: string
 }
 export type InputRequestedPayload = {
+  requestId: string
   messageId: string
   generationId: string
   inputId: string
@@ -184,6 +203,7 @@ export type InputRequestedPayload = {
   question?: JSONValue
 }
 export type InputRespondedPayload = {
+  requestId: string
   messageId: string
   generationId: string
   inputId: string
@@ -208,6 +228,9 @@ export type CompactionRequestedPayload = {
   generationId: string
   throughMessageId: string
   throughIndex?: number
+  inputTokenEstimate?: number
+  inputTokens?: number
+  thresholdTokens?: number
 }
 export type CompactionCompletedPayload<M extends UIMessage> = {
   generationId: string
@@ -335,9 +358,14 @@ export type AIState<M extends UIMessage = UIMessage> = {
   status: AgentStatus
   activeGeneration: ActiveGeneration | null
   activeRequestId: string | null
+  activeRequestReason: GenerationReason | null
+  activeRequestSourceGenerationId: string | null
   activeResponseMessageId: string | null
   activeProjection: ActiveGenerationProjection<M> | null
-  responseGenerationIds: Record<string, string>
+  responseOwners: Record<
+    string,
+    Pick<GenerationStartedPayload, 'generationId' | 'requestId' | 'attempt'>
+  >
   terminalRequestIds?: Record<string, true>
   terminalGenerations?: Record<
     string,
@@ -357,7 +385,7 @@ export type AgentDefinition<
   D extends EventDefs = AIEventDefs<M>,
 > = {
   readonly contract: Contract<D>
-  readonly reducer: Reducer<D, AIState<M>>
+  readonly reducer: AIStoredReducer<D, AIState<M>>
 }
 
 export type AgentToolEventDefs = EventDefs &
@@ -625,9 +653,14 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
         const parsed = validateMessage(value['message'], messageSchema)
         if (parsed.issues) return parsed
         if (
-          (action === 'edit' &&
-            (!Number.isSafeInteger(value['expectedRevision']) ||
-              Number(value['expectedRevision']) < 0)) ||
+          (action === 'send' || action === 'steer') &&
+          (parsed.value.id.length === 0 ||
+            parsed.value.id.startsWith(AI_EVENT_ID_PREFIX))
+        )
+          return issue(
+            'AI input IDs must be nonempty and outside the reserved internal prefix',
+          )
+        if (
           (action === 'steer' &&
             (!stringField(value, 'turnId') ||
               !optionalIndexField(value, 'lastSeenIndex'))) ||
@@ -762,7 +795,10 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
       'ai.generation.requested',
       (value) =>
         stringField(value, 'messageId') &&
-        optionalStringField(value, 'responseMessageId') &&
+        stringField(value, 'responseMessageId') &&
+        optionalStringField(value, 'sourceGenerationId') &&
+        (value['reason'] !== 'tool' ||
+          stringField(value, 'sourceGenerationId')) &&
         (value['reason'] === 'message' ||
           value['reason'] === 'tool' ||
           value['reason'] === 'input' ||
@@ -844,6 +880,7 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
     'ai.approval.requested': jsonObject(
       'ai.approval.requested',
       (value) =>
+        stringField(value, 'requestId') &&
         stringField(value, 'messageId') &&
         stringField(value, 'generationId') &&
         stringField(value, 'approvalId') &&
@@ -855,6 +892,7 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
     'ai.approval.responded': jsonObject(
       'ai.approval.responded',
       (value) =>
+        stringField(value, 'requestId') &&
         stringField(value, 'messageId') &&
         stringField(value, 'generationId') &&
         stringField(value, 'approvalId') &&
@@ -864,6 +902,7 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
     'ai.input.requested': jsonObject(
       'ai.input.requested',
       (value) =>
+        stringField(value, 'requestId') &&
         stringField(value, 'messageId') &&
         stringField(value, 'generationId') &&
         stringField(value, 'inputId') &&
@@ -872,6 +911,7 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
     'ai.input.responded': jsonObject(
       'ai.input.responded',
       (value) =>
+        stringField(value, 'requestId') &&
         stringField(value, 'messageId') &&
         stringField(value, 'generationId') &&
         stringField(value, 'inputId') &&
@@ -882,7 +922,10 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
       (value) =>
         stringField(value, 'generationId') &&
         stringField(value, 'throughMessageId') &&
-        optionalIndexField(value, 'throughIndex'),
+        optionalIndexField(value, 'throughIndex') &&
+        optionalIndexField(value, 'inputTokenEstimate') &&
+        optionalIndexField(value, 'inputTokens') &&
+        optionalIndexField(value, 'thresholdTokens'),
     ),
     'ai.compaction.completed': compactionPayloadSchema(messageSchema),
     'ai.model.metadata.requested': jsonObject(
@@ -919,9 +962,11 @@ const initialState = <M extends UIMessage>(): AIState<M> => ({
   status: 'idle',
   activeGeneration: null,
   activeRequestId: null,
+  activeRequestReason: null,
+  activeRequestSourceGenerationId: null,
   activeResponseMessageId: null,
   activeProjection: null,
-  responseGenerationIds: {},
+  responseOwners: {},
   pendingApprovals: [],
   pendingInputs: [],
   tools: [],
@@ -931,62 +976,6 @@ const initialState = <M extends UIMessage>(): AIState<M> => ({
   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 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,
-): 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,
@@ -998,90 +987,6 @@ const projectMessageChunks = <M extends UIMessage>(
       : 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'
@@ -1101,49 +1006,16 @@ const lifecycleStatus = <M extends UIMessage>(
 const isStaleGeneration = <M extends UIMessage>(
   state: AIState<M>,
   messageId: string,
-  generationId: string,
+  owner: { generationId: string; requestId: string },
 ): boolean => {
-  const current = state.responseGenerationIds[messageId]
-  return current !== undefined && current !== generationId
-}
-
-const projectedMessageStream = <M extends UIMessage>(
-  projection: ActiveGenerationProjection<M>,
-): { message: M; cursor: UIMessageCursor } => {
-  const projectionEvents = flattenProgressBatches(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 projectUIMessageStream({
-    messageId: projection.responseMessageId,
-    ...(projection.baseMessage === undefined
-      ? {}
-      : { base: projection.baseMessage }),
-    chunks: projectionEvents
-      .toSorted((left, right) => left.index - right.index)
-      .flatMap((event) => event.chunks),
-  })
+  const current = state.responseOwners[messageId]
+  return (
+    current !== undefined &&
+    (current.generationId !== owner.generationId ||
+      current.requestId !== owner.requestId)
+  )
 }
 
-const projectedMessage = <M extends UIMessage>(
-  projection: ActiveGenerationProjection<M>,
-): M => projectedMessageStream(projection).message
-
 const projectedResponse = <M extends UIMessage>(options: {
   state: AIState<M>
   projection: ActiveGenerationProjection<M>
@@ -1159,220 +1031,13 @@ const projectedResponse = <M extends UIMessage>(options: {
     : message
 }
 
-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: progressBatches(
-      flattenProgressBatches(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(
-          state.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,
+  requestId?: 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'
@@ -1429,6 +1094,8 @@ export function reduceAIState<M extends UIMessage>(
         ...visible,
         status,
         activeRequestId: commit.view.requestId,
+        activeRequestReason: commit.view.requestReason,
+        activeRequestSourceGenerationId: commit.view.sourceGenerationId,
         activeResponseMessageId: commit.view.responseMessageId,
       }
     }
@@ -1548,6 +1215,8 @@ export function reduceAIState<M extends UIMessage>(
             : state.compaction,
         activeGeneration: null,
         activeRequestId: null,
+        activeRequestReason: null,
+        activeRequestSourceGenerationId: null,
         activeResponseMessageId: null,
         activeProjection: null,
       }
@@ -1564,17 +1233,20 @@ export function reduceAIState<M extends UIMessage>(
     }
     case 'ai.message.completed': {
       const { messageId } = event.payload as MessageCompletedPayload
-      const generationId = state.responseGenerationIds[messageId]
+      const owner = state.responseOwners[messageId]
+      const generationId = owner?.generationId
+      const requestId = owner?.requestId
       if (
         generationId === undefined ||
-        terminalGenerationReason(state, generationId) !== undefined
+        terminalGenerationReason(state, generationId, requestId) !== undefined
       ) {
         return state
       }
-      const requestId = generationRequestId(generationId)
       return {
         ...state,
         activeRequestId: null,
+        activeRequestReason: null,
+        activeRequestSourceGenerationId: null,
         activeResponseMessageId: null,
         activeProjection:
           state.activeProjection?.generationId === generationId
@@ -1597,18 +1269,10 @@ export function reduceAIState<M extends UIMessage>(
     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 responseMessageId = payload.responseMessageId
+      const owner = state.responseOwners[responseMessageId]
+      const ownerGenerationId = owner?.generationId
+      const ownerRequestId = owner?.requestId
       const requestId = (event as { id?: string }).id
       if (
         requestId !== undefined &&
@@ -1621,30 +1285,35 @@ export function reduceAIState<M extends UIMessage>(
         payload.reason === 'message' ||
         (payload.reason === 'tool'
           ? ownerGenerationId !== undefined &&
-            requestId === `ai.generate:tools:${ownerGenerationId}` &&
+            payload.sourceGenerationId === 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 &&
+              payload.sourceGenerationId === ownerGenerationId &&
               state.activeRequestId === ownerRequestId &&
               state.status !== 'failed')
       if (!acceptsRequest) return state
+      const base =
+        payload.reason === 'retry' &&
+        state.activeProjection &&
+        state.activeProjection.completionIndex === undefined
+          ? restoreGenerationBase(
+              state,
+              state.activeProjection.generationId,
+              responseMessageId,
+              'failed',
+            )
+          : state
       return {
-        ...state,
-        messages:
-          payload.reason === 'retry' && payload.responseMessageId !== undefined
-            ? state.messages.filter(
-                (message) => message.id !== payload.responseMessageId,
-              )
-            : state.messages,
+        ...base,
         activeProjection:
           payload.reason === 'retry' ? null : state.activeProjection,
         status: schedulesToolContinuation ? state.status : 'generating',
         activeRequestId: requestId ?? null,
+        activeRequestReason: payload.reason,
+        activeRequestSourceGenerationId: payload.sourceGenerationId ?? null,
         activeResponseMessageId: responseMessageId,
         error: schedulesToolContinuation ? state.error : null,
       }
@@ -1653,7 +1322,11 @@ export function reduceAIState<M extends UIMessage>(
       const payload = event.payload as GenerationStartedPayload
       if (
         state.terminalRequestIds?.[payload.requestId] === true ||
-        terminalGenerationReason(state, payload.generationId) !== undefined
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
       ) {
         return state
       }
@@ -1663,31 +1336,26 @@ export function reduceAIState<M extends UIMessage>(
       ) {
         return state
       }
-      const currentOwner =
-        state.responseGenerationIds[payload.responseMessageId]
+      const currentOwner = state.responseOwners[payload.responseMessageId]
       const currentTerminal =
         currentOwner === undefined
           ? undefined
-          : terminalGenerationReason(state, currentOwner)
+          : terminalGenerationReason(
+              state,
+              currentOwner.generationId,
+              currentOwner.requestId,
+            )
       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
-      ) {
+        currentOwner?.requestId === payload.requestId &&
+        currentOwner.attempt >= payload.attempt
+      )
         return state
-      }
       const activeGeneration = state.activeGeneration
       const replacesActiveAttempt =
         activeGeneration !== null &&
@@ -1711,9 +1379,13 @@ export function reduceAIState<M extends UIMessage>(
         status: 'generating',
         activeGeneration: payload,
         activeResponseMessageId: null,
-        responseGenerationIds: {
-          ...baseState.responseGenerationIds,
-          [payload.responseMessageId]: payload.generationId,
+        responseOwners: {
+          ...baseState.responseOwners,
+          [payload.responseMessageId]: {
+            generationId: payload.generationId,
+            requestId: payload.requestId,
+            attempt: payload.attempt,
+          },
         },
         activeProjection: {
           generationId: payload.generationId,
@@ -1732,83 +1404,61 @@ 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) {
+      if (
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
+      ) {
         return state
       }
       const activeProjection = state.activeProjection
       if (
         state.activeGeneration?.generationId !== payload.generationId ||
+        state.activeGeneration.requestId !== payload.requestId ||
         activeProjection?.generationId !== payload.generationId
       ) {
         return state
       }
-      const ordered =
-        event.index > (activeProjection.batches.tail.at(-1)?.index ?? -1)
-      const projection: ActiveGenerationProjection<M> = {
-        ...activeProjection,
-        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,
-      }
+      const progress = advanceProjection({
+        projection: activeProjection,
+        base: state.messages.find(
+          (message) => message.id === activeProjection.responseMessageId,
+        ),
+        chunks: payload.chunks,
+        index: event.index,
+      })
       return {
         ...state,
         messages: upsertResponse({
           messages: state.messages,
-          message: projected.message,
+          message: progress.message,
           inputId: state.activeGeneration.messageId,
         }),
-        activeProjection: projection,
+        activeProjection: {
+          ...activeProjection,
+          batches: progress.batches,
+          cursor: progress.cursor,
+        },
         status: 'generating',
       }
     }
     case 'ai.generation.completed': {
       const payload = event.payload as GenerationCompletedPayload
-      if (terminalGenerationReason(state, payload.generationId) !== undefined) {
+      if (
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
+      ) {
         return state
       }
-      if (state.activeGeneration?.generationId !== payload.generationId) {
+      if (
+        state.activeGeneration?.generationId !== payload.generationId ||
+        state.activeGeneration.requestId !== payload.requestId
+      ) {
         return state
       }
       const usage = payload.usage
@@ -1832,17 +1482,26 @@ export function reduceAIState<M extends UIMessage>(
     }
     case 'ai.generation.failed': {
       const payload = event.payload as GenerationFailedPayload
-      if (terminalGenerationReason(state, payload.generationId) !== undefined) {
+      if (
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
+      ) {
         return state
       }
       if (
-        state.activeGeneration?.generationId !== payload.generationId &&
+        (state.activeGeneration?.generationId !== payload.generationId ||
+          state.activeGeneration.requestId !== payload.requestId) &&
         !(
           payload.stepLimit === true &&
-          state.responseGenerationIds[payload.responseMessageId] ===
+          state.activeRequestReason === 'tool' &&
+          state.responseOwners[payload.responseMessageId]?.requestId ===
+            payload.requestId &&
+          state.responseOwners[payload.responseMessageId]?.generationId ===
             payload.generationId &&
-          state.activeRequestId ===
-            `ai.generate:tools:${payload.generationId}` &&
+          state.activeRequestSourceGenerationId === payload.generationId &&
           state.activeResponseMessageId === payload.responseMessageId
         )
       ) {
@@ -1912,9 +1571,13 @@ export function reduceAIState<M extends UIMessage>(
     }
     case 'ai.message.interrupted': {
       const payload = event.payload as MessageInterruptedPayload
-      const currentGenerationId = state.responseGenerationIds[payload.messageId]
+      const owner = state.responseOwners[payload.messageId]
+      const currentGenerationId = owner?.generationId
       const targetRequestId =
-        payload.requestId ?? generationRequestId(payload.generationId)
+        payload.requestId ??
+        (owner?.generationId === payload.generationId
+          ? owner?.requestId
+          : undefined)
       const ownsRequestPhase =
         payload.requestId !== undefined &&
         state.activeGeneration === null &&
@@ -1925,7 +1588,7 @@ export function reduceAIState<M extends UIMessage>(
         !ownsRequestPhase &&
         (payload.generationId !== undefined
           ? currentGenerationId !== payload.generationId
-          : generationRequestId(currentGenerationId) !== payload.requestId)
+          : owner?.requestId !== payload.requestId)
       ) {
         return state
       }
@@ -1940,7 +1603,11 @@ export function reduceAIState<M extends UIMessage>(
       const currentTerminal =
         currentGenerationId === undefined
           ? undefined
-          : terminalGenerationReason(state, currentGenerationId)
+          : terminalGenerationReason(
+              state,
+              currentGenerationId,
+              owner?.requestId,
+            )
       if (
         currentTerminal === 'completed' ||
         currentTerminal === 'interrupted' ||
@@ -1970,6 +1637,8 @@ export function reduceAIState<M extends UIMessage>(
             ...baseState,
             status: 'idle',
             activeRequestId: null,
+            activeRequestReason: null,
+            activeRequestSourceGenerationId: null,
             activeResponseMessageId: null,
             activeProjection: null,
             terminalRequestIds: {
@@ -1990,7 +1659,7 @@ export function reduceAIState<M extends UIMessage>(
         const targetGenerationId =
           payload.generationId ??
           (currentGenerationId !== undefined &&
-          generationRequestId(currentGenerationId) === payload.requestId
+          owner?.requestId === payload.requestId
             ? currentGenerationId
             : undefined)
         if (targetGenerationId === undefined) return state
@@ -2019,7 +1688,8 @@ export function reduceAIState<M extends UIMessage>(
           currentGenerationId === targetGenerationId &&
           (state.activeRequestId === null ||
             state.activeRequestId === targetRequestId ||
-            state.activeRequestId === `ai.generate:tools:${targetGenerationId}`)
+            (state.activeRequestReason === 'tool' &&
+              state.activeRequestSourceGenerationId === targetGenerationId))
         const resolvesCurrentGeneration =
           currentGenerationId === targetGenerationId &&
           currentTerminal === undefined
@@ -2041,6 +1711,8 @@ export function reduceAIState<M extends UIMessage>(
         const next = {
           ...baseState,
           activeRequestId: null,
+          activeRequestReason: null,
+          activeRequestSourceGenerationId: null,
           activeResponseMessageId: null,
           activeProjection: null,
           messages:
@@ -2114,6 +1786,8 @@ export function reduceAIState<M extends UIMessage>(
       const next = {
         ...baseState,
         activeRequestId: null,
+        activeRequestReason: null,
+        activeRequestSourceGenerationId: null,
         activeResponseMessageId: null,
         terminalRequestIds: {
           ...baseState.terminalRequestIds,
@@ -2135,8 +1809,12 @@ export function reduceAIState<M extends UIMessage>(
       const payload = event.payload as ToolCalledPayload
       const projection = state.activeProjection
       if (
-        isStaleGeneration(state, payload.messageId, payload.generationId) ||
-        terminalGenerationReason(state, payload.generationId) !== undefined
+        isStaleGeneration(state, payload.messageId, payload) ||
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
       ) {
         return state
       }
@@ -2181,6 +1859,7 @@ export function reduceAIState<M extends UIMessage>(
       const terminalReason = terminalGenerationReason(
         state,
         payload.generationId,
+        payload.requestId,
       )
       const completesFailedTool =
         terminalReason === 'failed' &&
@@ -2194,7 +1873,7 @@ export function reduceAIState<M extends UIMessage>(
             activity.status === 'running',
         )
       if (
-        isStaleGeneration(state, payload.messageId, payload.generationId) ||
+        isStaleGeneration(state, payload.messageId, payload) ||
         (terminalReason !== undefined && !completesFailedTool)
       ) {
         return state
@@ -2266,8 +1945,12 @@ export function reduceAIState<M extends UIMessage>(
       const payload = event.payload as ApprovalRequestedPayload
       const projection = state.activeProjection
       if (
-        isStaleGeneration(state, payload.messageId, payload.generationId) ||
-        terminalGenerationReason(state, payload.generationId) !== undefined
+        isStaleGeneration(state, payload.messageId, payload) ||
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
       ) {
         return state
       }
@@ -2329,8 +2012,12 @@ export function reduceAIState<M extends UIMessage>(
     case 'ai.approval.responded': {
       const payload = event.payload as ApprovalRespondedPayload
       if (
-        isStaleGeneration(state, payload.messageId, payload.generationId) ||
-        terminalGenerationReason(state, payload.generationId) !== undefined
+        isStaleGeneration(state, payload.messageId, payload) ||
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
       ) {
         return state
       }
@@ -2398,8 +2085,12 @@ export function reduceAIState<M extends UIMessage>(
     case 'ai.input.requested': {
       const payload = event.payload as InputRequestedPayload
       if (
-        isStaleGeneration(state, payload.messageId, payload.generationId) ||
-        terminalGenerationReason(state, payload.generationId) !== undefined
+        isStaleGeneration(state, payload.messageId, payload) ||
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
       ) {
         return state
       }
@@ -2420,8 +2111,12 @@ export function reduceAIState<M extends UIMessage>(
     case 'ai.input.responded': {
       const payload = event.payload as InputRespondedPayload
       if (
-        isStaleGeneration(state, payload.messageId, payload.generationId) ||
-        terminalGenerationReason(state, payload.generationId) !== undefined
+        isStaleGeneration(state, payload.messageId, payload) ||
+        terminalGenerationReason(
+          state,
+          payload.generationId,
+          payload.requestId,
+        ) !== undefined
       ) {
         return state
       }
@@ -2455,13 +2150,18 @@ export function reduceAIState<M extends UIMessage>(
       ) {
         return state
       }
+      const compaction: CompactionState<M> = {
+        status: 'completed',
+        ...payload,
+        throughIndex:
+          payload.throughIndex ?? state.activeProjection!.promptThroughIndex!,
+      }
       return {
         ...state,
-        compaction: {
-          status: 'completed',
-          ...payload,
-          throughIndex:
-            payload.throughIndex ?? state.activeProjection!.promptThroughIndex!,
+        compaction,
+        activeProjection: state.activeProjection && {
+          ...state.activeProjection,
+          baseCompaction: compaction,
         },
       }
     }
@@ -2483,22 +2183,31 @@ export function deriveUIMessages<M extends UIMessage = UIMessage>(
 export function createReducer<D extends EventDefs>(options: {
   contract: Contract<D>
   name?: string
-}): Reducer<D, AIState<AIMessageFrom<D>>> {
+  fold?: Reducer<D, AIState<AIMessageFrom<D>>>['fold']
+  stateSchema?: StandardSchemaV1<unknown, AIState<AIMessageFrom<D>>>
+}): AIStoredReducer<D, AIState<AIMessageFrom<D>>> {
   type M = AIMessageFrom<D>
   const reducer = options.contract
     .reducer({
-      name: options.name ?? 'a2.ai.state.v17',
+      name: options.name ?? 'a2.ai.state.v22',
       initialState: initialState<M>(),
+      ...(options.stateSchema === undefined
+        ? {}
+        : { stateSchema: options.stateSchema }),
     })
-    .fold((state, event) => reduceAIState(state, event))
-  return { ...reducer, [clientStateProjection]: projectAIClientState }
+    .fold(options.fold ?? ((state, event) => reduceAIState(state, event)))
+  return storedAIReducer({
+    ...reducer,
+    [clientStateProjection]: projectAIClientState,
+  })
 }
 
 type CommandInput<C> = { type: 'ai.control.requested'; payload: C; id?: string }
 type ControlInput<
   M extends UIMessage,
   A extends ControlCommand<M>['action'],
-> = CommandInput<Extract<ControlCommand<M>, { action: A }>>
+> = CommandInput<Extract<ControlCommand<M>, { action: A }>> &
+  (A extends 'send' | 'steer' ? { id: string } : {})
 
 export type AIInputs = {
   message<M extends UIMessage>(
@@ -2515,7 +2224,6 @@ export type AIInputs = {
     }): ControlInput<UIMessage, 'send-now'>[]
     edit<M extends UIMessage>(options: {
       message: M
-      expectedRevision: number
     }): ControlInput<M, 'edit'>[]
     move(options: {
       inputId: string
@@ -2560,7 +2268,7 @@ const messageInput = <M extends UIMessage>(
 ): ControlInput<M, 'send'>[] => [
   {
     type: 'ai.control.requested',
-    id: `ai.message:${message.id}`,
+    id: message.id,
     payload: {
       action: 'send',
       message,
@@ -2579,7 +2287,13 @@ export const inputs: AIInputs = {
     move: (options) => commandInput({ action: 'move', ...options }),
   },
   stop: (options) => commandInput({ action: 'stop', ...options }),
-  steer: (options) => commandInput({ action: 'steer', ...options }),
+  steer: (options) => [
+    {
+      type: 'ai.control.requested',
+      id: options.message.id,
+      payload: { action: 'steer', ...options },
+    },
+  ],
   pause: (options) => commandInput({ action: 'pause', ...options }),
   resume: () => commandInput({ action: 'resume' }),
   message<M extends UIMessage>(message: M, options?: { generate?: boolean }) {
@@ -2591,19 +2305,19 @@ export const inputs: AIInputs = {
   approval: (response) => [
     {
       ...commandInput({ action: 'approval', response })[0]!,
-      id: `ai.approval:${encodeURIComponent(response.messageId)}:${encodeURIComponent(response.generationId)}:${encodeURIComponent(response.approvalId)}:response`,
+      id: `ai.approval:${encodeURIComponent(response.requestId)}:${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`,
+      id: `ai.input:${encodeURIComponent(response.requestId)}:${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`,
+      id: `ai.input:${encodeURIComponent(request.requestId)}:${encodeURIComponent(request.messageId)}:${encodeURIComponent(request.generationId)}:${encodeURIComponent(request.inputId)}:request`,
     },
   ],
   retry: (request) => [
@@ -2674,14 +2388,17 @@ export function agent<
   })
   const reducer = agentContract
     .reducer({
-      name: options.reducerName ?? 'a2.ai.state.v17',
+      name: options.reducerName ?? 'a2.ai.state.v22',
       initialState: initialState<M>(),
     })
     .fold((state, event) => reduceAIState(state, event))
 
   const definition = {
     contract: agentContract,
-    reducer: { ...reducer, [clientStateProjection]: projectAIClientState<M> },
+    reducer: storedAIReducer({
+      ...reducer,
+      [clientStateProjection]: projectAIClientState<M>,
+    }),
   }
   return definition as unknown as AgentDefinition<M, AIEventDefs<M> & E>
 }