experimental-a2 0.16.1 → 0.17.0

package/examples/playground/app/agent/telemetry.test.ts

@@ -0,0 +1,168 @@
+import { context, trace } from '@opentelemetry/api'
+import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks'
+import {
+  BasicTracerProvider,
+  InMemorySpanExporter,
+  SimpleSpanProcessor,
+} from '@opentelemetry/sdk-trace-base'
+import { simulateReadableStream } from 'ai'
+import { MockLanguageModelV4 } from 'ai/test'
+import { inputs } from 'experimental-a2/ai'
+import { agent } from 'experimental-a2/ai/server'
+import { memory } from 'experimental-a2/store-memory'
+import { otel } from 'experimental-a2/otel'
+import { afterAll, beforeAll, beforeEach, expect, it, vi } from 'vitest'
+import { demoAgent } from './model'
+import { modelTelemetry, traceMetadata } from './telemetry'
+import { traceUrl } from './trace'
+
+const exporter = new InMemorySpanExporter()
+const openSpans = new Set<string>()
+const provider = new BasicTracerProvider({
+  spanProcessors: [
+    new SimpleSpanProcessor(exporter),
+    {
+      onStart: (span) => {
+        openSpans.add(span.spanContext().spanId)
+      },
+      onEnd: (span) => {
+        openSpans.delete(span.spanContext().spanId)
+      },
+      forceFlush: async () => {},
+      shutdown: async () => {},
+    },
+  ],
+})
+beforeAll(() => {
+  context.setGlobalContextManager(
+    new AsyncLocalStorageContextManager().enable(),
+  )
+  trace.setGlobalTracerProvider(provider)
+})
+beforeEach(() => {
+  exporter.reset()
+})
+afterAll(async () => {
+  await provider.shutdown()
+  trace.disable()
+  context.disable()
+})
+
+it.each(['complete', 'stop'] as const)(
+  'correlates standard SDK traces and releases spans when generation ends with %s',
+  async (outcome) => {
+    const server = agent({
+      contract: demoAgent,
+      store: memory(),
+      telemetry: otel(),
+      generation: { telemetry: modelTelemetry() },
+      messageMetadata: traceMetadata,
+      model: new MockLanguageModelV4({
+        modelId: 'test-model',
+        doStream: async ({ abortSignal }) => ({
+          stream:
+            outcome === 'complete'
+              ? simulateReadableStream({
+                  chunks: [
+                    { type: 'stream-start', warnings: [] },
+                    { type: 'text-start', id: 'text' },
+                    { type: 'text-delta', id: 'text', delta: 'PRIVATE_OUTPUT' },
+                    { type: 'text-end', id: 'text' },
+                    {
+                      type: 'finish',
+                      finishReason: { unified: 'stop', raw: 'stop' },
+                      usage: {
+                        inputTokens: {
+                          total: 5,
+                          noCache: 5,
+                          cacheRead: 0,
+                          cacheWrite: 0,
+                        },
+                        outputTokens: { total: 2, text: 2, reasoning: 0 },
+                      },
+                      providerMetadata: {
+                        test: { secret: 'PRIVATE_METADATA' },
+                      },
+                    },
+                  ],
+                })
+              : new ReadableStream({
+                  start(controller) {
+                    controller.enqueue({ type: 'stream-start', warnings: [] })
+                    controller.enqueue({ type: 'text-start', id: 'text' })
+                    controller.enqueue({
+                      type: 'text-delta',
+                      id: 'text',
+                      delta: 'PRIVATE_PARTIAL',
+                    })
+                    abortSignal?.addEventListener(
+                      'abort',
+                      () => controller.close(),
+                      { once: true },
+                    )
+                  },
+                }),
+        }),
+      }),
+    })
+    const session = server.session(`trace-${outcome}`)
+    await session.append(
+      ...inputs.message({
+        id: 'user',
+        role: 'user',
+        parts: [{ type: 'text', text: 'PRIVATE_PROMPT' }],
+      }),
+    )
+    if (outcome === 'stop') {
+      await vi.waitUntil(async () => {
+        const { state } = await session.state(demoAgent.reducer)
+        return state.messages
+          .at(-1)
+          ?.parts.some(
+            (part) => part.type === 'text' && part.text === 'PRIVATE_PARTIAL',
+          )
+      })
+      const { state } = await session.state(demoAgent.reducer)
+      await session.append(
+        ...inputs.stop({
+          turnId: state.active!.turnId,
+        }),
+      )
+    }
+    await vi.waitFor(async () => {
+      const { state } = await session.state(demoAgent.reducer)
+      const view = state
+      expect(view.status).toBe('idle')
+      expect(view.messages.at(-1)?.role).toBe('assistant')
+      expect(openSpans.size).toBe(0)
+    })
+    const { state } = await session.state(demoAgent.reducer)
+    const metadata = state.messages.at(-1)?.metadata
+    const spans = exporter.getFinishedSpans()
+    const generation = spans.find((span) => span.name === 'a2.ai.generate')!
+    expect(metadata).toEqual({ traceId: generation.spanContext().traceId })
+    expect(traceUrl(metadata)).toContain(generation.spanContext().traceId)
+    const models = spans.filter(
+      (span) => span.instrumentationScope.name === 'gen_ai',
+    )
+    expect(models.length).toBeGreaterThan(0)
+    expect(
+      models.every(
+        (span) =>
+          span.spanContext().traceId === generation.spanContext().traceId,
+      ),
+    ).toBe(true)
+    expect(spans.some((span) => span.name.startsWith('a2.model'))).toBe(false)
+    expect(
+      JSON.stringify(
+        spans.map(({ attributes, events }) => ({ attributes, events })),
+      ),
+    ).not.toContain('PRIVATE_')
+  },
+)
+
+it('clears an earlier trace when no active trace is available', () => {
+  expect(
+    traceMetadata({ generationId: 'next', part: { type: 'start' } }),
+  ).toEqual({ traceId: null })
+})

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

@@ -0,0 +1,30 @@
+import { OpenTelemetry } from '@ai-sdk/otel'
+import { isSpanContextValid, trace } from '@opentelemetry/api'
+import type { TelemetryOptions, TextStreamPart, ToolSet } from 'ai'
+
+export function modelTelemetry(): TelemetryOptions {
+  return {
+    recordInputs: false,
+    recordOutputs: false,
+    integrations: [new OpenTelemetry()],
+  }
+}
+
+export function traceMetadata({
+  generationId,
+  part,
+}: {
+  generationId: string
+  part: TextStreamPart<ToolSet>
+}): { traceId: string | null } | undefined {
+  if (part.type !== 'start') return
+  const span = trace.getActiveSpan()
+  span?.setAttribute('a2.generation_id', generationId)
+  const spanContext = span?.spanContext()
+  return {
+    traceId:
+      spanContext && isSpanContextValid(spanContext)
+        ? spanContext.traceId
+        : null,
+  }
+}

package/examples/playground/app/agent/trace.test.ts

@@ -0,0 +1,23 @@
+import { expect, it } from 'vitest'
+import { traceUrl } from './trace'
+
+const traceId = '1234567890abcdef1234567890abcdef'
+
+it('links current and previously recorded trace metadata to Vercel', () => {
+  const current = traceUrl({ traceId })!
+  expect(traceUrl({ timing: { traceId } })).toBe(current)
+  const url = new URL(current)
+  expect(url.origin).toBe('https://vercel.com')
+  expect(url.searchParams.get('search')).toBe(`traceId:${traceId}`)
+})
+
+it('does not reuse an older trace after the latest generation clears it', () => {
+  expect(traceUrl({ traceId: null, timing: { traceId } })).toBeUndefined()
+})
+
+it.each([undefined, null, {}, { traceId: 'https://example.com' }])(
+  'omits the link for invalid metadata %j',
+  (metadata) => {
+    expect(traceUrl(metadata)).toBeUndefined()
+  },
+)

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

@@ -0,0 +1,18 @@
+import { z } from 'zod'
+
+const traceId = z.string().regex(/^[0-9a-f]{32}$/)
+const metadataSchema = z.object({
+  traceId: traceId.nullable().optional(),
+  timing: z.object({ traceId: traceId.optional() }).optional(),
+})
+
+export function traceUrl(metadata: unknown): string | undefined {
+  const parsed = metadataSchema.safeParse(metadata)
+  if (!parsed.success) return
+  const id =
+    parsed.data.traceId === undefined
+      ? parsed.data.timing?.traceId
+      : parsed.data.traceId
+  if (id)
+    return `https://vercel.com/vercel-labs/a2-playground/logs?${new URLSearchParams({ search: `traceId:${id}`, timeline: 'past2Weeks' })}`
+}

package/examples/playground/app/canvas/[canvasId]/canvas-client.tsx

@@ -1,4 +1,5 @@
 'use client'
+import type { StateSnapshot } from 'experimental-a2'
 import type {
   CSSProperties,
   PointerEvent as ReactPointerEvent,
@@ -30,11 +31,11 @@ const ADJECTIVES = ['Bright', 'Calm', 'Curious', 'Kind', 'Swift'] as const
 const ANIMALS = ['Fox', 'Heron', 'Otter', 'Panda', 'Wren'] as const
 const COLORS = ['teal', 'amber', 'blue', 'violet', 'rose'] as const
 const COLOR_HEX: Record<PaletteColor, string> = {
-  teal: '#16836a',
-  amber: '#bd7917',
-  blue: '#3874c8',
-  violet: '#7956b6',
-  rose: '#bd5269',
+  teal: '#3ccf9f',
+  amber: '#f2b54b',
+  blue: '#6ea8fe',
+  violet: '#a98bff',
+  rose: '#ff7a93',
 }
 
 /** The tab's identity: an opaque participant id, plus the name and
@@ -191,13 +192,12 @@ function Smoothed({
 
 export function CanvasSession({
   canvasId,
-  initialState,
-  initialIndex,
+  initial,
 }: {
   canvasId: string
-  initialState: CanvasState
-  initialIndex: number
+  initial: StateSnapshot<CanvasState>
 }): ReactNode {
+  const { state: initialState, index: initialIndex } = initial
   // The server snapshot is null so hydration matches exactly. The
   // browser snapshot then mints one stable tab identity before presence
   // becomes interactive.
@@ -205,8 +205,8 @@ export function CanvasSession({
   return (
     <SessionProvider
       sessionId={canvasId}
-      initialState={initialState}
-      initialIndex={initialIndex}
+      initial={initial}
+
       {...(identity ? { participant: identity.id } : {})}
     >
       <CanvasClient

package/examples/playground/app/canvas/[canvasId]/page.tsx

@@ -15,8 +15,7 @@ export default async function CanvasSessionPage({
   return (
     <CanvasSession
       canvasId={canvasId}
-      initialState={state}
-      initialIndex={index}
+      initial={{ state: state, index: index }}
     />
   )
 }

package/examples/playground/app/canvas/[canvasId]/session.ts

@@ -1,14 +1,14 @@
 'use client'
 import { createClient } from 'experimental-a2/client'
 import { createReact } from 'experimental-a2/react'
-import { canvasReducer } from '../model'
+import { canvas } from '../model'
 
 // One route, two wires: on Vercel (production and previews), pushes
 // and presence ride one WebSocket. Anywhere else, including `next
 // dev` (which drops the upgrade handshake before the route runs) and
 // a local `next start`, the same route serves the SSE + POST pair.
 export const canvasClient = createClient({
-  reducer: canvasReducer,
+  contract: canvas,
   api: process.env.NEXT_PUBLIC_VERCEL_ENV
     ? { type: 'ws', url: '/canvas/events' }
     : '/canvas/events',

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

@@ -5,6 +5,24 @@ const paletteSchema = z.enum(['teal', 'amber', 'blue', 'violet', 'rose'])
 
 export type PaletteColor = z.output<typeof paletteSchema>
 
+const canvasShapeSchema = z.object({
+  id: z.string(),
+  kind: z.enum(['rect', 'ellipse']),
+  x: z.number(),
+  y: z.number(),
+  w: z.number(),
+  h: z.number(),
+  color: paletteSchema,
+  movedAt: z.number(),
+})
+
+const canvasState = z.object({
+  shapes: z.record(z.string(), canvasShapeSchema),
+})
+
+export type CanvasState = z.output<typeof canvasState>
+export type CanvasShape = CanvasState['shapes'][string]
+
 export const canvas = a2.contract({
   name: 'canvas',
   events: {
@@ -32,63 +50,45 @@ export const canvas = a2.contract({
     cursor: z.object({ x: z.number(), y: z.number() }),
     dragging: z.object({ shapeId: z.string(), x: z.number(), y: z.number() }),
   },
-})
 
-const canvasShapeSchema = z.object({
-  id: z.string(),
-  kind: z.enum(['rect', 'ellipse']),
-  x: z.number(),
-  y: z.number(),
-  w: z.number(),
-  h: z.number(),
-  color: paletteSchema,
-  movedAt: z.number(),
-})
-
-const canvasState = z.object({
-  shapes: z.record(z.string(), canvasShapeSchema),
-})
-
-export type CanvasState = z.output<typeof canvasState>
-export type CanvasShape = CanvasState['shapes'][string]
-
-export const canvasReducer = canvas
-  .reducer({
+  reducer: {
     name: 'canvas-scene',
     initialState: { shapes: {} },
     stateSchema: canvasState,
-  })
-  .fold((state, event) => {
-    switch (event.type) {
-      case 'shape.created':
-        return {
-          shapes: {
-            ...state.shapes,
-            [event.payload.id]: { ...event.payload, movedAt: event.index },
-          },
-        }
-      case 'shape.moved': {
-        const shape = state.shapes[event.payload.id]
-        if (!shape) return state
-        return {
-          shapes: {
-            ...state.shapes,
-            [shape.id]: {
-              ...shape,
-              x: event.payload.x,
-              y: event.payload.y,
-              movedAt: event.index,
+    fold: (state, event) => {
+      switch (event.type) {
+        case 'shape.created':
+          return {
+            shapes: {
+              ...state.shapes,
+              [event.payload.id]: { ...event.payload, movedAt: event.index },
             },
-          },
+          }
+        case 'shape.moved': {
+          const shape = state.shapes[event.payload.id]
+          if (!shape) return state
+          return {
+            shapes: {
+              ...state.shapes,
+              [shape.id]: {
+                ...shape,
+                x: event.payload.x,
+                y: event.payload.y,
+                movedAt: event.index,
+              },
+            },
+          }
         }
+        case 'shape.deleted': {
+          if (!state.shapes[event.payload.id]) return state
+          const shapes = { ...state.shapes }
+          delete shapes[event.payload.id]
+          return { shapes }
+        }
+        default:
+          return state
       }
-      case 'shape.deleted': {
-        if (!state.shapes[event.payload.id]) return state
-        const shapes = { ...state.shapes }
-        delete shapes[event.payload.id]
-        return { shapes }
-      }
-      default:
-        return state
-    }
-  })
+    },
+  },
+})
+export const canvasReducer = canvas.reducer

package/examples/playground/app/chat/[chatId]/agent-state.test.ts

@@ -0,0 +1,176 @@
+import { expect, it } from 'vitest'
+import { reduceAIState, type AIState } from 'experimental-a2/ai'
+import { chatAgent, type ChatAgentMessage } from '../model'
+import { agentRetry, agentStreamStatus } from './chat-view'
+
+const generation = {
+  generationId: 'generation',
+  requestId: 'request',
+  responseMessageId: 'response',
+  messageId: 'input',
+  model: 'test/model',
+  attempt: 1,
+}
+
+function running() {
+  let index = 0
+  let state: AIState<ChatAgentMessage> = {
+    ...chatAgent.reducer.initialState,
+    active: {
+      turnId: 'turn',
+      phase: 'generating',
+      input: {
+        id: 'input',
+        generate: true,
+        message: { id: 'input', role: 'user', parts: [] },
+      },
+    },
+  }
+  const append = (event: { type: string; payload: unknown; id?: string }) => {
+    state = reduceAIState(state, { ...event, index: ++index })
+  }
+  append({
+    id: 'request',
+    type: 'ai.generation.requested',
+    payload: { ...generation, reason: 'message' },
+  })
+  append({ type: 'ai.generation.started', payload: generation })
+  const snapshot = () => state
+  return { append, snapshot }
+}
+
+it.each([false, true])(
+  'can retry from a restored snapshot before and after the control commit (step limit: %s)',
+  (stepLimit) => {
+    const log = running()
+    if (stepLimit) {
+      log.append({
+        type: 'ai.generation.completed',
+        payload: { ...generation, finishReason: 'tool-calls' },
+      })
+      log.append({
+        id: 'continuation',
+        type: 'ai.generation.requested',
+        payload: {
+          messageId: 'response',
+          responseMessageId: 'response',
+          reason: 'tool',
+          sourceGenerationId: 'generation',
+        },
+      })
+    }
+    log.append({
+      type: 'ai.generation.failed',
+      payload: {
+        ...generation,
+        error: 'failed',
+        ...(stepLimit ? { stepLimit: true } : {}),
+      },
+    })
+    expect(log.snapshot().status).toBe('failed')
+    expect(log.snapshot().activeResponseMessageId).toBeNull()
+    expect(agentRetry(log.snapshot())).toEqual({
+      messageId: 'input',
+      responseMessageId: 'response',
+    })
+    expect(
+      agentStreamStatus({
+        state: log.snapshot(),
+        responseMessageId: 'response',
+        streaming: false,
+      }),
+    ).toBe('failed')
+
+    log.append({
+      type: 'ai.control.committed',
+      payload: {
+        changes: [],
+        view: {
+          turnId: 'turn',
+          version: 1,
+          phase: 'failed',
+          paused: false,
+          closed: false,
+          requestId: null,
+          requestReason: null,
+          sourceGenerationId: null,
+          responseMessageId: 'response',
+        },
+      },
+    })
+    expect(log.snapshot().activeResponseMessageId).toBe('response')
+    expect(agentRetry(log.snapshot())).toEqual({
+      messageId: 'input',
+      responseMessageId: 'response',
+    })
+  },
+)
+
+it('uses the current response owner after retry rather than an older failure', () => {
+  const log = running()
+  log.append({
+    type: 'ai.generation.failed',
+    payload: { ...generation, error: 'failed' },
+  })
+  const retryGeneration = {
+    ...generation,
+    requestId: 'retry-request',
+    generationId: 'retry-generation',
+  }
+  log.append({
+    id: retryGeneration.requestId,
+    type: 'ai.generation.requested',
+    payload: { ...retryGeneration, reason: 'retry' },
+  })
+  log.append({ type: 'ai.generation.started', payload: retryGeneration })
+  expect(agentRetry(log.snapshot())).toBeNull()
+  expect(
+    agentStreamStatus({
+      state: log.snapshot(),
+      responseMessageId: 'response',
+      streaming: true,
+    }),
+  ).toBe('streaming')
+  log.append({ type: 'ai.generation.completed', payload: retryGeneration })
+  log.append({
+    type: 'ai.message.completed',
+    payload: { messageId: 'response' },
+  })
+  log.append({
+    type: 'ai.generation.failed',
+    payload: { ...generation, error: 'stale failure', superseded: true },
+  })
+  expect(
+    agentStreamStatus({
+      state: log.snapshot(),
+      responseMessageId: 'response',
+      streaming: false,
+    }),
+  ).toBe('completed')
+})
+
+it('shows interruption from the snapshot and does not invent missing retry identifiers', () => {
+  const log = running()
+  log.append({
+    type: 'ai.message.interrupted',
+    payload: { ...generation, messageId: 'response' },
+  })
+  expect(
+    agentStreamStatus({
+      state: log.snapshot(),
+      responseMessageId: 'response',
+      streaming: false,
+    }),
+  ).toBe('interrupted')
+  expect(
+    agentRetry({
+      ...log.snapshot(),
+      status: 'failed',
+      activeProjection: null,
+      activeResponseMessageId: null,
+    }),
+  ).toBeNull()
+  expect(
+    agentRetry({ ...log.snapshot(), status: 'failed', active: null }),
+  ).toBeNull()
+})

package/examples/playground/app/chat/[chatId]/agent-stream-drawer.tsx

@@ -2,6 +2,7 @@
 
 import type { ReactNode } from 'react'
 import { useEffect, useRef, useState } from 'react'
+import { isCompactionPart, type TranscriptMessage } from 'experimental-a2/ai'
 import type { ChatAgentMessage, ChatAgentStep } from '../model'
 import type { AgentStreamSelection } from './chat-view'
 
@@ -36,9 +37,23 @@ function StreamPart({
   part,
   position,
 }: {
-  part: ChatAgentMessage['parts'][number]
+  part: TranscriptMessage<ChatAgentMessage>['parts'][number]
   position: number
 }): ReactNode {
+  if (isCompactionPart(part)) {
+    return (
+      <div className="agent-stream-step" role="separator">
+        <span>
+          {part.data.status === 'completed'
+            ? 'Context compacted'
+            : part.data.status === 'running'
+              ? 'Compacting…'
+              : 'Compaction did not complete'}
+        </span>
+      </div>
+    )
+  }
+
   if (part.type === 'step-start') {
     return (
       <div className="agent-stream-step" role="separator">