experimental-a2 0.16.0 → 0.17.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,22 +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.v4',
-      initialState: initialControlState<M>(),
-    })
-    .fold((state, event) =>
-      event.type === 'ai.control.committed'
-        ? applyControlChanges({
-            state,
-            changes: (event.payload as ControlCommit).changes,
-          })
-        : state,
-    )
-  const coordinator = aiCoordinatorReducer(options.agent.contract)
+  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.contract)
 
   const handler = async (
     ctx: HandlerContext<
@@ -419,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,
                 ),
@@ -649,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,
           ),
@@ -665,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,
         )
@@ -717,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(
@@ -789,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-control.ts

@@ -77,6 +77,7 @@ export type ControlWork = {
 }
 
 export type AIControlState<M extends UIMessage = UIMessage> = {
+  pendingStops: Record<string, string>
   inbox: { paused: boolean; items: InboxItem<M>[] }
   active:
     | (ActiveTurn<M> & {
@@ -109,6 +110,7 @@ export type ToolExecutionRequest = {
 export const initialControlState = <
   M extends UIMessage,
 >(): AIControlState<M> => ({
+  pendingStops: {},
   inbox: { paused: false, items: [] },
   active: null,
   coordinator: { closed: false, queued: [] },

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-message-base.ts

@@ -0,0 +1,209 @@
+import type { UIMessage } from 'ai'
+import { sameValue } from './equal-json.ts'
+
+type FieldUndo = { value?: unknown } | { remove: true } | { length: number }
+type FieldsUndo = Record<string, FieldUndo>
+type PartUndo = { value: UIMessage['parts'][number] } | { fields: FieldsUndo }
+
+export type MessageBase = {
+  partCount: number
+  fields?: FieldsUndo
+  parts?: Record<string, PartUndo>
+}
+
+const record = (value: unknown): value is Record<string, unknown> =>
+  typeof value === 'object' && value !== null && !Array.isArray(value)
+
+const assertFields = (fields: unknown): void => {
+  if (fields === undefined) return
+  if (!record(fields)) throw new TypeError('invalid AI message base fields')
+  for (const [key, undo] of Object.entries(fields)) {
+    if (
+      !record(undo) ||
+      Object.keys(undo).length > 1 ||
+      !(
+        Object.keys(undo).length === 0 ||
+        'value' in undo ||
+        undo['remove'] === true ||
+        (key === 'text' &&
+          Number.isSafeInteger(undo['length']) &&
+          (undo['length'] as number) >= 0)
+      )
+    )
+      throw new TypeError('invalid AI message base field')
+  }
+}
+
+/** Checks decoded recovery data; `restoreMessageBase` trusts this shape. */
+export function assertMessageBase(
+  value: unknown,
+): asserts value is MessageBase {
+  if (
+    !record(value) ||
+    !Number.isSafeInteger(value['partCount']) ||
+    (value['partCount'] as number) < 0 ||
+    (value['partCount'] as number) > 0xffffffff - 1 ||
+    Object.keys(value).some(
+      (key) => !['partCount', 'fields', 'parts'].includes(key),
+    ) ||
+    (value['parts'] !== undefined && !record(value['parts']))
+  )
+    throw new TypeError('invalid AI message base')
+  assertFields(value['fields'])
+  const partCount = value['partCount'] as number
+  for (const [key, undo] of Object.entries(value['parts'] ?? {})) {
+    const index = Number(key)
+    if (
+      !Number.isSafeInteger(index) ||
+      String(index) !== key ||
+      index < 0 ||
+      index >= partCount ||
+      !record(undo) ||
+      Object.keys(undo).length !== 1 ||
+      !(('value' in undo && record(undo['value'])) || 'fields' in undo)
+    )
+      throw new TypeError('invalid AI message base part')
+    if ('fields' in undo) assertFields(undo['fields'])
+  }
+}
+
+const restoreFields = (
+  current: Record<string, unknown>,
+  fields: FieldsUndo | undefined,
+): Record<string, unknown> | undefined => {
+  if (fields === undefined) return current
+  const restored = { ...current }
+  for (const [key, undo] of Object.entries(fields)) {
+    if ('remove' in undo) delete restored[key]
+    else if ('length' in undo) {
+      const text = current[key]
+      if (typeof text !== 'string' || undo.length > text.length)
+        return undefined
+      restored[key] = text.slice(0, undo.length)
+    } else
+      Object.defineProperty(restored, key, {
+        value: undo.value,
+        enumerable: true,
+        writable: true,
+        configurable: true,
+      })
+  }
+  return restored
+}
+
+/** `undefined` when `base` no longer describes `current`, including a removed message. */
+export const restoreMessageBase = <M extends UIMessage>(options: {
+  current: M | undefined
+  base: MessageBase | undefined
+}): M | undefined => {
+  const { current, base } = options
+  if (base === undefined || current === undefined) return undefined
+  const parts = current.parts.slice(0, base.partCount)
+  for (const [key, undo] of Object.entries(base.parts ?? {})) {
+    const index = Number(key)
+    if ('value' in undo) {
+      parts[index] = undo.value as M['parts'][number]
+      continue
+    }
+    const part = parts[index]
+    const restored =
+      part === undefined
+        ? undefined
+        : restoreFields(part as unknown as Record<string, unknown>, undo.fields)
+    if (!restored) return undefined
+    parts[index] = restored as M['parts'][number]
+  }
+  if (
+    parts.length !== base.partCount ||
+    Object.keys(parts).length !== parts.length
+  )
+    return undefined
+  const { parts: _parts, ...fields } = current
+  const restored = restoreFields(fields, base.fields)
+  return restored && ({ ...restored, parts } as M)
+}
+
+const updateFields = (options: {
+  before: Record<string, unknown>
+  after: Record<string, unknown>
+  undo?: FieldsUndo
+}): FieldsUndo | undefined => {
+  const { before, after, undo } = options
+  const original = restoreFields(before, undo)
+  if (!original) return undefined
+  return Object.fromEntries(
+    [...new Set([...Object.keys(original), ...Object.keys(after)])].flatMap(
+      (key) => {
+        if (
+          Object.hasOwn(original, key) === Object.hasOwn(after, key) &&
+          sameValue(original[key], after[key])
+        )
+          return []
+        const value = original[key]
+        const next = after[key]
+        const change: FieldUndo = !Object.hasOwn(original, key)
+          ? { remove: true }
+          : key === 'text' &&
+              typeof value === 'string' &&
+              typeof next === 'string' &&
+              next.startsWith(value)
+            ? { length: value.length }
+            : value === undefined
+              ? {}
+              : { value }
+        return [[key, change]]
+      },
+    ),
+  )
+}
+
+/** `undefined` when `before` no longer matches `base`, so the recovery must be dropped. */
+export const updateMessageBase = (options: {
+  base: MessageBase
+  before: UIMessage
+  after: UIMessage
+}): MessageBase | undefined => {
+  const { base, before, after } = options
+  if (before === after) return base
+  const { parts: beforeParts, ...beforeFields } = before
+  const { parts: afterParts, ...afterFields } = after
+  const fields = updateFields({
+    before: beforeFields,
+    after: afterFields,
+    ...(base.fields && { undo: base.fields }),
+  })
+  if (!fields) return undefined
+  const parts: Array<[string, PartUndo]> = []
+  for (let index = 0; index < base.partCount; index += 1) {
+    const previous = beforeParts[index]
+    const current = afterParts[index]
+    const undo = base.parts?.[index]
+    const original = undo && 'value' in undo ? undo.value : previous
+    if (!original) return undefined
+    if (previous === current && undo === undefined) continue
+    const originalFields = original as unknown as Record<string, unknown>
+    const undoFields = undo && 'fields' in undo ? undo.fields : undefined
+    if (!current) {
+      const value = restoreFields(originalFields, undoFields)
+      if (!value) return undefined
+      parts.push([
+        String(index),
+        { value: value as UIMessage['parts'][number] },
+      ])
+      continue
+    }
+    const changed = updateFields({
+      before: originalFields,
+      after: current as unknown as Record<string, unknown>,
+      ...(undoFields && { undo: undoFields }),
+    })
+    if (!changed) return undefined
+    if (Object.keys(changed).length > 0)
+      parts.push([String(index), { fields: changed }])
+  }
+  return {
+    partCount: base.partCount,
+    ...(Object.keys(fields).length > 0 && { fields }),
+    ...(parts.length > 0 && { parts: Object.fromEntries(parts) }),
+  }
+}