experimental-a2 0.14.1 → 0.16.0

package/src/ai-stored-state.ts

@@ -0,0 +1,484 @@
+import type { UIMessage } from 'ai'
+import type { AIState } from './ai.ts'
+import type { MessageProjection } from './ai-message-projection.ts'
+import type { EventDefs } from './contract.ts'
+import { clientStateProjection, type Reducer } from './reducer.ts'
+import { validateSync } from './validate.ts'
+
+type Value =
+  null | boolean | number | string | Value[] | { [key: string]: Value }
+type Payload = { value?: unknown }
+type Undo =
+  | { value: Value }
+  | { length: number; items: Record<string, Undo> }
+  | { fields: Record<string, Undo>; remove: string[] }
+
+type Projection = {
+  fields: Value
+  hasMessage: boolean
+  message?: Undo
+  tools?: Undo
+}
+
+export type AIStoredState = {
+  version: 2
+  fields: Value
+  projection: Projection | null
+  values: Payload[]
+}
+
+export type AIStoredView = {
+  messages: UIMessage[]
+  activeProjection:
+    (MessageProjection<UIMessage> & { baseTools?: AIState['tools'] }) | null
+  tools?: AIState['tools']
+  usage?: AIState['usage']
+}
+
+export const aiStateReducer: unique symbol = Symbol('a2.ai.stateReducer')
+
+export type AIStoredReducer<
+  D extends EventDefs,
+  S extends AIStoredView,
+> = Reducer<D, AIStoredState, Record<never, never>, S> & {
+  readonly [aiStateReducer]: Reducer<D, S>
+  view(state: AIStoredState): S
+  fromView(state: S): AIStoredState
+}
+
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null && !Array.isArray(value)
+
+const hasKeys = (
+  value: unknown,
+  { required, optional = [] }: { required: string[]; optional?: string[] },
+): value is Record<string, unknown> =>
+  isRecord(value) &&
+  required.every((key) => Object.hasOwn(value, key)) &&
+  Object.keys(value).every(
+    (key) => required.includes(key) || optional.includes(key),
+  )
+
+const reference = (value: Value): number | undefined => {
+  if (!isRecord(value) || !Object.hasOwn(value, '$a2Value')) return undefined
+  const id = value['$a2Value']
+  if (
+    typeof id !== 'number' ||
+    !Number.isSafeInteger(id) ||
+    id < 0 ||
+    Object.keys(value).length !== 1
+  )
+    throw new TypeError('invalid AI payload reference')
+  return id
+}
+
+const payloadFields = new Set([
+  'input',
+  'output',
+  'rawInput',
+  'metadata',
+  'providerMetadata',
+  'callProviderMetadata',
+  'resultProviderMetadata',
+  'toolMetadata',
+  'data',
+])
+
+const keyFor = (value: unknown): string =>
+  value === undefined ? 'undefined' : JSON.stringify(value)
+
+const sameValue = (left: unknown, right: unknown): boolean => {
+  if (Object.is(left, right)) return true
+  if (
+    typeof left !== 'object' ||
+    left === null ||
+    typeof right !== 'object' ||
+    right === null
+  )
+    return false
+  const prototype = Object.getPrototypeOf(left)
+  if (
+    prototype !== Object.getPrototypeOf(right) ||
+    (prototype !== Object.prototype &&
+      prototype !== Array.prototype &&
+      prototype !== null)
+  )
+    return false
+  if (
+    Array.isArray(left) &&
+    (!Array.isArray(right) || left.length !== right.length)
+  )
+    return false
+  const keys = Reflect.ownKeys(left)
+  return (
+    keys.length === Reflect.ownKeys(right).length &&
+    keys.every(
+      (key) =>
+        Object.hasOwn(right, key) &&
+        sameValue(Reflect.get(left, key), Reflect.get(right, key)),
+    )
+  )
+}
+
+const restore = (current: Value, undo: Undo | undefined): Value => {
+  if (undo === undefined) return current
+  if (!isRecord(undo)) throw new TypeError('invalid AI recovery change')
+  if ('value' in undo) {
+    if (!hasKeys(undo, { required: ['value'] }))
+      throw new TypeError('invalid AI replacement recovery')
+    return undo.value
+  }
+  if ('length' in undo) {
+    if (
+      !hasKeys(undo, { required: ['length', 'items'] }) ||
+      !Number.isSafeInteger(undo.length) ||
+      undo.length < 0 ||
+      undo.length > 0xffffffff ||
+      !isRecord(undo.items) ||
+      !Array.isArray(current)
+    )
+      throw new TypeError('invalid AI array recovery')
+    const items = Object.entries(undo.items)
+    let added = 0
+    for (const [key, patch] of items) {
+      const index = Number(key)
+      if (
+        !Number.isSafeInteger(index) ||
+        index < 0 ||
+        index >= undo.length ||
+        String(index) !== key ||
+        patch === undefined
+      )
+        throw new TypeError('invalid AI array recovery index')
+      if (index >= current.length) added += 1
+    }
+    if (added !== Math.max(0, undo.length - current.length))
+      throw new TypeError('incomplete AI array recovery')
+    const result = current.slice(0, undo.length)
+    result.length = undo.length
+    for (const [index, patch] of items)
+      result[Number(index)] = restore(current[Number(index)] ?? null, patch)
+    return result
+  }
+  if (
+    !hasKeys(undo, { required: ['fields', 'remove'] }) ||
+    !isRecord(undo.fields) ||
+    !Array.isArray(undo.remove) ||
+    Object.keys(undo.remove).length !== undo.remove.length ||
+    Array.from(undo.remove).some((key) => typeof key !== 'string') ||
+    new Set(undo.remove).size !== undo.remove.length ||
+    !isRecord(current) ||
+    reference(current) !== undefined
+  )
+    throw new TypeError('invalid AI object recovery')
+  const result = { ...current } as Record<string, Value>
+  for (const key of undo.remove) delete result[key]
+  for (const [key, patch] of Object.entries(undo.fields)) {
+    if (patch === undefined || undo.remove.includes(key))
+      throw new TypeError('invalid AI object recovery field')
+    Object.defineProperty(result, key, {
+      value: restore(Object.hasOwn(result, key) ? result[key]! : null, patch),
+      enumerable: true,
+      configurable: true,
+      writable: true,
+    })
+  }
+  return result
+}
+
+const difference = (current: Value, before: Value): Undo | undefined => {
+  if (current === before) return undefined
+  if (reference(current) !== undefined || reference(before) !== undefined)
+    return reference(current) === reference(before)
+      ? undefined
+      : { value: before }
+  if (Array.isArray(current) && Array.isArray(before)) {
+    const items: Record<string, Undo> = {}
+    for (const [index, value] of before.entries()) {
+      const patch =
+        index < current.length ? difference(current[index]!, value) : { value }
+      if (patch !== undefined) items[index] = patch
+    }
+    return current.length === before.length && Object.keys(items).length === 0
+      ? undefined
+      : { length: before.length, items }
+  }
+  if (isRecord(current) && isRecord(before)) {
+    const fields: Record<string, Undo> = {}
+    const remove = Object.keys(current).filter(
+      (key) => !Object.hasOwn(before, key),
+    )
+    for (const [key, value] of Object.entries(before)) {
+      const patch = Object.hasOwn(current, key)
+        ? difference(current[key] as Value, value as Value)
+        : { value: value as Value }
+      if (patch !== undefined)
+        Object.defineProperty(fields, key, {
+          value: patch,
+          enumerable: true,
+          configurable: true,
+          writable: true,
+        })
+    }
+    return remove.length === 0 && Object.keys(fields).length === 0
+      ? undefined
+      : { fields, remove }
+  }
+  return { value: before }
+}
+
+const storeState = <S extends AIStoredView>(state: S): AIStoredState => {
+  const packed = new WeakMap<object, Value>()
+  const values: Payload[] = []
+  const ids = new Map<unknown, number>()
+  const contents = new Map<string, number[]>()
+  const intern = (value: unknown): Value => {
+    let id = ids.get(value)
+    if (id !== undefined) return { $a2Value: id }
+    const key =
+      typeof value === 'object' && value !== null ? keyFor(value) : undefined
+    const candidates = key === undefined ? undefined : contents.get(key)
+    if (candidates !== undefined)
+      id = candidates.find((candidate) =>
+        sameValue(values[candidate]!.value, value),
+      )
+    if (id === undefined) {
+      id = values.length
+      values.push(value === undefined ? {} : { value })
+    }
+    ids.set(value, id)
+    if (key !== undefined && !candidates?.includes(id))
+      contents.set(key, [...(candidates ?? []), id])
+    return { $a2Value: id }
+  }
+  const pack = (value: unknown, field?: string): Value => {
+    if (
+      value === undefined ||
+      (typeof value === 'string' && value.length > 80) ||
+      (field !== undefined &&
+        payloadFields.has(field) &&
+        !(field === 'input' && isRecord(value) && 'message' in value)) ||
+      (isRecord(value) && Object.hasOwn(value, '$a2Value'))
+    )
+      return intern(value)
+    if (
+      value === null ||
+      typeof value === 'string' ||
+      typeof value === 'boolean' ||
+      typeof value === 'number'
+    )
+      return value
+    if (typeof value !== 'object')
+      throw new TypeError('AI state must contain JSON values')
+    const cached = packed.get(value)
+    if (cached !== undefined) return cached
+    const result: Value = Array.isArray(value)
+      ? value.map((item) => pack(item))
+      : Object.fromEntries(
+          Object.entries(value)
+            .toSorted(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0))
+            .map(([key, item]) => [key, pack(item, key)]),
+        )
+    packed.set(value, result)
+    return result
+  }
+  const saveProjection = (
+    projection: NonNullable<AIStoredView['activeProjection']>,
+    messages: UIMessage[],
+    tools: AIStoredView['tools'],
+  ): Projection => {
+    const { baseMessage, baseTools, ...fields } = projection
+    const current = messages.find(
+      (message) => message.id === projection.responseMessageId,
+    )
+    const message =
+      baseMessage === undefined
+        ? undefined
+        : difference(pack(current), pack(baseMessage))
+    const toolPatch =
+      baseTools === undefined
+        ? undefined
+        : difference(pack(tools), pack(baseTools))
+    return {
+      fields: pack(fields),
+      hasMessage: baseMessage !== undefined,
+      ...(message === undefined ? {} : { message }),
+      ...(toolPatch === undefined ? {} : { tools: toolPatch }),
+    }
+  }
+  const { activeProjection, ...fields } = state
+  const projection =
+    activeProjection === null
+      ? null
+      : saveProjection(activeProjection, state.messages, state.tools)
+  return {
+    version: 2,
+    fields: pack(fields),
+    projection,
+    values,
+  }
+}
+
+const loadState = <S extends AIStoredView>(stored: AIStoredState): S => {
+  const unpack = (value: Value): unknown => {
+    const id = reference(value)
+    if (id !== undefined) {
+      if (!Object.hasOwn(stored.values, id))
+        throw new TypeError('missing AI payload')
+      return stored.values[id]!.value
+    }
+    if (Array.isArray(value)) {
+      if (Object.keys(value).length !== value.length)
+        throw new TypeError('invalid AI stored array')
+      return Array.from(value, unpack)
+    }
+    if (isRecord(value))
+      return Object.fromEntries(
+        Object.entries(value).map(([key, item]) => [
+          key,
+          unpack(item as Value),
+        ]),
+      )
+    if (
+      value === null ||
+      typeof value === 'boolean' ||
+      typeof value === 'string' ||
+      (typeof value === 'number' && Number.isFinite(value))
+    )
+      return value
+    throw new TypeError('invalid AI stored value')
+  }
+  const fields = stored.fields as Record<string, Value>
+  const state = unpack(stored.fields) as S
+  const loadProjection = (
+    saved: Projection,
+    messages: Value,
+    tools: Value | undefined,
+  ): AIStoredView['activeProjection'] => {
+    if (
+      !hasKeys(saved, {
+        required: ['fields', 'hasMessage'],
+        optional: ['message', 'tools'],
+      }) ||
+      typeof saved.hasMessage !== 'boolean' ||
+      !isRecord(saved.fields) ||
+      ('message' in saved &&
+        (!saved.hasMessage || saved.message === undefined)) ||
+      ('tools' in saved && saved.tools === undefined)
+    )
+      throw new TypeError('invalid AI projection recovery')
+    const projection = unpack(saved.fields) as NonNullable<
+      AIStoredView['activeProjection']
+    >
+    const currentMessages = unpack(messages) as UIMessage[]
+    const index = currentMessages.findIndex(
+      (message) => message.id === projection.responseMessageId,
+    )
+    if (saved.hasMessage) {
+      const current = index === -1 ? null : (messages as Value[])[index]!
+      projection.baseMessage = unpack(
+        restore(current, saved.message),
+      ) as UIMessage
+    }
+    if (tools !== undefined)
+      projection.baseTools = unpack(
+        restore(tools, saved.tools),
+      ) as AIState['tools']
+    else if (saved.tools !== undefined)
+      throw new TypeError('missing AI recovery tools')
+    return projection
+  }
+  state.activeProjection =
+    stored.projection === null
+      ? null
+      : loadProjection(stored.projection, fields['messages']!, fields['tools']!)
+  return state
+}
+
+export const storedAIReducer = <D extends EventDefs, S extends AIStoredView>(
+  source: Reducer<D, S>,
+): AIStoredReducer<D, S> => {
+  const views = new WeakMap<AIStoredState, S>()
+  const view = (state: AIStoredState): S => {
+    let result = views.get(state)
+    if (result === undefined) {
+      result = loadState<S>(state)
+      views.set(state, result)
+    }
+    return result
+  }
+  const defer = (next: S): AIStoredState => {
+    let stored: AIStoredState | undefined
+    const read = (): AIStoredState => (stored ??= storeState(next))
+    const state: AIStoredState = {
+      version: 2,
+      get fields() {
+        return read().fields
+      },
+      get projection() {
+        return read().projection
+      },
+      get values() {
+        return read().values
+      },
+    }
+    views.set(state, next)
+    return state
+  }
+  return {
+    ...source,
+    name: `${source.name}:records-v2`,
+    initialState: storeState(source.initialState),
+    [aiStateReducer]: source,
+    view,
+    fromView: (state) => storeState(state),
+    [clientStateProjection]: (state) =>
+      source[clientStateProjection]?.(view(state)) ?? view(state),
+    stateSchema: {
+      '~standard': {
+        version: 1,
+        vendor: 'a2-ai',
+        validate(value) {
+          try {
+            if (
+              !hasKeys(value, {
+                required: ['version', 'fields', 'projection', 'values'],
+              }) ||
+              value['version'] !== 2 ||
+              !Array.isArray(value['values']) ||
+              !isRecord(value['fields']) ||
+              (value['projection'] !== null && !isRecord(value['projection']))
+            )
+              return { issues: [{ message: 'invalid AI stored state' }] }
+            const state = value as AIStoredState
+            if (Object.keys(state.values).length !== state.values.length)
+              return { issues: [{ message: 'invalid AI payload table' }] }
+            for (const payload of state.values) {
+              if (!hasKeys(payload, { required: [], optional: ['value'] }))
+                return { issues: [{ message: 'invalid AI payload table' }] }
+            }
+            const decoded = loadState<S>(state)
+            if (source.stateSchema) {
+              const result = validateSync(
+                source.stateSchema,
+                decoded,
+                'AI view stateSchema',
+              )
+              if (result.issues) return result
+              return { value: storeState(result.value) }
+            }
+            return { value: state }
+          } catch {
+            return { issues: [{ message: 'invalid AI stored state' }] }
+          }
+        },
+      },
+    },
+    fold(state, event) {
+      const current = view(state)
+      const next = source.fold(current, event)
+      if (next === current) return state
+      return defer(next)
+    },
+  }
+}