experimental-a2 0.16.1 → 0.18.0

package/src/ai.ts

@@ -1,14 +1,17 @@
+import type { MessageBase } from './ai-message-base.ts'
+import { reducer as makeReducer } from './reducer.ts'
 /**
- * experimental-a2/ai — an isomorphic, extensible agent definition.
+ * experimental-a2/ai — the isomorphic AI protocol.
  *
- * `agent()` assembles a normal A2 contract, a reducer over the built-in
- * AI protocol, and pure builders for the events applications append.
- * Model calls and tool implementations stay in experimental-a2/ai/server.
+ * Event schemas, the `AIState` reducer, `agentContract()` for extending
+ * both, and pure builders for the events applications append. Model calls
+ * and tool implementations stay in experimental-a2/ai/server.
  */
 
 import type {
   FinishReason,
   JSONValue,
+  LanguageModelCallOptions,
   LanguageModelUsage,
   UIMessage,
   UIMessageChunk,
@@ -38,35 +41,55 @@ export type {
   InboxItem,
 } from './ai-control.ts'
 import { contract as createContract } from './contract.ts'
-import type { Contract, EventDefs } from './contract.ts'
+import type { Contract, ContractEvent, EventDefs } from './contract.ts'
 import {
   interruptUIMessage,
   pauseUIMessage,
   projectUIMessage,
+  projectUIMessageStream,
+  hasToolCall,
+  type ProjectionChunk,
   type UIMessageCursor,
 } from './ai-projector.ts'
 import { ambientToolScopeStorage } from './internal.ts'
 import {
   upsertMessage,
+  retainMessageBase,
+  findMessage,
   upsertResponse,
   answerApproval,
   approvalRequestedChunk,
   toolResultChunk,
+  toolCalledChunk,
   advanceProjection,
   projectedTools,
   restoreGenerationBase,
   stopGenerationTools,
   interruptGenerationProjection,
   projectedMessage,
+  projectedMessageStream,
   reduceToolActivity,
+  addCompactionPart,
+  isCompactionPart,
+  settleAbandonedCompaction,
+  settleCompactionPart,
 } from './ai-message-projection.ts'
-import { clientStateProjection, type Reducer } from './reducer.ts'
-import { storedAIReducer, type AIStoredReducer } from './ai-stored-state.ts'
-export type { AIStoredState, AIStoredReducer } from './ai-stored-state.ts'
+import type { Reducer } from './reducer.ts'
+import { withReducerRuntime } from './reducer-runtime.ts'
+import { withAIStorage } from './ai-stored-state.ts'
+import {
+  appSlice,
+  isAIStateKey,
+  registerAgentReducer,
+  withAppSlice,
+  type AppReducer,
+} from './ai-app-state.ts'
 import {
   flattenProgressBatches,
   progressBatches,
+  progressOwnedBy,
   type ProgressBatches,
+  type StoredProgressPayload,
 } from './ai-progress-batches.ts'
 import type { HandlerContext } from './server.ts'
 import type { StandardSchemaV1 } from './standard-schema.ts'
@@ -77,6 +100,43 @@ export type AgentStatus =
 
 export type GenerationReason = 'message' | 'tool' | 'input' | 'retry' | 'resume'
 
+export type ReasoningLevel = NonNullable<LanguageModelCallOptions['reasoning']>
+
+export const reasoningLevels: readonly ReasoningLevel[] = [
+  'provider-default',
+  'none',
+  'minimal',
+  'low',
+  'medium',
+  'high',
+  'xhigh',
+]
+
+export type CompactionPartData = {
+  status: 'running' | 'completed' | 'failed' | 'interrupted'
+  inputTokens?: number
+  thresholdTokens?: number
+  summary?: string
+  usage?: LanguageModelUsage
+}
+
+/** The transcript part A2 adds to a response where its context was compacted. */
+export type CompactionPart = {
+  type: 'data-a2-compaction'
+  id: string
+  data: CompactionPartData
+}
+
+export { isCompactionPart }
+
+/**
+ * A transcript message: `M` plus the compaction part A2 adds, so
+ * `isCompactionPart` narrows even when `M` declares its own data parts.
+ */
+export type TranscriptMessage<M extends UIMessage> = Omit<M, 'parts'> & {
+  parts: Array<M['parts'][number] | CompactionPart>
+}
+
 export type SessionCreatedPayload = { metadata?: JSONValue }
 export type SessionClosedPayload = { reason?: string }
 export type MessageCreatedPayload<M extends UIMessage> = {
@@ -115,15 +175,13 @@ export type GenerationStartedPayload = {
   generationId: string
   responseMessageId: string
   attempt: number
-  model: string
+  /** Absent when the attempt failed before its model resolved. */
+  model?: string
   promptThroughIndex?: number
 }
 export type GenerationProgressPayload = {
-  requestId: string
-  messageId: string
-  generationId: string
-  responseMessageId: string
-  sequence: number
+  /** Log index of the generation's `ai.generation.started`. */
+  startIndex: number
   chunks: UIMessageChunk[]
 }
 export type GenerationCompletedPayload = {
@@ -275,17 +333,20 @@ export type AIEventDefs<M extends UIMessage = UIMessage> = {
   'ai.model.metadata.resolved': StandardSchemaV1<ModelMetadataResolvedPayload>
 }
 
-export type ActiveGeneration = GenerationStartedPayload
+export type ActiveGeneration = GenerationStartedPayload & {
+  /** Log index of this generation's `ai.generation.started`. */
+  startIndex: number
+}
 
-export type ActiveGenerationProjection<M extends UIMessage = UIMessage> = {
+export type ActiveGenerationProjection = {
   generationId: string
   responseMessageId: string
   /** Log index of the completed model step while its response remains interruptible. */
   completionIndex?: number
-  baseMessage?: M
+  base?: MessageBase
   promptThroughIndex?: number
   cursor?: { throughIndex: number; stream: UIMessageCursor } | null
-  baseCompaction?: CompactionState<M> | null
+  baseCompaction?: CompactionState | null
   batches: ProgressBatches
   baseTools: ToolActivity[]
   toolEvents: Array<
@@ -315,28 +376,20 @@ export type ToolActivity = {
   generationId: string
   toolCallId: string
   toolName?: string
-  input?: unknown
-  rawInput?: unknown
-  output?: unknown
   error?: string
   preliminary?: boolean
   phase?: 'input' | 'execution'
   dynamic?: boolean
   providerExecuted?: boolean
-  callProviderMetadata?: unknown
-  resultProviderMetadata?: unknown
-  toolMetadata?: unknown
   title?: string
   status: 'running' | 'completed' | 'failed' | 'denied'
 }
 
-export type CompactionState<M extends UIMessage> = {
+export type CompactionState = {
   status: 'running' | 'completed'
   generationId: string
   throughMessageId: string
   throughIndex?: number
-  messages?: M[]
-  retainedMessageIds?: string[]
   summary?: string
   usage?: LanguageModelUsage
 }
@@ -354,14 +407,14 @@ export type AIState<M extends UIMessage = UIMessage> = {
   receipt: ControlReceipt | null
   pausedGenerationIds?: Record<string, true>
   session: { created: boolean; metadata?: JSONValue }
-  messages: M[]
+  messages: TranscriptMessage<M>[]
   status: AgentStatus
   activeGeneration: ActiveGeneration | null
   activeRequestId: string | null
   activeRequestReason: GenerationReason | null
   activeRequestSourceGenerationId: string | null
   activeResponseMessageId: string | null
-  activeProjection: ActiveGenerationProjection<M> | null
+  activeProjection: ActiveGenerationProjection | null
   responseOwners: Record<
     string,
     Pick<GenerationStartedPayload, 'generationId' | 'requestId' | 'attempt'>
@@ -374,19 +427,29 @@ export type AIState<M extends UIMessage = UIMessage> = {
   pendingApprovals: PendingApproval[]
   pendingInputs: PendingInput[]
   tools: ToolActivity[]
-  compaction: CompactionState<M> | null
+  compaction: CompactionState | null
   modelMetadata: Record<string, ModelMetadataState>
   usage: GenerationUsage[]
   error: string | null
 }
 
-export type AgentDefinition<
+/** The `UIMessage` type a contract's `ai.message.created` schema carries. */
+export type AIMessageOf<D> = D extends {
+  readonly 'ai.message.created': infer S extends StandardSchemaV1
+}
+  ? StandardSchemaV1.InferOutput<S> extends MessageCreatedPayload<
+      infer M extends UIMessage
+    >
+    ? M
+    : UIMessage
+  : UIMessage
+
+/** What `agentContract()` returns: the AI protocol, with `X` as application state. */
+export type AgentContract<
   M extends UIMessage = UIMessage,
   D extends EventDefs = AIEventDefs<M>,
-> = {
-  readonly contract: Contract<D>
-  readonly reducer: AIStoredReducer<D, AIState<M>>
-}
+  X extends object = Record<never, never>,
+> = Contract<D, Record<never, never>, AgentReducer<M, D, X>>
 
 export type AgentToolEventDefs = EventDefs &
   Pick<AIEventDefs, 'ai.tool.execution.requested'>
@@ -397,24 +460,50 @@ export type AgentToolContext<D extends AgentToolEventDefs> = HandlerContext<
   'ai.tool.execution.requested'
 >
 
-export type AgentOptions<M extends UIMessage, E extends EventDefs> = {
+/** Application state folded after the AI protocol, on the server and in clients. */
+/** Rejects application state keys that would shadow `AIState`'s. */
+type OwnKeys<X> = { [K in Extract<keyof X, keyof AIState>]: never }
+/** Rejects a stateless application reducer, whose name would not be built in. */
+type NonEmpty<X> = keyof X extends never ? never : unknown
+
+/**
+ * The reducer `agentContract()` builds. Without application state it carries
+ * the built-in name, which is what lets the contract-free client fold it.
+ * `~appState` is type-only, absent at runtime: it carries the application
+ * state's type to `agent()`, and marks the reducer as `agentContract()`'s.
+ */
+export type AgentReducer<
+  M extends UIMessage,
+  D extends EventDefs,
+  X extends object,
+> = (keyof X extends never
+  ? Reducer<D, AIState<M>> & { readonly name: typeof AI_STATE_REDUCER }
+  : Reducer<D, AIState<M> & X>) & { readonly '~appState': X }
+
+export type AgentContractOptions<
+  M extends UIMessage,
+  E extends EventDefs,
+  X extends object,
+> = {
   name: string
   events?: E
   /** Optional stricter schema for custom UIMessage metadata/data/tool types. */
   messageSchema?: StandardSchemaV1<M>
-  /** Snapshot identity. Bump when application-level projection semantics change. */
-  reducerName?: string
+  /**
+   * Application state, folded from events beside `AIState` and merged into
+   * it. `fold` sees only this state, never the transcript.
+   */
+  reducer?: {
+    /** Snapshot identity of the application state. Rename it when `fold` changes. */
+    name: string
+    initialState: X & OwnKeys<X> & NonEmpty<X>
+    fold: (
+      state: NoInfer<X>,
+      event: ContractEvent<AIEventDefs<M> & NoInfer<E>>,
+    ) => NoInfer<X>
+  }
 }
 
-type AIMessageFrom<D extends EventDefs> =
-  D extends Record<'ai.message.created', infer S extends StandardSchemaV1>
-    ? StandardSchemaV1.InferOutput<S> extends MessageCreatedPayload<
-        infer M extends UIMessage
-      >
-      ? M
-      : UIMessage
-    : UIMessage
-
 type SchemaResult<T> = StandardSchemaV1.Result<T>
 
 const issue = <T>(message: string): SchemaResult<T> => ({
@@ -540,13 +629,8 @@ const progressPayloadSchema: StandardSchemaV1<GenerationProgressPayload> =
     if (
       !isRecord(value) ||
       !isJSONCompatible(value) ||
-      !stringField(value, 'requestId') ||
-      !stringField(value, 'messageId') ||
-      !stringField(value, 'generationId') ||
-      !stringField(value, 'responseMessageId') ||
-      typeof value['sequence'] !== 'number' ||
-      !Number.isInteger(value['sequence']) ||
-      value['sequence'] < 0 ||
+      typeof value['startIndex'] !== 'number' ||
+      !optionalIndexField(value, 'startIndex') ||
       !Array.isArray(value['chunks']) ||
       !value['chunks'].every(
         (chunk) => isRecord(chunk) && typeof chunk['type'] === 'string',
@@ -822,7 +906,7 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
         typeof value['attempt'] === 'number' &&
         Number.isInteger(value['attempt']) &&
         value['attempt'] > 0 &&
-        stringField(value, 'model') &&
+        (value['model'] === undefined || stringField(value, 'model')) &&
         optionalIndexField(value, 'promptThroughIndex'),
     ),
     'ai.generation.progress': progressPayloadSchema,
@@ -979,13 +1063,18 @@ const initialState = <M extends UIMessage>(): AIState<M> => ({
 const projectMessageChunks = <M extends UIMessage>(
   messages: M[],
   messageId: string,
-  chunks: UIMessageChunk[],
-): M[] =>
-  messages.map((message) =>
-    message.id === messageId
-      ? projectUIMessage({ messageId, base: message, chunks })
-      : message,
+  chunks: ProjectionChunk[],
+): M[] => {
+  const base = findMessage(messages, messageId)
+  return upsertMessage(
+    messages,
+    projectUIMessage<M>({
+      messageId,
+      ...(base === undefined ? {} : { base }),
+      chunks,
+    }),
   )
+}
 
 const waitingStatus = <M extends UIMessage>(state: AIState<M>): AgentStatus => {
   if (state.pendingApprovals.length > 0 || state.pendingInputs.length > 0) {
@@ -1018,9 +1107,15 @@ const isStaleGeneration = <M extends UIMessage>(
 
 const projectedResponse = <M extends UIMessage>(options: {
   state: AIState<M>
-  projection: ActiveGenerationProjection<M>
-}): M => {
-  const message = projectedMessage(options.projection)
+  projection: ActiveGenerationProjection
+}): TranscriptMessage<M> => {
+  const message = projectedMessage({
+    projection: options.projection,
+    current: findMessage(
+      options.state.messages,
+      options.projection.responseMessageId,
+    ),
+  })
   return options.state.pausedGenerationIds?.[options.projection.generationId]
     ? pauseUIMessage({
         message,
@@ -1048,6 +1143,19 @@ const terminalGenerationReason = <M extends UIMessage>(
 export function reduceAIState<M extends UIMessage>(
   state: AIState<M>,
   event: { type: string; payload: unknown; index: number; id?: string },
+): AIState<M> {
+  return settleAbandonedCompaction(
+    state,
+    retainMessageBase({
+      before: state,
+      after: reduceAIEvent(state, event),
+    }),
+  )
+}
+
+function reduceAIEvent<M extends UIMessage>(
+  state: AIState<M>,
+  event: { type: string; payload: unknown; index: number; id?: string },
 ): AIState<M> {
   if (state.status === 'closed') return state
   switch (event.type) {
@@ -1371,13 +1479,14 @@ export function reduceAIState<M extends UIMessage>(
             'superseded',
           )
         : state
-      const baseMessage = baseState.messages.find(
-        (message) => message.id === payload.responseMessageId,
+      const baseMessage = findMessage(
+        baseState.messages,
+        payload.responseMessageId,
       )
       return {
         ...baseState,
         status: 'generating',
-        activeGeneration: payload,
+        activeGeneration: { ...payload, startIndex: event.index },
         activeResponseMessageId: null,
         responseOwners: {
           ...baseState.responseOwners,
@@ -1391,7 +1500,9 @@ export function reduceAIState<M extends UIMessage>(
           generationId: payload.generationId,
           responseMessageId: payload.responseMessageId,
           promptThroughIndex: payload.promptThroughIndex ?? event.index - 1,
-          ...(baseMessage === undefined ? {} : { baseMessage }),
+          ...(baseMessage === undefined
+            ? {}
+            : { base: { partCount: baseMessage.parts.length } }),
           baseCompaction: baseState.compaction,
           batches: progressBatches(),
           baseTools: baseState.tools,
@@ -1402,30 +1513,29 @@ export function reduceAIState<M extends UIMessage>(
       }
     }
     case 'ai.generation.progress': {
-      const payload = event.payload as GenerationProgressPayload
-      if (state.pausedGenerationIds?.[payload.generationId]) return state
+      const payload = event.payload as StoredProgressPayload
+      const generation = state.activeGeneration
+      const activeProjection = state.activeProjection
       if (
+        generation === null ||
+        !progressOwnedBy({
+          payload,
+          owner: generation,
+          startIndex: generation.startIndex,
+        }) ||
+        state.pausedGenerationIds?.[generation.generationId] ||
         terminalGenerationReason(
           state,
-          payload.generationId,
-          payload.requestId,
-        ) !== undefined
-      ) {
-        return state
-      }
-      const activeProjection = state.activeProjection
-      if (
-        state.activeGeneration?.generationId !== payload.generationId ||
-        state.activeGeneration.requestId !== payload.requestId ||
-        activeProjection?.generationId !== payload.generationId
+          generation.generationId,
+          generation.requestId,
+        ) !== undefined ||
+        activeProjection?.generationId !== generation.generationId
       ) {
         return state
       }
       const progress = advanceProjection({
         projection: activeProjection,
-        base: state.messages.find(
-          (message) => message.id === activeProjection.responseMessageId,
-        ),
+        base: findMessage(state.messages, activeProjection.responseMessageId),
         chunks: payload.chunks,
         index: event.index,
       })
@@ -1434,7 +1544,7 @@ export function reduceAIState<M extends UIMessage>(
         messages: upsertResponse({
           messages: state.messages,
           message: progress.message,
-          inputId: state.activeGeneration.messageId,
+          inputId: generation.messageId,
         }),
         activeProjection: {
           ...activeProjection,
@@ -1521,14 +1631,20 @@ export function reduceAIState<M extends UIMessage>(
           : null
       const shouldProject =
         projection !== null &&
-        (projection.baseMessage !== undefined || projection.batches.length > 0)
+        (projection.base !== undefined || projection.batches.length > 0)
       return {
         ...state,
         messages: shouldProject
           ? upsertMessage(
               state.messages,
               interruptUIMessage(
-                projectedMessage(projection),
+                projectedMessage({
+                  projection,
+                  current: findMessage(
+                    state.messages,
+                    projection.responseMessageId,
+                  ),
+                }),
                 `Generation failed before tool completion: ${payload.error}`,
               ),
             )
@@ -1832,6 +1948,9 @@ export function reduceAIState<M extends UIMessage>(
       if (projection?.generationId !== payload.generationId) {
         const next = {
           ...state,
+          messages: projectMessageChunks(state.messages, payload.messageId, [
+            toolCalledChunk(payload),
+          ]),
           tools: reduceToolActivity(state.tools, {
             type: 'ai.tool.called',
             payload,
@@ -1839,16 +1958,51 @@ export function reduceAIState<M extends UIMessage>(
         }
         return { ...next, status: lifecycleStatus(state, next) }
       }
-      const activeProjection: ActiveGenerationProjection<M> = {
+      const activeProjection: ActiveGenerationProjection = {
         ...projection,
+        cursor: null,
         toolEvents: [
           ...projection.toolEvents,
           { index: event.index, type: 'ai.tool.called', payload },
         ],
       }
+      const base = findMessage(state.messages, payload.messageId)
+      const paused = state.pausedGenerationIds?.[payload.generationId]
+      const cursor = projection.cursor
+      const incremental =
+        base && cursor && !paused && event.index > cursor.throughIndex
+      const projected = incremental
+        ? {
+            ...(hasToolCall({ message: base, call: payload })
+              ? { message: base, cursor: cursor.stream }
+              : projectUIMessageStream({
+                  messageId: payload.messageId,
+                  base,
+                  cursor: cursor.stream,
+                  chunks: [toolCalledChunk(payload)],
+                })),
+            throughIndex: event.index,
+          }
+        : projectedMessageStream({
+            projection: activeProjection,
+            current: base,
+          })
+      activeProjection.cursor = paused
+        ? null
+        : { throughIndex: projected.throughIndex, stream: projected.cursor }
       const next = {
         ...state,
         activeProjection,
+        messages: upsertResponse({
+          messages: state.messages,
+          message: paused
+            ? projectedResponse({ state, projection: activeProjection })
+            : projected.message,
+          inputId:
+            state.activeGeneration?.messageId ??
+            state.active?.input.id ??
+            payload.messageId,
+        }),
         tools: projectedTools(activeProjection),
       }
       return { ...next, status: lifecycleStatus(state, next) }
@@ -1903,7 +2057,7 @@ export function reduceAIState<M extends UIMessage>(
         }
         return { ...next, status: lifecycleStatus(state, next) }
       }
-      const activeProjection: ActiveGenerationProjection<M> = {
+      const activeProjection: ActiveGenerationProjection = {
         ...projection,
         cursor: null,
         toolEvents: [
@@ -2136,10 +2290,10 @@ export function reduceAIState<M extends UIMessage>(
       if (state.activeGeneration?.generationId !== payload.generationId) {
         return state
       }
-      return {
-        ...state,
-        compaction: { status: 'running', ...payload },
-      }
+      return addCompactionPart(
+        { ...state, compaction: { status: 'running', ...payload } },
+        payload,
+      )
     }
     case 'ai.compaction.completed': {
       const payload = event.payload as CompactionCompletedPayload<M>
@@ -2150,20 +2304,41 @@ export function reduceAIState<M extends UIMessage>(
       ) {
         return state
       }
-      const compaction: CompactionState<M> = {
+      const {
+        messages: _messages,
+        retainedMessageIds: _retained,
+        ...summary
+      } = payload
+      const compaction: CompactionState = {
         status: 'completed',
-        ...payload,
+        ...summary,
         throughIndex:
           payload.throughIndex ?? state.activeProjection!.promptThroughIndex!,
       }
-      return {
-        ...state,
-        compaction,
-        activeProjection: state.activeProjection && {
-          ...state.activeProjection,
-          baseCompaction: compaction,
+      const projection = state.activeProjection
+      return settleCompactionPart({
+        state: {
+          ...state,
+          compaction,
+          activeProjection: projection && {
+            ...projection,
+            baseCompaction: compaction,
+            // A new base identity records the completed part as the step's
+            // starting point, so a later cutoff cannot revert it to running.
+            ...(projection.base === undefined
+              ? {}
+              : { base: { ...projection.base } }),
+          },
         },
-      }
+        generationId: payload.generationId,
+        data: {
+          status: 'completed',
+          ...(summary.summary === undefined
+            ? {}
+            : { summary: summary.summary }),
+          ...(summary.usage === undefined ? {} : { usage: summary.usage }),
+        },
+      })
     }
     default:
       return state
@@ -2173,35 +2348,12 @@ export function reduceAIState<M extends UIMessage>(
 /** Derive AI SDK UI messages directly from an A2 event history. */
 export function deriveUIMessages<M extends UIMessage = UIMessage>(
   history: readonly { type: string; payload: unknown; index: number }[],
-): M[] {
+): TranscriptMessage<M>[] {
   let state = initialState<M>()
   for (const event of history) state = reduceAIState(state, event)
   return state.messages
 }
 
-/** Build the standard AI state projection for any contract containing the protocol. */
-export function createReducer<D extends EventDefs>(options: {
-  contract: Contract<D>
-  name?: string
-  fold?: Reducer<D, AIState<AIMessageFrom<D>>>['fold']
-  stateSchema?: StandardSchemaV1<unknown, AIState<AIMessageFrom<D>>>
-}): AIStoredReducer<D, AIState<AIMessageFrom<D>>> {
-  type M = AIMessageFrom<D>
-  const reducer = options.contract
-    .reducer({
-      name: options.name ?? 'a2.ai.state.v22',
-      initialState: initialState<M>(),
-      ...(options.stateSchema === undefined
-        ? {}
-        : { stateSchema: options.stateSchema }),
-    })
-    .fold(options.fold ?? ((state, event) => reduceAIState(state, event)))
-  return storedAIReducer({
-    ...reducer,
-    [clientStateProjection]: projectAIClientState,
-  })
-}
-
 type CommandInput<C> = { type: 'ai.control.requested'; payload: C; id?: string }
 type ControlInput<
   M extends UIMessage,
@@ -2338,22 +2490,24 @@ export const inputs: AIInputs = {
  * Read the durable handler context ambient to the current AI tool
  * execution. A2 runs each local tool's `execute` inside this scope, so
  * plain AI SDK `tool()` definitions reach their `event`, `attempt`,
- * `session`, and `signal` without an A2-specific wrapper. The agent (or
- * contract) argument carries the event types and is verified against
- * the executing server's contract.
+ * `session`, and `signal` without an A2-specific wrapper. Pass the
+ * agent's contract to type application events; it is verified against
+ * the executing agent.
  */
-export function handlerContext<D extends AgentToolEventDefs>(
-  agentOrContract: Contract<D> | { readonly contract: Contract<D> },
+export function handlerContext<D extends AgentToolEventDefs = AIEventDefs>(
+  contract?: Contract<D>,
 ): AgentToolContext<D> {
-  const contract =
-    'contract' in agentOrContract ? agentOrContract.contract : agentOrContract
   const scope = ambientToolScopeStorage()?.getStore()
   if (scope === undefined) {
     throw new TypeError(
-      `handlerContext('${contract.name}') requires an ambient A2 tool execution: call it inside a local tool's execute, or wrap direct tool tests with runWithHandlerContext`,
+      "handlerContext() requires an ambient A2 tool execution: call it inside a local tool's execute, or wrap direct tool tests with runWithHandlerContext",
     )
   }
-  if (scope.contract !== undefined && scope.contract.name !== contract.name) {
+  if (
+    contract !== undefined &&
+    scope.contract !== undefined &&
+    scope.contract.name !== contract.name
+  ) {
     throw new TypeError(
       `handlerContext('${contract.name}') was called during a tool execution for agent '${scope.contract.name}'`,
     )
@@ -2361,14 +2515,23 @@ export function handlerContext<D extends AgentToolEventDefs>(
   return scope.context as AgentToolContext<D>
 }
 
+/** The built-in AI reducer's name, and so its checkpoint identity. */
+export const AI_STATE_REDUCER = 'a2.ai.state.v26'
+
 /**
- * Define an agent: built-in AI protocol + application events and one state
- * reducer. The result is isomorphic.
+ * The built-in AI protocol plus application events, as a core contract whose
+ * reducer folds `AIState`. Isomorphic: share it when a client needs
+ * application events or extended state; otherwise `agent({ name })` on the
+ * server builds one and clients use the built-in protocol.
  */
-export function agent<
+export function agentContract<
   M extends UIMessage = UIMessage,
   E extends EventDefs = Record<never, never>,
->(options: AgentOptions<M, E>): AgentDefinition<M, AIEventDefs<M> & E> {
+  X extends object = Record<never, never>,
+>(
+  options: AgentContractOptions<M, E, X>,
+): AgentContract<M, AIEventDefs<M> & E, X> {
+  type D = AIEventDefs<M> & E
   const builtIns = createEvents(
     options.messageSchema === undefined
       ? undefined
@@ -2382,25 +2545,46 @@ export function agent<
       )
     }
   }
-  const agentContract = createContract({
+  const vocabulary = createContract({
     name: options.name,
-    events: { ...builtIns, ...extensions },
+    events: { ...builtIns, ...extensions } as D,
   })
-  const reducer = agentContract
-    .reducer({
-      name: options.reducerName ?? 'a2.ai.state.v22',
-      initialState: initialState<M>(),
-    })
-    .fold((state, event) => reduceAIState(state, event))
-
-  const definition = {
-    contract: agentContract,
-    reducer: storedAIReducer({
-      ...reducer,
-      [clientStateProjection]: projectAIClientState<M>,
-    }),
+  const app = options.reducer as AppReducer<D, X> | undefined
+  const aiState = initialState<M>()
+  for (const key of Object.keys(app?.initialState ?? {})) {
+    if (isAIStateKey(key))
+      throw new TypeError(`agent state '${key}' conflicts with AIState`)
   }
-  return definition as unknown as AgentDefinition<M, AIEventDefs<M> & E>
+  const reducer = makeReducer({
+    contract: vocabulary,
+    // Application checkpoints embed AI state, so they advance with it.
+    name:
+      app === undefined ? AI_STATE_REDUCER : `${app.name}+${AI_STATE_REDUCER}`,
+    initialState: { ...aiState, ...app?.initialState } as AIState<M> & X,
+  }).fold((state, event) => {
+    const next = reduceAIState(state, event) as AIState<M> & X
+    if (app === undefined) return next
+    const before = appSlice<X>(state)
+    const after = app.fold(before, event as ContractEvent<D>)
+    return after === before ? next : withAppSlice(next, after)
+  })
+  const shared = createContract({
+    name: options.name,
+    events: vocabulary.events,
+    reducer: withAIStorage(
+      withReducerRuntime({
+        reducer,
+        project: projectAIClientState<M> as unknown as (
+          state: AIState<M> & X,
+        ) => AIState<M> & X,
+      }),
+    ),
+  })
+  registerAgentReducer(
+    shared.reducer,
+    app as unknown as AppReducer<EventDefs, object> | undefined,
+  )
+  return shared as unknown as AgentContract<M, AIEventDefs<M> & E, X>
 }
 
 export type {