experimental-a2 0.4.0 → 0.5.1

package/src/ai-coordinator.ts

@@ -0,0 +1,358 @@
+import type { UIMessage } from 'ai'
+import type {
+  AIEventDefs,
+  ApprovalRequestedPayload,
+  ApprovalRespondedPayload,
+  GenerationCompletedPayload,
+  GenerationFailedPayload,
+  GenerationRequestedPayload,
+  GenerationStartedPayload,
+  InputRequestedPayload,
+  MessageCreatedPayload,
+  MessageInterruptedPayload,
+  ToolCalledPayload,
+  ToolResultPayload,
+} from './ai.ts'
+import type { Contract, ContractEvent, EventDefs } from './contract.ts'
+import type { Reducer } from './reducer.ts'
+
+export type CoordinatorCall = {
+  index: number
+  call: ToolCalledPayload
+  approval?: ApprovalRequestedPayload
+  response?: ApprovalRespondedPayload
+  responseIndex?: number
+  terminal: boolean
+}
+
+export type CoordinatorResponse = {
+  rootMessageId: string
+  responseMessageId: string
+  status: 'requested' | 'generating' | 'waiting' | 'failed'
+  activeRequestId?: string
+  generation?: GenerationStartedPayload
+  completion?: GenerationCompletedPayload
+  failure?: GenerationFailedPayload
+  calls: CoordinatorCall[]
+  inputs: InputRequestedPayload[]
+  inputResponse?: {
+    index: number
+    generationId: string
+    inputId: string
+    name: string
+  }
+}
+
+export type AICoordinatorState = {
+  closed: boolean
+  queued: Array<{ index: number; messageId: string; generate: boolean }>
+  response?: CoordinatorResponse
+}
+
+const updateResponse = (
+  state: AICoordinatorState,
+  update: (response: CoordinatorResponse) => CoordinatorResponse,
+): AICoordinatorState =>
+  state.response === undefined
+    ? state
+    : { ...state, response: update(state.response) }
+
+const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
+  state: AICoordinatorState,
+  event: ContractEvent<D>,
+): AICoordinatorState => {
+  switch (event.type) {
+    case 'ai.session.closed':
+      return { closed: true, queued: [] }
+    case 'ai.message.created': {
+      const payload = event.payload as MessageCreatedPayload<M>
+      const message = payload.message
+      if (message.role !== 'user') return state
+      return {
+        ...state,
+        queued: [
+          ...state.queued.filter((item) => item.messageId !== message.id),
+          {
+            index: event.index,
+            messageId: message.id,
+            generate: payload.generate !== false,
+          },
+        ].toSorted((left, right) => left.index - right.index),
+      }
+    }
+    case 'ai.generation.requested': {
+      const request = event.payload as GenerationRequestedPayload
+      if (request.reason === 'message') {
+        const responseMessageId =
+          request.responseMessageId ?? `${request.messageId}:assistant`
+        const requested = state.queued.find(
+          (item) => item.messageId === request.messageId,
+        )
+        return {
+          ...state,
+          queued:
+            requested === undefined
+              ? state.queued.filter(
+                  (item) => item.messageId !== request.messageId,
+                )
+              : state.queued.filter((item) => item.index > requested.index),
+          response: {
+            rootMessageId: request.messageId,
+            responseMessageId,
+            status: 'requested',
+            activeRequestId: event.id,
+            calls: [],
+            inputs: [],
+          },
+        }
+      }
+      if (
+        state.response === undefined ||
+        request.responseMessageId !== state.response.responseMessageId
+      ) {
+        return state
+      }
+      if (
+        (request.reason === 'tool' &&
+          (!continuationReady(state) ||
+            event.id !==
+              `ai.generate:tools:${state.response.generation?.generationId}`)) ||
+        (request.reason === 'retry' && state.response.status !== 'failed') ||
+        (request.reason === 'input' &&
+          state.response.inputResponse === undefined)
+      ) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        rootMessageId: response.rootMessageId,
+        responseMessageId: response.responseMessageId,
+        status: 'requested',
+        activeRequestId: event.id,
+        calls: [],
+        inputs: [],
+      }))
+    }
+    case 'ai.generation.started': {
+      const generation = event.payload as GenerationStartedPayload
+      if (
+        state.response?.responseMessageId !== generation.responseMessageId ||
+        state.response.activeRequestId !== generation.requestId ||
+        (state.response.generation?.requestId === generation.requestId &&
+          state.response.generation.attempt >= generation.attempt)
+      ) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        rootMessageId: response.rootMessageId,
+        responseMessageId: response.responseMessageId,
+        status: 'generating',
+        generation,
+        activeRequestId: generation.requestId,
+        calls: [],
+        inputs: [],
+      }))
+    }
+    case 'ai.generation.completed': {
+      const completion = event.payload as GenerationCompletedPayload
+      if (
+        state.response?.generation?.generationId !== completion.generationId
+      ) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        ...response,
+        status: 'waiting',
+        completion,
+      }))
+    }
+    case 'ai.generation.failed': {
+      const failure = event.payload as GenerationFailedPayload
+      const response = state.response
+      const ownsActiveGeneration =
+        response?.generation?.generationId === failure.generationId
+      const ownsActiveStepLimit =
+        failure.stepLimit === true &&
+        response?.activeRequestId ===
+          `ai.generate:tools:${failure.generationId}` &&
+        response.responseMessageId === failure.responseMessageId
+      if (!ownsActiveGeneration && !ownsActiveStepLimit) {
+        return state
+      }
+      return updateResponse(state, (current) => ({
+        ...current,
+        status: 'failed',
+        failure,
+      }))
+    }
+    case 'ai.message.completed': {
+      const messageId = (event.payload as { messageId: string }).messageId
+      return state.response?.responseMessageId === messageId
+        ? { closed: state.closed, queued: state.queued }
+        : state
+    }
+    case 'ai.message.interrupted': {
+      const interruption = event.payload as MessageInterruptedPayload
+      const response = state.response
+      return response?.responseMessageId === interruption.messageId &&
+        (interruption.generationId === response.generation?.generationId ||
+          interruption.requestId === response.activeRequestId)
+        ? { closed: state.closed, queued: state.queued }
+        : state
+    }
+    case 'ai.tool.called': {
+      const call = event.payload as ToolCalledPayload
+      if (state.response?.generation?.generationId !== call.generationId) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        ...response,
+        calls: [
+          ...response.calls.filter(
+            (candidate) => candidate.call.toolCallId !== call.toolCallId,
+          ),
+          { index: event.index, call, terminal: false },
+        ],
+      }))
+    }
+    case 'ai.approval.requested': {
+      const approval = event.payload as ApprovalRequestedPayload
+      if (state.response?.generation?.generationId !== approval.generationId) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        ...response,
+        calls: response.calls.map((candidate) =>
+          candidate.call.toolCallId === approval.toolCallId
+            ? { ...candidate, approval }
+            : candidate,
+        ),
+      }))
+    }
+    case 'ai.approval.responded': {
+      const approval = event.payload as ApprovalRespondedPayload
+      if (
+        state.response?.responseMessageId !== approval.messageId ||
+        state.response.generation?.generationId !== approval.generationId
+      ) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        ...response,
+        calls: response.calls.map((candidate) =>
+          candidate.approval?.approvalId === approval.approvalId
+            ? { ...candidate, response: approval, responseIndex: event.index }
+            : candidate,
+        ),
+      }))
+    }
+    case 'ai.tool.result': {
+      const result = event.payload as ToolResultPayload
+      if (
+        result.preliminary === true ||
+        state.response?.generation?.generationId !== result.generationId
+      ) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        ...response,
+        calls: response.calls.map((candidate) =>
+          candidate.call.toolCallId === result.toolCallId
+            ? { ...candidate, terminal: true }
+            : candidate,
+        ),
+      }))
+    }
+    case 'ai.input.requested': {
+      const input = event.payload as InputRequestedPayload
+      if (
+        state.response?.responseMessageId !== input.messageId ||
+        state.response.generation?.generationId !== input.generationId
+      ) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        ...response,
+        inputs: [
+          ...response.inputs.filter(
+            (candidate) =>
+              candidate.messageId !== input.messageId ||
+              candidate.generationId !== input.generationId ||
+              candidate.inputId !== input.inputId,
+          ),
+          input,
+        ],
+      }))
+    }
+    case 'ai.input.responded': {
+      const input = event.payload as {
+        messageId: string
+        generationId: string
+        inputId: string
+        name: string
+      }
+      const requested = state.response?.inputs.find(
+        (candidate) =>
+          candidate.messageId === input.messageId &&
+          candidate.generationId === input.generationId &&
+          candidate.inputId === input.inputId &&
+          candidate.name === input.name,
+      )
+      if (
+        state.response?.responseMessageId !== input.messageId ||
+        state.response.generation?.generationId !== input.generationId ||
+        requested === undefined
+      ) {
+        return state
+      }
+      return updateResponse(state, (response) => ({
+        ...response,
+        inputs: response.inputs.filter(
+          (candidate) => candidate.inputId !== input.inputId,
+        ),
+        inputResponse: {
+          index: event.index,
+          generationId: input.generationId,
+          inputId: input.inputId,
+          name: requested.name,
+        },
+      }))
+    }
+    default:
+      return state
+  }
+}
+
+export const aiCoordinatorReducer = <
+  M extends UIMessage,
+  D extends AIEventDefs<M> & EventDefs,
+>(
+  contract: Contract<D>,
+): Reducer<D, AICoordinatorState> =>
+  contract
+    .reducer({
+      name: 'a2.ai.coordinator.v3',
+      initialState: {
+        closed: false,
+        queued: [],
+      } as AICoordinatorState,
+    })
+    .fold((state, event) => foldCoordinator<M, D>(state, event))
+
+export const continuationReady = (state: AICoordinatorState): boolean => {
+  const response = state.response
+  return (
+    !state.closed &&
+    response?.completion?.finishReason === 'tool-calls' &&
+    response.generation?.generationId === response.completion.generationId &&
+    response.failure === undefined &&
+    response.calls.length > 0 &&
+    response.calls.every(
+      (call) =>
+        call.terminal ||
+        (call.call.providerExecuted === true &&
+          call.call.supportsDeferredResults !== true &&
+          call.approval !== undefined &&
+          call.response !== undefined),
+    )
+  )
+}