experimental-a2 0.11.0 → 0.13.0

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

@@ -0,0 +1,337 @@
+import type { UIMessageChunk } from 'ai'
+import { reduceAIState } from 'experimental-a2/ai'
+import { expect, it } from 'vitest'
+import { compactionTimeline, type TimelineEntry } from './compaction-timeline'
+import { demoAgent, type AgentEvent } from './model'
+
+type EventInput = {
+  [K in AgentEvent['type']]: Pick<
+    Extract<AgentEvent, { type: K }>,
+    'type' | 'payload'
+  > & { id?: string }
+}[AgentEvent['type']]
+type Generation = Extract<
+  AgentEvent,
+  { type: 'ai.generation.started' }
+>['payload']
+
+function transcript() {
+  const events: AgentEvent[] = []
+  const append = (input: EventInput): void => {
+    events.push({
+      ...input,
+      id: input.id ?? `event:${events.length + 1}`,
+      index: events.length + 1,
+      sessionId: 'timeline',
+      createdAt: new Date('2026-09-14T12:00:00Z'),
+    })
+  }
+  const user = (id: string): void =>
+    append({
+      type: 'ai.message.created',
+      payload: {
+        message: { id, role: 'user', parts: [{ type: 'text', text: id }] },
+      },
+    })
+  const start = ({
+    requestId = 'request',
+    reason = 'message',
+    attempt = 1,
+  }: {
+    requestId?: string
+    reason?: 'message' | 'tool' | 'retry'
+    attempt?: number
+  } = {}): Generation => {
+    const generation = {
+      requestId,
+      messageId: reason === 'tool' ? 'assistant' : 'user',
+      responseMessageId: 'assistant',
+      generationId: `${requestId}:generation:${attempt}`,
+      attempt,
+      model: 'test/model',
+    }
+    if (attempt === 1)
+      append({
+        id: requestId,
+        type: 'ai.generation.requested',
+        payload: {
+          messageId: generation.messageId,
+          responseMessageId: 'assistant',
+          reason,
+        },
+      })
+    append({ type: 'ai.generation.started', payload: generation })
+    return generation
+  }
+  const progress = (generation: Generation, chunks: UIMessageChunk[]): void =>
+    append({
+      type: 'ai.generation.progress',
+      payload: { ...generation, sequence: events.length, chunks },
+    })
+  const text = (generation: Generation, value: string, step = true): void =>
+    progress(generation, [
+      { type: 'start', messageId: generation.responseMessageId },
+      ...(step ? [{ type: 'start-step' as const }] : []),
+      { type: 'text-start', id: generation.generationId },
+      { type: 'text-delta', id: generation.generationId, delta: value },
+      { type: 'text-end', id: generation.generationId },
+    ])
+  const compact = (generation: Generation, completed = true): void => {
+    const payload = {
+      generationId: generation.generationId,
+      throughMessageId: generation.messageId,
+      throughIndex: events.length,
+    }
+    append({ type: 'ai.compaction.requested', payload })
+    if (completed)
+      append({
+        type: 'ai.compaction.completed',
+        payload: {
+          ...payload,
+          messages: [],
+          summary: `Summary ${generation.generationId}`,
+        },
+      })
+  }
+  const state = () =>
+    events.reduce(reduceAIState, demoAgent.reducer.initialState)
+  const timeline = () => compactionTimeline({ state: state(), events })
+  return {
+    events,
+    append,
+    user,
+    start,
+    progress,
+    text,
+    compact,
+    state,
+    timeline,
+  }
+}
+
+const labels = (entries: TimelineEntry[]): string[] =>
+  entries.map((entry) =>
+    entry.type === 'compaction'
+      ? entry.id
+      : entry.message.parts
+          .flatMap((part) => (part.type === 'text' ? [part.text] : []))
+          .join(''),
+  )
+
+it('keeps a stable marker from running to completed, including after replay', () => {
+  const log = transcript()
+  log.user('user')
+  const generation = log.start()
+  log.compact(generation, false)
+  const running = log.timeline().find((entry) => entry.type === 'compaction')
+  expect(running?.running).toBe(true)
+  expect(labels(log.timeline())).toEqual([
+    'user',
+    `compaction:${generation.generationId}`,
+  ])
+  log.append({
+    type: 'ai.compaction.completed',
+    payload: {
+      generationId: generation.generationId,
+      throughMessageId: 'user',
+      messages: [],
+      summary: 'Summary',
+    },
+  })
+  log.text(generation, 'Answer')
+  log.append({ type: 'ai.generation.completed', payload: generation })
+  log.append({
+    type: 'ai.message.completed',
+    payload: { messageId: 'assistant' },
+  })
+  const completed = log.timeline().find((entry) => entry.type === 'compaction')
+  expect(completed).toMatchObject({
+    id: running?.id,
+    running: false,
+    event: { type: 'ai.compaction.completed' },
+  })
+  expect(labels(log.timeline())).toEqual([
+    'user',
+    `compaction:${generation.generationId}`,
+    'Answer',
+  ])
+})
+
+it('places compaction between tool steps without moving a queued user into the response', () => {
+  const log = transcript()
+  log.user('user')
+  const first = log.start()
+  log.text(first, 'Before tool')
+  log.append({
+    type: 'ai.generation.completed',
+    payload: { ...first, finishReason: 'tool-calls' },
+  })
+  log.user('queued')
+  const second = log.start({
+    requestId: `ai.generate:tools:${first.generationId}`,
+    reason: 'tool',
+  })
+  log.compact(second, false)
+  expect(log.state().compaction?.status).toBe('running')
+  expect(labels(log.timeline())).toEqual([
+    'user',
+    'Before tool',
+    `compaction:${second.generationId}`,
+    'queued',
+  ])
+  log.append({
+    type: 'ai.compaction.completed',
+    payload: {
+      generationId: second.generationId,
+      throughMessageId: 'assistant',
+      messages: [],
+    },
+  })
+  log.text(second, 'After tool')
+  expect(labels(log.timeline())).toEqual([
+    'user',
+    'Before tool',
+    `compaction:${second.generationId}`,
+    'After tool',
+    'queued',
+  ])
+  expect(
+    log
+      .timeline()
+      .flatMap((entry) =>
+        entry.type === 'message' && entry.message.id === 'assistant'
+          ? entry.message.parts
+          : [],
+      ),
+  ).toEqual(
+    log.state().messages.find((message) => message.id === 'assistant')?.parts,
+  )
+})
+
+it('preserves compaction history but drops superseded attempt content', () => {
+  const log = transcript()
+  log.user('user')
+  const first = log.start()
+  log.compact(first)
+  log.text(first, 'Superseded')
+  const second = log.start({ attempt: 2 })
+  log.compact(second)
+  log.text(second, 'Replacement')
+  expect(log.state().terminalGenerations?.[first.generationId]).toBe(
+    'superseded',
+  )
+  expect(labels(log.timeline())).toEqual([
+    'user',
+    `compaction:${first.generationId}`,
+    `compaction:${second.generationId}`,
+    'Replacement',
+  ])
+  expect(
+    log
+      .timeline()
+      .find((entry) => entry.id === `compaction:${first.generationId}`),
+  ).toMatchObject({ outcome: 'superseded', running: false })
+})
+
+it('keeps old checkpoints before an explicit retry without restoring discarded parts', () => {
+  const log = transcript()
+  log.user('user')
+  const first = log.start()
+  log.compact(first)
+  log.text(first, 'Failed response')
+  log.append({
+    type: 'ai.generation.failed',
+    payload: { ...first, error: 'Test failure' },
+  })
+  const retry = log.start({ requestId: 'retry', reason: 'retry' })
+  log.compact(retry)
+  log.text(retry, 'Retried response')
+  expect(labels(log.timeline())).toEqual([
+    'user',
+    `compaction:${first.generationId}`,
+    `compaction:${retry.generationId}`,
+    'Retried response',
+  ])
+  expect(
+    log
+      .timeline()
+      .find((entry) => entry.id === `compaction:${first.generationId}`),
+  ).toMatchObject({ outcome: 'failed' })
+})
+
+it('does not leave failed compaction spinning', () => {
+  const log = transcript()
+  log.user('user')
+  const generation = log.start()
+  log.compact(generation, false)
+  log.append({
+    type: 'ai.generation.failed',
+    payload: { ...generation, error: 'Compaction failed' },
+  })
+  expect(
+    log.timeline().find((entry) => entry.type === 'compaction'),
+  ).toMatchObject({
+    running: false,
+    outcome: 'failed',
+    event: { type: 'ai.compaction.requested' },
+  })
+})
+
+it.each([false, true])(
+  'preserves complete messages without step boundaries, delta only=%s',
+  (deltaOnly) => {
+    const log = transcript()
+    log.user('user')
+    const first = log.start()
+    log.text(first, 'Before')
+    log.append({
+      type: 'ai.generation.completed',
+      payload: { ...first, finishReason: 'tool-calls' },
+    })
+    const second = log.start({
+      requestId: `ai.generate:tools:${first.generationId}`,
+      reason: 'tool',
+    })
+    log.compact(second)
+    if (deltaOnly)
+      log.progress(second, [
+        {
+          type: 'text-delta',
+          id: 'after',
+          delta: 'After without a step boundary',
+        },
+      ])
+    else log.text(second, 'After without a step boundary', false)
+    const messages = log
+      .timeline()
+      .flatMap((entry) => (entry.type === 'message' ? [entry.message] : []))
+    expect(messages).toEqual(log.state().messages)
+    expect(
+      log.timeline().filter((entry) => entry.type === 'compaction'),
+    ).toHaveLength(1)
+    expect(labels(log.timeline())[1]).toBe(`compaction:${second.generationId}`)
+  },
+)
+
+it('ignores a rejected later generation when mapping canonical parts', () => {
+  const log = transcript()
+  log.user('user')
+  const first = log.start()
+  log.compact(first)
+  log.text(first, 'Kept')
+  log.append({ type: 'ai.generation.completed', payload: first })
+  log.append({
+    type: 'ai.message.completed',
+    payload: { messageId: 'assistant' },
+  })
+  const rejected = log.start({ requestId: 'stale-retry', reason: 'retry' })
+  log.text(rejected, 'Ignored')
+  expect(log.state().responseGenerationIds['assistant']).toBe(
+    first.generationId,
+  )
+  expect(labels(log.timeline())).toEqual([
+    'user',
+    `compaction:${first.generationId}`,
+    'Kept',
+  ])
+})

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,49 @@
 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 function compactionOptions(
+  history: readonly AgentEvent[],
+): AutomaticCompaction | false {
+  const { settings } = latestCompactionSettings(history)
+  return settings.mode === 'disabled'
+    ? false
+    : settings.mode === 'manual'
+      ? { thresholdTokens: settings.thresholdTokens }
+      : {}
+}

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

@@ -8,7 +8,7 @@ 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, demoAgent } from './model'
 
 const maxReminderDelaySeconds = 6 * 24 * 60 * 60
 const bashTimeoutMs = 120_000
@@ -131,7 +131,8 @@ export const agentServer = createAgentServer({
   ...(store ? { store } : {}),
   ...(agentScheduler ? { scheduler: agentScheduler } : {}),
   agent: demoAgent,
-  model: 'openai/gpt-5.6-luna',
+  model: agentModel,
+  compaction: ({ history }) => compactionOptions(history),
   tools,
   instructions: `You are an incident-response agent for the playground.
 When responding to an incident, inspect checkout-api before taking action. Its