experimental-a2 0.12.0 → 0.13.0

package/src/ai.ts

@@ -64,6 +64,7 @@ export type GenerationStartedPayload = {
   responseMessageId: string
   attempt: number
   model: string
+  promptThroughIndex?: number
 }
 export type GenerationProgressPayload = {
   requestId: string
@@ -80,6 +81,7 @@ export type GenerationCompletedPayload = {
   responseMessageId: string
   finishReason?: FinishReason
   usage?: LanguageModelUsage
+  inputTokenEstimate?: number
 }
 export type GenerationFailedPayload = {
   requestId: string
@@ -152,15 +154,33 @@ export type InputRespondedPayload = {
   name: string
   value: JSONValue
 }
+export type ModelLimits = {
+  contextWindow: number
+  maxOutputTokens: number
+}
+export type ModelMetadataRequestedPayload = { modelId: string }
+export type ModelMetadataResolvedPayload = {
+  modelId: string
+  limits: ModelLimits | null
+}
+export type ModelMetadataState =
+  | { status: 'pending' }
+  | { status: 'resolved'; limits: ModelLimits }
+  | { status: 'unavailable' }
+
 export type CompactionRequestedPayload = {
   generationId: string
   throughMessageId: string
+  throughIndex?: number
 }
 export type CompactionCompletedPayload<M extends UIMessage> = {
   generationId: string
   throughMessageId: string
+  throughIndex?: number
   messages: M[]
   retainedMessageIds?: string[]
+  summary?: string
+  usage?: LanguageModelUsage
 }
 
 export type AIEventDefs<M extends UIMessage = UIMessage> = {
@@ -183,6 +203,8 @@ export type AIEventDefs<M extends UIMessage = UIMessage> = {
   'ai.input.responded': StandardSchemaV1<InputRespondedPayload>
   'ai.compaction.requested': StandardSchemaV1<CompactionRequestedPayload>
   'ai.compaction.completed': StandardSchemaV1<CompactionCompletedPayload<M>>
+  'ai.model.metadata.requested': StandardSchemaV1<ModelMetadataRequestedPayload>
+  'ai.model.metadata.resolved': StandardSchemaV1<ModelMetadataResolvedPayload>
 }
 
 export type ActiveGeneration = GenerationStartedPayload
@@ -242,8 +264,11 @@ export type CompactionState<M extends UIMessage> = {
   status: 'running' | 'completed'
   generationId: string
   throughMessageId: string
+  throughIndex?: number
   messages?: M[]
   retainedMessageIds?: string[]
+  summary?: string
+  usage?: LanguageModelUsage
 }
 
 export type GenerationUsage = {
@@ -269,6 +294,7 @@ export type AIState<M extends UIMessage = UIMessage> = {
   pendingInputs: PendingInput[]
   tools: ToolActivity[]
   compaction: CompactionState<M> | null
+  modelMetadata: Record<string, ModelMetadataState>
   usage: GenerationUsage[]
   error: string | null
 }
@@ -373,6 +399,15 @@ const jsonObject = <T extends object>(
 const stringField = (value: Record<string, unknown>, key: string): boolean =>
   typeof value[key] === 'string' && value[key].length > 0
 
+const optionalIndexField = (
+  value: Record<string, unknown>,
+  key: string,
+): boolean =>
+  value[key] === undefined ||
+  (typeof value[key] === 'number' &&
+    Number.isSafeInteger(value[key]) &&
+    value[key] >= 0)
+
 const optionalStringField = (
   value: Record<string, unknown>,
   key: string,
@@ -449,6 +484,7 @@ const completedPayloadSchema: StandardSchemaV1<GenerationCompletedPayload> =
       !stringField(value, 'messageId') ||
       !stringField(value, 'generationId') ||
       !stringField(value, 'responseMessageId') ||
+      !optionalIndexField(value, 'inputTokenEstimate') ||
       !optionalStringField(value, 'finishReason')
     ) {
       return issue('invalid ai.generation.completed payload')
@@ -465,6 +501,8 @@ const compactionPayloadSchema = <M extends UIMessage>(
       !isJSONCompatible(value) ||
       !stringField(value, 'generationId') ||
       !stringField(value, 'throughMessageId') ||
+      !optionalIndexField(value, 'throughIndex') ||
+      !optionalStringField(value, 'summary') ||
       !Array.isArray(value['messages']) ||
       (value['retainedMessageIds'] !== undefined &&
         (!Array.isArray(value['retainedMessageIds']) ||
@@ -484,7 +522,16 @@ const compactionPayloadSchema = <M extends UIMessage>(
       value: {
         generationId: value['generationId'] as string,
         throughMessageId: value['throughMessageId'] as string,
+        ...(value['throughIndex'] === undefined
+          ? {}
+          : { throughIndex: value['throughIndex'] as number }),
         messages,
+        ...(value['summary'] === undefined
+          ? {}
+          : { summary: value['summary'] as string }),
+        ...(value['usage'] === undefined
+          ? {}
+          : { usage: value['usage'] as LanguageModelUsage }),
         ...(value['retainedMessageIds'] === undefined
           ? {}
           : { retainedMessageIds: value['retainedMessageIds'] as string[] }),
@@ -547,7 +594,8 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
         typeof value['attempt'] === 'number' &&
         Number.isInteger(value['attempt']) &&
         value['attempt'] > 0 &&
-        stringField(value, 'model'),
+        stringField(value, 'model') &&
+        optionalIndexField(value, 'promptThroughIndex'),
     ),
     'ai.generation.progress': progressPayloadSchema,
     'ai.generation.completed': completedPayloadSchema,
@@ -641,9 +689,27 @@ export function createEvents<M extends UIMessage = UIMessage>(options?: {
       'ai.compaction.requested',
       (value) =>
         stringField(value, 'generationId') &&
-        stringField(value, 'throughMessageId'),
+        stringField(value, 'throughMessageId') &&
+        optionalIndexField(value, 'throughIndex'),
     ),
     'ai.compaction.completed': compactionPayloadSchema(messageSchema),
+    'ai.model.metadata.requested': jsonObject(
+      'ai.model.metadata.requested',
+      (value) => stringField(value, 'modelId'),
+    ),
+    'ai.model.metadata.resolved': jsonObject(
+      'ai.model.metadata.resolved',
+      (value) =>
+        stringField(value, 'modelId') &&
+        (value['limits'] === null ||
+          (isRecord(value['limits']) &&
+            typeof value['limits']['contextWindow'] === 'number' &&
+            Number.isSafeInteger(value['limits']['contextWindow']) &&
+            value['limits']['contextWindow'] > 0 &&
+            typeof value['limits']['maxOutputTokens'] === 'number' &&
+            Number.isSafeInteger(value['limits']['maxOutputTokens']) &&
+            value['limits']['maxOutputTokens'] > 0)),
+    ),
   }
   return definition
 }
@@ -663,6 +729,7 @@ const initialState = <M extends UIMessage>(): AIState<M> => ({
   pendingInputs: [],
   tools: [],
   compaction: null,
+  modelMetadata: {},
   usage: [],
   error: null,
 })
@@ -1084,6 +1151,30 @@ export function reduceAIState<M extends UIMessage>(
 ): AIState<M> {
   if (state.status === 'closed') return state
   switch (event.type) {
+    case 'ai.model.metadata.requested': {
+      const { modelId } = event.payload as ModelMetadataRequestedPayload
+      if (Object.hasOwn(state.modelMetadata, modelId)) return state
+      return {
+        ...state,
+        modelMetadata: {
+          ...state.modelMetadata,
+          [modelId]: { status: 'pending' },
+        },
+      }
+    }
+    case 'ai.model.metadata.resolved': {
+      const { modelId, limits } = event.payload as ModelMetadataResolvedPayload
+      return {
+        ...state,
+        modelMetadata: {
+          ...state.modelMetadata,
+          [modelId]:
+            limits === null
+              ? { status: 'unavailable' }
+              : { status: 'resolved', limits },
+        },
+      }
+    }
     case 'ai.session.created': {
       const payload = event.payload as SessionCreatedPayload
       return {
@@ -1102,6 +1193,10 @@ export function reduceAIState<M extends UIMessage>(
       return {
         ...state,
         status: 'closed',
+        compaction:
+          state.compaction?.status === 'running'
+            ? (state.activeProjection?.baseCompaction ?? null)
+            : state.compaction,
         activeGeneration: null,
         activeRequestId: null,
         activeResponseMessageId: null,
@@ -1968,7 +2063,7 @@ export function createReducer<D extends EventDefs>(options: {
   type M = AIMessageFrom<D>
   return options.contract
     .reducer({
-      name: options.name ?? 'a2.ai.state.v8',
+      name: options.name ?? 'a2.ai.state.v10',
       initialState: initialState<M>(),
     })
     .fold((state, event) => reduceAIState(state, event))
@@ -2134,7 +2229,7 @@ export function agent<
   })
   const reducer = agentContract
     .reducer({
-      name: options.reducerName ?? 'a2.ai.state.v8',
+      name: options.reducerName ?? 'a2.ai.state.v10',
       initialState: initialState<M>(),
     })
     .fold((state, event) => reduceAIState(state, event))