experimental-a2 0.15.0 → 0.16.1

package/src/ai-control-server.ts

@@ -28,6 +28,7 @@ import type { AICoordinatorState, CoordinatorCall } from './ai-coordinator.ts'
 import { validateSync } from './validate.ts'
 import type { StandardSchemaV1 } from './standard-schema.ts'
 import { A2Error } from './errors.ts'
+import { aiEventId } from './ai-id.ts'
 
 export class ControlCancelled extends Error {}
 
@@ -58,7 +59,7 @@ export type AIControlRuntime<
   settled(options: {
     ctx: HandlerContext<D>
     events: void | AppendInput<D> | readonly AppendInput<D>[]
-  }): AppendInput<D>
+  }): Promise<AppendInput<D>>
   cancelled(options: { error: unknown; signal: AbortSignal }): boolean
   reducer: Reducer<D, AIControlState<M>>
   coordinator: Reducer<D, AICoordinatorState>
@@ -72,17 +73,36 @@ export const createControlRuntime = <
 }): AIControlRuntime<M, D> => {
   const reducer = options.agent.contract
     .reducer({
-      name: 'a2.ai.control.v1',
+      name: 'a2.ai.control.v5',
       initialState: initialControlState<M>(),
     })
-    .fold((state, event) =>
-      event.type === 'ai.control.committed'
+    .fold((state, event) => {
+      if (event.type === 'ai.control.requested') {
+        const command = event.payload as ControlCommand<M>
+        if (command.action === 'stop')
+          return {
+            ...state,
+            pendingStops: {
+              ...state.pendingStops,
+              [event.id]: command.turnId,
+            },
+          }
+      }
+      if (event.type === 'ai.control.decided') {
+        const receipt = event.payload as ControlReceipt
+        if (Object.hasOwn(state.pendingStops, receipt.commandId)) {
+          const pendingStops = { ...state.pendingStops }
+          delete pendingStops[receipt.commandId]
+          return { ...state, pendingStops }
+        }
+      }
+      return event.type === 'ai.control.committed'
         ? applyControlChanges({
             state,
             changes: (event.payload as ControlCommit).changes,
           })
-        : state,
-    )
+        : state
+    })
   const coordinator = aiCoordinatorReducer(options.agent.contract)
 
   const handler = async (
@@ -97,16 +117,22 @@ export const createControlRuntime = <
     const state = structuredClone(snapshot.state)
     const facts: WorkFact[] = []
     const jobs: AppendInput<D>[] = []
-    const emit = ({
+    const emit = async ({
       type,
       payload,
-      id = `${ctx.event.id}:fact:${facts.length}`,
+      id,
     }: {
       type: string
       payload: unknown
       id?: string
-    }): void => {
-      const fact = { type, payload, id }
+    }): Promise<void> => {
+      const fact = {
+        type,
+        payload,
+        id:
+          id ??
+          (await aiEventId('a2.ai.fact', ctx.event.id, String(facts.length))),
+      }
       facts.push(fact)
       state.coordinator = foldCoordinator<M, D>(state.coordinator, {
         ...ctx.event,
@@ -114,20 +140,20 @@ export const createControlRuntime = <
       } as ContractEvent<D>)
     }
     let reason: ControlReceipt['reason']
-    const end = ({
+    const end = async ({
       lastSeenIndex,
       reason: interruptionReason,
     }: {
       lastSeenIndex: number | undefined
       reason: string | undefined
-    }): void => {
+    }): Promise<void> => {
       const response = state.coordinator.response
       if (response === undefined) return
       const owner =
         response.generation === undefined
           ? { requestId: response.activeRequestId! }
           : { generationId: response.generation.generationId }
-      emit({
+      await emit({
         type: 'ai.message.interrupted',
         payload: {
           messageId: response.responseMessageId,
@@ -140,24 +166,6 @@ export const createControlRuntime = <
       })
       state.active = null
     }
-    const knownInput = (id: string): boolean =>
-      state.inbox.items.some((candidate) => candidate.id === id) ||
-      Object.hasOwn(state.reservedIds, id)
-    const missingInputReason = (id: string): ControlReceipt['reason'] =>
-      state.reservedIds[id] === 'removed' || !knownInput(id)
-        ? 'not-found'
-        : 'already-active'
-    const inputTaken = ({
-      message,
-      generate = true,
-    }: {
-      message: M
-      generate?: boolean
-    }): boolean =>
-      knownInput(message.id) ||
-      (generate &&
-        message.role === 'user' &&
-        knownInput(`${message.id}:assistant`))
     const insert = ({
       message,
       generate = true,
@@ -171,18 +179,10 @@ export const createControlRuntime = <
     }): void => {
       const candidate = {
         id: message.id,
-        revision: 0,
         message,
         generate: generate && message.role === 'user',
         ...(afterStepOf === undefined ? {} : { afterStepOf }),
       }
-      state.reservedIds = {
-        ...state.reservedIds,
-        [candidate.id]: true,
-        ...(candidate.generate
-          ? { [`${candidate.id}:assistant`]: true as const }
-          : {}),
-      }
       if (first) state.inbox.items.unshift(candidate)
       else state.inbox.items.push(candidate)
     }
@@ -230,61 +230,41 @@ export const createControlRuntime = <
       else
         switch (command.action) {
           case 'send':
-            if (
-              inputTaken({
-                message: command.message,
-                generate: command.generate !== false,
-              })
-            )
-              reason = 'duplicate-input'
-            else
-              insert({
-                message: command.message,
-                generate: command.generate !== false,
-              })
+            insert({
+              message: command.message,
+              generate: command.generate !== false,
+            })
             break
           case 'edit': {
             const item = state.inbox.items.find(
               (candidate) => candidate.id === command.message.id,
             )
-            if (item === undefined)
-              reason = missingInputReason(command.message.id)
-            else if (item.revision !== command.expectedRevision)
-              reason = 'revision-conflict'
+            if (item === undefined) reason = 'not-pending'
             else if (item.message.role !== command.message.role)
-              reason = 'revision-conflict'
-            else {
-              item.message = command.message
-              item.revision += 1
-            }
+              reason = 'role-mismatch'
+            else item.message = command.message
             break
           }
           case 'remove': {
             const position = state.inbox.items.findIndex(
               (item) => item.id === command.inputId,
             )
-            if (position < 0) reason = missingInputReason(command.inputId)
-            else {
-              state.inbox.items.splice(position, 1)
-              state.reservedIds = {
-                ...state.reservedIds,
-                [command.inputId]: 'removed',
-              }
-            }
+            if (position < 0) reason = 'not-pending'
+            else state.inbox.items.splice(position, 1)
             break
           }
           case 'move': {
             const index = state.inbox.items.findIndex(
               (candidate) => candidate.id === command.inputId,
             )
-            if (index < 0) reason = missingInputReason(command.inputId)
+            if (index < 0) reason = 'not-pending'
             else if (
               command.beforeId !== null &&
               !state.inbox.items.some(
                 (candidate) => candidate.id === command.beforeId,
               )
             )
-              reason = 'not-found'
+              reason = 'not-pending'
             else {
               const destination =
                 command.beforeId === null
@@ -304,23 +284,12 @@ export const createControlRuntime = <
               (item) => item.id === command.inputId,
             )
             const selected = state.inbox.items[position]
-            if (!selected) reason = missingInputReason(command.inputId)
+            if (!selected) reason = 'not-pending'
             else if ((state.active?.turnId ?? null) !== command.turnId)
               reason = 'stale-turn'
             else if (selected.message.role !== 'user') reason = 'not-user-input'
-            else if (
-              !selected.generate &&
-              knownInput(`${selected.id}:assistant`)
-            )
-              reason = 'duplicate-input'
             else {
-              if (!selected.generate) {
-                state.reservedIds = {
-                  ...state.reservedIds,
-                  [`${selected.id}:assistant`]: true,
-                }
-                selected.generate = true
-              }
+              selected.generate = true
               if (state.active) selected.afterStepOf = state.active.turnId
               else delete selected.afterStepOf
               moveInput({ from: position, to: 0 })
@@ -330,19 +299,18 @@ export const createControlRuntime = <
           case 'stop':
             if (state.active?.turnId !== command.turnId) reason = 'stale-turn'
             else
-              end({ lastSeenIndex: command.lastSeenIndex, reason: undefined })
+              await end({
+                lastSeenIndex: command.lastSeenIndex,
+                reason: undefined,
+              })
             break
           case 'steer':
-            if (inputTaken({ message: command.message }))
-              reason = 'duplicate-input'
-            else {
-              insert({
-                message: command.message,
-                generate: true,
-                first: true,
-                afterStepOf: command.turnId,
-              })
-            }
+            insert({
+              message: command.message,
+              generate: true,
+              first: true,
+              afterStepOf: command.turnId,
+            })
             break
           case 'interrupt': {
             const response = state.coordinator.response
@@ -355,7 +323,7 @@ export const createControlRuntime = <
             )
               reason = 'stale-turn'
             else
-              end({
+              await end({
                 lastSeenIndex: command.request.lastSeenIndex,
                 reason: command.request.reason,
               })
@@ -373,7 +341,7 @@ export const createControlRuntime = <
               active.version += 1
               const generation = state.coordinator.response?.generation
               if (generation)
-                emit({
+                await emit({
                   type: 'ai.generation.paused',
                   payload: {
                     generationId: generation.generationId,
@@ -389,6 +357,7 @@ export const createControlRuntime = <
           case 'tool-result': {
             const call = state.coordinator.response?.calls.find(
               (candidate) =>
+                candidate.call.requestId === command.result.requestId &&
                 candidate.call.generationId === command.result.generationId &&
                 candidate.call.messageId === command.result.messageId &&
                 candidate.call.toolCallId === command.result.toolCallId &&
@@ -396,15 +365,15 @@ export const createControlRuntime = <
                 !candidate.terminal,
             )
             if (!call) reason = 'not-found'
-            else emit({ type: 'ai.tool.result', payload: command.result })
+            else await emit({ type: 'ai.tool.result', payload: command.result })
             break
           }
           case 'approval': {
             const response = state.coordinator.response
             const call = response?.calls.find(
               (candidate) =>
-                candidate.approval?.approvalId ===
-                  command.response.approvalId &&
+                candidate.approval?.requestId === command.response.requestId &&
+                candidate.approval.approvalId === command.response.approvalId &&
                 candidate.approval.generationId ===
                   command.response.generationId &&
                 candidate.approval.messageId === command.response.messageId,
@@ -413,7 +382,10 @@ export const createControlRuntime = <
             else if (!call || call.response || call.terminal)
               reason = 'not-found'
             else
-              emit({ type: 'ai.approval.responded', payload: command.response })
+              await emit({
+                type: 'ai.approval.responded',
+                payload: command.response,
+              })
             break
           }
           case 'request-input':
@@ -423,7 +395,8 @@ export const createControlRuntime = <
             const response = state.coordinator.response
             if (
               response?.failure !== undefined ||
-              response?.generation?.generationId !== payload.generationId ||
+              response?.generation?.requestId !== payload.requestId ||
+              response.generation.generationId !== payload.generationId ||
               response.responseMessageId !== payload.messageId ||
               (command.action === 'input' &&
                 !response.inputs.some(
@@ -434,7 +407,7 @@ export const createControlRuntime = <
             )
               reason = 'not-found'
             else
-              emit({
+              await emit({
                 type:
                   command.action === 'input'
                     ? 'ai.input.responded'
@@ -453,14 +426,22 @@ export const createControlRuntime = <
             )
               reason = 'stale-turn'
             else if (!state.inbox.paused && !state.active.suspended) {
-              emit({ type: 'ai.retry.requested', payload: command.request })
+              await emit({
+                type: 'ai.retry.requested',
+                payload: command.request,
+              })
               requestModel({
                 request: {
                   messageId: response.rootMessageId,
                   responseMessageId: response.responseMessageId,
                   reason: 'retry',
                 },
-                id: `ai.generate:retry:${command.request.retryId}`,
+                id: await aiEventId(
+                  'a2.ai.request.retry',
+                  options.agent.contract.name,
+                  ctx.event.sessionId,
+                  command.request.retryId,
+                ),
               })
             } else reason = 'stale-turn'
             break
@@ -468,6 +449,15 @@ export const createControlRuntime = <
         }
     } else {
       const report = ctx.event.payload as WorkReport
+      if (
+        report.kind === 'tool' &&
+        report.settled &&
+        state.coordinator.unsettledTools
+      ) {
+        delete state.coordinator.unsettledTools[report.workId]
+        if (Object.keys(state.coordinator.unsettledTools).length === 0)
+          delete state.coordinator.unsettledTools
+      }
       const active = state.active
       const response = state.coordinator.response
       const call = response?.calls.find(
@@ -492,11 +482,18 @@ export const createControlRuntime = <
             authorized ||
             (report.kind === 'tool' && event.type === 'ai.tool.result')
           ) {
-            emit({ type: event.type, payload: event.payload, id: event.id })
+            await emit({
+              type: event.type,
+              payload: event.payload,
+              id: event.id,
+            })
             work.facts = [...(work.facts ?? []), event.id]
           }
         }
-        if (report.settled) work.settled = true
+        if (report.settled) {
+          work.settled = true
+          delete work.facts
+        }
         if (!authorized) reason = 'stale-work'
       }
     }
@@ -534,7 +531,7 @@ export const createControlRuntime = <
           )
             continue
           if (call.response?.approved === false) {
-            emit({
+            await emit({
               type: 'ai.tool.result',
               payload: {
                 ...call.call,
@@ -544,8 +541,17 @@ export const createControlRuntime = <
             })
             continue
           }
-          const id = `${ctx.event.id}:tool:${call.call.toolCallId}`
+          const id = await aiEventId(
+            'a2.ai.tool.execution',
+            ctx.event.id,
+            call.call.generationId,
+            call.call.toolCallId,
+          )
           call.work = { id, version: active.version, settled: false }
+          state.coordinator.unsettledTools = {
+            ...state.coordinator.unsettledTools,
+            [id]: true,
+          }
           jobs.push({
             type: 'ai.tool.execution.requested',
             id,
@@ -593,7 +599,7 @@ export const createControlRuntime = <
         response.completion !== undefined &&
         response.completion.finishReason !== 'tool-calls'
       if (handoffReady && !naturalCompletion) {
-        end({ lastSeenIndex: undefined, reason: 'steered after step' })
+        await end({ lastSeenIndex: undefined, reason: 'steered after step' })
       } else if (readyForModel && (!handoff || handoffReady)) {
         const resume =
           active.model !== null && active.model.version !== active.version
@@ -605,20 +611,29 @@ export const createControlRuntime = <
             (response.completion && input))
         ) {
           const nextReason = resume ? 'resume' : input ? 'input' : 'tool'
+          const sourceGenerationId = (
+            response.generation ?? response.source?.generation
+          )?.generationId
           const id =
             nextReason === 'tool'
-              ? `ai.generate:tools:${response.generation!.generationId}`
-              : `${ctx.event.id}:generate:${nextReason}`
+              ? await aiEventId(
+                  'a2.ai.request.tool',
+                  response.generation!.generationId,
+                )
+              : await aiEventId('a2.ai.request', ctx.event.id, nextReason)
           requestModel({
             request: {
               messageId: response.responseMessageId,
               responseMessageId: response.responseMessageId,
               reason: nextReason,
+              ...(sourceGenerationId === undefined
+                ? {}
+                : { sourceGenerationId }),
             },
             id,
           })
         } else if (naturalCompletion) {
-          emit({
+          await emit({
             type: 'ai.message.completed',
             payload: {
               messageId: response.responseMessageId,
@@ -648,10 +663,15 @@ export const createControlRuntime = <
       while (state.inbox.items.length > 0) {
         const item = state.inbox.items.shift()!
         delete item.afterStepOf
-        emit({
+        await emit({
           type: 'ai.message.created',
           payload: { message: item.message, generate: item.generate },
-          id: `ai.admitted:${item.id}`,
+          id: await aiEventId(
+            'a2.ai.admitted',
+            options.agent.contract.name,
+            ctx.event.sessionId,
+            item.id,
+          ),
         })
         if (!item.generate) continue
         state.active = {
@@ -662,9 +682,19 @@ export const createControlRuntime = <
           suspended: false,
           model: null,
         }
+        const requestId = await aiEventId(
+          'a2.ai.request.message',
+          options.agent.contract.name,
+          ctx.event.sessionId,
+          item.id,
+        )
         requestModel({
-          request: { messageId: item.id, reason: 'message' },
-          id: `ai.generate:message:${item.id}`,
+          request: {
+            messageId: item.id,
+            responseMessageId: requestId,
+            reason: 'message',
+          },
+          id: requestId,
         })
         break
       }
@@ -679,14 +709,19 @@ export const createControlRuntime = <
       ...(facts as AppendInput<D>[]),
       {
         type: 'ai.control.committed',
+        id: await aiEventId('a2.ai.control.committed', ctx.event.id),
         payload: controlCommit({ before, after: state, moves }),
       } as AppendInput<D>,
       ...jobs,
-      { type: 'ai.control.decided', payload: receipt } as AppendInput<D>,
+      {
+        type: 'ai.control.decided',
+        id: await aiEventId('a2.ai.control.decided', ctx.event.id),
+        payload: receipt,
+      } as AppendInput<D>,
     ]
   }
 
-  const report = ({
+  const report = async ({
     ctx,
     events,
     settled,
@@ -696,7 +731,7 @@ export const createControlRuntime = <
     events: readonly AppendInput<D>[]
     settled: boolean
     name: string
-  }): AppendInput<D> => {
+  }): Promise<AppendInput<D>> => {
     const payload = ctx.event.payload as
       GenerationRequestedPayload | { turnId: string; version: number }
     const { turnId, version } = 'turnId' in payload ? payload : payload.control!
@@ -716,7 +751,12 @@ export const createControlRuntime = <
     }
     return {
       type: 'ai.work.reported',
-      id: `${ctx.event.id}:report:${ctx.attempt}:${name}`,
+      id: await aiEventId(
+        'a2.ai.report',
+        ctx.event.id,
+        String(ctx.attempt),
+        name,
+      ),
       payload: {
         workId: ctx.event.id,
         attempt: ctx.attempt,
@@ -724,19 +764,27 @@ export const createControlRuntime = <
         version,
         kind: ctx.event.type === 'ai.generation.requested' ? 'model' : 'tool',
         settled,
-        events: events
-          .filter(
-            (event) =>
-              event.type !== 'ai.message.completed' &&
-              event.type !== 'ai.message.interrupted',
-          )
-          .map((event, index) => ({
-            type: event.type,
-            payload: event.payload,
-            id:
-              event.id ??
-              `${ctx.event.id}:fact:${ctx.attempt}:${name}:${index}`,
-          })),
+        events: await Promise.all(
+          events
+            .filter(
+              (event) =>
+                event.type !== 'ai.message.completed' &&
+                event.type !== 'ai.message.interrupted',
+            )
+            .map(async (event, index) => ({
+              type: event.type,
+              payload: event.payload,
+              id:
+                event.id ??
+                (await aiEventId(
+                  'a2.ai.work.fact',
+                  ctx.event.id,
+                  String(ctx.attempt),
+                  name,
+                  String(index),
+                )),
+            })),
+        ),
       },
     } as AppendInput<D>
   }
@@ -757,7 +805,7 @@ export const createControlRuntime = <
       (event) => event.type !== 'ai.generation.progress',
     )
     if (lifecycle.length === 0) return ctx.session.append(name, ...progress)
-    const input = report({ ctx, events: lifecycle, settled: false, name })
+    const input = await report({ ctx, events: lifecycle, settled: false, name })
     const persisted = await ctx.session.append(name, ...progress, input)
     const index = persisted.at(-1)!.index
     const feed = ctx.session.stream({ startAfter: index })
@@ -803,7 +851,7 @@ export const createControlRuntime = <
   }: {
     ctx: HandlerContext<D>
     events: void | AppendInput<D> | readonly AppendInput<D>[]
-  }): AppendInput<D> =>
+  }): Promise<AppendInput<D>> =>
     report({
       ctx,
       events:

package/src/ai-control-state.ts

@@ -58,6 +58,10 @@ export const controlChanges = ({
       !Array.isArray(left) &&
       !Array.isArray(right)
     ) {
+      if (Object.keys(right).length === 0) {
+        changes.push({ path, value: right })
+        return
+      }
       for (const key of Object.keys(left))
         if (!Object.hasOwn(right, key))
           changes.push({ path: [...path, key], remove: true })
@@ -146,6 +150,8 @@ export const controlCommit = <M extends UIMessage>(options: {
       paused: options.after.inbox.paused,
       closed: options.after.coordinator.closed,
       requestId: response?.activeRequestId ?? null,
+      requestReason: response?.requestReason ?? null,
+      sourceGenerationId: response?.sourceGenerationId ?? null,
       responseMessageId: response?.responseMessageId ?? null,
     },
   }