experimental-a2 0.3.0 → 0.5.0

package/src/ai-projector.ts

@@ -0,0 +1,524 @@
+import type { UIMessage, UIMessageChunk } from 'ai'
+
+type MutablePart = Record<string, unknown>
+
+type PartialToolCall = {
+  text: string
+  toolName: string
+  dynamic: boolean
+  providerExecuted?: boolean
+  providerMetadata?: unknown
+  toolMetadata?: unknown
+  title?: string
+}
+
+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
+
+const currentStepParts = (parts: MutablePart[]): MutablePart[] => {
+  let start = parts.length - 1
+  while (start >= 0 && parts[start]?.['type'] !== 'step-start') start -= 1
+  return parts.slice(start + 1)
+}
+
+const isToolPart = (part: MutablePart): boolean =>
+  typeof part['toolCallId'] === 'string' &&
+  (part['type'] === 'dynamic-tool' ||
+    (typeof part['type'] === 'string' && part['type'].startsWith('tool-')))
+
+const findToolPart = (
+  parts: MutablePart[],
+  toolCallId: string,
+): MutablePart | undefined => {
+  const current = currentStepParts(parts).find(
+    (part) => isToolPart(part) && part['toolCallId'] === toolCallId,
+  )
+  if (current) return current
+  return parts.findLast(
+    (part) => isToolPart(part) && part['toolCallId'] === toolCallId,
+  )
+}
+
+const findApprovalPart = (
+  parts: MutablePart[],
+  approvalId: string,
+): MutablePart | undefined =>
+  parts.find((part) => {
+    const approval = part['approval']
+    return (
+      isToolPart(part) && isRecord(approval) && approval['id'] === approvalId
+    )
+  })
+
+const setOptional = (
+  target: MutablePart,
+  key: string,
+  value: unknown,
+): void => {
+  if (value !== undefined) target[key] = value
+}
+
+const setToolPart = (options: {
+  parts: MutablePart[]
+  toolCallId: string
+  toolName: string
+  dynamic: boolean
+  state: string
+  input?: unknown
+  rawInput?: unknown
+  output?: unknown
+  errorText?: string
+  preliminary?: boolean | undefined
+  providerExecuted?: boolean | undefined
+  providerMetadata?: unknown
+  toolMetadata?: unknown
+  title?: string | undefined
+}): MutablePart => {
+  let part = findToolPart(options.parts, options.toolCallId)
+  if (!part) {
+    part = {
+      type: options.dynamic ? 'dynamic-tool' : `tool-${options.toolName}`,
+      ...(options.dynamic ? { toolName: options.toolName } : {}),
+      toolCallId: options.toolCallId,
+    }
+    options.parts.push(part)
+  }
+
+  part['state'] = options.state
+  part['input'] = options.input
+  delete part['output']
+  delete part['rawInput']
+  delete part['errorText']
+  delete part['preliminary']
+  setOptional(part, 'rawInput', options.rawInput)
+  setOptional(part, 'output', options.output)
+  setOptional(part, 'errorText', options.errorText)
+  setOptional(part, 'preliminary', options.preliminary)
+  setOptional(part, 'providerExecuted', options.providerExecuted)
+  setOptional(part, 'title', options.title)
+  setOptional(part, 'toolMetadata', options.toolMetadata)
+  if (options.dynamic) part['toolName'] = options.toolName
+  if (options.providerMetadata !== undefined) {
+    part[
+      options.state === 'output-available' || options.state === 'output-error'
+        ? 'resultProviderMetadata'
+        : 'callProviderMetadata'
+    ] = options.providerMetadata
+  }
+  return part
+}
+
+const mergeMetadata = (current: unknown, update: unknown): unknown => {
+  if (update === undefined || update === null) return current
+  return isRecord(current) && isRecord(update)
+    ? { ...current, ...update }
+    : update
+}
+
+const parseCompleteJSON = (value: string): unknown => {
+  try {
+    return JSON.parse(value) as unknown
+  } catch {
+    return undefined
+  }
+}
+
+/** 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 {
+  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>()
+
+  for (const chunk of options.chunks) {
+    switch (chunk.type) {
+      case 'text-start': {
+        const part: MutablePart = {
+          type: 'text',
+          text: '',
+          state: 'streaming',
+          ...(chunk.providerMetadata === undefined
+            ? {}
+            : { providerMetadata: chunk.providerMetadata }),
+        }
+        activeText.set(chunk.id, part)
+        parts.push(part)
+        break
+      }
+      case 'text-delta': {
+        let part = activeText.get(chunk.id)
+        if (!part) {
+          part = { type: 'text', text: '', state: 'streaming' }
+          activeText.set(chunk.id, part)
+          parts.push(part)
+        }
+        part['text'] =
+          `${typeof part['text'] === 'string' ? part['text'] : ''}${chunk.delta}`
+        setOptional(part, 'providerMetadata', chunk.providerMetadata)
+        break
+      }
+      case 'text-end': {
+        const part = activeText.get(chunk.id)
+        if (part) {
+          part['state'] = 'done'
+          setOptional(part, 'providerMetadata', chunk.providerMetadata)
+          activeText.delete(chunk.id)
+        }
+        break
+      }
+      case 'reasoning-start': {
+        const part: MutablePart = {
+          type: 'reasoning',
+          text: '',
+          state: 'streaming',
+          ...(chunk.providerMetadata === undefined
+            ? {}
+            : { providerMetadata: chunk.providerMetadata }),
+        }
+        activeReasoning.set(chunk.id, part)
+        parts.push(part)
+        break
+      }
+      case 'reasoning-delta': {
+        let part = activeReasoning.get(chunk.id)
+        if (!part) {
+          part = { type: 'reasoning', text: '', state: 'streaming' }
+          activeReasoning.set(chunk.id, part)
+          parts.push(part)
+        }
+        part['text'] =
+          `${typeof part['text'] === 'string' ? part['text'] : ''}${chunk.delta}`
+        setOptional(part, 'providerMetadata', chunk.providerMetadata)
+        break
+      }
+      case 'reasoning-end': {
+        const part = activeReasoning.get(chunk.id)
+        if (part) {
+          part['state'] = 'done'
+          setOptional(part, 'providerMetadata', chunk.providerMetadata)
+          activeReasoning.delete(chunk.id)
+        }
+        break
+      }
+      case 'custom':
+        parts.push({
+          type: 'custom',
+          kind: chunk.kind,
+          ...(chunk.providerMetadata === undefined
+            ? {}
+            : { providerMetadata: chunk.providerMetadata }),
+        })
+        break
+      case 'file':
+      case 'reasoning-file':
+        parts.push({
+          type: chunk.type,
+          mediaType: chunk.mediaType,
+          url: chunk.url,
+          ...(chunk.providerMetadata === undefined
+            ? {}
+            : { providerMetadata: chunk.providerMetadata }),
+        })
+        break
+      case 'source-url':
+        parts.push({
+          type: chunk.type,
+          sourceId: chunk.sourceId,
+          url: chunk.url,
+          ...(chunk.title === undefined ? {} : { title: chunk.title }),
+          ...(chunk.providerMetadata === undefined
+            ? {}
+            : { providerMetadata: chunk.providerMetadata }),
+        })
+        break
+      case 'source-document':
+        parts.push({
+          type: chunk.type,
+          sourceId: chunk.sourceId,
+          mediaType: chunk.mediaType,
+          title: chunk.title,
+          ...(chunk.filename === undefined ? {} : { filename: chunk.filename }),
+          ...(chunk.providerMetadata === undefined
+            ? {}
+            : { providerMetadata: chunk.providerMetadata }),
+        })
+        break
+      case 'tool-input-start': {
+        partialTools.set(chunk.toolCallId, {
+          text: '',
+          toolName: chunk.toolName,
+          dynamic: chunk.dynamic === true,
+          ...(chunk.providerExecuted === undefined
+            ? {}
+            : { providerExecuted: chunk.providerExecuted }),
+          ...(chunk.providerMetadata === undefined
+            ? {}
+            : { providerMetadata: chunk.providerMetadata }),
+          ...(chunk.toolMetadata === undefined
+            ? {}
+            : { toolMetadata: chunk.toolMetadata }),
+          ...(chunk.title === undefined ? {} : { title: chunk.title }),
+        })
+        setToolPart({
+          parts,
+          toolCallId: chunk.toolCallId,
+          toolName: chunk.toolName,
+          dynamic: chunk.dynamic === true,
+          state: 'input-streaming',
+          providerExecuted: chunk.providerExecuted,
+          providerMetadata: chunk.providerMetadata,
+          toolMetadata: chunk.toolMetadata,
+          title: chunk.title,
+        })
+        break
+      }
+      case 'tool-input-delta': {
+        const partial = partialTools.get(chunk.toolCallId)
+        if (!partial) break
+        partial.text += chunk.inputTextDelta
+        setToolPart({
+          parts,
+          toolCallId: chunk.toolCallId,
+          toolName: partial.toolName,
+          dynamic: partial.dynamic,
+          state: 'input-streaming',
+          input: parseCompleteJSON(partial.text),
+          providerExecuted: partial.providerExecuted,
+          providerMetadata: partial.providerMetadata,
+          toolMetadata: partial.toolMetadata,
+          title: partial.title,
+        })
+        break
+      }
+      case 'tool-input-available':
+        setToolPart({
+          parts,
+          toolCallId: chunk.toolCallId,
+          toolName: chunk.toolName,
+          dynamic: chunk.dynamic === true,
+          state: 'input-available',
+          input: chunk.input,
+          providerExecuted: chunk.providerExecuted,
+          providerMetadata: chunk.providerMetadata,
+          toolMetadata: chunk.toolMetadata,
+          title: chunk.title,
+        })
+        partialTools.delete(chunk.toolCallId)
+        break
+      case 'tool-input-error': {
+        const existing = findToolPart(parts, chunk.toolCallId)
+        const dynamic = existing?.['type'] === 'dynamic-tool' || chunk.dynamic
+        setToolPart({
+          parts,
+          toolCallId: chunk.toolCallId,
+          toolName: chunk.toolName,
+          dynamic: dynamic === true,
+          state: 'output-error',
+          input: dynamic ? chunk.input : undefined,
+          rawInput: dynamic ? undefined : chunk.input,
+          errorText: chunk.errorText,
+          providerExecuted: chunk.providerExecuted,
+          providerMetadata: chunk.providerMetadata,
+          toolMetadata: chunk.toolMetadata,
+          title: chunk.title,
+        })
+        partialTools.delete(chunk.toolCallId)
+        break
+      }
+      case 'tool-approval-request': {
+        const part = findToolPart(parts, chunk.toolCallId)
+        if (!part) break
+        part['state'] = 'approval-requested'
+        part['approval'] = {
+          id: chunk.approvalId,
+          ...(chunk.isAutomatic === true ? { isAutomatic: true } : {}),
+          ...(chunk.signature === undefined
+            ? {}
+            : { signature: chunk.signature }),
+        }
+        break
+      }
+      case 'tool-approval-response': {
+        const part = findApprovalPart(parts, chunk.approvalId)
+        if (!part) break
+        const approval = isRecord(part['approval']) ? part['approval'] : {}
+        part['state'] = 'approval-responded'
+        part['approval'] = {
+          id: chunk.approvalId,
+          approved: chunk.approved,
+          ...(chunk.reason === undefined ? {} : { reason: chunk.reason }),
+          ...(approval['isAutomatic'] === true ? { isAutomatic: true } : {}),
+          ...(typeof approval['signature'] === 'string'
+            ? { signature: approval['signature'] }
+            : {}),
+        }
+        setOptional(part, 'providerExecuted', chunk.providerExecuted)
+        setOptional(part, 'callProviderMetadata', chunk.providerMetadata)
+        break
+      }
+      case 'tool-output-available': {
+        const part = findToolPart(parts, chunk.toolCallId)
+        if (!part) break
+        setToolPart({
+          parts,
+          toolCallId: chunk.toolCallId,
+          toolName:
+            part['type'] === 'dynamic-tool' &&
+            typeof part['toolName'] === 'string'
+              ? part['toolName']
+              : String(part['type']).slice('tool-'.length),
+          dynamic: part['type'] === 'dynamic-tool',
+          state: 'output-available',
+          input: part['input'],
+          output: chunk.output,
+          preliminary: chunk.preliminary,
+          providerExecuted: chunk.providerExecuted,
+          providerMetadata: chunk.providerMetadata,
+          toolMetadata: chunk.toolMetadata ?? part['toolMetadata'],
+          title: typeof part['title'] === 'string' ? part['title'] : undefined,
+        })
+        break
+      }
+      case 'tool-output-error': {
+        const part = findToolPart(parts, chunk.toolCallId)
+        if (!part) break
+        setToolPart({
+          parts,
+          toolCallId: chunk.toolCallId,
+          toolName:
+            part['type'] === 'dynamic-tool' &&
+            typeof part['toolName'] === 'string'
+              ? part['toolName']
+              : String(part['type']).slice('tool-'.length),
+          dynamic: part['type'] === 'dynamic-tool',
+          state: 'output-error',
+          input: part['input'],
+          rawInput: part['rawInput'],
+          errorText: chunk.errorText,
+          providerExecuted: chunk.providerExecuted,
+          providerMetadata: chunk.providerMetadata,
+          toolMetadata: chunk.toolMetadata ?? part['toolMetadata'],
+          title: typeof part['title'] === 'string' ? part['title'] : undefined,
+        })
+        break
+      }
+      case 'tool-output-denied': {
+        const part = findToolPart(parts, chunk.toolCallId)
+        if (part) part['state'] = 'output-denied'
+        break
+      }
+      case 'start-step':
+        parts.push({ type: 'step-start' })
+        break
+      case 'finish-step':
+        activeText.clear()
+        activeReasoning.clear()
+        break
+      case 'start':
+        if (chunk.messageId !== undefined) record['id'] = chunk.messageId
+        if (
+          chunk.messageMetadata !== undefined &&
+          chunk.messageMetadata !== null
+        ) {
+          record['metadata'] = mergeMetadata(
+            record['metadata'],
+            chunk.messageMetadata,
+          )
+        }
+        break
+      case 'finish':
+        if (
+          chunk.messageMetadata !== undefined &&
+          chunk.messageMetadata !== null
+        ) {
+          record['metadata'] = mergeMetadata(
+            record['metadata'],
+            chunk.messageMetadata,
+          )
+        }
+        break
+      case 'message-metadata':
+        if (
+          chunk.messageMetadata !== undefined &&
+          chunk.messageMetadata !== null
+        ) {
+          record['metadata'] = mergeMetadata(
+            record['metadata'],
+            chunk.messageMetadata,
+          )
+        }
+        break
+      case 'abort':
+      case 'error':
+        break
+      default: {
+        if (!chunk.type.startsWith('data-') || chunk.transient) break
+        const existing =
+          chunk.id === undefined
+            ? undefined
+            : parts.find(
+                (part) =>
+                  part['type'] === chunk.type && part['id'] === chunk.id,
+              )
+        if (existing) existing['data'] = chunk.data
+        else
+          parts.push({
+            type: chunk.type,
+            ...(chunk.id === undefined ? {} : { id: chunk.id }),
+            data: chunk.data,
+          })
+      }
+    }
+  }
+
+  return message
+}
+
+export function interruptUIMessage<M extends UIMessage>(
+  message: M,
+  toolError = 'Tool execution was interrupted.',
+): M {
+  const next = structuredClone(message)
+  for (const value of next.parts as unknown as MutablePart[]) {
+    if (value['type'] === 'text' || value['type'] === 'reasoning') {
+      if (value['state'] === 'streaming') value['state'] = 'done'
+      continue
+    }
+    if (!isToolPart(value)) continue
+    const state = value['state']
+    if (state === 'approval-responded') {
+      const approval = value['approval']
+      if (isRecord(approval) && approval['approved'] === false) {
+        value['state'] = 'output-denied'
+        continue
+      }
+    }
+    if (state === 'output-available' && value['preliminary'] === true) {
+      delete value['output']
+      delete value['preliminary']
+      value['state'] = 'output-error'
+      value['errorText'] = toolError
+      continue
+    }
+    if (
+      state === 'input-streaming' ||
+      state === 'input-available' ||
+      state === 'approval-requested' ||
+      state === 'approval-responded'
+    ) {
+      if (state === 'approval-requested') delete value['approval']
+      value['state'] = 'output-error'
+      value['errorText'] = toolError
+    }
+  }
+  return next
+}