experimental-a2 0.16.1 → 0.18.0

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

@@ -1,377 +0,0 @@
-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,
-    sourceGenerationId,
-  }: {
-    requestId?: string
-    reason?: 'message' | 'tool' | 'retry'
-    attempt?: number
-    sourceGenerationId?: string
-  } = {}): Generation => {
-    const generation = {
-      requestId,
-      messageId: reason === 'tool' ? 'assistant' : 'user',
-      responseMessageId: 'assistant',
-      generationId: `generation-${events.length}-${attempt}`,
-      attempt,
-      model: 'test/model',
-    }
-    if (attempt === 1)
-      append({
-        id: requestId,
-        type: 'ai.generation.requested',
-        payload: {
-          messageId: generation.messageId,
-          responseMessageId: 'assistant',
-          reason,
-          ...(sourceGenerationId === undefined ? {} : { sourceGenerationId }),
-        },
-      })
-    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.view(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: 'continuation-request',
-    sourceGenerationId: 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: undefined, 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: undefined })
-})
-
-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('keeps completed steps around committed compactions when retrying a later step', () => {
-  const log = transcript()
-  log.user('user')
-  const first = log.start()
-  log.compact(first)
-  log.text(first, 'Completed step')
-  log.append({
-    type: 'ai.generation.completed',
-    payload: { ...first, finishReason: 'tool-calls' },
-  })
-  const second = log.start({
-    requestId: 'second',
-    reason: 'tool',
-    sourceGenerationId: first.generationId,
-  })
-  log.compact(second)
-  log.text(second, 'Discarded partial step')
-  log.append({
-    type: 'ai.generation.failed',
-    payload: { ...second, error: 'Failed' },
-  })
-  const retry = log.start({ requestId: 'retry', reason: 'retry' })
-  log.text(retry, 'Retried step')
-  expect(labels(log.timeline())).toEqual([
-    'user',
-    `compaction:${first.generationId}`,
-    'Completed step',
-    `compaction:${second.generationId}`,
-    'Retried step',
-  ])
-})
-
-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: 'continuation-request',
-      sourceGenerationId: 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().responseOwners['assistant']?.generationId).toBe(
-    first.generationId,
-  )
-  expect(labels(log.timeline())).toEqual([
-    'user',
-    `compaction:${first.generationId}`,
-    'Kept',
-  ])
-})

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

@@ -1,208 +0,0 @@
-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 completedGenerations = 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
-      case 'ai.generation.completed':
-        completedGenerations.add(event.payload.generationId)
-        break
-    }
-  }
-
-  const contributions = new Map<string, Generation[]>()
-  for (const generation of generations.values()) {
-    const { generationId, responseMessageId, requestId } = generation.payload
-    const ownerId = state.responseOwners[responseMessageId]?.generationId
-    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 = steps.filter((step) =>
-        completedGenerations.has(step.payload.generationId),
-      )
-    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:
-        event.type === 'ai.compaction.completed' || 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]
-}