experimental-a2 0.13.0 → 0.14.1

package/src/ai-control-state.ts

@@ -0,0 +1,152 @@
+import type {
+  AIControlState,
+  ControlCommit,
+  ControlChange,
+} from './ai-control.ts'
+import type { UIMessage } from 'ai'
+
+type RecordValue = Record<string, unknown> | unknown[]
+const container = (value: unknown): value is RecordValue =>
+  typeof value === 'object' && value !== null
+const same = (left: unknown, right: unknown): boolean =>
+  JSON.stringify(left) === JSON.stringify(right)
+
+export const controlChanges = ({
+  before,
+  after,
+}: {
+  before: unknown
+  after: unknown
+}): ControlChange[] => {
+  const changes: ControlChange[] = []
+  const visit = (
+    left: unknown,
+    right: unknown,
+    path: (string | number)[],
+  ): void => {
+    if (same(left, right)) return
+    if (Array.isArray(left) && Array.isArray(right)) {
+      let prefix = 0
+      while (
+        prefix < Math.min(left.length, right.length) &&
+        same(left[prefix], right[prefix])
+      )
+        prefix += 1
+      let suffix = 0
+      while (
+        suffix < Math.min(left.length, right.length) - prefix &&
+        same(left[left.length - 1 - suffix], right[right.length - 1 - suffix])
+      )
+        suffix += 1
+      if (left.length === right.length) {
+        for (let index = prefix; index < left.length - suffix; index += 1)
+          visit(left[index], right[index], [...path, index])
+      } else
+        changes.push({
+          path,
+          splice: {
+            index: prefix,
+            deleteCount: left.length - prefix - suffix,
+            items: right.slice(prefix, right.length - suffix),
+          },
+        })
+      return
+    }
+    if (
+      container(left) &&
+      container(right) &&
+      !Array.isArray(left) &&
+      !Array.isArray(right)
+    ) {
+      for (const key of Object.keys(left))
+        if (!Object.hasOwn(right, key))
+          changes.push({ path: [...path, key], remove: true })
+      for (const key of Object.keys(right))
+        visit(Object.hasOwn(left, key) ? left[key] : undefined, right[key], [
+          ...path,
+          key,
+        ])
+      return
+    }
+    changes.push({ path, value: right })
+  }
+  visit(before, after, [])
+  return changes
+}
+
+export const applyControlChanges = <T>({
+  state,
+  changes,
+}: {
+  state: T
+  changes: readonly ControlChange[]
+}): T => {
+  if (changes.length === 0) return state
+  let next = structuredClone(state)
+  for (const change of changes) {
+    const path = change.path
+    if (path.length === 0) {
+      next = structuredClone(change.value) as T
+      continue
+    }
+    let target: unknown = next
+    for (const key of path.slice(0, -1)) {
+      if (!container(target) || !Object.hasOwn(target, key))
+        throw new TypeError('Invalid AI control change path')
+      target = (target as Record<string | number, unknown>)[key]
+    }
+    if (!container(target))
+      throw new TypeError('Invalid AI control change target')
+    const parent = target as Record<string | number, unknown>
+    const key = path.at(-1)!
+    if (change.remove) delete parent[key]
+    else if (change.splice) {
+      const items = parent[key] as unknown[]
+      Object.defineProperty(parent, key, {
+        value: [
+          ...items.slice(0, change.splice.index),
+          ...structuredClone(change.splice.items),
+          ...items.slice(change.splice.index + change.splice.deleteCount),
+        ],
+        enumerable: true,
+        configurable: true,
+        writable: true,
+      })
+    } else if (change.move) {
+      const items = parent[key] as unknown[]
+      const [item] = items.splice(change.move.from, 1)
+      items.splice(change.move.to, 0, item)
+    } else
+      Object.defineProperty(parent, key, {
+        value: structuredClone(change.value),
+        enumerable: true,
+        configurable: true,
+        writable: true,
+      })
+  }
+  return next
+}
+
+export const controlCommit = <M extends UIMessage>(options: {
+  before: AIControlState<M>
+  after: AIControlState<M>
+  moves: ControlChange[]
+}): ControlCommit => {
+  const active = options.after.active
+  const response = options.after.coordinator.response
+  return {
+    changes: [
+      ...options.moves,
+      ...controlChanges({ before: options.before, after: options.after }),
+    ],
+    view: {
+      turnId: active?.turnId ?? null,
+      version: active?.version ?? null,
+      phase: active?.phase ?? null,
+      paused: options.after.inbox.paused,
+      closed: options.after.coordinator.closed,
+      requestId: response?.activeRequestId ?? null,
+      responseMessageId: response?.responseMessageId ?? null,
+    },
+  }
+}

package/src/ai-control.ts

@@ -0,0 +1,139 @@
+import type { UIMessage } from 'ai'
+import type { AICoordinatorState } from './ai-coordinator.ts'
+import type {
+  ApprovalRespondedPayload,
+  GenerationStartedPayload,
+  InputRequestedPayload,
+  InputRespondedPayload,
+  MessageInterruptedPayload,
+  RetryRequestedPayload,
+  ToolCalledPayload,
+  ToolResultPayload,
+} from './ai.ts'
+
+export type InboxItem<M extends UIMessage = UIMessage> = {
+  id: string
+  revision: number
+  message: M
+  generate: boolean
+  afterStepOf?: string
+}
+
+export type ControlCommand<M extends UIMessage = UIMessage> =
+  | { action: 'send'; message: M; generate?: boolean }
+  | { action: 'edit'; message: M; expectedRevision: number }
+  | { action: 'move'; inputId: string; beforeId: string | null }
+  | { action: 'remove'; inputId: string }
+  | {
+      action: 'send-now'
+      inputId: string
+      turnId: string | null
+      lastSeenIndex?: number
+    }
+  | { action: 'stop'; turnId: string; lastSeenIndex?: number }
+  | { action: 'steer'; turnId: string; message: M; lastSeenIndex?: number }
+  | { action: 'pause'; when: 'now' | 'after-turn' }
+  | { action: 'resume' }
+  | { action: 'tool-result'; result: ToolResultPayload }
+  | { action: 'approval'; response: ApprovalRespondedPayload }
+  | { action: 'input'; response: InputRespondedPayload }
+  | { action: 'request-input'; request: InputRequestedPayload }
+  | { action: 'retry'; request: RetryRequestedPayload }
+  | { action: 'interrupt'; request: MessageInterruptedPayload }
+
+export type ControlReceipt = {
+  commandId: string
+  work?: true
+  outcome: 'applied' | 'rejected'
+  reason?:
+    | 'already-active'
+    | 'revision-conflict'
+    | 'stale-turn'
+    | 'not-found'
+    | 'duplicate-input'
+    | 'closed'
+    | 'stale-work'
+    | 'not-user-input'
+}
+
+export type ActiveTurn<M extends UIMessage = UIMessage> = {
+  turnId: string
+  input: InboxItem<M>
+  phase:
+    | 'starting'
+    | 'generating'
+    | 'tools'
+    | 'waiting'
+    | 'failed'
+    | 'pausing'
+    | 'paused'
+}
+
+export type ControlWork = {
+  id: string
+  version: number
+  settled: boolean
+  attempt?: number
+  facts?: string[]
+}
+
+export type AIControlState<M extends UIMessage = UIMessage> = {
+  reservedIds: Record<string, true | 'removed'>
+  inbox: { paused: boolean; items: InboxItem<M>[] }
+  active:
+    | (ActiveTurn<M> & {
+        version: number
+        suspended: boolean
+        model: ControlWork | null
+      })
+    | null
+  coordinator: AICoordinatorState
+}
+
+export type WorkFact = { type: string; id: string; payload: unknown }
+export type WorkReport = {
+  attempt: number
+  workId: string
+  turnId: string
+  version: number
+  kind: 'model' | 'tool'
+  settled: boolean
+  events: WorkFact[]
+}
+
+export type ToolExecutionRequest = {
+  turnId: string
+  version: number
+  call: ToolCalledPayload
+  generation: GenerationStartedPayload
+}
+
+export const initialControlState = <
+  M extends UIMessage,
+>(): AIControlState<M> => ({
+  reservedIds: {},
+  inbox: { paused: false, items: [] },
+  active: null,
+  coordinator: { closed: false, queued: [] },
+})
+
+export type ControlChange = {
+  path: (string | number)[]
+  value?: unknown
+  remove?: true
+  splice?: { index: number; deleteCount: number; items: unknown[] }
+  move?: { from: number; to: number }
+}
+
+export type ControlCommit = {
+  changes: ControlChange[]
+  view: {
+    turnId: string | null
+    version: number | null
+    phase: ActiveTurn['phase'] | null
+    paused: boolean
+    closed: boolean
+    requestId: string | null
+    responseMessageId: string | null
+  }
+}

package/src/ai-coordinator.ts

@@ -1,4 +1,6 @@
 import type { UIMessage } from 'ai'
+import { applyControlChanges } from './ai-control-state.ts'
+import type { ControlCommit, ControlWork } from './ai-control.ts'
 import type {
   AIEventDefs,
   ApprovalRequestedPayload,
@@ -17,26 +19,32 @@ 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
+  work?: ControlWork
 }
 
 export type CoordinatorResponse = {
   rootMessageId: string
   responseMessageId: string
   status: 'requested' | 'generating' | 'waiting' | 'failed'
+  stepCount: number
   activeRequestId?: string
   generation?: GenerationStartedPayload
+  promptThroughIndex?: number
+  source?: {
+    generation: GenerationStartedPayload
+    promptThroughIndex: number
+    failed: boolean
+    canCompact: boolean
+  }
   completion?: GenerationCompletedPayload
   failure?: GenerationFailedPayload
   calls: CoordinatorCall[]
   inputs: InputRequestedPayload[]
   inputResponse?: {
-    index: number
     generationId: string
     inputId: string
     name: string
@@ -45,8 +53,9 @@ export type CoordinatorResponse = {
 
 export type AICoordinatorState = {
   closed: boolean
-  queued: Array<{ index: number; messageId: string; generate: boolean }>
+  queued: Array<{ messageId: string; generate: boolean }>
   response?: CoordinatorResponse
+  calibration?: { model: string; inputTokens: number; estimate: number }
 }
 
 const updateResponse = (
@@ -57,13 +66,18 @@ const updateResponse = (
     ? state
     : { ...state, response: update(state.response) }
 
-const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
+export 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.retry.requested':
+    case 'ai.compaction.completed': {
+      const { calibration: _, ...rest } = state
+      return rest
+    }
     case 'ai.message.created': {
       const payload = event.payload as MessageCreatedPayload<M>
       const message = payload.message
@@ -73,11 +87,10 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
         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': {
@@ -85,21 +98,22 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       if (request.reason === 'message') {
         const responseMessageId =
           request.responseMessageId ?? `${request.messageId}:assistant`
-        const requested = state.queued.find(
+        const requested = state.queued.findIndex(
           (item) => item.messageId === request.messageId,
         )
         return {
           ...state,
           queued:
-            requested === undefined
+            requested < 0
               ? state.queued.filter(
                   (item) => item.messageId !== request.messageId,
                 )
-              : state.queued.filter((item) => item.index > requested.index),
+              : state.queued.slice(requested + 1),
           response: {
             rootMessageId: request.messageId,
             responseMessageId,
             status: 'requested',
+            stepCount: 0,
             activeRequestId: event.id,
             calls: [],
             inputs: [],
@@ -123,14 +137,31 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       ) {
         return state
       }
-      return updateResponse(state, (response) => ({
-        rootMessageId: response.rootMessageId,
-        responseMessageId: response.responseMessageId,
-        status: 'requested',
-        activeRequestId: event.id,
-        calls: [],
-        inputs: [],
-      }))
+      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
+            ? {}
+            : {
+                source: {
+                  generation,
+                  promptThroughIndex,
+                  failed:
+                    response.failure?.generationId === generation.generationId,
+                  canCompact: response.calls.every((call) => call.terminal),
+                },
+              }),
+          calls: [],
+          inputs: [],
+        }
+      })
     }
     case 'ai.generation.started': {
       const generation = event.payload as GenerationStartedPayload
@@ -146,7 +177,12 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
         rootMessageId: response.rootMessageId,
         responseMessageId: response.responseMessageId,
         status: 'generating',
+        stepCount: response.stepCount,
         generation,
+        ...(generation.promptThroughIndex === undefined
+          ? {}
+          : { promptThroughIndex: generation.promptThroughIndex }),
+        ...(response.source === undefined ? {} : { source: response.source }),
         activeRequestId: generation.requestId,
         calls: [],
         inputs: [],
@@ -159,11 +195,29 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       ) {
         return state
       }
-      return updateResponse(state, (response) => ({
+      const next = updateResponse(state, (response) => ({
         ...response,
         status: 'waiting',
         completion,
+        stepCount:
+          response.stepCount + (response.completion === undefined ? 1 : 0),
       }))
+      const { calibration: _, ...uncalibrated } = next
+      const inputTokens = completion.usage?.inputTokens
+      const estimate = completion.inputTokenEstimate
+      return inputTokens !== undefined &&
+        Number.isFinite(inputTokens) &&
+        inputTokens >= 0 &&
+        estimate !== undefined
+        ? {
+            ...uncalibrated,
+            calibration: {
+              model: state.response.generation.model,
+              inputTokens,
+              estimate,
+            },
+          }
+        : uncalibrated
     }
     case 'ai.generation.failed': {
       const failure = event.payload as GenerationFailedPayload
@@ -186,18 +240,21 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
     }
     case 'ai.message.completed': {
       const messageId = (event.payload as { messageId: string }).messageId
-      return state.response?.responseMessageId === messageId
-        ? { closed: state.closed, queued: state.queued }
-        : state
+      if (state.response?.responseMessageId !== messageId) return state
+      const { response: _, ...rest } = state
+      return rest
     }
     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
+      if (
+        response?.responseMessageId !== interruption.messageId ||
+        (interruption.generationId !== response.generation?.generationId &&
+          interruption.requestId !== response.activeRequestId)
+      )
+        return state
+      const { response: _, ...rest } = state
+      return rest
     }
     case 'ai.tool.called': {
       const call = event.payload as ToolCalledPayload
@@ -210,7 +267,7 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
           ...response.calls.filter(
             (candidate) => candidate.call.toolCallId !== call.toolCallId,
           ),
-          { index: event.index, call, terminal: false },
+          { call, terminal: false },
         ],
       }))
     }
@@ -240,7 +297,7 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
         ...response,
         calls: response.calls.map((candidate) =>
           candidate.approval?.approvalId === approval.approvalId
-            ? { ...candidate, response: approval, responseIndex: event.index }
+            ? { ...candidate, response: approval }
             : candidate,
         ),
       }))
@@ -310,7 +367,6 @@ const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
           (candidate) => candidate.inputId !== input.inputId,
         ),
         inputResponse: {
-          index: event.index,
           generationId: input.generationId,
           inputId: input.inputId,
           name: requested.name,
@@ -330,13 +386,24 @@ export const aiCoordinatorReducer = <
 ): Reducer<D, AICoordinatorState> =>
   contract
     .reducer({
-      name: 'a2.ai.coordinator.v3',
+      name: 'a2.ai.coordinator.v5',
       initialState: {
         closed: false,
         queued: [],
       } as AICoordinatorState,
     })
-    .fold((state, event) => foldCoordinator<M, D>(state, event))
+    .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-progress-batches.ts

@@ -0,0 +1,68 @@
+import type { UIMessageChunk } from 'ai'
+
+export type ProgressBatch = { index: number; chunks: UIMessageChunk[] }
+type BatchBlock = ProgressBatch[] | { left: BatchBlock; right: BatchBlock }
+
+export type ProgressBatches = {
+  length: number
+  tail: ProgressBatch[]
+  blocks: Array<BatchBlock | null>
+}
+
+const PAGE_SIZE = 64
+
+export const appendProgressBatch = ({
+  batches,
+  batch,
+}: {
+  batches: ProgressBatches
+  batch: ProgressBatch
+}): ProgressBatches => {
+  if (batches.tail.length < PAGE_SIZE)
+    return {
+      ...batches,
+      length: batches.length + 1,
+      tail: [...batches.tail, batch],
+    }
+
+  const blocks = [...batches.blocks]
+  let block: BatchBlock = batches.tail
+  let level = 0
+  for (;;) {
+    const left = blocks[level]
+    if (left == null) break
+    block = { left, right: block }
+    blocks[level] = null
+    level += 1
+  }
+  blocks[level] = block
+  return { length: batches.length + 1, tail: [batch], blocks }
+}
+
+export const progressBatches = (
+  items: readonly ProgressBatch[] = [],
+): ProgressBatches => {
+  let batches: ProgressBatches = { length: 0, tail: [], blocks: [] }
+  for (const batch of items) batches = appendProgressBatch({ batches, batch })
+  return batches
+}
+
+export const flattenProgressBatches = (
+  batches: ProgressBatches,
+): ProgressBatch[] => {
+  const items: ProgressBatch[] = []
+  const visit = (block: BatchBlock): void => {
+    if (Array.isArray(block)) items.push(...block)
+    else {
+      visit(block.left)
+      visit(block.right)
+    }
+  }
+  // Level n holds 2^n completed pages; higher levels contain older batches.
+  for (let level = batches.blocks.length - 1; level >= 0; level -= 1) {
+    const block = batches.blocks[level]
+    if (block) visit(block)
+  }
+  items.push(...batches.tail)
+  return items
+}