experimental-a2 0.18.0 → 0.19.0

package/src/ai-context.ts

@@ -19,7 +19,12 @@ import type {
   ToolCalledPayload,
   ToolResultPayload,
 } from './ai.ts'
-import type { ControlCommit } from './ai-control.ts'
+import {
+  type ControlCommand,
+  type ControlCommit,
+  type ControlReceipt,
+  unexplainedAdmission,
+} from './ai-control.ts'
 import { contextSchema } from './ai-context-schema.ts'
 import type { Contract, ContractEvent, EventDefs } from './contract.ts'
 import {
@@ -49,8 +54,22 @@ export type AIModelContext<M extends UIMessage = UIMessage> = {
   modelMetadata: Record<string, ModelMetadataState>
 }
 
+/**
+ * Inputs' content between their command and their admission: each input's
+ * content is stored only in the command that brought it in. Lists, not
+ * records: clients choose the ids, and a record would read an id such as
+ * `constructor` or `__proto__` from its prototype.
+ */
+export type AIContextInputs<M extends UIMessage> = {
+  /** Undecided input commands; a remove holds no message. */
+  commands: { commandId: string; inputId: string; message?: M }[]
+  /** Queued inputs' current content. */
+  waiting: { inputId: string; message: M }[]
+}
+
 export type AIContextState<M extends UIMessage> = AIModelContext<M> & {
   closed: boolean
+  inputs: AIContextInputs<M>
   authority: Pick<
     ControlCommit['view'],
     | 'requestId'
@@ -211,6 +230,7 @@ export const createContextReducer = <M extends UIMessage, D extends EventDefs>({
     compaction: null,
     modelMetadata: {},
     closed: false,
+    inputs: { commands: [], waiting: [] },
     authority: null,
     generation: null,
     activeProjection: null,
@@ -222,7 +242,7 @@ export const createContextReducer = <M extends UIMessage, D extends EventDefs>({
     // The application state rides in this checkpoint so every generation
     // reads it without another state read.
     name:
-      app === undefined ? 'a2.ai.context.v6' : `${app.name}+a2.ai.context.v6`,
+      app === undefined ? 'a2.ai.context.v7' : `${app.name}+a2.ai.context.v7`,
     initialState,
     stateSchema: contextSchema<M>(app),
   }).fold((state, event) => {
@@ -265,19 +285,79 @@ const reduceContext = <M extends UIMessage, D extends EventDefs>(
       }
       return view.closed ? endGeneration(next) : next
     }
+    case 'ai.control.requested': {
+      const command = event.payload as ControlCommand<M>
+      const held =
+        command.action === 'send' ||
+        command.action === 'steer' ||
+        command.action === 'edit'
+          ? {
+              commandId: event.id,
+              inputId: command.message.id,
+              message: command.message,
+            }
+          : command.action === 'remove'
+            ? { commandId: event.id, inputId: command.inputId }
+            : undefined
+      if (held === undefined) return state
+      return {
+        ...state,
+        inputs: {
+          ...state.inputs,
+          commands: [...state.inputs.commands, held],
+        },
+      }
+    }
+    case 'ai.control.decided': {
+      const receipt = event.payload as ControlReceipt
+      const held = receipt.work
+        ? undefined
+        : state.inputs.commands.find(
+            (command) => command.commandId === receipt.commandId,
+          )
+      if (held === undefined) return state
+      const commands = state.inputs.commands.filter(
+        (command) => command !== held,
+      )
+      if (receipt.outcome !== 'applied')
+        return { ...state, inputs: { ...state.inputs, commands } }
+      const waiting = state.inputs.waiting.filter(
+        (input) => input.inputId !== held.inputId,
+      )
+      if (held.message !== undefined)
+        waiting.push({ inputId: held.inputId, message: held.message })
+      return { ...state, inputs: { commands, waiting } }
+    }
     case 'ai.session.closed':
       return endGeneration({
         ...state,
         closed: true,
+        inputs: { commands: [], waiting: [] },
         compaction:
           state.compaction?.status === 'running'
             ? (state.activeProjection?.baseCompaction ?? null)
             : state.compaction,
       })
     case 'ai.message.created': {
-      const { message } = event.payload as MessageCreatedPayload<M>
+      const { messageId } = event.payload as MessageCreatedPayload
+      // A send or steer admitted with its own receipt is still undecided
+      // here; its command id is the input's.
+      const message =
+        state.inputs.waiting.find((input) => input.inputId === messageId)
+          ?.message ??
+        state.inputs.commands.find((command) => command.commandId === messageId)
+          ?.message
+      if (message === undefined) throw unexplainedAdmission(messageId)
       return {
         ...state,
+        inputs: {
+          commands: state.inputs.commands.filter(
+            (command) => command.commandId !== messageId,
+          ),
+          waiting: state.inputs.waiting.filter(
+            (input) => input.inputId !== messageId,
+          ),
+        },
         messages: upsertMessage(state.messages, message),
         activeProjection:
           state.activeProjection?.responseMessageId === message.id

package/src/ai-control-server.ts

@@ -6,6 +6,7 @@ import type { UIMessage } from 'ai'
 import type {
   AIEventDefs,
   GenerationRequestedPayload,
+  MessageCreatedPayload,
   MessageInterruptedPayload,
 } from './ai.ts'
 import {
@@ -15,6 +16,9 @@ import {
 } from './ai-coordinator.ts'
 import {
   initialControlState,
+  STEERED_REASON,
+  steeringEnd,
+  steersOn,
   type AIControlState,
   type ControlCommand,
   type ControlReceipt,
@@ -67,7 +71,7 @@ export type AIControlRuntime<
     events: void | AppendInput<D> | readonly AppendInput<D>[]
   }): Promise<AppendInput<D>>
   cancelled(options: { error: unknown; signal: AbortSignal }): boolean
-  reducer: Reducer<D, AIControlState<M>>
+  reducer: Reducer<D, AIControlState>
   coordinator: Reducer<D, AICoordinatorState>
 }
 
@@ -81,8 +85,8 @@ export const createControlRuntime = <
   const telemetry = options.telemetry ?? CONSOLE_TELEMETRY
   const reducer = makeReducer({
     contract: options.contract,
-    name: 'a2.ai.control.v5',
-    initialState: initialControlState<M>(),
+    name: 'a2.ai.control.v6',
+    initialState: initialControlState(),
   }).fold((state, event) => {
     if (event.type === 'ai.control.requested') {
       const command = event.payload as ControlCommand<M>
@@ -141,7 +145,7 @@ export const createControlRuntime = <
           (await aiEventId('a2.ai.fact', ctx.event.id, String(facts.length))),
       }
       facts.push(fact)
-      state.coordinator = foldCoordinator<M, D>(state.coordinator, {
+      state.coordinator = foldCoordinator<D>(state.coordinator, {
         ...ctx.event,
         ...fact,
       } as ContractEvent<D>)
@@ -173,25 +177,52 @@ export const createControlRuntime = <
       })
       state.active = null
     }
+    const decide = async (): Promise<AppendInput<D>[]> => {
+      const receipt: ControlReceipt = {
+        commandId: ctx.event.id,
+        ...(ctx.event.type === 'ai.work.reported'
+          ? { work: true as const }
+          : {}),
+        outcome: reason === undefined ? 'applied' : 'rejected',
+        ...(reason === undefined ? {} : { reason }),
+      }
+      return [
+        ...(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',
+          id: await aiEventId('a2.ai.control.decided', ctx.event.id),
+          payload: receipt,
+        } as AppendInput<D>,
+      ]
+    }
     const insert = ({
       message,
       generate = true,
-      first = false,
       afterStepOf,
     }: {
       message: M
       generate?: boolean
-      first?: boolean
       afterStepOf?: string
     }): void => {
       const candidate = {
         id: message.id,
-        message,
+        role: message.role,
         generate: generate && message.role === 'user',
         ...(afterStepOf === undefined ? {} : { afterStepOf }),
       }
-      if (first) state.inbox.items.unshift(candidate)
-      else state.inbox.items.push(candidate)
+      if (afterStepOf === undefined) state.inbox.items.push(candidate)
+      else
+        state.inbox.items.splice(
+          steeringEnd(state.inbox.items, afterStepOf),
+          0,
+          candidate,
+        )
     }
     const moveInput = ({ from, to }: { from: number; to: number }): void => {
       if (from === to) return
@@ -213,7 +244,7 @@ export const createControlRuntime = <
         ...request,
         control: { turnId: active.turnId, version: active.version },
       }
-      state.coordinator = foldCoordinator<M, D>(state.coordinator, {
+      state.coordinator = foldCoordinator<D>(state.coordinator, {
         ...ctx.event,
         type: 'ai.generation.requested',
         id,
@@ -247,10 +278,11 @@ export const createControlRuntime = <
               (candidate) => candidate.id === command.message.id,
             )
             if (item === undefined) reason = 'not-pending'
-            else if (item.message.role !== command.message.role)
+            else if (item.role !== command.message.role)
               reason = 'role-mismatch'
-            else item.message = command.message
-            break
+            // Its content takes effect at its receipt. An edit never runs
+            // admission, so no input's admission shares a batch with it.
+            return decide()
           }
           case 'remove': {
             const position = state.inbox.items.findIndex(
@@ -294,12 +326,18 @@ export const createControlRuntime = <
             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.role !== 'user') reason = 'not-user-input'
             else {
               selected.generate = true
               if (state.active) selected.afterStepOf = state.active.turnId
               else delete selected.afterStepOf
-              moveInput({ from: position, to: 0 })
+              const others = state.inbox.items.filter(
+                (item) => item !== selected,
+              )
+              moveInput({
+                from: position,
+                to: state.active ? steeringEnd(others, state.active.turnId) : 0,
+              })
             }
             break
           }
@@ -315,7 +353,6 @@ export const createControlRuntime = <
             insert({
               message: command.message,
               generate: true,
-              first: true,
               afterStepOf: command.turnId,
             })
             break
@@ -586,15 +623,19 @@ export const createControlRuntime = <
           !call.terminal &&
           !approvalBlocksExecution(call),
       )
+      // A paused turn keeps its steering until Resume, so every steer sent
+      // while it's paused targets it and they enter together.
       const handoffReady =
         handoff &&
         workSettled &&
         !externalPending &&
-        (active.suspended ||
-          response.failure !== undefined ||
-          response.calls.every(
-            (call) => callReadyForModel(call) || approvalBlocksExecution(call),
-          ))
+        (active.suspended
+          ? !state.inbox.paused
+          : response.failure !== undefined ||
+            response.calls.every(
+              (call) =>
+                callReadyForModel(call) || approvalBlocksExecution(call),
+            ))
       const readyForModel =
         !active.suspended &&
         response.failure === undefined &&
@@ -606,7 +647,7 @@ export const createControlRuntime = <
         response.completion !== undefined &&
         response.completion.finishReason !== 'tool-calls'
       if (handoffReady && !naturalCompletion) {
-        await end({ lastSeenIndex: undefined, reason: 'steered after step' })
+        await end({ lastSeenIndex: undefined, reason: STEERED_REASON })
       } else if (readyForModel && (!handoff || handoffReady)) {
         const resume =
           active.model !== null && active.model.version !== active.version
@@ -669,10 +710,15 @@ export const createControlRuntime = <
     ) {
       while (state.inbox.items.length > 0) {
         const item = state.inbox.items.shift()!
+        if (steersOn(item, state.inbox.items)) item.generate = false
         delete item.afterStepOf
         await emit({
           type: 'ai.message.created',
-          payload: { message: item.message, generate: item.generate },
+          payload: {
+            messageId: item.id,
+            role: item.role,
+            ...(item.generate ? {} : { generate: false as const }),
+          } satisfies MessageCreatedPayload,
           id: await aiEventId(
             'a2.ai.admitted',
             options.contract.name,
@@ -706,26 +752,7 @@ export const createControlRuntime = <
         break
       }
     }
-    const receipt: ControlReceipt = {
-      commandId: ctx.event.id,
-      ...(ctx.event.type === 'ai.work.reported' ? { work: true as const } : {}),
-      outcome: reason === undefined ? 'applied' : 'rejected',
-      ...(reason === undefined ? {} : { reason }),
-    }
-    return [
-      ...(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',
-        id: await aiEventId('a2.ai.control.decided', ctx.event.id),
-        payload: receipt,
-      } as AppendInput<D>,
-    ]
+    return decide()
   }
 
   const report = async ({

package/src/ai-control-state.ts

@@ -3,7 +3,6 @@ import type {
   ControlCommit,
   ControlChange,
 } from './ai-control.ts'
-import type { UIMessage } from 'ai'
 
 type RecordValue = Record<string, unknown> | unknown[]
 const container = (value: unknown): value is RecordValue =>
@@ -131,9 +130,9 @@ export const applyControlChanges = <T>({
   return next
 }
 
-export const controlCommit = <M extends UIMessage>(options: {
-  before: AIControlState<M>
-  after: AIControlState<M>
+export const controlCommit = (options: {
+  before: AIControlState
+  after: AIControlState
   moves: ControlChange[]
 }): ControlCommit => {
   const active = options.after.active

package/src/ai-control.ts

@@ -12,13 +12,22 @@ import type {
   ToolResultPayload,
 } from './ai.ts'
 
-export type InboxItem<M extends UIMessage = UIMessage> = {
+/**
+ * An input in the queue or the active turn. Its content stays in the
+ * `ai.control.requested` command that brought it, and in no other event.
+ */
+export type InputRef = {
   id: string
-  message: M
+  role: UIMessage['role']
   generate: boolean
   afterStepOf?: string
 }
 
+/** A queued input with its content, as `AIState` shows it. */
+export type InboxItem<M extends UIMessage = UIMessage> = InputRef & {
+  message: M
+}
+
 export type ControlCommand<M extends UIMessage = UIMessage> =
   | { action: 'send'; message: M; generate?: boolean }
   | { action: 'edit'; message: M }
@@ -55,9 +64,9 @@ export type ControlReceipt = {
     | 'not-user-input'
 }
 
-export type ActiveTurn<M extends UIMessage = UIMessage> = {
+export type ActiveTurn = {
   turnId: string
-  input: InboxItem<M>
+  input: InputRef
   phase:
     | 'starting'
     | 'generating'
@@ -76,11 +85,11 @@ export type ControlWork = {
   facts?: string[]
 }
 
-export type AIControlState<M extends UIMessage = UIMessage> = {
+export type AIControlState = {
   pendingStops: Record<string, string>
-  inbox: { paused: boolean; items: InboxItem<M>[] }
+  inbox: { paused: boolean; items: InputRef[] }
   active:
-    | (ActiveTurn<M> & {
+    | (ActiveTurn & {
         version: number
         suspended: boolean
         model: ControlWork | null
@@ -107,9 +116,49 @@ export type ToolExecutionRequest = {
   generation: GenerationStartedPayload
 }
 
-export const initialControlState = <
-  M extends UIMessage,
->(): AIControlState<M> => ({
+/**
+ * The `ai.message.interrupted` reason A2 writes when steering ends a response
+ * after a step; a harness's replay continues the turn on it. Reserved, so a
+ * client's interruption can't carry it.
+ */
+export const STEERED_REASON = 'steered after step'
+
+/**
+ * A fold met an admission no command explains. Only A2's control lane writes
+ * `ai.message.created`, after the command that brought its input in.
+ */
+export const unexplainedAdmission = (messageId: string): TypeError =>
+  new TypeError(
+    `ai.message.created for ${messageId} follows no command that brought it in; only A2's control lane writes it`,
+  )
+
+/**
+ * Where new steering for a turn goes in the inbox: right after the turn's
+ * earlier steering, or first, so it hands off in the order it was sent.
+ */
+export const steeringEnd = (
+  items: readonly InputRef[],
+  turnId: string,
+): number => items.findLastIndex((item) => item.afterStepOf === turnId) + 1
+
+/**
+ * Whether an input steering a turn enters as context for later steering: the
+ * run of inputs steering one turn is admitted together, and only its last
+ * generating input starts the next turn.
+ */
+export const steersOn = (
+  item: InputRef,
+  rest: readonly InputRef[],
+): boolean => {
+  if (item.afterStepOf === undefined) return false
+  for (const next of rest) {
+    if (next.afterStepOf !== item.afterStepOf) return false
+    if (next.generate) return true
+  }
+  return false
+}
+
+export const initialControlState = (): AIControlState => ({
   pendingStops: {},
   inbox: { paused: false, items: [] },
   active: null,

package/src/ai-coordinator.ts

@@ -107,7 +107,7 @@ const updateResponse = (
     ? state
     : { ...state, response: update(state.response) }
 
-export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
+export const foldCoordinator = <D extends EventDefs>(
   state: AICoordinatorState,
   event: ContractEvent<D>,
 ): AICoordinatorState => {
@@ -120,15 +120,16 @@ export const foldCoordinator = <M extends UIMessage, D extends EventDefs>(
       return rest
     }
     case 'ai.message.created': {
-      const payload = event.payload as MessageCreatedPayload<M>
-      const message = payload.message
-      if (message.role !== 'user') return state
+      const payload = event.payload as MessageCreatedPayload
+      if (payload.role !== 'user') return state
       return {
         ...state,
         queued: [
-          ...state.queued.filter((item) => item.messageId !== message.id),
+          ...state.queued.filter(
+            (item) => item.messageId !== payload.messageId,
+          ),
           {
-            messageId: message.id,
+            messageId: payload.messageId,
             generate: payload.generate !== false,
           },
         ],