experimental-a2 0.14.1 → 0.16.0

package/src/ai-coordinator.ts

@@ -8,8 +8,10 @@ import type {
   GenerationCompletedPayload,
   GenerationFailedPayload,
   GenerationRequestedPayload,
+  GenerationReason,
   GenerationStartedPayload,
   InputRequestedPayload,
+  InputRespondedPayload,
   MessageCreatedPayload,
   MessageInterruptedPayload,
   ToolCalledPayload,
@@ -18,12 +20,51 @@ import type {
 import type { Contract, ContractEvent, EventDefs } from './contract.ts'
 import type { Reducer } from './reducer.ts'
 
+type CompletedCall = Pick<
+  ToolCalledPayload,
+  | 'requestId'
+  | 'generationId'
+  | 'messageId'
+  | 'toolCallId'
+  | 'toolName'
+  | 'providerExecuted'
+  | 'supportsDeferredResults'
+>
+
 export type CoordinatorCall = {
-  call: ToolCalledPayload
   approval?: ApprovalRequestedPayload
   response?: ApprovalRespondedPayload
-  terminal: boolean
   work?: ControlWork
+} & (
+  | { call: ToolCalledPayload; terminal: false }
+  | { call: CompletedCall; terminal: true }
+)
+
+const completeCall = (candidate: CoordinatorCall): CoordinatorCall => {
+  const {
+    requestId,
+    generationId,
+    messageId,
+    toolCallId,
+    toolName,
+    providerExecuted,
+    supportsDeferredResults,
+  } = candidate.call
+  return {
+    terminal: true,
+    call: {
+      requestId,
+      generationId,
+      messageId,
+      toolCallId,
+      toolName,
+      ...(providerExecuted === undefined ? {} : { providerExecuted }),
+      ...(supportsDeferredResults === undefined
+        ? {}
+        : { supportsDeferredResults }),
+    },
+    ...(candidate.work === undefined ? {} : { work: candidate.work }),
+  }
 }
 
 export type CoordinatorResponse = {
@@ -32,13 +73,11 @@ export type CoordinatorResponse = {
   status: 'requested' | 'generating' | 'waiting' | 'failed'
   stepCount: number
   activeRequestId?: string
+  requestReason: GenerationReason
+  sourceGenerationId?: string
   generation?: GenerationStartedPayload
-  promptThroughIndex?: number
   source?: {
     generation: GenerationStartedPayload
-    promptThroughIndex: number
-    failed: boolean
-    canCompact: boolean
   }
   completion?: GenerationCompletedPayload
   failure?: GenerationFailedPayload
@@ -53,6 +92,7 @@ export type CoordinatorResponse = {
 
 export type AICoordinatorState = {
   closed: boolean
+  unsettledTools?: Record<string, true>
   queued: Array<{ messageId: string; generate: boolean }>
   response?: CoordinatorResponse
   calibration?: { model: string; inputTokens: number; estimate: number }
@@ -96,8 +136,7 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
     case 'ai.generation.requested': {
       const request = event.payload as GenerationRequestedPayload
       if (request.reason === 'message') {
-        const responseMessageId =
-          request.responseMessageId ?? `${request.messageId}:assistant`
+        const responseMessageId = request.responseMessageId
         const requested = state.queued.findIndex(
           (item) => item.messageId === request.messageId,
         )
@@ -115,6 +154,10 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
             status: 'requested',
             stepCount: 0,
             activeRequestId: event.id,
+            requestReason: request.reason,
+            ...(request.sourceGenerationId === undefined
+              ? {}
+              : { sourceGenerationId: request.sourceGenerationId }),
             calls: [],
             inputs: [],
           },
@@ -129,8 +172,12 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       if (
         (request.reason === 'tool' &&
           (!continuationReady(state) ||
-            event.id !==
-              `ai.generate:tools:${state.response.generation?.generationId}`)) ||
+            request.sourceGenerationId !==
+              state.response.generation?.generationId)) ||
+        (request.sourceGenerationId !== undefined &&
+          request.sourceGenerationId !==
+            (state.response.generation ?? state.response.source?.generation)
+              ?.generationId) ||
         (request.reason === 'retry' && state.response.status !== 'failed') ||
         (request.reason === 'input' &&
           state.response.inputResponse === undefined)
@@ -139,23 +186,21 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       }
       return updateResponse(state, (response) => {
         const generation = response.generation ?? response.source?.generation
-        const promptThroughIndex =
-          response.promptThroughIndex ?? response.source?.promptThroughIndex
         return {
           rootMessageId: response.rootMessageId,
           responseMessageId: response.responseMessageId,
           status: 'requested',
           stepCount: response.stepCount,
           activeRequestId: event.id,
-          ...(generation === undefined || promptThroughIndex === undefined
+          requestReason: request.reason,
+          ...(request.sourceGenerationId === undefined
+            ? {}
+            : { sourceGenerationId: request.sourceGenerationId }),
+          ...(generation === undefined
             ? {}
             : {
                 source: {
                   generation,
-                  promptThroughIndex,
-                  failed:
-                    response.failure?.generationId === generation.generationId,
-                  canCompact: response.calls.every((call) => call.terminal),
                 },
               }),
           calls: [],
@@ -179,11 +224,12 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
         status: 'generating',
         stepCount: response.stepCount,
         generation,
-        ...(generation.promptThroughIndex === undefined
-          ? {}
-          : { promptThroughIndex: generation.promptThroughIndex }),
         ...(response.source === undefined ? {} : { source: response.source }),
         activeRequestId: generation.requestId,
+        requestReason: response.requestReason,
+        ...(response.sourceGenerationId === undefined
+          ? {}
+          : { sourceGenerationId: response.sourceGenerationId }),
         calls: [],
         inputs: [],
       }))
@@ -191,7 +237,8 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
     case 'ai.generation.completed': {
       const completion = event.payload as GenerationCompletedPayload
       if (
-        state.response?.generation?.generationId !== completion.generationId
+        state.response?.generation?.requestId !== completion.requestId ||
+        state.response.generation.generationId !== completion.generationId
       ) {
         return state
       }
@@ -223,11 +270,13 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       const failure = event.payload as GenerationFailedPayload
       const response = state.response
       const ownsActiveGeneration =
-        response?.generation?.generationId === failure.generationId
+        response?.generation?.generationId === failure.generationId &&
+        response.generation.requestId === failure.requestId
       const ownsActiveStepLimit =
         failure.stepLimit === true &&
-        response?.activeRequestId ===
-          `ai.generate:tools:${failure.generationId}` &&
+        response?.requestReason === 'tool' &&
+        response.sourceGenerationId === failure.generationId &&
+        response.source?.generation.requestId === failure.requestId &&
         response.responseMessageId === failure.responseMessageId
       if (!ownsActiveGeneration && !ownsActiveStepLimit) {
         return state
@@ -258,7 +307,10 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
     }
     case 'ai.tool.called': {
       const call = event.payload as ToolCalledPayload
-      if (state.response?.generation?.generationId !== call.generationId) {
+      if (
+        state.response?.generation?.requestId !== call.requestId ||
+        state.response.generation.generationId !== call.generationId
+      ) {
         return state
       }
       return updateResponse(state, (response) => ({
@@ -273,7 +325,10 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
     }
     case 'ai.approval.requested': {
       const approval = event.payload as ApprovalRequestedPayload
-      if (state.response?.generation?.generationId !== approval.generationId) {
+      if (
+        state.response?.generation?.requestId !== approval.requestId ||
+        state.response.generation.generationId !== approval.generationId
+      ) {
         return state
       }
       return updateResponse(state, (response) => ({
@@ -289,7 +344,8 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       const approval = event.payload as ApprovalRespondedPayload
       if (
         state.response?.responseMessageId !== approval.messageId ||
-        state.response.generation?.generationId !== approval.generationId
+        state.response.generation?.requestId !== approval.requestId ||
+        state.response.generation.generationId !== approval.generationId
       ) {
         return state
       }
@@ -306,7 +362,8 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       const result = event.payload as ToolResultPayload
       if (
         result.preliminary === true ||
-        state.response?.generation?.generationId !== result.generationId
+        state.response?.generation?.requestId !== result.requestId ||
+        state.response.generation.generationId !== result.generationId
       ) {
         return state
       }
@@ -314,7 +371,7 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
         ...response,
         calls: response.calls.map((candidate) =>
           candidate.call.toolCallId === result.toolCallId
-            ? { ...candidate, terminal: true }
+            ? completeCall(candidate)
             : candidate,
         ),
       }))
@@ -323,7 +380,8 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       const input = event.payload as InputRequestedPayload
       if (
         state.response?.responseMessageId !== input.messageId ||
-        state.response.generation?.generationId !== input.generationId
+        state.response.generation?.requestId !== input.requestId ||
+        state.response.generation.generationId !== input.generationId
       ) {
         return state
       }
@@ -341,12 +399,7 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       }))
     }
     case 'ai.input.responded': {
-      const input = event.payload as {
-        messageId: string
-        generationId: string
-        inputId: string
-        name: string
-      }
+      const input = event.payload as InputRespondedPayload
       const requested = state.response?.inputs.find(
         (candidate) =>
           candidate.messageId === input.messageId &&
@@ -356,7 +409,8 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       )
       if (
         state.response?.responseMessageId !== input.messageId ||
-        state.response.generation?.generationId !== input.generationId ||
+        state.response.generation?.requestId !== input.requestId ||
+        state.response.generation.generationId !== input.generationId ||
         requested === undefined
       ) {
         return state
@@ -386,7 +440,7 @@ export const aiCoordinatorReducer = <
 ): Reducer<D, AICoordinatorState> =>
   contract
     .reducer({
-      name: 'a2.ai.coordinator.v5',
+      name: 'a2.ai.coordinator.v6',
       initialState: {
         closed: false,
         queued: [],

package/src/ai-id.ts

@@ -0,0 +1,6 @@
+import { idempotentId } from './idempotent-id.ts'
+
+export const AI_EVENT_ID_PREFIX = 'a2.ai:'
+
+export const aiEventId = async (...parts: string[]): Promise<string> =>
+  `${AI_EVENT_ID_PREFIX}${await idempotentId(...parts)}`

package/src/ai-input-tokens.ts

@@ -0,0 +1,75 @@
+import { Buffer } from 'node:buffer'
+import { asSchema } from 'ai'
+import type { Instructions, ModelMessage, ToolSet } from 'ai'
+import type { AICoordinatorState } from './ai-coordinator.ts'
+
+const withoutProviderOptions = <T extends object>(value: T): T => {
+  const result = { ...value }
+  Reflect.deleteProperty(result, 'providerOptions')
+  return result
+}
+
+const estimatedMessage = (message: ModelMessage): unknown => ({
+  ...withoutProviderOptions(message),
+  content:
+    typeof message.content === 'string'
+      ? message.content
+      : message.content.map((part) =>
+          part.type === 'tool-result'
+            ? {
+                ...withoutProviderOptions(part),
+                output: {
+                  ...withoutProviderOptions(part.output),
+                  ...(part.output.type === 'content'
+                    ? {
+                        value: part.output.value.map(withoutProviderOptions),
+                      }
+                    : {}),
+                },
+              }
+            : withoutProviderOptions(part),
+        ),
+})
+
+export const estimateInputTokens = async (options: {
+  messages: ModelMessage[]
+  instructions: Instructions | undefined
+  tools: ToolSet
+}): Promise<number> => {
+  const tools = await Promise.all(
+    Object.entries(options.tools).map(async ([name, tool]) => ({
+      name,
+      description: tool.description,
+      inputSchema: await asSchema(tool.inputSchema).jsonSchema,
+    })),
+  )
+  const instructions = options.instructions
+  return Math.ceil(
+    Buffer.byteLength(
+      JSON.stringify({
+        messages: options.messages.map(estimatedMessage),
+        instructions:
+          instructions === undefined || typeof instructions === 'string'
+            ? instructions
+            : Array.isArray(instructions)
+              ? instructions.map(estimatedMessage)
+              : estimatedMessage(instructions),
+        tools,
+      }),
+      'utf8',
+    ) / 4,
+  )
+}
+
+export const measuredInputTokens = (options: {
+  estimate: number
+  model: string
+  calibration: AICoordinatorState['calibration']
+}): number => {
+  const previous = options.calibration
+  return previous !== undefined &&
+    previous.model === options.model &&
+    options.estimate >= previous.estimate
+    ? Math.ceil(previous.inputTokens + options.estimate - previous.estimate)
+    : options.estimate
+}