experimental-a2 0.15.0 → 0.16.0

package/src/ai-context.ts

@@ -0,0 +1,680 @@
+import type { UIMessage } from 'ai'
+import type {
+  AgentDefinition,
+  ApprovalRequestedPayload,
+  ApprovalRespondedPayload,
+  CompactionCompletedPayload,
+  CompactionRequestedPayload,
+  CompactionState,
+  GenerationCompletedPayload,
+  GenerationFailedPayload,
+  GenerationProgressPayload,
+  GenerationRequestedPayload,
+  GenerationStartedPayload,
+  MessageCreatedPayload,
+  MessageInterruptedPayload,
+  ModelMetadataRequestedPayload,
+  ModelMetadataResolvedPayload,
+  ModelMetadataState,
+  ToolCalledPayload,
+  ToolResultPayload,
+} from './ai.ts'
+import type { ControlCommit } from './ai-control.ts'
+import { contextSchema } from './ai-context-schema.ts'
+import type { ContractEvent, EventDefs } from './contract.ts'
+import {
+  projectedMessage,
+  advanceProjection,
+  reduceToolActivity,
+  upsertMessage,
+  upsertResponse,
+  type MessageProjection,
+} from './ai-message-projection.ts'
+import { interruptUIMessage, pauseUIMessage } from './ai-projector.ts'
+import {
+  flattenProgressBatches,
+  progressBatches,
+} from './ai-progress-batches.ts'
+import { storedAIReducer, type AIStoredReducer } from './ai-stored-state.ts'
+
+export type AIModelContext<M extends UIMessage = UIMessage> = {
+  messages: M[]
+  compaction: Omit<
+    CompactionState<M>,
+    'messages' | 'retainedMessageIds' | 'usage'
+  > | null
+  modelMetadata: Record<string, ModelMetadataState>
+}
+
+export type AIContextState<M extends UIMessage> = AIModelContext<M> & {
+  closed: boolean
+  authority: Pick<
+    ControlCommit['view'],
+    | 'requestId'
+    | 'requestReason'
+    | 'sourceGenerationId'
+    | 'responseMessageId'
+    | 'phase'
+  > | null
+  generation: {
+    owner: GenerationStartedPayload
+    phase: 'streaming' | 'completed' | 'paused' | 'failed' | 'superseded'
+    error?: string
+    responseIndex?: number
+  } | null
+  activeProjection: MessageProjection<M> | null
+  messageIndexes: Record<string, number>
+}
+
+const replaceResponse = <M extends UIMessage>({
+  state,
+  message,
+  index,
+}: {
+  state: AIContextState<M>
+  message: M
+  index: number
+}): AIContextState<M> => {
+  let messages = upsertResponse({
+    messages: state.messages,
+    message,
+    inputId: state.generation!.owner.messageId,
+  })
+  if (
+    state.generation?.responseIndex !== undefined &&
+    !state.messages.some((item) => item.id === message.id)
+  ) {
+    messages = messages.filter((item) => item.id !== message.id)
+    messages.splice(state.generation.responseIndex, 0, message)
+  }
+  return {
+    ...state,
+    messages,
+    messageIndexes: Object.hasOwn(state.messageIndexes, message.id)
+      ? state.messageIndexes
+      : { ...state.messageIndexes, [message.id]: index },
+  }
+}
+
+const calledTools = <M extends UIMessage>(
+  projection: MessageProjection<M>,
+): Set<string> =>
+  new Set([
+    ...(projection.baseMessage?.parts.flatMap((part) =>
+      'toolCallId' in part ? [part.toolCallId] : [],
+    ) ?? []),
+    ...projection.toolEvents
+      .filter((event) => event.type === 'ai.tool.called')
+      .map((event) => event.payload.toolCallId),
+  ])
+
+const renderProjection = <M extends UIMessage>({
+  state,
+  index,
+}: {
+  state: AIContextState<M>
+  index: number
+}): AIContextState<M> => {
+  const projection = state.activeProjection!
+  let message = projectedMessage(projection)
+  if (
+    projection.baseMessage === undefined &&
+    projection.batches.length === 0 &&
+    message.parts.length === 0
+  )
+    return state
+  if (state.generation?.phase === 'paused')
+    message = pauseUIMessage({ message, toolCallIds: calledTools(projection) })
+  if (state.generation?.phase === 'failed')
+    message = interruptUIMessage(
+      message,
+      `Generation failed before tool completion: ${state.generation.error}`,
+    )
+  return replaceResponse({
+    state: { ...state, activeProjection: { ...projection, cursor: null } },
+    message,
+    index,
+  })
+}
+
+export const generationBase = <M extends UIMessage>(
+  state: AIContextState<M>,
+): AIContextState<M> => {
+  const projection = state.activeProjection
+  if (!projection) return state
+  return {
+    ...state,
+    messages:
+      projection.baseMessage === undefined
+        ? state.messages.filter(
+            (message) => message.id !== projection.responseMessageId,
+          )
+        : upsertMessage(state.messages, projection.baseMessage),
+    compaction: projection.baseCompaction ?? null,
+  }
+}
+
+export const contextForAttempt = <M extends UIMessage>(
+  state: AIContextState<M>,
+): AIContextState<M> => ({
+  ...(state.activeProjection?.completionIndex === undefined
+    ? generationBase(state)
+    : state),
+  generation: state.generation && { ...state.generation, phase: 'superseded' },
+  activeProjection: null,
+})
+
+const endGeneration = <M extends UIMessage>(
+  state: AIContextState<M>,
+): AIContextState<M> => ({
+  ...state,
+  authority: null,
+  generation: null,
+  activeProjection: null,
+})
+
+const ownsGeneration = <M extends UIMessage>(
+  state: AIContextState<M>,
+  owner: { generationId: string; requestId: string },
+): boolean =>
+  state.generation?.owner.generationId === owner.generationId &&
+  state.generation.owner.requestId === owner.requestId &&
+  state.generation.phase !== 'superseded'
+
+export const createContextReducer = <M extends UIMessage, D extends EventDefs>({
+  agent,
+}: {
+  agent: AgentDefinition<M, D>
+}): AIStoredReducer<D, AIContextState<M>> => {
+  const initialState: AIContextState<M> = {
+    messages: [],
+    compaction: null,
+    modelMetadata: {},
+    closed: false,
+    authority: null,
+    generation: null,
+    activeProjection: null,
+    messageIndexes: {},
+  }
+  const reducer = agent.contract
+    .reducer({
+      name: 'a2.ai.context.v3',
+      initialState,
+      stateSchema: contextSchema<M>(),
+    })
+    .fold((state, event) => reduceContext(state, event))
+  return storedAIReducer(reducer)
+}
+
+const reduceContext = <M extends UIMessage, D extends EventDefs>(
+  state: AIContextState<M>,
+  event: ContractEvent<D>,
+): AIContextState<M> => {
+  if (state.closed) return state
+  switch (event.type) {
+    case 'ai.control.committed': {
+      const { view } = event.payload as ControlCommit
+      const {
+        requestId,
+        requestReason,
+        sourceGenerationId,
+        responseMessageId,
+        phase,
+      } = view
+      const next = {
+        ...state,
+        authority: {
+          requestId,
+          requestReason,
+          sourceGenerationId,
+          responseMessageId,
+          phase,
+        },
+        closed: view.closed,
+      }
+      return view.closed ? endGeneration(next) : next
+    }
+    case 'ai.session.closed':
+      return endGeneration({
+        ...state,
+        closed: true,
+        compaction:
+          state.compaction?.status === 'running'
+            ? (state.activeProjection?.baseCompaction ?? null)
+            : state.compaction,
+      })
+    case 'ai.message.created': {
+      const { message } = event.payload as MessageCreatedPayload<M>
+      return {
+        ...state,
+        messages: upsertMessage(state.messages, message),
+        activeProjection:
+          state.activeProjection?.responseMessageId === message.id
+            ? { ...state.activeProjection, cursor: null }
+            : state.activeProjection,
+        messageIndexes: { ...state.messageIndexes, [message.id]: event.index },
+      }
+    }
+    case 'ai.generation.requested': {
+      const request = event.payload as GenerationRequestedPayload
+      if (
+        state.authority?.requestId !== event.id ||
+        request.responseMessageId !== state.authority.responseMessageId ||
+        request.reason !== state.authority.requestReason ||
+        (request.sourceGenerationId ?? null) !==
+          state.authority.sourceGenerationId ||
+        request.reason !== 'retry'
+      )
+        return state
+      return contextForAttempt(state)
+    }
+    case 'ai.generation.started': {
+      const owner = event.payload as GenerationStartedPayload
+      if (
+        state.authority?.requestId !== owner.requestId ||
+        state.authority.responseMessageId !== owner.responseMessageId ||
+        state.authority.phase === 'paused' ||
+        state.authority.phase === 'pausing'
+      )
+        return state
+      const previous = state.generation
+      if (
+        previous?.owner.requestId === owner.requestId &&
+        (previous.owner.attempt >= owner.attempt ||
+          previous.phase === 'failed' ||
+          previous.phase === 'completed' ||
+          state.activeProjection?.completionIndex !== undefined)
+      )
+        return state
+      const base =
+        previous?.owner.requestId === owner.requestId
+          ? generationBase(state)
+          : state
+      const baseMessage = base.messages.find(
+        (message) => message.id === owner.responseMessageId,
+      )
+      return {
+        ...base,
+        generation: {
+          owner,
+          phase: 'streaming',
+          ...(previous?.owner.responseMessageId === owner.responseMessageId &&
+          previous.responseIndex !== undefined
+            ? { responseIndex: previous.responseIndex }
+            : {}),
+        },
+        activeProjection: {
+          generationId: owner.generationId,
+          responseMessageId: owner.responseMessageId,
+          promptThroughIndex: owner.promptThroughIndex ?? event.index - 1,
+          ...(baseMessage === undefined ? {} : { baseMessage }),
+          baseCompaction: base.compaction,
+          batches: progressBatches(),
+          toolEvents: [],
+          approvalEvents: [],
+        },
+      }
+    }
+    case 'ai.generation.progress': {
+      const payload = event.payload as GenerationProgressPayload
+      if (
+        !ownsGeneration(state, payload) ||
+        state.generation!.owner.responseMessageId !==
+          payload.responseMessageId ||
+        state.generation!.phase !== 'streaming' ||
+        !state.activeProjection
+      )
+        return state
+      const projection = state.activeProjection
+      const progress = advanceProjection({
+        projection,
+        base: state.messages.find(
+          (message) => message.id === projection.responseMessageId,
+        ),
+        chunks: payload.chunks,
+        index: event.index,
+      })
+      return replaceResponse({
+        state: {
+          ...state,
+          activeProjection: {
+            ...projection,
+            batches: progress.batches,
+            cursor: progress.cursor,
+          },
+        },
+        message: progress.message,
+        index: event.index,
+      })
+    }
+    case 'ai.generation.completed': {
+      const payload = event.payload as GenerationCompletedPayload
+      if (
+        !ownsGeneration(state, payload) ||
+        state.generation!.phase !== 'streaming'
+      )
+        return state
+      return {
+        ...state,
+        generation: { ...state.generation!, phase: 'completed' },
+        activeProjection: {
+          ...state.activeProjection!,
+          completionIndex: event.index,
+          cursor: null,
+        },
+      }
+    }
+    case 'ai.generation.failed': {
+      const payload = event.payload as GenerationFailedPayload
+      if (!ownsGeneration(state, payload)) return state
+      if (
+        state.generation!.phase !== 'streaming' &&
+        !(
+          payload.stepLimit &&
+          state.authority?.requestReason === 'tool' &&
+          state.authority.sourceGenerationId === payload.generationId
+        )
+      )
+        return state
+      if (payload.superseded) {
+        const next = generationBase(state)
+        return {
+          ...next,
+          generation: { ...state.generation!, phase: 'superseded' },
+          activeProjection: null,
+        }
+      }
+      return renderProjection({
+        state: {
+          ...state,
+          compaction: state.activeProjection?.baseCompaction ?? null,
+          generation: {
+            ...state.generation!,
+            phase: 'failed',
+            error: payload.error,
+          },
+        },
+        index: event.index,
+      })
+    }
+    case 'ai.message.interrupted': {
+      const payload = event.payload as MessageInterruptedPayload
+      const owner = state.generation?.owner
+      const ownsRequest =
+        payload.requestId !== undefined &&
+        payload.requestId === state.authority?.requestId &&
+        payload.messageId === state.authority.responseMessageId
+      const ownsOwner =
+        owner?.responseMessageId === payload.messageId &&
+        (payload.generationId === owner.generationId ||
+          payload.requestId === owner.requestId)
+      if (!ownsRequest && !ownsOwner) return state
+      if (
+        ownsOwner &&
+        state.authority?.requestId !== owner!.requestId &&
+        !ownsRequest
+      )
+        return state
+      const projection = state.activeProjection
+      if (!projection) return endGeneration(state)
+      const cutoff = payload.lastSeenIndex ?? Number.POSITIVE_INFINITY
+      const retained = {
+        ...projection,
+        batches: progressBatches(
+          flattenProgressBatches(projection.batches).filter(
+            (batch) => batch.index <= cutoff,
+          ),
+        ),
+        toolEvents: projection.toolEvents.filter(
+          (item) => item.index <= cutoff,
+        ),
+        approvalEvents: projection.approvalEvents.filter(
+          (item) => item.index <= cutoff,
+        ),
+      }
+      let next = generationBase(state)
+      if (retained.baseMessage !== undefined || retained.batches.length > 0)
+        next = replaceResponse({
+          state: next,
+          message: interruptUIMessage(projectedMessage(retained)),
+          index: event.index,
+        })
+      return endGeneration(next)
+    }
+    case 'ai.message.completed': {
+      const { messageId } = event.payload as { messageId: string }
+      return state.generation?.owner.responseMessageId === messageId &&
+        state.generation.phase !== 'failed' &&
+        state.generation.phase !== 'superseded'
+        ? endGeneration(state)
+        : state
+    }
+    case 'ai.generation.paused': {
+      const { generationId } = event.payload as { generationId: string }
+      if (
+        state.generation?.owner.generationId !== generationId ||
+        state.generation.phase === 'failed' ||
+        state.generation.phase === 'superseded' ||
+        !state.activeProjection
+      )
+        return state
+      const toolCallIds = calledTools(state.activeProjection)
+      return renderProjection({
+        state: {
+          ...state,
+          generation: { ...state.generation, phase: 'paused' },
+          compaction:
+            state.compaction?.status === 'running'
+              ? (state.activeProjection.baseCompaction ?? null)
+              : state.compaction,
+          activeProjection: {
+            ...state.activeProjection,
+            cursor: null,
+            batches: progressBatches(
+              flattenProgressBatches(state.activeProjection.batches).map(
+                (batch) => ({
+                  index: batch.index,
+                  chunks: batch.chunks.filter(
+                    (chunk) =>
+                      !('toolCallId' in chunk) ||
+                      toolCallIds.has(chunk.toolCallId),
+                  ),
+                }),
+              ),
+            ),
+          },
+        },
+        index: event.index,
+      })
+    }
+    case 'ai.tool.called':
+    case 'ai.tool.result': {
+      const payload = event.payload as ToolCalledPayload | ToolResultPayload
+      if (
+        !ownsGeneration(state, payload) ||
+        !state.activeProjection ||
+        state.generation!.owner.responseMessageId !== payload.messageId
+      )
+        return state
+      if (state.generation!.phase === 'failed') {
+        const result = payload as ToolResultPayload
+        const running = state.activeProjection.toolEvents
+          .reduce(
+            (tools, item) => reduceToolActivity(tools, item),
+            [] as ReturnType<typeof reduceToolActivity>,
+          )
+          .some(
+            (tool) =>
+              tool.toolCallId === payload.toolCallId &&
+              tool.status === 'running',
+          )
+        if (event.type !== 'ai.tool.result' || result.preliminary || !running)
+          return state
+      }
+      const toolEvent =
+        event.type === 'ai.tool.called'
+          ? {
+              index: event.index,
+              type: 'ai.tool.called' as const,
+              payload: payload as ToolCalledPayload,
+            }
+          : {
+              index: event.index,
+              type: 'ai.tool.result' as const,
+              payload: payload as ToolResultPayload,
+            }
+      const next = {
+        ...state,
+        activeProjection: {
+          ...state.activeProjection,
+          toolEvents: [...state.activeProjection.toolEvents, toolEvent],
+        },
+      }
+      if (event.type === 'ai.tool.called') return next
+      next.activeProjection.cursor = null
+      return renderProjection({ state: next, index: event.index })
+    }
+    case 'ai.approval.requested':
+    case 'ai.approval.responded': {
+      const payload = event.payload as
+        ApprovalRequestedPayload | ApprovalRespondedPayload
+      if (
+        !ownsGeneration(state, payload) ||
+        !state.activeProjection ||
+        state.generation!.phase === 'failed' ||
+        state.generation!.owner.responseMessageId !== payload.messageId
+      )
+        return state
+      const approvalEvent =
+        event.type === 'ai.approval.requested'
+          ? {
+              index: event.index,
+              type: 'ai.approval.requested' as const,
+              payload: payload as ApprovalRequestedPayload,
+            }
+          : {
+              index: event.index,
+              type: 'ai.approval.responded' as const,
+              payload: payload as ApprovalRespondedPayload,
+            }
+      return renderProjection({
+        state: {
+          ...state,
+          activeProjection: {
+            ...state.activeProjection,
+            cursor: null,
+            approvalEvents: [
+              ...state.activeProjection.approvalEvents,
+              approvalEvent,
+            ],
+          },
+        },
+        index: event.index,
+      })
+    }
+    case 'ai.compaction.requested': {
+      const payload = event.payload as CompactionRequestedPayload
+      if (
+        state.generation?.owner.generationId !== payload.generationId ||
+        state.generation.phase !== 'streaming'
+      )
+        return state
+      return {
+        ...state,
+        compaction: {
+          status: 'running',
+          generationId: payload.generationId,
+          throughMessageId: payload.throughMessageId,
+          throughIndex:
+            payload.throughIndex ?? state.activeProjection!.promptThroughIndex!,
+        },
+      }
+    }
+    case 'ai.compaction.completed': {
+      const payload = event.payload as CompactionCompletedPayload<M>
+      const projection = state.activeProjection
+      if (
+        state.generation?.owner.generationId !== payload.generationId ||
+        state.generation.phase !== 'streaming' ||
+        state.compaction?.status !== 'running' ||
+        state.compaction.generationId !== payload.generationId ||
+        !projection
+      )
+        return state
+      const throughIndex =
+        payload.throughIndex ?? projection.promptThroughIndex!
+      const replacements = new Set(
+        payload.messages.map((message) => message.id),
+      )
+      const retained = new Set(payload.retainedMessageIds ?? [])
+      const messages = [
+        ...payload.messages,
+        ...state.messages.filter(
+          (message) =>
+            !replacements.has(message.id) &&
+            (retained.has(message.id) ||
+              (state.messageIndexes[message.id] ?? 0) > throughIndex),
+        ),
+      ]
+      const activeProjection = { ...projection, cursor: null }
+      delete activeProjection.baseMessage
+      const baseMessage = messages.find(
+        (message) => message.id === projection.responseMessageId,
+      )
+      if (baseMessage !== undefined) activeProjection.baseMessage = baseMessage
+      const compaction: AIModelContext<M>['compaction'] = {
+        status: 'completed',
+        generationId: payload.generationId,
+        throughMessageId: payload.throughMessageId,
+        throughIndex,
+        ...(payload.summary === undefined ? {} : { summary: payload.summary }),
+      }
+      activeProjection.baseCompaction = compaction
+      return {
+        ...state,
+        messages,
+        compaction,
+        generation: {
+          ...state.generation,
+          responseIndex: payload.messages.length,
+        },
+        activeProjection,
+        messageIndexes: Object.fromEntries(
+          messages.map((message) => [
+            message.id,
+            replacements.has(message.id)
+              ? throughIndex
+              : state.messageIndexes[message.id]!,
+          ]),
+        ),
+      }
+    }
+    case 'ai.model.metadata.requested': {
+      const { modelId } = event.payload as ModelMetadataRequestedPayload
+      return Object.hasOwn(state.modelMetadata, modelId)
+        ? state
+        : {
+            ...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 },
+        },
+      }
+    }
+    default:
+      return state
+  }
+}