experimental-a2 0.13.0 → 0.14.1

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

@@ -244,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