experimental-a2 0.15.0 → 0.16.0

package/examples/playground/app/agent/[agentId]/agent-client.tsx

@@ -6,7 +6,11 @@ import type {
 } from 'react'
 import { useEffect, useRef, useState } from 'react'
 import type { UIMessage } from 'ai'
-import { inputs, type ControlCommand } from 'experimental-a2/ai'
+import {
+  inputs,
+  type ControlCommand,
+  type PendingApproval,
+} from 'experimental-a2/ai'
 import { ActivityFeed } from '@/app/components/activity-feed'
 import { ConnectionPill } from '@/app/components/connection-pill'
 import { useSession } from '../session'
@@ -215,20 +219,22 @@ export function AgentClient({ agentId }: { agentId: string }): ReactNode {
     }
   }
 
-  const answerApproval = async (
-    messageId: string,
-    generationId: string,
-    approvalId: string,
-    approved: boolean,
-  ): Promise<void> => {
+  const answerApproval = async ({
+    approval,
+    approved,
+  }: {
+    approval: PendingApproval
+    approved: boolean
+  }): Promise<void> => {
     setResponding(true)
     setError(null)
     try {
       await push(
         ...inputs.approval({
-          messageId,
-          generationId,
-          approvalId,
+          requestId: approval.requestId,
+          messageId: approval.messageId,
+          generationId: approval.generationId,
+          approvalId: approval.approvalId,
           approved,
           ...(!approved ? { reason: 'operator denied the restart' } : {}),
         }),
@@ -349,12 +355,7 @@ export function AgentClient({ agentId }: { agentId: string }): ReactNode {
                 type="button"
                 disabled={responding}
                 onClick={() =>
-                  void answerApproval(
-                    approval.messageId,
-                    approval.generationId,
-                    approval.approvalId,
-                    true,
-                  )
+                  void answerApproval({ approval, approved: true })
                 }
               >
                 Approve restart
@@ -364,12 +365,7 @@ export function AgentClient({ agentId }: { agentId: string }): ReactNode {
                 className="secondary"
                 disabled={responding}
                 onClick={() =>
-                  void answerApproval(
-                    approval.messageId,
-                    approval.generationId,
-                    approval.approvalId,
-                    false,
-                  )
+                  void answerApproval({ approval, approved: false })
                 }
               >
                 Deny
@@ -387,10 +383,6 @@ export function AgentClient({ agentId }: { agentId: string }): ReactNode {
         <AgentQueue
           items={inbox.items}
           paused={inbox.paused}
-          pendingEdits={(state.pendingQueueCommands ?? []).flatMap(
-            ({ command: pending }) =>
-              pending.action === 'edit' ? [pending.message.id] : [],
-          )}
           remove={async (inputId) => {
             await command(inputs.queue.remove({ inputId }))
           }}
@@ -402,8 +394,8 @@ export function AgentClient({ agentId }: { agentId: string }): ReactNode {
               }),
             )
           }}
-          edit={async (message, expectedRevision) => {
-            await push(...inputs.queue.edit({ message, expectedRevision }))
+          edit={async (message) => {
+            await push(...inputs.queue.edit({ message }))
           }}
           move={(inputId, beforeId) =>
             void command(inputs.queue.move({ inputId, beforeId }))

package/examples/playground/app/agent/[agentId]/agent-queue.test.tsx

@@ -0,0 +1,123 @@
+// @vitest-environment jsdom
+import { act, cleanup, fireEvent, render, screen } from '@testing-library/react'
+import { afterEach, expect, it, vi } from 'vitest'
+import type { UIMessage } from 'ai'
+import { AgentQueue } from './agent-queue'
+
+afterEach(cleanup)
+
+type Edit = (message: UIMessage) => Promise<void>
+
+const queue = (edit: Edit) =>
+  render(
+    <AgentQueue
+      items={[
+        {
+          id: 'queued',
+          message: {
+            id: 'queued',
+            role: 'user',
+            parts: [{ type: 'text', text: 'Original' }],
+          },
+          generate: true,
+        },
+      ]}
+      paused
+      edit={edit}
+      move={vi.fn<() => void>()}
+      remove={vi.fn<() => Promise<void>>().mockResolvedValue(undefined)}
+      sendNow={vi.fn<() => Promise<void>>().mockResolvedValue(undefined)}
+    />,
+  )
+
+const open = (): void => {
+  fireEvent.click(screen.getByRole('button', { name: 'Edit queued message 1' }))
+}
+const save = (text: string): void => {
+  fireEvent.change(screen.getByRole('textbox'), { target: { value: text } })
+  fireEvent.click(screen.getByRole('button', { name: 'Save' }))
+}
+
+it('closes immediately and permits another edit while saves remain in flight', async () => {
+  const first = Promise.withResolvers<void>()
+  const second = Promise.withResolvers<void>()
+  const edit = vi
+    .fn<Edit>()
+    .mockImplementationOnce(() => first.promise)
+    .mockImplementationOnce(() => second.promise)
+  queue(edit)
+  open()
+  save('First')
+  expect(screen.queryByRole('textbox')).toBeNull()
+  expect(
+    (
+      screen.getByRole('button', {
+        name: 'Edit queued message 1',
+      }) as HTMLButtonElement
+    ).disabled,
+  ).toBe(false)
+  expect(
+    (
+      screen.getByRole('button', {
+        name: 'Send queued message 1 after current step',
+      }) as HTMLButtonElement
+    ).disabled,
+  ).toBe(false)
+  open()
+  save('Second')
+  expect(edit.mock.calls.map(([message]) => message.parts)).toEqual([
+    [{ type: 'text', text: 'First' }],
+    [{ type: 'text', text: 'Second' }],
+  ])
+  await act(async () => {
+    first.reject(new Error('Earlier save failed'))
+    await first.promise.catch(() => {})
+  })
+  expect(screen.queryByRole('textbox')).toBeNull()
+  expect(screen.queryByRole('alert')).toBeNull()
+  await act(async () => {
+    second.resolve()
+    await second.promise
+  })
+})
+
+it('restores the latest submitted draft on transport failure', async () => {
+  const pending = Promise.withResolvers<void>()
+  queue(() => pending.promise)
+  open()
+  save('Keep this draft')
+  expect(screen.queryByRole('textbox')).toBeNull()
+  await act(async () => {
+    pending.reject(new Error('Save failed'))
+    await pending.promise.catch(() => {})
+  })
+  expect((screen.getByRole('textbox') as HTMLTextAreaElement).value).toBe(
+    'Keep this draft',
+  )
+  expect(screen.getByRole('alert').textContent).toBe('Save failed')
+})
+
+it.each(['open', 'cancel'] as const)(
+  'does not restore an old failed save after a newer draft is %s',
+  async (action) => {
+    const pending = Promise.withResolvers<void>()
+    queue(() => pending.promise)
+    open()
+    save('Old draft')
+    open()
+    fireEvent.change(screen.getByRole('textbox'), {
+      target: { value: 'New draft' },
+    })
+    if (action === 'cancel')
+      fireEvent.click(screen.getByRole('button', { name: 'Cancel' }))
+    await act(async () => {
+      pending.reject(new Error('Old save failed'))
+      await pending.promise.catch(() => {})
+    })
+    expect(
+      (screen.queryByRole('textbox') as HTMLTextAreaElement | null)?.value ??
+        null,
+    ).toBe(action === 'open' ? 'New draft' : null)
+    expect(screen.queryByRole('alert')).toBeNull()
+  },
+)

package/examples/playground/app/agent/[agentId]/agent-queue.tsx

@@ -7,8 +7,7 @@ import type { InboxItem } from 'experimental-a2/ai'
 type QueueProps = {
   items: InboxItem[]
   paused: boolean
-  pendingEdits: readonly string[]
-  edit(message: UIMessage, revision: number): Promise<void>
+  edit(message: UIMessage): Promise<void>
   move(inputId: string, beforeId: string | null): void
   remove(inputId: string): Promise<void>
   sendNow(inputId: string): Promise<void>
@@ -50,7 +49,6 @@ function QueuedInput({
   remove,
   sendNow,
   paused,
-  pendingEdit,
 }: {
   item: InboxItem
   position: number
@@ -61,26 +59,25 @@ function QueuedInput({
   remove: QueueProps['remove']
   sendNow: QueueProps['sendNow']
   paused: boolean
-  pendingEdit: boolean
 }): ReactNode {
-  const [draft, setDraft] = useState<{ text: string; revision: number } | null>(
-    null,
-  )
+  const [draft, setDraft] = useState<string | null>(null)
+  const editAttempt = useRef(0)
   const [pending, setPending] = useState(false)
   const [error, setError] = useState<string | null>(null)
   const editButton = useRef<HTMLButtonElement>(null)
   const restoreFocus = useRef(false)
   useEffect(() => {
-    if (draft === null && !pending && !pendingEdit && restoreFocus.current) {
+    if (draft === null && !pending && restoreFocus.current) {
       if (document.activeElement === document.body) editButton.current?.focus()
       restoreFocus.current = false
     }
-  }, [draft, pending, pendingEdit])
+  }, [draft, pending])
   const text = item.message.parts
     .filter((part) => part.type === 'text')
     .map((part) => part.text)
     .join('\n')
   const cancel = (): void => {
+    editAttempt.current += 1
     restoreFocus.current = true
     setDraft(null)
     setError(null)
@@ -98,17 +95,23 @@ function QueuedInput({
     }
   }
   const save = async (): Promise<void> => {
-    if (!draft || pendingEdit || !draft.text.trim()) return
-    await perform(async () => {
-      await edit(
-        { ...item.message, parts: [{ type: 'text', text: draft.text }] },
-        draft.revision,
-      )
-      cancel()
-    })
+    if (draft === null || !draft.trim()) return
+    const submitted = draft
+    cancel()
+    const attempt = editAttempt.current
+    try {
+      await edit({
+        ...item.message,
+        parts: [{ type: 'text', text: submitted }],
+      })
+    } catch (failure) {
+      if (editAttempt.current !== attempt) return
+      setDraft(submitted)
+      setError(failure instanceof Error ? failure.message : String(failure))
+    }
   }
   return (
-    <li className={`agent-queue-item${draft ? ' editing' : ''}`}>
+    <li className={`agent-queue-item${draft !== null ? ' editing' : ''}`}>
       <div className="agent-queue-row">
         <span
           className="agent-queue-number"
@@ -129,12 +132,7 @@ function QueuedInput({
                 : 'Send after current step'
             }
             aria-label={`Send queued message ${position + 1} after current step`}
-            disabled={
-              pending ||
-              pendingEdit ||
-              draft !== null ||
-              item.message.role !== 'user'
-            }
+            disabled={pending || draft !== null || item.message.role !== 'user'}
             onClick={() => void perform(() => sendNow(item.id))}
           >
             <QueueIcon name="send" />
@@ -142,15 +140,18 @@ function QueuedInput({
           <button
             ref={editButton}
             type="button"
-            title={
-              pendingEdit ? 'Waiting for edit confirmation' : 'Edit message'
-            }
+            title="Edit message"
             aria-label={`Edit queued message ${position + 1}`}
             aria-expanded={draft !== null}
-            disabled={pending || pendingEdit}
-            onClick={() =>
-              draft ? cancel() : setDraft({ text, revision: item.revision })
-            }
+            disabled={pending}
+            onClick={() => {
+              if (draft !== null) cancel()
+              else {
+                editAttempt.current += 1
+                setDraft(text)
+                setError(null)
+              }
+            }}
           >
             <QueueIcon name="edit" />
           </button>
@@ -189,7 +190,7 @@ function QueuedInput({
           {error}
         </p>
       ) : null}
-      {draft ? (
+      {draft !== null ? (
         <form
           className="agent-queue-editor"
           onSubmit={(event) => {
@@ -201,11 +202,9 @@ function QueuedInput({
             autoFocus
             aria-label={`Edit text for queued message ${position + 1}`}
             rows={2}
-            value={draft.text}
+            value={draft}
             disabled={pending}
-            onChange={(event) =>
-              setDraft({ ...draft, text: event.target.value })
-            }
+            onChange={(event) => setDraft(event.target.value)}
             onKeyDown={(event) => {
               if (event.nativeEvent.isComposing) return
               if (event.key === 'Escape' && !pending) {
@@ -228,11 +227,8 @@ function QueuedInput({
             >
               Cancel
             </button>
-            <button
-              type="submit"
-              disabled={pending || pendingEdit || !draft.text.trim()}
-            >
-              {pending ? 'Saving…' : 'Save'}
+            <button type="submit" disabled={pending || !draft.trim()}>
+              Save
             </button>
           </div>
         </form>
@@ -244,7 +240,6 @@ function QueuedInput({
 export function AgentQueue({
   items,
   paused,
-  pendingEdits,
   edit,
   move,
   remove,
@@ -272,7 +267,6 @@ export function AgentQueue({
             remove={remove}
             sendNow={sendNow}
             paused={paused}
-            pendingEdit={pendingEdits.includes(item.id)}
             moveUp={() => {
               const previous = items[position - 1]
               if (previous) move(item.id, previous.id)

package/examples/playground/app/agent/[agentId]/compaction-panel.tsx

@@ -183,6 +183,11 @@ export function CompactionPanel({
   const completed = events
     .filter((event) => event.type === 'ai.compaction.completed')
     .toReversed()
+  const requested = new Map(
+    events
+      .filter((event) => event.type === 'ai.compaction.requested')
+      .map((event) => [event.payload.generationId, event.payload]),
+  )
   const compacting = state.compaction?.status === 'running'
 
   return (
@@ -232,10 +237,15 @@ export function CompactionPanel({
             <dd>{tokens(lastStep?.payload.usage?.inputTokens)}</dd>
           </div>
           <div>
-            <dt>Last step input · estimated</dt>
+            <dt>Last step input · raw estimate</dt>
             <dd>{tokens(lastStep?.payload.inputTokenEstimate)}</dd>
           </div>
         </dl>
+        <p className="hint">
+          The raw estimate excludes provider metadata. Compaction adjusts it
+          using the previous step’s reported usage when the model matches and
+          the prompt has not shrunk.
+        </p>
         {limits === undefined ? (
           <p className="hint">
             {metadata?.status === 'unavailable'
@@ -262,29 +272,39 @@ export function CompactionPanel({
             </p>
           ) : (
             <ol>
-              {completed.map((event) => (
-                <li key={event.id}>
-                  <div className="compaction-marker">
-                    <strong>Context compacted · #{event.index}</strong>
-                    <time dateTime={event.createdAt.toISOString()}>
-                      {formatActivityTime(event.createdAt)}
-                    </time>
-                  </div>
-                  <p className="hint">
-                    {event.payload.throughIndex === undefined
-                      ? `Through message ${event.payload.throughMessageId}`
-                      : `Summarized through event #${event.payload.throughIndex}`}{' '}
-                    · {tokens(event.payload.usage?.inputTokens)} in ·{' '}
-                    {tokens(event.payload.usage?.outputTokens)} out
-                  </p>
-                  {event.payload.summary ? (
-                    <details>
-                      <summary>View summary</summary>
-                      <pre>{event.payload.summary}</pre>
-                    </details>
-                  ) : null}
-                </li>
-              ))}
+              {completed.map((event) => {
+                const trigger = requested.get(event.payload.generationId)
+                return (
+                  <li key={event.id}>
+                    <div className="compaction-marker">
+                      <strong>Context compacted · #{event.index}</strong>
+                      <time dateTime={event.createdAt.toISOString()}>
+                        {formatActivityTime(event.createdAt)}
+                      </time>
+                    </div>
+                    <p className="hint">
+                      {event.payload.throughIndex === undefined
+                        ? `Through message ${event.payload.throughMessageId}`
+                        : `Summarized through event #${event.payload.throughIndex}`}{' '}
+                      · {tokens(event.payload.usage?.inputTokens)} in ·{' '}
+                      {tokens(event.payload.usage?.outputTokens)} out
+                    </p>
+                    {trigger?.inputTokens === undefined ? null : (
+                      <p className="hint">
+                        Trigger: {tokens(trigger.inputTokens)} against a{' '}
+                        {tokens(trigger.thresholdTokens)} threshold. Raw
+                        estimate: {tokens(trigger.inputTokenEstimate)}.
+                      </p>
+                    )}
+                    {event.payload.summary ? (
+                      <details>
+                        <summary>View summary</summary>
+                        <pre>{event.payload.summary}</pre>
+                      </details>
+                    ) : null}
+                  </li>
+                )
+              })}
             </ol>
           )}
         </details>

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

@@ -37,16 +37,18 @@ function transcript() {
     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: `${requestId}:generation:${attempt}`,
+      generationId: `generation-${events.length}-${attempt}`,
       attempt,
       model: 'test/model',
     }
@@ -58,6 +60,7 @@ function transcript() {
           messageId: generation.messageId,
           responseMessageId: 'assistant',
           reason,
+          ...(sourceGenerationId === undefined ? {} : { sourceGenerationId }),
         },
       })
     append({ type: 'ai.generation.started', payload: generation })
@@ -94,7 +97,10 @@ function transcript() {
       })
   }
   const state = () =>
-    events.reduce(reduceAIState, demoAgent.reducer.initialState)
+    events.reduce(
+      reduceAIState,
+      demoAgent.reducer.view(demoAgent.reducer.initialState),
+    )
   const timeline = () => compactionTimeline({ state: state(), events })
   return {
     events,
@@ -168,7 +174,8 @@ it('places compaction between tool steps without moving a queued user into the r
   })
   log.user('queued')
   const second = log.start({
-    requestId: `ai.generate:tools:${first.generationId}`,
+    requestId: 'continuation-request',
+    sourceGenerationId: first.generationId,
     reason: 'tool',
   })
   log.compact(second, false)
@@ -230,7 +237,7 @@ it('preserves compaction history but drops superseded attempt content', () => {
     log
       .timeline()
       .find((entry) => entry.id === `compaction:${first.generationId}`),
-  ).toMatchObject({ outcome: 'superseded', running: false })
+  ).toMatchObject({ outcome: undefined, running: false })
 })
 
 it('keeps old checkpoints before an explicit retry without restoring discarded parts', () => {
@@ -256,7 +263,7 @@ it('keeps old checkpoints before an explicit retry without restoring discarded p
     log
       .timeline()
       .find((entry) => entry.id === `compaction:${first.generationId}`),
-  ).toMatchObject({ outcome: 'failed' })
+  ).toMatchObject({ outcome: undefined })
 })
 
 it('does not leave failed compaction spinning', () => {
@@ -277,6 +284,38 @@ it('does not leave failed compaction spinning', () => {
   })
 })
 
+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) => {
@@ -289,7 +328,8 @@ it.each([false, true])(
       payload: { ...first, finishReason: 'tool-calls' },
     })
     const second = log.start({
-      requestId: `ai.generate:tools:${first.generationId}`,
+      requestId: 'continuation-request',
+      sourceGenerationId: first.generationId,
       reason: 'tool',
     })
     log.compact(second)
@@ -326,7 +366,7 @@ it('ignores a rejected later generation when mapping canonical parts', () => {
   })
   const rejected = log.start({ requestId: 'stale-retry', reason: 'retry' })
   log.text(rejected, 'Ignored')
-  expect(log.state().responseGenerationIds['assistant']).toBe(
+  expect(log.state().responseOwners['assistant']?.generationId).toBe(
     first.generationId,
   )
   expect(labels(log.timeline())).toEqual([

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

@@ -31,6 +31,7 @@ export function compactionTimeline({
   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) {
@@ -72,13 +73,16 @@ export function compactionTimeline({
       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.responseGenerationIds[responseMessageId]
+    const ownerId = state.responseOwners[responseMessageId]?.generationId
     const owner = ownerId === undefined ? undefined : generations.get(ownerId)
     if (
       owner === undefined ||
@@ -87,7 +91,10 @@ export function compactionTimeline({
     )
       continue
     let steps = contributions.get(responseMessageId) ?? []
-    if (requests.get(requestId) === 'retry') steps = []
+    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)
@@ -122,7 +129,10 @@ export function compactionTimeline({
       running:
         state.compaction?.generationId === generationId &&
         state.compaction.status === 'running',
-      outcome: terminal === 'completed' ? undefined : terminal,
+      outcome:
+        event.type === 'ai.compaction.completed' || terminal === 'completed'
+          ? undefined
+          : terminal,
     }
     const responseId = generations.get(generationId)?.payload.responseMessageId
     const anchor =

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

@@ -376,11 +376,7 @@ export function ChatClient({ chatId }: { chatId: string }): ReactNode {
         )
       : undefined
   const failedRootMessageId = failedGeneration
-    ? state.messages.find(
-        (message) =>
-          `${message.id}:assistant` ===
-          failedGeneration.payload.responseMessageId,
-      )?.id
+    ? state.ai.active?.input.id
     : undefined
 
   const sendTypingPing = useCallback((): void => {