experimental-a2 0.16.1 → 0.18.0

package/src/ai-app-state.ts

@@ -0,0 +1,91 @@
+import type { AIState } from './ai.ts'
+import type { ContractEvent, EventDefs } from './contract.ts'
+
+/** An application reducer folded beside the AI protocol, from events alone. */
+export type AppReducer<D extends EventDefs, X extends object> = {
+  /** Snapshot identity of the application state. Rename it when `fold` changes. */
+  name: string
+  initialState: X
+  fold: (state: X, event: ContractEvent<D>) => X
+}
+
+// Every AIState key, so application state is exactly the rest of a merged state.
+const aiStateKeys: Record<keyof AIState, true> = {
+  inbox: true,
+  pendingQueueCommands: true,
+  active: true,
+  starting: true,
+  receipt: true,
+  pausedGenerationIds: true,
+  session: true,
+  messages: true,
+  status: true,
+  activeGeneration: true,
+  activeRequestId: true,
+  activeRequestReason: true,
+  activeRequestSourceGenerationId: true,
+  activeResponseMessageId: true,
+  activeProjection: true,
+  responseOwners: true,
+  terminalRequestIds: true,
+  terminalGenerations: true,
+  pendingApprovals: true,
+  pendingInputs: true,
+  tools: true,
+  compaction: true,
+  modelMetadata: true,
+  usage: true,
+  error: true,
+}
+
+export const isAIStateKey = (key: string): boolean =>
+  Object.hasOwn(aiStateKeys, key)
+
+const keep = (state: object, ai: boolean): Record<string, unknown> => {
+  const kept: Record<string, unknown> = {}
+  for (const [key, value] of Object.entries(state))
+    if (isAIStateKey(key) === ai) kept[key] = value
+  return kept
+}
+
+/** The application part of a merged state: every key `AIState` does not own. */
+export const appSlice = <X extends object>(state: object): X =>
+  keep(state, false) as X
+
+/** A merged state with its application part replaced. */
+export const withAppSlice = <S extends object>(state: S, app: object): S =>
+  ({ ...keep(state, true), ...keep(app, false) }) as S
+
+// agentContract() registers every reducer it builds here, keyed by the reducer
+// so spread contracts keep it: agent() serves only these, and folds their
+// application state into the model-context checkpoint every generation reads.
+const agentReducers = new WeakMap<
+  object,
+  AppReducer<EventDefs, object> | null
+>()
+
+export const registerAgentReducer = (
+  reducer: object,
+  app: AppReducer<EventDefs, object> | undefined,
+): void => {
+  agentReducers.set(reducer, app ?? null)
+}
+
+const reducerOf = (contract: { reducer?: unknown }): object | undefined =>
+  typeof contract.reducer === 'object' && contract.reducer !== null
+    ? contract.reducer
+    : undefined
+
+export const isAgentContract = (contract: { reducer?: unknown }): boolean => {
+  const reducer = reducerOf(contract)
+  return reducer !== undefined && agentReducers.has(reducer)
+}
+
+export const appReducerOf = (contract: {
+  reducer?: unknown
+}): AppReducer<EventDefs, object> | undefined => {
+  const reducer = reducerOf(contract)
+  return reducer === undefined
+    ? undefined
+    : (agentReducers.get(reducer) ?? undefined)
+}

package/src/ai-client-state.ts

@@ -1,7 +1,10 @@
 import type { UIMessage } from 'ai'
 import type { AIState } from './ai.ts'
 import type { ControlCommand, InboxItem } from './ai-control.ts'
-import { interruptGenerationProjection } from './ai-message-projection.ts'
+import {
+  interruptGenerationProjection,
+  settleAbandonedCompaction,
+} from './ai-message-projection.ts'
 
 export type QueueCommand<M extends UIMessage> = Extract<
   ControlCommand<M>,
@@ -200,7 +203,7 @@ export function projectAIClientState<M extends UIMessage>(
           projection.responseMessageId,
           stop.lastSeenIndex ?? Number.POSITIVE_INFINITY,
         ) ?? displayed)
-  return {
+  return settleAbandonedCompaction(displayed, {
     ...stopped,
     active: null,
     status: paused ? 'paused' : 'idle',
@@ -213,5 +216,5 @@ export function projectAIClientState<M extends UIMessage>(
     pendingApprovals: [],
     pendingInputs: [],
     error: null,
-  }
+  })
 }

package/src/ai-client.ts

@@ -0,0 +1,62 @@
+/**
+ * experimental-a2/ai/client — a client for an `agent()` route that needs no
+ * contract: the AI protocol and its `AIState` reducer are library code, and
+ * the agent's types arrive through a type-only `typeof` import.
+ *
+ * ```ts
+ * import type { assistant } from '@/server/assistant' // erased at build
+ * const client = createAgentClient<typeof assistant>({ api: '/api/agent' })
+ * await client.session(id).push(...inputs.message(message))
+ * ```
+ */
+
+import {
+  agentContract,
+  type AI_STATE_REDUCER,
+  type AIEventDefs,
+  type AIMessageOf,
+  type AIState,
+} from './ai.ts'
+import {
+  createClient,
+  type A2Client,
+  type CreateClientOptions,
+} from './client.ts'
+
+/**
+ * The structural shape of `typeof` an `agent({ name })` server: one whose
+ * state the built-in reducer folds. Agents built from an extended contract
+ * use `createClient({ contract })` instead.
+ */
+export type AnyAgent = {
+  readonly contract: {
+    readonly events: object
+    readonly reducer: { readonly name: typeof AI_STATE_REDUCER }
+  }
+}
+
+export type AgentMessage<A extends AnyAgent> = AIMessageOf<
+  A['contract']['events']
+>
+
+export type AgentClient<A extends AnyAgent = AnyAgent> = A2Client<
+  AIEventDefs<AgentMessage<A>>,
+  AIState<AgentMessage<A>>
+>
+
+export type CreateAgentClientOptions = Omit<
+  CreateClientOptions<AIEventDefs, AIState>,
+  'contract'
+>
+
+const protocol = agentContract({ name: 'a2.ai' })
+
+/**
+ * A client for the built-in AI protocol. Agents built from a shared
+ * `agentContract()` use `createClient({ contract })` instead.
+ */
+export function createAgentClient<A extends AnyAgent = AnyAgent>(
+  options: CreateAgentClientOptions,
+): AgentClient<A> {
+  return createClient({ ...options, contract: protocol }) as AgentClient<A>
+}

package/src/ai-context-schema.ts

@@ -1,3 +1,4 @@
+import { assertMessageBase, restoreMessageBase } from './ai-message-base.ts'
 import type { UIMessage } from 'ai'
 import { events } from './ai.ts'
 import type { AIContextState } from './ai-context.ts'
@@ -103,6 +104,16 @@ const indexedEvents = (
       ),
   )
 
+const restorableBase = (options: {
+  base: unknown
+  current: UIMessage | undefined
+}): boolean => {
+  assertMessageBase(options.base)
+  return isMessage(
+    restoreMessageBase({ current: options.current, base: options.base }),
+  )
+}
+
 const projection = (value: unknown, messages: UIMessage[]): boolean => {
   if (
     !record(value) ||
@@ -113,7 +124,7 @@ const projection = (value: unknown, messages: UIMessage[]): boolean => {
           'responseMessageId',
           'promptThroughIndex',
           'completionIndex',
-          'baseMessage',
+          'base',
           'baseCompaction',
           'batches',
           'toolEvents',
@@ -126,7 +137,13 @@ const projection = (value: unknown, messages: UIMessage[]): boolean => {
     !index(value['promptThroughIndex']) ||
     (value['completionIndex'] !== undefined &&
       !index(value['completionIndex'])) ||
-    (value['baseMessage'] !== undefined && !isMessage(value['baseMessage'])) ||
+    (value['base'] !== undefined &&
+      !restorableBase({
+        base: value['base'],
+        current: messages.find(
+          (message) => message.id === value['responseMessageId'],
+        ),
+      })) ||
     (value['baseCompaction'] !== undefined &&
       !compaction(value['baseCompaction'])) ||
     !progress(value['batches']) ||
@@ -200,9 +217,9 @@ const authority = (value: unknown): boolean =>
         'paused',
       ].includes(String(value['phase']))))
 
-export const contextSchema = <M extends UIMessage>(): StandardSchemaV1<
-  AIContextState<M>
-> => ({
+export const contextSchema = <M extends UIMessage>(
+  app?: object,
+): StandardSchemaV1<AIContextState<M>> => ({
   '~standard': {
     version: 1,
     vendor: 'a2-context',
@@ -217,6 +234,7 @@ export const contextSchema = <M extends UIMessage>(): StandardSchemaV1<
         'generation',
         'activeProjection',
         'messageIndexes',
+        ...(app === undefined ? [] : ['app']),
       ]
       try {
         if (
@@ -254,6 +272,7 @@ export const contextSchema = <M extends UIMessage>(): StandardSchemaV1<
           generation !== null &&
           (!record(generation) ||
             !matches(events['ai.generation.started'], generation['owner']) ||
+            !index(generation['startIndex']) ||
             ![
               'streaming',
               'completed',
@@ -287,6 +306,7 @@ export const contextSchema = <M extends UIMessage>(): StandardSchemaV1<
           return invalid
         if (value['closed'] && (generation !== null || current !== null))
           return invalid
+        if (app !== undefined && !record(value['app'])) return invalid
         return { value: value as AIContextState<M> }
       } catch {
         return invalid

package/src/ai-context.ts

@@ -1,6 +1,7 @@
+import { restoreMessageBase } from './ai-message-base.ts'
+import { reducer as makeReducer } from './reducer.ts'
 import type { UIMessage } from 'ai'
 import type {
-  AgentDefinition,
   ApprovalRequestedPayload,
   ApprovalRespondedPayload,
   CompactionCompletedPayload,
@@ -8,7 +9,6 @@ import type {
   CompactionState,
   GenerationCompletedPayload,
   GenerationFailedPayload,
-  GenerationProgressPayload,
   GenerationRequestedPayload,
   GenerationStartedPayload,
   MessageCreatedPayload,
@@ -21,9 +21,11 @@ import type {
 } from './ai.ts'
 import type { ControlCommit } from './ai-control.ts'
 import { contextSchema } from './ai-context-schema.ts'
-import type { ContractEvent, EventDefs } from './contract.ts'
+import type { Contract, ContractEvent, EventDefs } from './contract.ts'
 import {
   projectedMessage,
+  retainMessageBase,
+  findMessage,
   advanceProjection,
   reduceToolActivity,
   upsertMessage,
@@ -34,15 +36,16 @@ import { interruptUIMessage, pauseUIMessage } from './ai-projector.ts'
 import {
   flattenProgressBatches,
   progressBatches,
+  progressOwnedBy,
+  type StoredProgressPayload,
 } from './ai-progress-batches.ts'
-import { storedAIReducer, type AIStoredReducer } from './ai-stored-state.ts'
+import { appReducerOf, appSlice } from './ai-app-state.ts'
+import { withAIStorage } from './ai-stored-state.ts'
+import type { Reducer } from './reducer.ts'
 
 export type AIModelContext<M extends UIMessage = UIMessage> = {
   messages: M[]
-  compaction: Omit<
-    CompactionState<M>,
-    'messages' | 'retainedMessageIds' | 'usage'
-  > | null
+  compaction: Omit<CompactionState, 'usage'> | null
   modelMetadata: Record<string, ModelMetadataState>
 }
 
@@ -58,12 +61,16 @@ export type AIContextState<M extends UIMessage> = AIModelContext<M> & {
   > | null
   generation: {
     owner: GenerationStartedPayload
+    /** Log index of the owner's `ai.generation.started`. */
+    startIndex: number
     phase: 'streaming' | 'completed' | 'paused' | 'failed' | 'superseded'
     error?: string
     responseIndex?: number
   } | null
-  activeProjection: MessageProjection<M> | null
+  activeProjection: MessageProjection | null
   messageIndexes: Record<string, number>
+  /** The contract's application state, when `agentContract()` has a reducer. */
+  app?: object
 }
 
 const replaceResponse = <M extends UIMessage>({
@@ -96,12 +103,16 @@ const replaceResponse = <M extends UIMessage>({
   }
 }
 
-const calledTools = <M extends UIMessage>(
-  projection: MessageProjection<M>,
-): Set<string> =>
+const calledTools = <M extends UIMessage>({
+  projection,
+  current,
+}: {
+  projection: MessageProjection
+  current: M | undefined
+}): Set<string> =>
   new Set([
-    ...(projection.baseMessage?.parts.flatMap((part) =>
-      'toolCallId' in part ? [part.toolCallId] : [],
+    ...(restoreMessageBase({ current, base: projection.base })?.parts.flatMap(
+      (part) => ('toolCallId' in part ? [part.toolCallId] : []),
     ) ?? []),
     ...projection.toolEvents
       .filter((event) => event.type === 'ai.tool.called')
@@ -116,15 +127,19 @@ const renderProjection = <M extends UIMessage>({
   index: number
 }): AIContextState<M> => {
   const projection = state.activeProjection!
-  let message = projectedMessage(projection)
+  const current = findMessage(state.messages, projection.responseMessageId)
+  let message = projectedMessage({ projection, current })
   if (
-    projection.baseMessage === undefined &&
+    projection.base === undefined &&
     projection.batches.length === 0 &&
     message.parts.length === 0
   )
     return state
   if (state.generation?.phase === 'paused')
-    message = pauseUIMessage({ message, toolCallIds: calledTools(projection) })
+    message = pauseUIMessage({
+      message,
+      toolCallIds: calledTools({ projection, current }),
+    })
   if (state.generation?.phase === 'failed')
     message = interruptUIMessage(
       message,
@@ -142,14 +157,18 @@ export const generationBase = <M extends UIMessage>(
 ): AIContextState<M> => {
   const projection = state.activeProjection
   if (!projection) return state
+  const base = restoreMessageBase({
+    current: findMessage(state.messages, projection.responseMessageId),
+    base: projection.base,
+  })
   return {
     ...state,
     messages:
-      projection.baseMessage === undefined
+      base === undefined
         ? state.messages.filter(
             (message) => message.id !== projection.responseMessageId,
           )
-        : upsertMessage(state.messages, projection.baseMessage),
+        : upsertMessage(state.messages, base),
     compaction: projection.baseCompaction ?? null,
   }
 }
@@ -182,10 +201,11 @@ const ownsGeneration = <M extends UIMessage>(
   state.generation.phase !== 'superseded'
 
 export const createContextReducer = <M extends UIMessage, D extends EventDefs>({
-  agent,
+  contract,
 }: {
-  agent: AgentDefinition<M, D>
-}): AIStoredReducer<D, AIContextState<M>> => {
+  contract: Contract<D, Record<never, never>, unknown>
+}): Reducer<D, AIContextState<M>> => {
+  const app = appReducerOf(contract)
   const initialState: AIContextState<M> = {
     messages: [],
     compaction: null,
@@ -195,15 +215,26 @@ export const createContextReducer = <M extends UIMessage, D extends EventDefs>({
     generation: null,
     activeProjection: null,
     messageIndexes: {},
+    ...(app === undefined ? {} : { app: app.initialState }),
   }
-  const reducer = agent.contract
-    .reducer({
-      name: 'a2.ai.context.v3',
-      initialState,
-      stateSchema: contextSchema<M>(),
+  const reducer = makeReducer({
+    contract,
+    // The application state rides in this checkpoint so every generation
+    // reads it without another state read.
+    name:
+      app === undefined ? 'a2.ai.context.v6' : `${app.name}+a2.ai.context.v6`,
+    initialState,
+    stateSchema: contextSchema<M>(app),
+  }).fold((state, event) => {
+    const after = retainMessageBase({
+      before: state,
+      after: reduceContext(state, event),
     })
-    .fold((state, event) => reduceContext(state, event))
-  return storedAIReducer(reducer)
+    if (app === undefined) return after
+    const next = app.fold(state.app!, event as ContractEvent<EventDefs>)
+    return next === state.app ? after : { ...after, app: appSlice(next) }
+  })
+  return withAIStorage(reducer)
 }
 
 const reduceContext = <M extends UIMessage, D extends EventDefs>(
@@ -297,6 +328,7 @@ const reduceContext = <M extends UIMessage, D extends EventDefs>(
         ...base,
         generation: {
           owner,
+          startIndex: event.index,
           phase: 'streaming',
           ...(previous?.owner.responseMessageId === owner.responseMessageId &&
           previous.responseIndex !== undefined
@@ -307,7 +339,9 @@ const reduceContext = <M extends UIMessage, D extends EventDefs>(
           generationId: owner.generationId,
           responseMessageId: owner.responseMessageId,
           promptThroughIndex: owner.promptThroughIndex ?? event.index - 1,
-          ...(baseMessage === undefined ? {} : { baseMessage }),
+          ...(baseMessage === undefined
+            ? {}
+            : { base: { partCount: baseMessage.parts.length } }),
           baseCompaction: base.compaction,
           batches: progressBatches(),
           toolEvents: [],
@@ -316,12 +350,16 @@ const reduceContext = <M extends UIMessage, D extends EventDefs>(
       }
     }
     case 'ai.generation.progress': {
-      const payload = event.payload as GenerationProgressPayload
+      const payload = event.payload as StoredProgressPayload
+      const generation = state.generation
       if (
-        !ownsGeneration(state, payload) ||
-        state.generation!.owner.responseMessageId !==
-          payload.responseMessageId ||
-        state.generation!.phase !== 'streaming' ||
+        generation === null ||
+        generation.phase !== 'streaming' ||
+        !progressOwnedBy({
+          payload,
+          owner: generation.owner,
+          startIndex: generation.startIndex,
+        }) ||
         !state.activeProjection
       )
         return state
@@ -433,10 +471,15 @@ const reduceContext = <M extends UIMessage, D extends EventDefs>(
         ),
       }
       let next = generationBase(state)
-      if (retained.baseMessage !== undefined || retained.batches.length > 0)
+      if (retained.base !== undefined || retained.batches.length > 0)
         next = replaceResponse({
           state: next,
-          message: interruptUIMessage(projectedMessage(retained)),
+          message: interruptUIMessage(
+            projectedMessage({
+              projection: retained,
+              current: findMessage(state.messages, retained.responseMessageId),
+            }),
+          ),
           index: event.index,
         })
       return endGeneration(next)
@@ -458,7 +501,13 @@ const reduceContext = <M extends UIMessage, D extends EventDefs>(
         !state.activeProjection
       )
         return state
-      const toolCallIds = calledTools(state.activeProjection)
+      const toolCallIds = calledTools({
+        projection: state.activeProjection,
+        current: findMessage(
+          state.messages,
+          state.activeProjection.responseMessageId,
+        ),
+      })
       return renderProjection({
         state: {
           ...state,
@@ -617,11 +666,12 @@ const reduceContext = <M extends UIMessage, D extends EventDefs>(
         ),
       ]
       const activeProjection = { ...projection, cursor: null }
-      delete activeProjection.baseMessage
+      delete activeProjection.base
       const baseMessage = messages.find(
         (message) => message.id === projection.responseMessageId,
       )
-      if (baseMessage !== undefined) activeProjection.baseMessage = baseMessage
+      if (baseMessage !== undefined)
+        activeProjection.base = { partCount: baseMessage.parts.length }
       const compaction: AIModelContext<M>['compaction'] = {
         status: 'completed',
         generationId: payload.generationId,