experimental-a2 0.12.0 → 0.14.0

package/examples/playground/app/agent/compaction-timeline.ts

@@ -0,0 +1,198 @@
+import type { UIMessage } from 'ai'
+import type { AIState } from 'experimental-a2/ai'
+import type { AgentEvent } from './model'
+
+type Generation = Extract<AgentEvent, { type: 'ai.generation.started' }>
+type CompactionEvent = Extract<
+  AgentEvent,
+  { type: 'ai.compaction.requested' | 'ai.compaction.completed' }
+>
+
+export type CompactionEntry = {
+  type: 'compaction'
+  id: string
+  event: CompactionEvent
+  running: boolean
+  outcome: 'failed' | 'interrupted' | 'superseded' | undefined
+}
+
+export type TimelineEntry =
+  | { type: 'message'; id: string; message: UIMessage; lastSegment: boolean }
+  | CompactionEntry
+
+export function compactionTimeline({
+  state,
+  events,
+}: {
+  state: AIState
+  events: readonly AgentEvent[]
+}): TimelineEntry[] {
+  const generations = new Map<string, Generation>()
+  const requests = new Map<string, string>()
+  const stepCounts = new Map<string, number>()
+  const contentGenerations = new Set<string>()
+  const compactions = new Map<string, CompactionEvent>()
+  const roots = new Map<string, string>()
+  for (const event of events) {
+    switch (event.type) {
+      case 'ai.generation.requested':
+        requests.set(event.id, event.payload.reason)
+        break
+      case 'ai.generation.started': {
+        const { generationId, messageId, responseMessageId } = event.payload
+        generations.set(generationId, event)
+        if (messageId !== responseMessageId)
+          roots.set(responseMessageId, messageId)
+        break
+      }
+      case 'ai.generation.progress': {
+        const id = event.payload.generationId
+        const count = event.payload.chunks.filter(
+          (chunk) => chunk.type === 'start-step',
+        ).length
+        stepCounts.set(id, (stepCounts.get(id) ?? 0) + count)
+        if (
+          event.payload.chunks.some(
+            (chunk) =>
+              ![
+                'start',
+                'finish',
+                'start-step',
+                'finish-step',
+                'message-metadata',
+                'abort',
+                'error',
+              ].includes(chunk.type),
+          )
+        )
+          contentGenerations.add(id)
+        break
+      }
+      case 'ai.compaction.requested':
+      case 'ai.compaction.completed':
+        compactions.set(event.payload.generationId, event)
+        break
+    }
+  }
+
+  const contributions = new Map<string, Generation[]>()
+  for (const generation of generations.values()) {
+    const { generationId, responseMessageId, requestId } = generation.payload
+    const ownerId = state.responseGenerationIds[responseMessageId]
+    const owner = ownerId === undefined ? undefined : generations.get(ownerId)
+    if (
+      owner === undefined ||
+      generation.index > owner.index ||
+      state.terminalGenerations?.[generationId] === 'superseded'
+    )
+      continue
+    let steps = contributions.get(responseMessageId) ?? []
+    if (requests.get(requestId) === 'retry') steps = []
+    if (steps.at(-1)?.payload.requestId === requestId) steps.pop()
+    steps.push(generation)
+    contributions.set(responseMessageId, steps)
+  }
+
+  const included = new Set(
+    [...contributions.values()].flatMap((steps) =>
+      steps.map((step) => step.payload.generationId),
+    ),
+  )
+  const offsets = new Map<string, number>()
+  const counts = new Map<string, number>()
+  for (const { payload } of generations.values()) {
+    const count = counts.get(payload.responseMessageId) ?? 0
+    offsets.set(payload.generationId, count)
+    if (included.has(payload.generationId))
+      counts.set(
+        payload.responseMessageId,
+        count + (stepCounts.get(payload.generationId) ?? 0),
+      )
+  }
+
+  const markers = new Map<string, CompactionEntry[]>()
+  const orphaned: CompactionEntry[] = []
+  const messageIds = new Set(state.messages.map((message) => message.id))
+  for (const [generationId, event] of compactions) {
+    const terminal = state.terminalGenerations?.[generationId]
+    const entry: CompactionEntry = {
+      type: 'compaction',
+      id: `compaction:${generationId}`,
+      event,
+      running:
+        state.compaction?.generationId === generationId &&
+        state.compaction.status === 'running',
+      outcome: terminal === 'completed' ? undefined : terminal,
+    }
+    const responseId = generations.get(generationId)?.payload.responseMessageId
+    const anchor =
+      responseId !== undefined && messageIds.has(responseId)
+        ? responseId
+        : responseId === undefined
+          ? undefined
+          : roots.get(responseId)
+    if (anchor === undefined || !messageIds.has(anchor)) orphaned.push(entry)
+    else {
+      const entries = markers.get(anchor) ?? []
+      entries.push(entry)
+      markers.set(anchor, entries)
+    }
+  }
+
+  const timeline: TimelineEntry[] = []
+  for (const message of state.messages) {
+    const entries = markers.get(message.id) ?? []
+    if (entries.length === 0) {
+      timeline.push({
+        type: 'message',
+        id: `${message.id}:0`,
+        message,
+        lastSegment: true,
+      })
+      continue
+    }
+    const steps = contributions.get(message.id) ?? []
+    const boundaries = message.parts.flatMap((part, index) =>
+      part.type === 'step-start' ? [index] : [],
+    )
+    const mapped =
+      boundaries[0] === 0 &&
+      boundaries.length === counts.get(message.id) &&
+      steps.every(
+        ({ payload }) =>
+          !contentGenerations.has(payload.generationId) ||
+          (stepCounts.get(payload.generationId) ?? 0) > 0,
+      )
+    let offset = 0
+    const appendMessage = (end: number): void => {
+      if (end <= offset) return
+      timeline.push({
+        type: 'message',
+        id: `${message.id}:${offset}`,
+        message: { ...message, parts: message.parts.slice(offset, end) },
+        lastSegment: end === message.parts.length,
+      })
+      offset = end
+    }
+    for (const entry of entries) {
+      if (message.role === 'user') {
+        appendMessage(message.parts.length)
+      } else if (mapped) {
+        const stepIndex = offsets.get(entry.event.payload.generationId) ?? 0
+        appendMessage(boundaries[stepIndex] ?? message.parts.length)
+      }
+      timeline.push(entry)
+    }
+    if (message.parts.length === 0) {
+      timeline.push({
+        type: 'message',
+        id: message.id,
+        message,
+        lastSegment: true,
+      })
+    } else {
+      appendMessage(message.parts.length)
+    }
+  }
+  return [...timeline, ...orphaned]
+}

package/examples/playground/app/agent/model.ts

@@ -1,3 +1,58 @@
 import { agent } from 'experimental-a2/ai'
+import type { ContractEvent } from 'experimental-a2'
+import type { AutomaticCompaction } from 'experimental-a2/ai/server'
+import { z } from 'zod'
 
-export const demoAgent = agent({ name: 'agent' })
+export const agentModel = 'openai/gpt-5.6-luna'
+
+export const compactionSettingsSchema = z.discriminatedUnion('mode', [
+  z.strictObject({ mode: z.literal('automatic') }),
+  z.strictObject({
+    mode: z.literal('manual'),
+    thresholdTokens: z.number().int().min(1).max(Number.MAX_SAFE_INTEGER),
+  }),
+  z.strictObject({ mode: z.literal('disabled') }),
+])
+
+export type CompactionSettings = z.infer<typeof compactionSettingsSchema>
+
+export const demoAgent = agent({
+  name: 'agent',
+  events: { 'compaction.settings.updated': compactionSettingsSchema },
+})
+
+export type AgentEvent = ContractEvent<typeof demoAgent.contract.events>
+
+const automatic: CompactionSettings = { mode: 'automatic' }
+
+export function latestCompactionSettings(history: readonly AgentEvent[]): {
+  settings: CompactionSettings
+  id: string
+} {
+  const event = history.findLast(
+    (candidate) => candidate.type === 'compaction.settings.updated',
+  )
+  return event
+    ? { settings: event.payload, id: event.id }
+    : { settings: automatic, id: 'default' }
+}
+
+export const compactionSettings = demoAgent.contract
+  .reducer({
+    name: 'agent.compaction-settings.v1',
+    initialState: automatic,
+    stateSchema: compactionSettingsSchema,
+  })
+  .fold((state, event): CompactionSettings =>
+    event.type === 'compaction.settings.updated' ? event.payload : state,
+  )
+
+export function compactionOptions(
+  settings: CompactionSettings,
+): AutomaticCompaction | false {
+  return settings.mode === 'disabled'
+    ? false
+    : settings.mode === 'manual'
+      ? { thresholdTokens: settings.thresholdTokens }
+      : {}
+}

package/examples/playground/app/agent/server.ts

@@ -8,7 +8,12 @@ import { createAgentServer } from 'experimental-a2/ai/server'
 import { vercelQueues } from 'experimental-a2/scheduler-vercel'
 import { z } from 'zod'
 import { store } from '@/lib/store'
-import { demoAgent } from './model'
+import {
+  agentModel,
+  compactionOptions,
+  compactionSettings,
+  demoAgent,
+} from './model'
 
 const maxReminderDelaySeconds = 6 * 24 * 60 * 60
 const bashTimeoutMs = 120_000
@@ -131,7 +136,9 @@ export const agentServer = createAgentServer({
   ...(store ? { store } : {}),
   ...(agentScheduler ? { scheduler: agentScheduler } : {}),
   agent: demoAgent,
-  model: 'openai/gpt-5.6-luna',
+  model: agentModel,
+  compaction: async ({ session }) =>
+    compactionOptions((await session.state(compactionSettings)).state),
   tools,
   instructions: `You are an incident-response agent for the playground.
 When responding to an incident, inspect checkout-api before taking action. Its

package/examples/playground/app/chat/[chatId]/chat-client.tsx

@@ -539,22 +539,12 @@ export function ChatClient({ chatId }: { chatId: string }): ReactNode {
   }
 
   const stopAgent = async (): Promise<void> => {
-    if (!activeResponseId) return
-    const target = state.ai.activeGeneration
-      ? { generationId: state.ai.activeGeneration.generationId }
-      : runningTool
-        ? { generationId: runningTool.generationId }
-        : state.ai.activeRequestId
-          ? { requestId: state.ai.activeRequestId }
-          : null
-    if (target === null) return
+    if (!state.ai.active) return
     setError(null)
     try {
       await push(
-        ...inputs.interrupt({
-          messageId: activeResponseId,
-          ...target,
-          reason: 'a chat participant stopped the agent',
+        ...inputs.stop({
+          turnId: state.ai.active.turnId,
           lastSeenIndex: index,
         }),
       )

package/examples/playground/app/chat/model.ts

@@ -164,7 +164,7 @@ export const chatAgent = agent<ChatAgentMessage, typeof chatEvents>({
   name: 'chat',
   events: chatEvents,
   messageSchema: chatMessageSchema,
-  reducerName: 'a2.playground.chat.ai.v1',
+  reducerName: 'a2.playground.chat.ai.v3',
 })
 
 export const chat = chatAgent.contract
@@ -208,7 +208,7 @@ const initialChatState: ChatState = {
 
 export const chatReducer = chat
   .reducer({
-    name: 'a2.playground.chat.v1',
+    name: 'a2.playground.chat.v3',
     initialState: initialChatState,
   })
   .fold((state, event) => {

package/examples/playground/app/chat/server.ts

@@ -21,16 +21,13 @@ const tools = {
     inputSchema: chatPostInputSchema,
     execute: async ({ message }, { toolCallId }) => {
       const context = handlerContext(chatAgent)
-      if (context.event.type !== 'ai.tool.called') {
-        throw new Error('The chat post tool must execute automatically')
-      }
       context.signal.throwIfAborted()
       await context.session.append('post', {
         type: 'agent.message.posted',
         payload: {
           message,
-          generationId: context.event.payload.generationId,
-          responseMessageId: context.event.payload.messageId,
+          generationId: context.event.payload.call.generationId,
+          responseMessageId: context.event.payload.call.messageId,
           toolCallId,
         },
       })
@@ -118,23 +115,33 @@ export const authorizeChatOperation = (
       return false
     }
     return !(
-      event.type === 'ai.approval.responded' ||
-      event.type === 'ai.input.responded' ||
-      event.type === 'ai.message.interrupted' ||
-      event.type === 'ai.retry.requested'
+      event.type === 'ai.control.requested' &&
+      typeof event.payload === 'object' &&
+      event.payload !== null &&
+      'action' in event.payload &&
+      typeof event.payload.action === 'string' &&
+      [
+        'approval',
+        'input',
+        'interrupt',
+        'stop',
+        'retry',
+        'pause',
+        'resume',
+      ].includes(event.payload.action)
     )
   })
   return rejected === undefined
 }
 
-const toolCalledHandler = aiHandlers['ai.tool.called']
-if (toolCalledHandler === undefined) {
-  throw new Error('A2 AI handlers must include ai.tool.called')
+const toolExecutionHandler = aiHandlers['ai.tool.execution.requested']
+if (toolExecutionHandler === undefined) {
+  throw new Error('A2 AI handlers must include ai.tool.execution.requested')
 }
-const serializedToolCalledHandler =
-  typeof toolCalledHandler === 'function'
-    ? { handler: toolCalledHandler, lane: 'a2.chat.agent-tools' }
-    : { ...toolCalledHandler, lane: 'a2.chat.agent-tools' }
+const serializedToolExecutionHandler =
+  typeof toolExecutionHandler === 'function'
+    ? { handler: toolExecutionHandler, lane: 'a2.chat.agent-tools' }
+    : { ...toolExecutionHandler, lane: 'a2.chat.agent-tools' }
 
 export const chatScheduler =
   process.env.VERCEL || process.env.A2_QUEUES
@@ -147,7 +154,7 @@ export const chatServer = createServer({
   contract: chatAgent.contract,
   handlers: {
     ...aiHandlers,
-    'ai.tool.called': serializedToolCalledHandler,
+    'ai.tool.execution.requested': serializedToolExecutionHandler,
     'message.sent': {
       lane: 'a2.chat.ingress',
       handler: async (context) => {