experimental-a2 0.16.0 → 0.17.0

package/src/ai-message-projection.ts

@@ -1,13 +1,27 @@
+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,
   ToolResultPayload,
   ToolActivity,
 } from './ai.ts'
-import { type UIMessageCursor, projectUIMessageStream } from './ai-projector.ts'
+import {
+  type UIMessageCursor,
+  type ProjectionChunk,
+  projectUIMessageStream,
+  interruptUIMessage,
+} from './ai-projector.ts'
 import {
   flattenProgressBatches,
   appendProgressBatch,
@@ -15,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: {
@@ -40,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>(
@@ -103,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 {
@@ -166,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({
@@ -188,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[],
@@ -211,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 [
@@ -236,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
       ? {}
@@ -259,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'
@@ -289,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
@@ -323,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,
@@ -346,3 +372,301 @@ export const advanceProjection = <M extends UIMessage>({
     },
   }
 }
+
+export const projectedTools = (
+  projection: ActiveGenerationProjection,
+): 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
+  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:
+      base === undefined
+        ? state.messages.filter((message) => message.id !== responseMessageId)
+        : upsertMessage(state.messages, base),
+    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 = {
+    ...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.base !== undefined || retained.batches.length > 0
+  const restored = restoreGenerationBase(
+    state,
+    generationId,
+    responseMessageId,
+    'interrupted',
+  )
+  return {
+    ...restored,
+    messages: shouldProject
+      ? upsertMessage(
+          state.messages,
+          interruptUIMessage(
+            projectedMessage({
+              projection: retained,
+              current: findMessage(state.messages, responseMessageId),
+            }),
+          ),
+        )
+      : 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),
+  }
+}
+
+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' },
+  })
+}

package/src/ai-progress-batches.ts

@@ -1,4 +1,30 @@
 import type { UIMessageChunk } from 'ai'
+import type { GenerationProgressPayload } from './ai.ts'
+
+/** Progress already in a log may name its generation by ids instead of `startIndex`. */
+export type StoredProgressPayload =
+  | GenerationProgressPayload
+  | {
+      requestId: string
+      generationId: string
+      responseMessageId: string
+      chunks: UIMessageChunk[]
+    }
+
+export const progressOwnedBy = ({
+  payload,
+  owner,
+  startIndex,
+}: {
+  payload: StoredProgressPayload
+  owner: { requestId: string; generationId: string; responseMessageId: string }
+  startIndex: number
+}): boolean =>
+  'startIndex' in payload
+    ? payload.startIndex === startIndex
+    : payload.requestId === owner.requestId &&
+      payload.generationId === owner.generationId &&
+      payload.responseMessageId === owner.responseMessageId
 
 export type ProgressBatch = { index: number; chunks: UIMessageChunk[] }
 type BatchBlock = ProgressBatch[] | { left: BatchBlock; right: BatchBlock }