experimental-a2 0.16.1 → 0.17.0

package/src/ai-stored-state.ts

@@ -1,9 +1,11 @@
+import { assertMessageBase, restoreMessageBase } from './ai-message-base.ts'
+import { sameValue } from './equal-json.ts'
 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'
+import type { Reducer } from './reducer.ts'
+import { withReducerRuntime } from './reducer-runtime.ts'
 
 type Value =
   null | boolean | number | string | Value[] | { [key: string]: Value }
@@ -15,13 +17,11 @@ type Undo =
 
 type Projection = {
   fields: Value
-  hasMessage: boolean
-  message?: Undo
   tools?: Undo
 }
 
 export type AIStoredState = {
-  version: 2
+  version: 3
   fields: Value
   projection: Projection | null
   values: Payload[]
@@ -30,22 +30,11 @@ export type AIStoredState = {
 export type AIStoredView = {
   messages: UIMessage[]
   activeProjection:
-    (MessageProjection<UIMessage> & { baseTools?: AIState['tools'] }) | null
+    (MessageProjection & { 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)
 
@@ -87,39 +76,6 @@ const payloadFields = new Set([
 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')
@@ -284,25 +240,15 @@ const storeState = <S extends AIStoredView>(state: S): AIStoredState => {
   }
   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 { baseTools, ...fields } = projection
     const toolPatch =
       baseTools === undefined
         ? undefined
         : difference(pack(tools), pack(baseTools))
     return {
       fields: pack(fields),
-      hasMessage: baseMessage !== undefined,
-      ...(message === undefined ? {} : { message }),
       ...(toolPatch === undefined ? {} : { tools: toolPatch }),
     }
   }
@@ -310,9 +256,9 @@ const storeState = <S extends AIStoredView>(state: S): AIStoredState => {
   const projection =
     activeProjection === null
       ? null
-      : saveProjection(activeProjection, state.messages, state.tools)
+      : saveProjection(activeProjection, state.tools)
   return {
-    version: 2,
+    version: 3,
     fields: pack(fields),
     projection,
     values,
@@ -352,33 +298,30 @@ const loadState = <S extends AIStoredView>(stored: AIStoredState): S => {
   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'],
+        required: ['fields'],
+        optional: ['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 (projection.base !== undefined) {
+      assertMessageBase(projection.base)
+      const current = state.messages.find(
+        (message) => message.id === projection.responseMessageId,
+      )
+      if (
+        current !== undefined &&
+        restoreMessageBase({ current, base: projection.base }) === undefined
+      )
+        throw new TypeError('AI message base does not match its message')
     }
     if (tools !== undefined)
       projection.baseTools = unpack(
@@ -391,94 +334,36 @@ const loadState = <S extends AIStoredView>(stored: AIStoredState): S => {
   state.activeProjection =
     stored.projection === null
       ? null
-      : loadProjection(stored.projection, fields['messages']!, fields['tools']!)
+      : loadProjection(stored.projection, fields['tools']!)
   return state
 }
 
-export const storedAIReducer = <D extends EventDefs, S extends AIStoredView>(
+export const withAIStorage = <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' }] }
-          }
-        },
+): Reducer<D, S> =>
+  withReducerRuntime({
+    reducer: source,
+    storage: {
+      encode: (state: S) => storeState(state),
+      decode(value: unknown): S {
+        if (
+          !hasKeys(value, {
+            required: ['version', 'fields', 'projection', 'values'],
+          }) ||
+          value['version'] !== 3 ||
+          !Array.isArray(value['values']) ||
+          !isRecord(value['fields']) ||
+          (value['projection'] !== null && !isRecord(value['projection']))
+        )
+          throw new TypeError('invalid AI stored state')
+        const state = value as AIStoredState
+        if (Object.keys(state.values).length !== state.values.length)
+          throw new TypeError('invalid AI payload table')
+        for (const payload of state.values) {
+          if (!hasKeys(payload, { required: [], optional: ['value'] }))
+            throw new TypeError('invalid AI payload table')
+        }
+        return loadState<S>(state)
       },
     },
-    fold(state, event) {
-      const current = view(state)
-      const next = source.fold(current, event)
-      if (next === current) return state
-      return defer(next)
-    },
-  }
-}
+  })