experimental-a2 0.16.1 → 0.18.0

package/src/ai-control-server.ts

@@ -1,10 +1,10 @@
+import { reducer as makeReducer } from './reducer.ts'
 // oxlint-disable no-await-in-loop -- work acknowledgements follow the durable stream.
 import { applyControlChanges, controlCommit } from './ai-control-state.ts'
 import type { ControlChange, ControlCommit } from './ai-control.ts'
 import type { UIMessage } from 'ai'
 import type {
   AIEventDefs,
-  AgentDefinition,
   GenerationRequestedPayload,
   MessageInterruptedPayload,
 } from './ai.ts'
@@ -21,7 +21,12 @@ import {
   type WorkFact,
   type WorkReport,
 } from './ai-control.ts'
-import type { AppendInput, ContractEvent, EventDefs } from './contract.ts'
+import type {
+  AppendInput,
+  Contract,
+  ContractEvent,
+  EventDefs,
+} from './contract.ts'
 import type { HandlerContext } from './server.ts'
 import type { Reducer } from './reducer.ts'
 import type { AICoordinatorState, CoordinatorCall } from './ai-coordinator.ts'
@@ -29,6 +34,7 @@ import { validateSync } from './validate.ts'
 import type { StandardSchemaV1 } from './standard-schema.ts'
 import { A2Error } from './errors.ts'
 import { aiEventId } from './ai-id.ts'
+import { CONSOLE_TELEMETRY, type A2Telemetry } from './telemetry.ts'
 
 export class ControlCancelled extends Error {}
 
@@ -69,41 +75,42 @@ export const createControlRuntime = <
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
 >(options: {
-  agent: AgentDefinition<M, D>
+  contract: Contract<D, Record<never, never>, unknown>
+  telemetry?: A2Telemetry
 }): AIControlRuntime<M, D> => {
-  const reducer = options.agent.contract
-    .reducer({
-      name: 'a2.ai.control.v5',
-      initialState: initialControlState<M>(),
-    })
-    .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 }
+  const telemetry = options.telemetry ?? CONSOLE_TELEMETRY
+  const reducer = makeReducer({
+    contract: options.contract,
+    name: 'a2.ai.control.v5',
+    initialState: initialControlState<M>(),
+  }).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
-    })
-  const coordinator = aiCoordinatorReducer(options.agent.contract)
+    }
+    return event.type === 'ai.control.committed'
+      ? applyControlChanges({
+          state,
+          changes: (event.payload as ControlCommit).changes,
+        })
+      : state
+  })
+  const coordinator = aiCoordinatorReducer(options.contract)
 
   const handler = async (
     ctx: HandlerContext<
@@ -438,7 +445,7 @@ export const createControlRuntime = <
                 },
                 id: await aiEventId(
                   'a2.ai.request.retry',
-                  options.agent.contract.name,
+                  options.contract.name,
                   ctx.event.sessionId,
                   command.request.retryId,
                 ),
@@ -668,7 +675,7 @@ export const createControlRuntime = <
           payload: { message: item.message, generate: item.generate },
           id: await aiEventId(
             'a2.ai.admitted',
-            options.agent.contract.name,
+            options.contract.name,
             ctx.event.sessionId,
             item.id,
           ),
@@ -684,7 +691,7 @@ export const createControlRuntime = <
         }
         const requestId = await aiEventId(
           'a2.ai.request.message',
-          options.agent.contract.name,
+          options.contract.name,
           ctx.event.sessionId,
           item.id,
         )
@@ -736,7 +743,7 @@ export const createControlRuntime = <
       GenerationRequestedPayload | { turnId: string; version: number }
     const { turnId, version } = 'turnId' in payload ? payload : payload.control!
     for (const event of events) {
-      const schema = options.agent.contract.events[
+      const schema = options.contract.events[
         event.type
       ] as StandardSchemaV1<unknown>
       const result = validateSync(
@@ -808,41 +815,54 @@ export const createControlRuntime = <
     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 })
-    const stream = feed[Symbol.asyncIterator]()
-    const canonical: ContractEvent<D>[] = []
-    const wanted = new Set(lifecycle.map((event) => event.id))
-    let rejectAbort!: (error: unknown) => void
-    const aborted = new Promise<never>((_, reject) => {
-      rejectAbort = reject
-    })
-    const onAbort = (): void => {
-      rejectAbort(ctx.signal.reason)
-    }
-    ctx.signal.addEventListener('abort', onAbort, { once: true })
-    try {
-      ctx.signal.throwIfAborted()
-      for (;;) {
-        const item = await Promise.race([stream.next(), aborted])
-        if (item.done)
-          throw new Error('AI control acknowledgement stream ended')
-        if (wanted.has(item.value.id)) canonical.push(item.value)
-        if (
-          item.value.type === 'ai.control.decided' &&
-          (item.value.payload as ControlReceipt).commandId === input.id
-        ) {
-          if ((item.value.payload as ControlReceipt).outcome === 'rejected')
-            throw new ControlCancelled('stale AI work')
-          return [
-            ...persisted.filter((event) => event.type !== 'ai.work.reported'),
-            ...canonical,
-          ]
+    return telemetry.span(
+      'a2.ai.ack',
+      {
+        'a2.contract': options.contract.name,
+        'a2.session_id': ctx.event.sessionId,
+        'a2.ai.ack.request_id': ctx.event.id,
+        'a2.ai.ack.report_id': input.id!,
+      },
+      async () => {
+        const feed = ctx.session.stream({ startAfter: index })
+        const stream = feed[Symbol.asyncIterator]()
+        const canonical: ContractEvent<D>[] = []
+        const wanted = new Set(lifecycle.map((event) => event.id))
+        let rejectAbort!: (error: unknown) => void
+        const aborted = new Promise<never>((_, reject) => {
+          rejectAbort = reject
+        })
+        const onAbort = (): void => {
+          rejectAbort(ctx.signal.reason)
         }
-      }
-    } finally {
-      ctx.signal.removeEventListener('abort', onAbort)
-      await stream.return?.()
-    }
+        ctx.signal.addEventListener('abort', onAbort, { once: true })
+        try {
+          ctx.signal.throwIfAborted()
+          for (;;) {
+            const item = await Promise.race([stream.next(), aborted])
+            if (item.done)
+              throw new Error('AI control acknowledgement stream ended')
+            if (wanted.has(item.value.id)) canonical.push(item.value)
+            if (
+              item.value.type === 'ai.control.decided' &&
+              (item.value.payload as ControlReceipt).commandId === input.id
+            ) {
+              if ((item.value.payload as ControlReceipt).outcome === 'rejected')
+                throw new ControlCancelled('stale AI work')
+              return [
+                ...persisted.filter(
+                  (event) => event.type !== 'ai.work.reported',
+                ),
+                ...canonical,
+              ]
+            }
+          }
+        } finally {
+          ctx.signal.removeEventListener('abort', onAbort)
+          await stream.return?.()
+        }
+      },
+    )
   }
 
   const settled = ({

package/src/ai-coordinator.ts

@@ -1,3 +1,4 @@
+import { reducer } from './reducer.ts'
 import type { UIMessage } from 'ai'
 import { applyControlChanges } from './ai-control-state.ts'
 import type { ControlCommit, ControlWork } from './ai-control.ts'
@@ -255,7 +256,8 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       return inputTokens !== undefined &&
         Number.isFinite(inputTokens) &&
         inputTokens >= 0 &&
-        estimate !== undefined
+        estimate !== undefined &&
+        state.response.generation.model !== undefined
         ? {
             ...uncalibrated,
             calibration: {
@@ -436,28 +438,27 @@ export const aiCoordinatorReducer = <
   M extends UIMessage,
   D extends AIEventDefs<M> & EventDefs,
 >(
-  contract: Contract<D>,
+  contract: Contract<D, Record<never, never>, unknown>,
 ): Reducer<D, AICoordinatorState> =>
-  contract
-    .reducer({
-      name: 'a2.ai.coordinator.v6',
-      initialState: {
-        closed: false,
-        queued: [],
-      } as AICoordinatorState,
-    })
-    .fold((state, event) =>
-      event.type === 'ai.control.committed'
-        ? applyControlChanges({
-            state,
-            changes: (event.payload as ControlCommit).changes
-              .filter((change) => change.path[0] === 'coordinator')
-              .map((change) =>
-                Object.assign({}, change, { path: change.path.slice(1) }),
-              ),
-          })
-        : state,
-    )
+  reducer({
+    contract: contract,
+    name: 'a2.ai.coordinator.v6',
+    initialState: {
+      closed: false,
+      queued: [],
+    } as AICoordinatorState,
+  }).fold((state, event) =>
+    event.type === 'ai.control.committed'
+      ? applyControlChanges({
+          state,
+          changes: (event.payload as ControlCommit).changes
+            .filter((change) => change.path[0] === 'coordinator')
+            .map((change) =>
+              Object.assign({}, change, { path: change.path.slice(1) }),
+            ),
+        })
+      : state,
+  )
 
 export const continuationReady = (state: AICoordinatorState): boolean => {
   const response = state.response

package/src/ai-fx.browser.ts

@@ -0,0 +1,5 @@
+/** experimental-a2/ai/fx is server-only; browsers reach an agent through its route. */
+throw new Error(
+  'experimental-a2/ai/fx is server-only — a client bundle imported it. Keep agent({ harness: fx() }) in server code.',
+)
+export {}