experimental-a2 0.12.0 → 0.14.0

package/src/ai-progress-batches.ts

@@ -0,0 +1,68 @@
+import type { UIMessageChunk } from 'ai'
+
+export type ProgressBatch = { index: number; chunks: UIMessageChunk[] }
+type BatchBlock = ProgressBatch[] | { left: BatchBlock; right: BatchBlock }
+
+export type ProgressBatches = {
+  length: number
+  tail: ProgressBatch[]
+  blocks: Array<BatchBlock | null>
+}
+
+const PAGE_SIZE = 64
+
+export const appendProgressBatch = ({
+  batches,
+  batch,
+}: {
+  batches: ProgressBatches
+  batch: ProgressBatch
+}): ProgressBatches => {
+  if (batches.tail.length < PAGE_SIZE)
+    return {
+      ...batches,
+      length: batches.length + 1,
+      tail: [...batches.tail, batch],
+    }
+
+  const blocks = [...batches.blocks]
+  let block: BatchBlock = batches.tail
+  let level = 0
+  for (;;) {
+    const left = blocks[level]
+    if (left == null) break
+    block = { left, right: block }
+    blocks[level] = null
+    level += 1
+  }
+  blocks[level] = block
+  return { length: batches.length + 1, tail: [batch], blocks }
+}
+
+export const progressBatches = (
+  items: readonly ProgressBatch[] = [],
+): ProgressBatches => {
+  let batches: ProgressBatches = { length: 0, tail: [], blocks: [] }
+  for (const batch of items) batches = appendProgressBatch({ batches, batch })
+  return batches
+}
+
+export const flattenProgressBatches = (
+  batches: ProgressBatches,
+): ProgressBatch[] => {
+  const items: ProgressBatch[] = []
+  const visit = (block: BatchBlock): void => {
+    if (Array.isArray(block)) items.push(...block)
+    else {
+      visit(block.left)
+      visit(block.right)
+    }
+  }
+  // Level n holds 2^n completed pages; higher levels contain older batches.
+  for (let level = batches.blocks.length - 1; level >= 0; level -= 1) {
+    const block = batches.blocks[level]
+    if (block) visit(block)
+  }
+  items.push(...batches.tail)
+  return items
+}

package/src/ai-projector.ts

@@ -1,4 +1,5 @@
 import type { UIMessage, UIMessageChunk } from 'ai'
+import { parsePartialJSON } from './parse-partial-json.ts'
 
 type MutablePart = Record<string, unknown>
 
@@ -16,7 +17,13 @@ const isRecord = (value: unknown): value is Record<string, unknown> =>
   typeof value === 'object' && value !== null && !Array.isArray(value)
 
 const cloneMessage = <M extends UIMessage>(messageId: string, base?: M): M =>
-  structuredClone(base ?? { id: messageId, role: 'assistant', parts: [] }) as M
+  (base === undefined
+    ? { id: messageId, role: 'assistant', parts: [] }
+    : {
+        ...base,
+        id: messageId,
+        parts: base.parts.map((part) => ({ ...part })),
+      }) as M
 
 const currentStepParts = (parts: MutablePart[]): MutablePart[] => {
   let start = parts.length - 1
@@ -118,26 +125,47 @@ const mergeMetadata = (current: unknown, update: unknown): unknown => {
     : update
 }
 
-const parseCompleteJSON = (value: string): unknown => {
-  try {
-    return JSON.parse(value) as unknown
-  } catch {
-    return undefined
-  }
+export type UIMessageCursor = {
+  text: Record<string, number>
+  reasoning: Record<string, number>
+  tools: Record<string, PartialToolCall>
 }
 
-/** Internal synchronous counterpart of the AI SDK UI message stream fold. */
 export function projectUIMessage<M extends UIMessage>(options: {
   messageId: string
   base?: M
   chunks: readonly UIMessageChunk[]
 }): M {
+  return projectUIMessageStream(options).message
+}
+
+export function projectUIMessageStream<M extends UIMessage>(options: {
+  messageId: string
+  base?: M
+  cursor?: UIMessageCursor
+  chunks: readonly UIMessageChunk[]
+}): { message: M; cursor: UIMessageCursor } {
   const message = cloneMessage(options.messageId, options.base)
   const record = message as unknown as Record<string, unknown>
   const parts = message.parts as unknown as MutablePart[]
-  const activeText = new Map<string, MutablePart>()
-  const activeReasoning = new Map<string, MutablePart>()
-  const partialTools = new Map<string, PartialToolCall>()
+  const activeText = new Map(
+    Object.entries(options.cursor?.text ?? {}).map(([id, index]) => [
+      id,
+      parts[index]!,
+    ]),
+  )
+  const activeReasoning = new Map(
+    Object.entries(options.cursor?.reasoning ?? {}).map(([id, index]) => [
+      id,
+      parts[index]!,
+    ]),
+  )
+  const partialTools = new Map(
+    Object.entries(options.cursor?.tools ?? {}).map(([id, partial]) => [
+      id,
+      { ...partial },
+    ]),
+  )
 
   for (const chunk of options.chunks) {
     switch (chunk.type) {
@@ -291,7 +319,7 @@ export function projectUIMessage<M extends UIMessage>(options: {
           toolName: partial.toolName,
           dynamic: partial.dynamic,
           state: 'input-streaming',
-          input: parseCompleteJSON(partial.text),
+          input: parsePartialJSON(partial.text),
           providerExecuted: partial.providerExecuted,
           providerMetadata: partial.providerMetadata,
           toolMetadata: partial.toolMetadata,
@@ -424,7 +452,6 @@ export function projectUIMessage<M extends UIMessage>(options: {
         activeReasoning.clear()
         break
       case 'start':
-        if (chunk.messageId !== undefined) record['id'] = chunk.messageId
         if (
           chunk.messageMetadata !== undefined &&
           chunk.messageMetadata !== null
@@ -480,7 +507,38 @@ export function projectUIMessage<M extends UIMessage>(options: {
     }
   }
 
-  return message
+  const indexes = new Map(parts.map((part, index) => [part, index]))
+  return {
+    message,
+    cursor: {
+      text: Object.fromEntries(
+        [...activeText].map(([id, part]) => [id, indexes.get(part)!]),
+      ),
+      reasoning: Object.fromEntries(
+        [...activeReasoning].map(([id, part]) => [id, indexes.get(part)!]),
+      ),
+      tools: Object.fromEntries(partialTools),
+    },
+  }
+}
+
+export function pauseUIMessage<M extends UIMessage>(options: {
+  message: M
+  toolCallIds: ReadonlySet<string>
+}): M {
+  const next = structuredClone(options.message)
+  next.parts = next.parts.filter(
+    (part) =>
+      !('toolCallId' in part && !options.toolCallIds.has(part.toolCallId)),
+  )
+  for (const part of next.parts) {
+    if (
+      (part.type === 'text' || part.type === 'reasoning') &&
+      part.state === 'streaming'
+    )
+      part.state = 'done'
+  }
+  return next
 }
 
 export function interruptUIMessage<M extends UIMessage>(
@@ -488,6 +546,10 @@ export function interruptUIMessage<M extends UIMessage>(
   toolError = 'Tool execution was interrupted.',
 ): M {
   const next = structuredClone(message)
+  next.parts = next.parts.filter((value) => {
+    const part = value as unknown as MutablePart
+    return !isToolPart(part) || part['state'] !== 'input-streaming'
+  })
   for (const value of next.parts as unknown as MutablePart[]) {
     if (value['type'] === 'text' || value['type'] === 'reasoning') {
       if (value['state'] === 'streaming') value['state'] = 'done'
@@ -510,7 +572,6 @@ export function interruptUIMessage<M extends UIMessage>(
       continue
     }
     if (
-      state === 'input-streaming' ||
       state === 'input-available' ||
       state === 'approval-requested' ||
       state === 'approval-responded'

package/src/ai-sdk-step.ts

@@ -10,6 +10,7 @@ import type {
   Instructions,
   LanguageModel,
   LanguageModelUsage,
+  ModelMessage,
   TextStreamPart,
   ToolSet,
   UIMessage,
@@ -127,6 +128,7 @@ export type AISDKStepInput<M extends UIMessage, T extends ToolSet> = {
   model: LanguageModel
   tools: T
   messages: M[]
+  modelMessages?: ModelMessage[]
   responseMessageId: string
   abortSignal: AbortSignal
   instructions?: Instructions
@@ -201,7 +203,9 @@ export async function generateAISDKStep<M extends UIMessage, T extends ToolSet>(
 ): Promise<AISDKStepResult<M>> {
   const settings = safeSettings(input.settings)
   const tools = modelToolSet(input.tools)
-  const messages = await convertToModelMessages(input.messages, { tools })
+  const messages =
+    input.modelMessages ??
+    (await convertToModelMessages(input.messages, { tools }))
   const streamOptions = {
     ...settings,
     model: input.model,
@@ -240,7 +244,6 @@ export async function generateAISDKStep<M extends UIMessage, T extends ToolSet>(
     stream: toUIMessageStream<T, M>({
       stream: observed,
       tools,
-      originalMessages: input.messages,
       generateMessageId: () => input.responseMessageId,
       onError: errorMessage,
       ...(input.messageMetadata === undefined