experimental-a2 0.15.0 → 0.16.1

package/src/ai-message-projection.ts

@@ -0,0 +1,474 @@
+import type { UIMessage, UIMessageChunk } from 'ai'
+import type {
+  ActiveGenerationProjection,
+  AIState,
+  ApprovalRequestedPayload,
+  ApprovalRespondedPayload,
+  ToolCalledPayload,
+  ToolResultPayload,
+  ToolActivity,
+} from './ai.ts'
+import {
+  type UIMessageCursor,
+  projectUIMessageStream,
+  interruptUIMessage,
+} from './ai-projector.ts'
+import {
+  flattenProgressBatches,
+  appendProgressBatch,
+  progressBatches,
+  type ProgressBatches,
+} from './ai-progress-batches.ts'
+
+export type MessageProjection<M extends UIMessage> = Omit<
+  ActiveGenerationProjection<M>,
+  'baseTools'
+>
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null && !Array.isArray(value)
+
+export 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
+}
+
+export 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),
+  ]
+}
+
+export 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
+  })
+
+export 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 }),
+})
+
+export const approvalRespondedChunk = (
+  payload: ApprovalRespondedPayload,
+): UIMessageChunk => ({
+  type: 'tool-approval-response',
+  approvalId: payload.approvalId,
+  approved: payload.approved,
+  ...(payload.reason === undefined ? {} : { reason: payload.reason }),
+})
+
+export 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
+}
+
+export const projectedMessageStream = <M extends UIMessage>(
+  projection: MessageProjection<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),
+  })
+}
+
+export const projectedMessage = <M extends UIMessage>(
+  projection: MessageProjection<M>,
+): M => projectedMessageStream(projection).message
+
+export 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,
+  ]
+}
+
+export const advanceProjection = <M extends UIMessage>({
+  projection,
+  base,
+  chunks,
+  index,
+}: {
+  projection: MessageProjection<M>
+  base: M | undefined
+  chunks: UIMessageChunk[]
+  index: number
+}): {
+  message: M
+  batches: ProgressBatches
+  cursor: NonNullable<MessageProjection<M>['cursor']>
+} => {
+  const ordered = index > (projection.batches.tail.at(-1)?.index ?? -1)
+  const batches = ordered
+    ? appendProgressBatch({
+        batches: projection.batches,
+        batch: { index, chunks },
+      })
+    : progressBatches(
+        [
+          ...flattenProgressBatches(projection.batches).filter(
+            (batch) => batch.index !== index,
+          ),
+          { index, chunks },
+        ].toSorted((left, right) => left.index - right.index),
+      )
+  const incremental =
+    base !== undefined &&
+    projection.cursor != null &&
+    index > projection.cursor.throughIndex
+  const projected = incremental
+    ? projectUIMessageStream({
+        messageId: projection.responseMessageId,
+        base,
+        cursor: projection.cursor!.stream,
+        chunks,
+      })
+    : projectedMessageStream({ ...projection, batches })
+  return {
+    message: projected.message,
+    batches,
+    cursor: {
+      throughIndex: incremental
+        ? index
+        : Math.max(
+            index,
+            batches.tail.at(-1)?.index ?? 0,
+            projection.toolEvents.reduce(
+              (latest, event) => Math.max(latest, event.index),
+              0,
+            ),
+            projection.approvalEvents.reduce(
+              (latest, event) => Math.max(latest, event.index),
+              0,
+            ),
+          ),
+      stream: projected.cursor,
+    },
+  }
+}
+
+export const projectedTools = <M extends UIMessage>(
+  projection: ActiveGenerationProjection<M>,
+): ToolActivity[] =>
+  projection.toolEvents.reduce(
+    (tools, event) => reduceToolActivity(tools, event),
+    projection.baseTools,
+  )
+
+export const restoreGenerationBase = <M extends UIMessage>(
+  state: AIState<M>,
+  generationId: string,
+  responseMessageId: string,
+  reason: 'interrupted' | 'superseded' | 'failed',
+): 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),
+  }
+}
+
+export 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,
+  )
+
+export 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),
+  }
+}