experimental-a2 0.12.0 → 0.13.0

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

@@ -0,0 +1,128 @@
+import { expect, it } from 'vitest'
+import { createAgentServer } from 'experimental-a2/ai/server'
+import { memory } from 'experimental-a2/store-memory'
+import { compactionOptions, demoAgent, latestCompactionSettings } from './model'
+import { updateCompactionSettings } from './compaction-settings'
+
+const request = (
+  settings: unknown,
+  options: { origin?: string; contentType?: string; id?: string } = {},
+): Request =>
+  new Request('https://example.test/agent/session/compaction', {
+    method: 'POST',
+    headers: {
+      origin: options.origin ?? 'https://example.test',
+      'content-type': options.contentType ?? 'application/json',
+    },
+    body: JSON.stringify({ id: options.id ?? crypto.randomUUID(), settings }),
+  })
+
+it('uses the request Host when Next normalizes the internal URL', async () => {
+  let writes = 0
+  const response = await updateCompactionSettings({
+    request: new Request('http://localhost:3107/agent/s/compaction', {
+      method: 'POST',
+      headers: {
+        host: '127.0.0.1:3107',
+        origin: 'http://127.0.0.1:3107',
+        'content-type': 'application/json',
+      },
+      body: JSON.stringify({
+        id: crypto.randomUUID(),
+        settings: { mode: 'manual', thresholdTokens: 1 },
+      }),
+    }),
+    sessionId: 's',
+    append: async () => {
+      writes += 1
+    },
+  })
+  expect(response.status).toBe(200)
+  expect(writes).toBe(1)
+})
+
+it('saves session settings durably and idempotently without starting a turn', async () => {
+  let generations = 0
+  const server = createAgentServer({
+    agent: demoAgent,
+    store: memory(),
+    model: 'test/model',
+    compaction: ({ history }) => compactionOptions(history),
+    generate: () => {
+      generations += 1
+      return new ReadableStream({
+        start(controller) {
+          controller.close()
+        },
+      })
+    },
+  })
+  const session = server.session('settings-test')
+  const id = crypto.randomUUID()
+  const save = () =>
+    updateCompactionSettings({
+      request: request({ mode: 'manual', thresholdTokens: 1 }, { id }),
+      sessionId: session.id,
+      append: (event) => session.append(event),
+    })
+  expect((await save()).status).toBe(200)
+  expect((await save()).status).toBe(200)
+  await server.drain(session.id)
+  const history = await session.history()
+  expect(history.map((event) => event.type)).toEqual([
+    'compaction.settings.updated',
+  ])
+  expect(latestCompactionSettings(history).settings).toEqual({
+    mode: 'manual',
+    thresholdTokens: 1,
+  })
+  expect(compactionOptions(history)).toEqual({ thresholdTokens: 1 })
+  expect(compactionOptions(await server.session('other').history())).toEqual({})
+  expect(generations).toBe(0)
+
+  await updateCompactionSettings({
+    request: request({ mode: 'disabled' }),
+    sessionId: session.id,
+    append: (event) => session.append(event),
+  })
+  expect(compactionOptions(await session.history())).toBe(false)
+  await updateCompactionSettings({
+    request: request({ mode: 'automatic' }),
+    sessionId: session.id,
+    append: (event) => session.append(event),
+  })
+  expect(compactionOptions(await session.history())).toEqual({})
+})
+
+it.each([0, -1, 1.5, Number.MAX_SAFE_INTEGER + 1, '1', null])(
+  'rejects invalid threshold %s',
+  async (thresholdTokens) => {
+    let writes = 0
+    const response = await updateCompactionSettings({
+      request: request({ mode: 'manual', thresholdTokens }),
+      sessionId: 's',
+      append: async () => {
+        writes += 1
+      },
+    })
+    expect(response.status).toBe(400)
+    expect(writes).toBe(0)
+  },
+)
+
+it.each([
+  { origin: 'https://other.test', status: 403 },
+  { contentType: 'text/plain', status: 415 },
+  { id: 'not-an-id', status: 400 },
+])('rejects invalid request boundaries %j', async ({ status, ...options }) => {
+  let writes = 0
+  const response = await updateCompactionSettings({
+    request: request({ mode: 'automatic' }, options),
+    sessionId: 's',
+    append: async () => {
+      writes += 1
+    },
+  })
+  expect(response.status).toBe(status)
+  expect(writes).toBe(0)
+})

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

@@ -0,0 +1,49 @@
+import { z } from 'zod'
+import { compactionSettingsSchema, type CompactionSettings } from './model'
+
+const updateSchema = z.strictObject({
+  id: z.uuid(),
+  settings: compactionSettingsSchema,
+})
+
+export async function updateCompactionSettings(options: {
+  request: Request
+  sessionId: string
+  append: (event: {
+    id: string
+    type: 'compaction.settings.updated'
+    payload: CompactionSettings
+  }) => Promise<unknown>
+}): Promise<Response> {
+  const { request, sessionId } = options
+  const origin = request.headers.get('origin')
+  const requestUrl = new URL(request.url)
+  requestUrl.host = request.headers.get('host') ?? requestUrl.host
+  if (origin !== null && origin !== requestUrl.origin) {
+    return Response.json(
+      { error: 'Use this page to update compaction settings.' },
+      { status: 403 },
+    )
+  }
+  if (
+    request.headers.get('content-type')?.split(';')[0]?.trim() !==
+    'application/json'
+  ) {
+    return Response.json({ error: 'Expected JSON settings.' }, { status: 415 })
+  }
+  if (sessionId.length === 0)
+    return Response.json({ error: 'Session is required.' }, { status: 400 })
+  const parsed = updateSchema.safeParse(await request.json().catch(() => null))
+  if (!parsed.success) {
+    return Response.json(
+      { error: 'Choose a mode and a positive whole-number threshold.' },
+      { status: 400 },
+    )
+  }
+  await options.append({
+    id: `agent.compaction.settings:${encodeURIComponent(sessionId)}:${parsed.data.id}`,
+    type: 'compaction.settings.updated',
+    payload: parsed.data.settings,
+  })
+  return Response.json({ settings: parsed.data.settings })
+}

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]
+}