experimental-a2 0.16.1 → 0.18.0

package/src/ai-message-base.ts

@@ -0,0 +1,209 @@
+import type { UIMessage } from 'ai'
+import { sameValue } from './equal-json.ts'
+
+type FieldUndo = { value?: unknown } | { remove: true } | { length: number }
+type FieldsUndo = Record<string, FieldUndo>
+type PartUndo = { value: UIMessage['parts'][number] } | { fields: FieldsUndo }
+
+export type MessageBase = {
+  partCount: number
+  fields?: FieldsUndo
+  parts?: Record<string, PartUndo>
+}
+
+const record = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null && !Array.isArray(value)
+
+const assertFields = (fields: unknown): void => {
+  if (fields === undefined) return
+  if (!record(fields)) throw new TypeError('invalid AI message base fields')
+  for (const [key, undo] of Object.entries(fields)) {
+    if (
+      !record(undo) ||
+      Object.keys(undo).length > 1 ||
+      !(
+        Object.keys(undo).length === 0 ||
+        'value' in undo ||
+        undo['remove'] === true ||
+        (key === 'text' &&
+          Number.isSafeInteger(undo['length']) &&
+          (undo['length'] as number) >= 0)
+      )
+    )
+      throw new TypeError('invalid AI message base field')
+  }
+}
+
+/** Checks decoded recovery data; `restoreMessageBase` trusts this shape. */
+export function assertMessageBase(
+  value: unknown,
+): asserts value is MessageBase {
+  if (
+    !record(value) ||
+    !Number.isSafeInteger(value['partCount']) ||
+    (value['partCount'] as number) < 0 ||
+    (value['partCount'] as number) > 0xffffffff - 1 ||
+    Object.keys(value).some(
+      (key) => !['partCount', 'fields', 'parts'].includes(key),
+    ) ||
+    (value['parts'] !== undefined && !record(value['parts']))
+  )
+    throw new TypeError('invalid AI message base')
+  assertFields(value['fields'])
+  const partCount = value['partCount'] as number
+  for (const [key, undo] of Object.entries(value['parts'] ?? {})) {
+    const index = Number(key)
+    if (
+      !Number.isSafeInteger(index) ||
+      String(index) !== key ||
+      index < 0 ||
+      index >= partCount ||
+      !record(undo) ||
+      Object.keys(undo).length !== 1 ||
+      !(('value' in undo && record(undo['value'])) || 'fields' in undo)
+    )
+      throw new TypeError('invalid AI message base part')
+    if ('fields' in undo) assertFields(undo['fields'])
+  }
+}
+
+const restoreFields = (
+  current: Record<string, unknown>,
+  fields: FieldsUndo | undefined,
+): Record<string, unknown> | undefined => {
+  if (fields === undefined) return current
+  const restored = { ...current }
+  for (const [key, undo] of Object.entries(fields)) {
+    if ('remove' in undo) delete restored[key]
+    else if ('length' in undo) {
+      const text = current[key]
+      if (typeof text !== 'string' || undo.length > text.length)
+        return undefined
+      restored[key] = text.slice(0, undo.length)
+    } else
+      Object.defineProperty(restored, key, {
+        value: undo.value,
+        enumerable: true,
+        writable: true,
+        configurable: true,
+      })
+  }
+  return restored
+}
+
+/** `undefined` when `base` no longer describes `current`, including a removed message. */
+export const restoreMessageBase = <M extends UIMessage>(options: {
+  current: M | undefined
+  base: MessageBase | undefined
+}): M | undefined => {
+  const { current, base } = options
+  if (base === undefined || current === undefined) return undefined
+  const parts = current.parts.slice(0, base.partCount)
+  for (const [key, undo] of Object.entries(base.parts ?? {})) {
+    const index = Number(key)
+    if ('value' in undo) {
+      parts[index] = undo.value as M['parts'][number]
+      continue
+    }
+    const part = parts[index]
+    const restored =
+      part === undefined
+        ? undefined
+        : restoreFields(part as unknown as Record<string, unknown>, undo.fields)
+    if (!restored) return undefined
+    parts[index] = restored as M['parts'][number]
+  }
+  if (
+    parts.length !== base.partCount ||
+    Object.keys(parts).length !== parts.length
+  )
+    return undefined
+  const { parts: _parts, ...fields } = current
+  const restored = restoreFields(fields, base.fields)
+  return restored && ({ ...restored, parts } as M)
+}
+
+const updateFields = (options: {
+  before: Record<string, unknown>
+  after: Record<string, unknown>
+  undo?: FieldsUndo
+}): FieldsUndo | undefined => {
+  const { before, after, undo } = options
+  const original = restoreFields(before, undo)
+  if (!original) return undefined
+  return Object.fromEntries(
+    [...new Set([...Object.keys(original), ...Object.keys(after)])].flatMap(
+      (key) => {
+        if (
+          Object.hasOwn(original, key) === Object.hasOwn(after, key) &&
+          sameValue(original[key], after[key])
+        )
+          return []
+        const value = original[key]
+        const next = after[key]
+        const change: FieldUndo = !Object.hasOwn(original, key)
+          ? { remove: true }
+          : key === 'text' &&
+              typeof value === 'string' &&
+              typeof next === 'string' &&
+              next.startsWith(value)
+            ? { length: value.length }
+            : value === undefined
+              ? {}
+              : { value }
+        return [[key, change]]
+      },
+    ),
+  )
+}
+
+/** `undefined` when `before` no longer matches `base`, so the recovery must be dropped. */
+export const updateMessageBase = (options: {
+  base: MessageBase
+  before: UIMessage
+  after: UIMessage
+}): MessageBase | undefined => {
+  const { base, before, after } = options
+  if (before === after) return base
+  const { parts: beforeParts, ...beforeFields } = before
+  const { parts: afterParts, ...afterFields } = after
+  const fields = updateFields({
+    before: beforeFields,
+    after: afterFields,
+    ...(base.fields && { undo: base.fields }),
+  })
+  if (!fields) return undefined
+  const parts: Array<[string, PartUndo]> = []
+  for (let index = 0; index < base.partCount; index += 1) {
+    const previous = beforeParts[index]
+    const current = afterParts[index]
+    const undo = base.parts?.[index]
+    const original = undo && 'value' in undo ? undo.value : previous
+    if (!original) return undefined
+    if (previous === current && undo === undefined) continue
+    const originalFields = original as unknown as Record<string, unknown>
+    const undoFields = undo && 'fields' in undo ? undo.fields : undefined
+    if (!current) {
+      const value = restoreFields(originalFields, undoFields)
+      if (!value) return undefined
+      parts.push([
+        String(index),
+        { value: value as UIMessage['parts'][number] },
+      ])
+      continue
+    }
+    const changed = updateFields({
+      before: originalFields,
+      after: current as unknown as Record<string, unknown>,
+      ...(undoFields && { undo: undoFields }),
+    })
+    if (!changed) return undefined
+    if (Object.keys(changed).length > 0)
+      parts.push([String(index), { fields: changed }])
+  }
+  return {
+    partCount: base.partCount,
+    ...(Object.keys(fields).length > 0 && { fields }),
+    ...(parts.length > 0 && { parts: Object.fromEntries(parts) }),
+  }
+}

package/src/ai-message-projection.ts

@@ -1,7 +1,15 @@
+import {
+  restoreMessageBase,
+  updateMessageBase,
+  type MessageBase,
+} from './ai-message-base.ts'
 import type { UIMessage, UIMessageChunk } from 'ai'
 import type {
   ActiveGenerationProjection,
   AIState,
+  CompactionPart,
+  CompactionPartData,
+  CompactionRequestedPayload,
   ApprovalRequestedPayload,
   ApprovalRespondedPayload,
   ToolCalledPayload,
@@ -10,6 +18,7 @@ import type {
 } from './ai.ts'
 import {
   type UIMessageCursor,
+  type ProjectionChunk,
   projectUIMessageStream,
   interruptUIMessage,
 } from './ai-projector.ts'
@@ -20,23 +29,26 @@ import {
   type ProgressBatches,
 } from './ai-progress-batches.ts'
 
-export type MessageProjection<M extends UIMessage> = Omit<
-  ActiveGenerationProjection<M>,
-  'baseTools'
->
+export type MessageProjection = Omit<ActiveGenerationProjection, 'baseTools'>
 
 const isRecord = (value: unknown): value is Record<string, unknown> =>
   typeof value === 'object' && value !== null && !Array.isArray(value)
 
+const positionOf = (messages: readonly UIMessage[], id: string): number =>
+  messages.findLastIndex((message) => message.id === id)
+
+export const findMessage = <M extends UIMessage>(
+  messages: M[],
+  id: string,
+): M | undefined => messages.findLast((message) => message.id === id)
+
 export const upsertMessage = <M extends UIMessage>(
   messages: M[],
   message: M,
 ): M[] => {
-  const index = messages.findIndex((candidate) => candidate.id === message.id)
+  const index = positionOf(messages, message.id)
   if (index === -1) return [...messages, message]
-  const next = [...messages]
-  next[index] = message
-  return next
+  return messages[index] === message ? messages : messages.with(index, message)
 }
 
 export const upsertResponse = <M extends UIMessage>(options: {
@@ -45,15 +57,11 @@ export const upsertResponse = <M extends UIMessage>(options: {
   inputId: string
 }): M[] => {
   const { messages, message, inputId } = options
-  if (messages.some((item) => item.id === message.id))
+  if (positionOf(messages, message.id) !== -1)
     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 inputIndex = positionOf(messages, inputId)
+  const position = inputIndex === -1 ? messages.length : inputIndex + 1
+  return [...messages.slice(0, position), message, ...messages.slice(position)]
 }
 
 export const answerApproval = <M extends UIMessage>(
@@ -108,6 +116,10 @@ export const approvalRespondedChunk = (
   ...(payload.reason === undefined ? {} : { reason: payload.reason }),
 })
 
+export const toolCalledChunk = (
+  payload: ToolCalledPayload,
+): ProjectionChunk => ({ type: 'a2-tool-called', call: payload })
+
 export const toolResultChunk = (payload: ToolResultPayload): UIMessageChunk => {
   if (payload.phase === 'input') {
     return {
@@ -171,17 +183,27 @@ export const toolResultChunk = (payload: ToolResultPayload): UIMessageChunk => {
   } as UIMessageChunk
 }
 
-export const projectedMessageStream = <M extends UIMessage>(
-  projection: MessageProjection<M>,
-): { message: M; cursor: UIMessageCursor } => {
-  const projectionEvents = flattenProgressBatches(projection.batches)
+export const projectedMessageStream = <M extends UIMessage>({
+  projection,
+  current,
+}: {
+  projection: MessageProjection
+  current: M | undefined
+}): { message: M; cursor: UIMessageCursor; throughIndex: number } => {
+  const base = restoreMessageBase({ current, base: projection.base })
+  const projectionEvents: Array<{
+    index: number
+    chunks: readonly ProjectionChunk[]
+  }> = flattenProgressBatches(projection.batches)
   for (const event of projection.toolEvents) {
-    if (event.type === 'ai.tool.result') {
-      projectionEvents.push({
-        index: event.index,
-        chunks: [toolResultChunk(event.payload)],
-      })
-    }
+    projectionEvents.push({
+      index: event.index,
+      chunks: [
+        event.type === 'ai.tool.called'
+          ? toolCalledChunk(event.payload)
+          : toolResultChunk(event.payload),
+      ],
+    })
   }
   for (const event of projection.approvalEvents) {
     projectionEvents.push({
@@ -193,20 +215,23 @@ export const projectedMessageStream = <M extends UIMessage>(
       ],
     })
   }
-  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 ordered = projectionEvents.toSorted(
+    (left, right) => left.index - right.index,
+  )
+  return {
+    ...projectUIMessageStream<M>({
+      messageId: projection.responseMessageId,
+      ...(base === undefined ? {} : { base }),
+      chunks: ordered.flatMap((event) => event.chunks),
+    }),
+    throughIndex: ordered.at(-1)?.index ?? 0,
+  }
 }
 
-export const projectedMessage = <M extends UIMessage>(
-  projection: MessageProjection<M>,
-): M => projectedMessageStream(projection).message
+export const projectedMessage = <M extends UIMessage>(options: {
+  projection: MessageProjection
+  current: M | undefined
+}): M => projectedMessageStream(options).message
 
 export const reduceToolActivity = (
   tools: ToolActivity[],
@@ -216,12 +241,17 @@ export const reduceToolActivity = (
 ): ToolActivity[] => {
   if (event.type === 'ai.tool.called') {
     const payload = event.payload
-    const { providerMetadata, ...called } = payload
     const activity: ToolActivity = {
-      ...called,
-      ...(providerMetadata === undefined
+      requestId: payload.requestId,
+      messageId: payload.messageId,
+      generationId: payload.generationId,
+      toolCallId: payload.toolCallId,
+      toolName: payload.toolName,
+      ...(payload.dynamic === undefined ? {} : { dynamic: payload.dynamic }),
+      ...(payload.providerExecuted === undefined
         ? {}
-        : { callProviderMetadata: providerMetadata }),
+        : { providerExecuted: payload.providerExecuted }),
+      ...(payload.title === undefined ? {} : { title: payload.title }),
       status: 'running',
     }
     return [
@@ -241,22 +271,15 @@ export const reduceToolActivity = (
       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
       ? {}
@@ -264,11 +287,6 @@ export const reduceToolActivity = (
     ...(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'
@@ -294,14 +312,14 @@ export const advanceProjection = <M extends UIMessage>({
   chunks,
   index,
 }: {
-  projection: MessageProjection<M>
+  projection: MessageProjection
   base: M | undefined
   chunks: UIMessageChunk[]
   index: number
 }): {
   message: M
   batches: ProgressBatches
-  cursor: NonNullable<MessageProjection<M>['cursor']>
+  cursor: NonNullable<MessageProjection['cursor']>
 } => {
   const ordered = index > (projection.batches.tail.at(-1)?.index ?? -1)
   const batches = ordered
@@ -328,7 +346,10 @@ export const advanceProjection = <M extends UIMessage>({
         cursor: projection.cursor!.stream,
         chunks,
       })
-    : projectedMessageStream({ ...projection, batches })
+    : projectedMessageStream({
+        projection: { ...projection, batches },
+        current: base,
+      })
   return {
     message: projected.message,
     batches,
@@ -352,8 +373,8 @@ export const advanceProjection = <M extends UIMessage>({
   }
 }
 
-export const projectedTools = <M extends UIMessage>(
-  projection: ActiveGenerationProjection<M>,
+export const projectedTools = (
+  projection: ActiveGenerationProjection,
 ): ToolActivity[] =>
   projection.toolEvents.reduce(
     (tools, event) => reduceToolActivity(tools, event),
@@ -370,12 +391,22 @@ export const restoreGenerationBase = <M extends UIMessage>(
     state.activeProjection?.generationId === generationId
       ? state.activeProjection
       : null
+  const restored = restoreMessageBase({
+    current: findMessage(state.messages, responseMessageId),
+    base: projection?.base,
+  })
+  // A superseded attempt's unfinished compaction reruns in its replacement;
+  // a completed one stays in effect, and so does its part.
+  const base =
+    restored !== undefined && reason === 'superseded'
+      ? withoutRunningCompactionPart(restored, generationId)
+      : restored
   return {
     ...state,
     messages:
-      projection?.baseMessage === undefined
+      base === undefined
         ? state.messages.filter((message) => message.id !== responseMessageId)
-        : upsertMessage(state.messages, projection.baseMessage),
+        : upsertMessage(state.messages, base),
     tools:
       projection === null
         ? state.tools.filter(
@@ -429,7 +460,7 @@ export const interruptGenerationProjection = <M extends UIMessage>(
       ? state.activeProjection
       : null
   if (projection === null) return null
-  const retained: ActiveGenerationProjection<M> = {
+  const retained: ActiveGenerationProjection = {
     ...projection,
     batches: progressBatches(
       flattenProgressBatches(projection.batches).filter(
@@ -444,7 +475,7 @@ export const interruptGenerationProjection = <M extends UIMessage>(
     ),
   }
   const shouldProject =
-    retained.baseMessage !== undefined || retained.batches.length > 0
+    retained.base !== undefined || retained.batches.length > 0
   const restored = restoreGenerationBase(
     state,
     generationId,
@@ -456,7 +487,12 @@ export const interruptGenerationProjection = <M extends UIMessage>(
     messages: shouldProject
       ? upsertMessage(
           state.messages,
-          interruptUIMessage(projectedMessage(retained)),
+          interruptUIMessage(
+            projectedMessage({
+              projection: retained,
+              current: findMessage(state.messages, responseMessageId),
+            }),
+          ),
         )
       : restored.messages,
     tools: stopGenerationTools(
@@ -472,3 +508,165 @@ export const interruptGenerationProjection = <M extends UIMessage>(
         : restored.usage.filter((entry) => entry.generationId !== generationId),
   }
 }
+
+export const retainMessageBase = <
+  S extends {
+    messages: UIMessage[]
+    activeProjection: { responseMessageId: string; base?: MessageBase } | null
+  },
+>(options: {
+  before: S
+  after: S
+}): S => {
+  const { before, after } = options
+  const projection = after.activeProjection
+  if (
+    before.messages === after.messages ||
+    !projection?.base ||
+    projection.base !== before.activeProjection?.base
+  )
+    return after
+  const previous = findMessage(before.messages, projection.responseMessageId)
+  const current = findMessage(after.messages, projection.responseMessageId)
+  if (previous !== undefined && previous === current) return after
+  // A response removed or rewritten outside this tracking (for example, by an
+  // extension's own transition) cannot be restored, so a later re-created
+  // message must not inherit the stale recovery.
+  const base =
+    previous &&
+    current &&
+    updateMessageBase({
+      base: projection.base,
+      before: previous,
+      after: current,
+    })
+  const { base: _stale, ...unrecoverable } = projection
+  return {
+    ...after,
+    activeProjection: base ? { ...unrecoverable, base } : unrecoverable,
+  }
+}
+
+export const isCompactionPart = (part: {
+  type: string
+}): part is CompactionPart => part.type === 'data-a2-compaction'
+
+const compactionPartOf = (message: UIMessage, generationId: string): number =>
+  message.parts.findIndex(
+    (part) => isCompactionPart(part) && part.id === generationId,
+  )
+
+const withoutRunningCompactionPart = <M extends UIMessage>(
+  message: M,
+  generationId: string,
+): M => {
+  const index = compactionPartOf(message, generationId)
+  const part = message.parts[index]
+  return part !== undefined &&
+    isCompactionPart(part) &&
+    part.data.status === 'running'
+    ? { ...message, parts: message.parts.toSpliced(index, 1) }
+    : message
+}
+
+/**
+ * Compaction runs before its model step streams, so the part lands between
+ * the response's earlier steps and this one. Counting it in the step's base
+ * keeps it through Stop cutoffs, failures, and retries of that step.
+ */
+export const addCompactionPart = <M extends UIMessage>(
+  state: AIState<M>,
+  payload: CompactionRequestedPayload,
+): AIState<M> => {
+  const projection = state.activeProjection
+  if (projection?.generationId !== payload.generationId) return state
+  const current = findMessage(state.messages, projection.responseMessageId)
+  if (
+    current === undefined
+      ? projection.base !== undefined
+      : projection.base?.partCount !== current.parts.length
+  )
+    return state
+  const inputTokens = payload.inputTokens ?? payload.inputTokenEstimate
+  const part: CompactionPart = {
+    type: 'data-a2-compaction',
+    id: payload.generationId,
+    data: {
+      status: 'running',
+      ...(inputTokens === undefined ? {} : { inputTokens }),
+      ...(payload.thresholdTokens === undefined
+        ? {}
+        : { thresholdTokens: payload.thresholdTokens }),
+    },
+  }
+  const message = (
+    current === undefined
+      ? { id: projection.responseMessageId, role: 'assistant', parts: [part] }
+      : { ...current, parts: [...current.parts, part] }
+  ) as M
+  return {
+    ...state,
+    messages: upsertResponse({
+      messages: state.messages,
+      message,
+      inputId: state.activeGeneration?.messageId ?? message.id,
+    }),
+    activeProjection: {
+      ...projection,
+      cursor: null,
+      base:
+        projection.base === undefined
+          ? { partCount: 1 }
+          : { ...projection.base, partCount: projection.base.partCount + 1 },
+    },
+  }
+}
+
+export const settleCompactionPart = <M extends UIMessage>(options: {
+  state: AIState<M>
+  generationId: string
+  data: Partial<CompactionPartData>
+}): AIState<M> => {
+  const { state, generationId, data } = options
+  const position = state.messages.findLastIndex(
+    (message) => compactionPartOf(message, generationId) !== -1,
+  )
+  if (position === -1) return state
+  const message = state.messages[position]!
+  const index = compactionPartOf(message, generationId)
+  const part = message.parts[index] as unknown as CompactionPart
+  return {
+    ...state,
+    messages: state.messages.with(position, {
+      ...message,
+      parts: message.parts.with(index, {
+        ...part,
+        data: { ...part.data, ...data },
+      } as unknown as M['parts'][number]),
+    }),
+  }
+}
+
+/** A running compaction whose state reverted without completing was abandoned. */
+export const settleAbandonedCompaction = <M extends UIMessage>(
+  before: AIState<M>,
+  after: AIState<M>,
+): AIState<M> => {
+  const running = before.compaction
+  if (
+    running?.status !== 'running' ||
+    after.compaction === running ||
+    (after.compaction?.status === 'completed' &&
+      after.compaction.generationId === running.generationId)
+  )
+    return after
+  const interrupted =
+    after.status === 'closed' ||
+    after.pausedGenerationIds?.[running.generationId] === true ||
+    after.terminalGenerations?.[running.generationId] === 'interrupted'
+  return settleCompactionPart({
+    state: after,
+    generationId: running.generationId,
+    data: { status: interrupted ? 'interrupted' : 'failed' },
+  })
+}