experimental-a2 0.16.1 → 0.18.0

package/src/ai-agent-actions.ts

@@ -0,0 +1,219 @@
+import type { JSONValue, UIMessage } from 'ai'
+import {
+  inputs,
+  isCompactionPart,
+  type AIState,
+  type AgentStatus,
+  type PendingApproval,
+  type PendingInput,
+} from './ai.ts'
+import type { AppendInput, EventDefs } from './contract.ts'
+import type { PushResult } from './client.ts'
+import type { UseSessionResult } from './react.ts'
+
+/** The input message type of a state: what its inbox holds. */
+type InputOf<S extends AIState> = S['inbox']['items'][number]['message']
+
+/** What the UI shows while nothing new is visible in the response. */
+export type AgentIndicator = 'starting' | 'thinking' | 'compacting' | null
+
+/**
+ * Every action resolves like `push()`: to the persisted command events, whose
+ * `id` is the `commandId` of its `ai.control.decided` receipt. An action with
+ * nothing to address (Stop with no active turn) resolves to no events.
+ */
+export type AgentActions<D extends EventDefs, M extends UIMessage> = {
+  /** Queue a user message. A string becomes one text part. */
+  sendMessage: (
+    message: string | M,
+    options?: { generate?: boolean },
+  ) => PushResult<D>
+  /** Hand off after the active turn's current step, or send when idle. */
+  steer: (message: string | M) => PushResult<D>
+  /** Stop the active turn at the frontier this render showed. */
+  stop: () => PushResult<D>
+  approve: (approval: PendingApproval) => PushResult<D>
+  deny: (
+    approval: PendingApproval,
+    options?: { reason?: string },
+  ) => PushResult<D>
+  respond: (input: PendingInput, value: JSONValue) => PushResult<D>
+  /** Rerun the failed model step of the active turn. */
+  retry: () => PushResult<D>
+  pause: (options: { when: 'now' | 'after-turn' }) => PushResult<D>
+  resume: () => PushResult<D>
+  queue: {
+    edit: (message: M) => PushResult<D>
+    move: (options: {
+      inputId: string
+      beforeId: string | null
+    }) => PushResult<D>
+    remove: (inputId: string) => PushResult<D>
+    /** Prioritize a queued input and hand off after the current step. */
+    sendNow: (inputId: string) => PushResult<D>
+  }
+}
+
+export type UseAgentResult<
+  D extends EventDefs,
+  S extends AIState,
+> = UseSessionResult<D, S> &
+  AgentActions<D, InputOf<S>> & {
+    messages: S['messages']
+    status: AgentStatus
+    indicator: AgentIndicator
+  }
+
+const visible = (part: UIMessage['parts'][number]): boolean =>
+  part.type === 'text'
+    ? part.text.length > 0
+    : part.type !== 'reasoning' &&
+      part.type !== 'step-start' &&
+      !isCompactionPart(part)
+
+const indicatorOf = (state: AIState): AgentIndicator => {
+  if (state.starting) return 'starting'
+  if (state.status !== 'generating') return null
+  if (state.compaction?.status === 'running') return 'compacting'
+  const responseId =
+    state.activeGeneration?.responseMessageId ??
+    state.activeResponseMessageId ??
+    state.activeProjection?.responseMessageId
+  const response =
+    responseId === undefined || responseId === null
+      ? undefined
+      : state.messages.findLast((message) => message.id === responseId)
+  // A requested step that has not started has produced nothing yet.
+  if (
+    response === undefined ||
+    (state.activeGeneration === null && state.activeResponseMessageId !== null)
+  )
+    return 'thinking'
+  const projection = state.activeProjection
+  // A started step's own output begins after its base; before its first chunk
+  // the previous step's parts are not this step's progress.
+  const start =
+    state.activeGeneration !== null &&
+    projection?.generationId === state.activeGeneration.generationId &&
+    projection.base !== undefined
+      ? projection.base.partCount
+      : response.parts.findLastIndex((part) => part.type === 'step-start') + 1
+  return response.parts.slice(start).some(visible) ? null : 'thinking'
+}
+
+const userMessage = <M extends UIMessage>(message: string | M): M =>
+  typeof message === 'string'
+    ? ({
+        id: crypto.randomUUID(),
+        role: 'user',
+        parts: [{ type: 'text', text: message }],
+      } as UIMessage as M)
+    : message
+
+export function withAgentActions<D extends EventDefs, S extends AIState>(
+  session: UseSessionResult<D, S>,
+): UseAgentResult<D, S> {
+  type M = InputOf<S>
+  const { state, index, push } = session
+  const send = (
+    events: { type: 'ai.control.requested'; payload: unknown }[],
+  ): PushResult<D> =>
+    events.length === 0
+      ? Object.assign(Promise.resolve([]), {
+          confirmed: Promise.resolve([]),
+        })
+      : push(...(events as AppendInput<D>[]))
+  const active = state.active
+  // `index` is this render's frontier: Stop freezes the response the user
+  // saw, not whatever arrived before the click was handled.
+  const lastSeenIndex = index
+  const actions: AgentActions<D, M> = {
+    sendMessage: (message, options) =>
+      send(inputs.message(userMessage(message), options)),
+    steer: (message) =>
+      send(
+        active === null
+          ? inputs.message(userMessage(message))
+          : inputs.steer({
+              turnId: active.turnId,
+              message: userMessage(message),
+              lastSeenIndex,
+            }),
+      ),
+    stop: () =>
+      send(
+        active === null
+          ? []
+          : inputs.stop({ turnId: active.turnId, lastSeenIndex }),
+      ),
+    approve: (approval) =>
+      send(inputs.approval({ ...approvalOwner(approval), approved: true })),
+    deny: (approval, options) =>
+      send(
+        inputs.approval({
+          ...approvalOwner(approval),
+          approved: false,
+          ...(options?.reason === undefined ? {} : { reason: options.reason }),
+        }),
+      ),
+    respond: (input, value) =>
+      send(
+        inputs.input({
+          requestId: input.requestId,
+          messageId: input.messageId,
+          generationId: input.generationId,
+          inputId: input.inputId,
+          name: input.name,
+          value,
+        }),
+      ),
+    retry: () => {
+      const responseMessageId = state.activeProjection?.responseMessageId
+      return send(
+        state.status !== 'failed' ||
+          active === null ||
+          responseMessageId === undefined
+          ? []
+          : inputs.retry({
+              messageId: active.input.id,
+              responseMessageId,
+              retryId: crypto.randomUUID(),
+            }),
+      )
+    },
+    pause: (options) => send(inputs.pause(options)),
+    resume: () => send(inputs.resume()),
+    queue: {
+      edit: (message) => send(inputs.queue.edit({ message })),
+      move: (options) => send(inputs.queue.move(options)),
+      remove: (inputId) => send(inputs.queue.remove({ inputId })),
+      sendNow: (inputId) =>
+        send(
+          inputs.queue.sendNow({
+            inputId,
+            turnId: active?.turnId ?? null,
+            lastSeenIndex,
+          }),
+        ),
+    },
+  }
+  return {
+    ...session,
+    ...actions,
+    messages: state.messages,
+    status: state.status,
+    indicator: indicatorOf(state as AIState),
+  }
+}
+
+const approvalOwner = (
+  approval: PendingApproval,
+): Pick<
+  PendingApproval,
+  'requestId' | 'messageId' | 'generationId' | 'approvalId'
+> => ({
+  requestId: approval.requestId,
+  messageId: approval.messageId,
+  generationId: approval.generationId,
+  approvalId: approval.approvalId,
+})