experimental-a2 0.18.1 → 0.20.0

package/examples/playground/app/a0/chat.test.tsx

@@ -0,0 +1,595 @@
+// @vitest-environment jsdom
+import {
+  act,
+  cleanup,
+  createEvent,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+  within,
+} from '@testing-library/react'
+import { tool } from 'ai'
+import { handlerContext, inputs } from 'experimental-a2/ai'
+import { agent } from 'experimental-a2/ai/server'
+import type { A2Operation } from 'experimental-a2/server'
+import { memory } from 'experimental-a2/store-memory'
+import { afterEach, beforeEach, expect, it, vi } from 'vitest'
+import { z } from 'zod'
+import { say, testModel, toolCall, type ModelStep } from '@/test/model'
+import { authorizer, serve, type Served } from '@/test/serve'
+import { Chat } from './chat'
+import { appAgent, type AppMessage } from './model'
+import { SessionProvider } from './session'
+
+const blobStore = vi.hoisted(() => ({ failure: null as Error | null }))
+
+// The App Router's router only exists inside a Next.js app.
+vi.mock('next/navigation', () => ({
+  useRouter: () => ({ replace: () => {} }),
+  useSearchParams: () => new URLSearchParams(),
+}))
+// Vercel Blob has no local stand-in: an upload lands at its path.
+vi.mock('@vercel/blob/client', () => ({
+  upload: async (pathname: string, file: File) => {
+    if (blobStore.failure) throw blobStore.failure
+    return {
+      url: `https://store.public.blob.vercel-storage.com/${pathname}`,
+      contentType: file.type,
+    }
+  },
+}))
+
+/** Pi's tools, by name, without a sandbox to run in. */
+const tools = {
+  write: tool({
+    inputSchema: z.object({ path: z.string(), content: z.string() }),
+    execute: () => ({ content: [] }),
+  }),
+  edit: tool({
+    inputSchema: z.object({
+      path: z.string(),
+      oldText: z.string(),
+      newText: z.string(),
+    }),
+    execute: () => ({ content: [] }),
+  }),
+  bash: tool({
+    inputSchema: z.object({ command: z.string() }),
+    execute: async () => {
+      const { signal } = handlerContext()
+      await new Promise<void>((resolve, reject) => {
+        void bashDone.promise.then(resolve)
+        signal.addEventListener('abort', () => reject(signal.reason), {
+          once: true,
+        })
+      })
+      return { content: [] }
+    },
+  }),
+}
+
+/** An app's coding agent, on a memory store. */
+const appServer = (model: string) =>
+  agent({ contract: appAgent, store: memory(), model, tools })
+
+let appId: string
+let model: ReturnType<typeof testModel>
+let server: ReturnType<typeof appServer>
+let probe: ReturnType<typeof authorizer>
+let served: Served
+/** Lets the bash call that waits on it finish. */
+let bashDone: PromiseWithResolvers<void>
+
+beforeEach(async () => {
+  appId = crypto.randomUUID()
+  model = testModel()
+  bashDone = Promise.withResolvers()
+  server = appServer(model.id)
+  probe = authorizer()
+  // The app's client follows `/a0/events` on this origin.
+  served = await serve({
+    routes: {
+      '/a0/events': (request) =>
+        server.fetch(request, { authorize: probe.authorize }),
+    },
+  })
+  blobStore.failure = null
+  Object.assign(URL, {
+    createObjectURL: () => 'blob:thumbnail',
+    revokeObjectURL: () => {},
+  })
+  // jsdom lays nothing out, so the transcript has nothing to follow.
+  vi.stubGlobal(
+    'ResizeObserver',
+    class {
+      observe(): void {}
+      disconnect(): void {}
+    },
+  )
+})
+
+afterEach(async () => {
+  cleanup()
+  probe.release()
+  bashDone.resolve()
+  model.close()
+  // The agent's work ends with this test.
+  await vi.waitUntil(async () => (await server.drain(appId)).settled)
+  await served.close()
+  vi.unstubAllGlobals()
+})
+
+const user = (id: string, text: string): AppMessage => ({
+  id,
+  role: 'user',
+  parts: [{ type: 'text', text }],
+})
+
+const chat = async () =>
+  render(
+    <SessionProvider
+      sessionId={appId}
+      initial={await server.session(appId).state()}
+    >
+      <Chat appId={appId} harness="pi" onPrompt={() => {}} autoFocus={false} />
+    </SessionProvider>,
+  )
+
+/** A turn the model is in the middle of: its first step is open. */
+const generating = async (): Promise<{ turnId: string; step: ModelStep }> => {
+  await server
+    .session(appId)
+    .append(...inputs.message(user('u1', 'Build a blog.')))
+  const step = await model.step(0)
+  const { state } = await server.session(appId).state()
+  return { turnId: state.active!.turnId, step }
+}
+
+const write = (text: string): void => {
+  const box = screen.getByRole('textbox', { name: 'Message' })
+  fireEvent.change(box, { target: { value: text } })
+  fireEvent.keyDown(box, { key: 'Enter' })
+}
+
+/** The controls the app's log holds, in order. */
+const controls = async (): Promise<unknown[]> =>
+  (await server.session(appId).history()).flatMap((event) =>
+    event.type === 'ai.control.requested' ? [event.payload] : [],
+  )
+
+const isPush = (operation: A2Operation): boolean =>
+  operation.type === 'push' && operation.sessionId === appId
+
+const blob = (name: string): string =>
+  `https://store.public.blob.vercel-storage.com/a0/${appId}/${name}`
+
+const paused = () => screen.findByText('Paused')
+
+/** The queued message that reads `summary`. */
+const queued = (summary: string) =>
+  within(
+    within(screen.getByRole('list', { name: 'Queued messages' }))
+      .getByText(summary)
+      .closest('li')!,
+  )
+
+const chevron = () => screen.getByRole('button', { name: 'Send mode' })
+
+it('sends a message between turns', async () => {
+  await chat()
+  expect(screen.queryByRole('button', { name: 'Send mode' })).toBeNull()
+  write('Build a blog.')
+  await waitFor(async () =>
+    expect(await controls()).toEqual([
+      expect.objectContaining({
+        action: 'send',
+        message: expect.objectContaining({
+          parts: [{ type: 'text', text: 'Build a blog.' }],
+        }),
+      }),
+    ]),
+  )
+  expect(JSON.stringify((await model.step(0)).options.prompt)).toContain(
+    'Build a blog.',
+  )
+})
+
+it('steers the active turn by default, and queues behind it from the chevron', async () => {
+  const { turnId } = await generating()
+  await chat()
+  expect(screen.getByRole('button', { name: 'Steer' })).toBeTruthy()
+  write('Use a dark theme.')
+  await waitFor(async () =>
+    expect(await controls()).toContainEqual(
+      expect.objectContaining({
+        action: 'steer',
+        turnId,
+        message: expect.objectContaining({
+          parts: [{ type: 'text', text: 'Use a dark theme.' }],
+        }),
+      }),
+    ),
+  )
+
+  fireEvent.click(chevron())
+  expect(
+    screen
+      .getByRole('menuitemradio', { name: 'Steer' })
+      .getAttribute('aria-checked'),
+  ).toBe('true')
+  fireEvent.click(screen.getByRole('menuitemradio', { name: 'Queue' }))
+  expect(screen.getByRole('button', { name: 'Queue' })).toBeTruthy()
+  write('Then add an about page.')
+  await waitFor(async () =>
+    expect(await controls()).toContainEqual(
+      expect.objectContaining({
+        action: 'send',
+        message: expect.objectContaining({
+          parts: [{ type: 'text', text: 'Then add an about page.' }],
+        }),
+      }),
+    ),
+  )
+})
+
+it('pauses, resumes, and stops the turn like a player', async () => {
+  const { turnId } = await generating()
+  await chat()
+  fireEvent.click(screen.getByRole('button', { name: 'Pause' }))
+  await paused()
+  expect(screen.queryByText('Thinking')).toBeNull()
+  fireEvent.click(screen.getByRole('button', { name: 'Resume' }))
+  await waitFor(() => expect(screen.queryByText('Paused')).toBeNull())
+  fireEvent.click(screen.getByRole('button', { name: 'Stop' }))
+  await waitFor(async () =>
+    expect((await controls()).slice(1)).toEqual([
+      { action: 'pause', when: 'now' },
+      { action: 'resume' },
+      { action: 'stop', turnId },
+    ]),
+  )
+  await waitFor(() =>
+    expect(screen.queryByRole('group', { name: 'Turn' })).toBeNull(),
+  )
+})
+
+it('steers with a queued message, or removes it', async () => {
+  const { turnId } = await generating()
+  await server
+    .session(appId)
+    .append(
+      ...inputs.message(user('q1', 'Add tags.')),
+      ...inputs.message(user('q2', 'Add a footer.')),
+    )
+  await chat()
+  fireEvent.click(
+    queued('Add tags.').getByRole('button', { name: 'Steer with it now' }),
+  )
+  fireEvent.click(
+    queued('Add a footer.').getByRole('button', { name: 'Remove' }),
+  )
+  await waitFor(async () =>
+    expect((await controls()).slice(3)).toEqual([
+      { action: 'send-now', inputId: 'q1', turnId },
+      { action: 'remove', inputId: 'q2' },
+    ]),
+  )
+})
+
+it('shows what a tool call writes while it streams, folded once it ran', async () => {
+  const { step } = await generating()
+  const { container } = await chat()
+  const call = (name: string) =>
+    container.querySelector(`[data-tool="${name}"]`)
+  step.write(
+    { type: 'tool-input-start', id: 'w1', toolName: 'write' },
+    {
+      type: 'tool-input-delta',
+      id: 'w1',
+      delta: '{"path":"app/page.tsx","content":"export default function',
+    },
+  )
+  await waitFor(() =>
+    expect(call('write')?.querySelector('.a0-tool-text')?.textContent).toBe(
+      'export default function',
+    ),
+  )
+  expect(call('write')?.querySelector('details')?.hasAttribute('open')).toBe(
+    true,
+  )
+
+  step.write(
+    { type: 'tool-input-delta', id: 'w1', delta: ' Page() {}"}' },
+    { type: 'tool-input-end', id: 'w1' },
+    {
+      type: 'tool-call',
+      toolCallId: 'w1',
+      toolName: 'write',
+      input: JSON.stringify({
+        path: 'app/page.tsx',
+        content: 'export default function Page() {}',
+      }),
+    },
+  )
+  step.finish('tool-calls')
+  await waitFor(() =>
+    expect(call('write')?.querySelector('details')?.hasAttribute('open')).toBe(
+      false,
+    ),
+  )
+
+  // A one-line command fits its step's line.
+  ;(await model.step(1)).write(
+    ...toolCall({ id: 'b1', name: 'bash', chunks: ['{"command":"ls"}'] }),
+  )
+  await waitFor(() => expect(call('bash')?.textContent).toContain('ls'))
+  expect(call('bash')?.classList.contains('a0-step')).toBe(true)
+  expect(call('bash')?.querySelector('details')).toBeNull()
+})
+
+it('closes the send-mode menu from its chevron, and drops it when the turn ends', async () => {
+  const { step } = await generating()
+  await chat()
+  fireEvent.click(chevron())
+  // A press on the chevron is its own, not one outside the menu.
+  fireEvent.pointerDown(chevron())
+  expect(screen.getAllByRole('menuitemradio')).toHaveLength(2)
+  fireEvent.click(chevron())
+  expect(screen.queryByRole('menuitemradio')).toBeNull()
+
+  fireEvent.click(chevron())
+  step.write(...say({ id: 'done', text: 'Done.' }))
+  step.finish()
+  await waitFor(() =>
+    expect(screen.queryByRole('button', { name: 'Send mode' })).toBeNull(),
+  )
+  await server
+    .session(appId)
+    .append(...inputs.message(user('u2', 'Add a footer.')))
+  await screen.findByRole('button', { name: 'Send mode' })
+  expect(screen.queryByRole('menuitemradio')).toBeNull()
+})
+
+it('ends the pause with a stop, so the next message goes out', async () => {
+  const { turnId } = await generating()
+  await chat()
+  fireEvent.click(screen.getByRole('button', { name: 'Pause' }))
+  await paused()
+  fireEvent.click(screen.getByRole('button', { name: 'Stop' }))
+  await waitFor(async () =>
+    expect((await controls()).slice(2)).toEqual([
+      { action: 'stop', turnId },
+      { action: 'resume' },
+    ]),
+  )
+  await waitFor(() => expect(screen.queryByText('Paused')).toBeNull())
+  write('Start over.')
+  const next = await model.step(model.steps.length)
+  expect(JSON.stringify(next.options.prompt)).toContain('Start over.')
+})
+
+it('offers to steer with a queued message only while a turn runs', async () => {
+  await chat()
+  // The first message is on its way, and its turn has no id yet.
+  const sending = probe.hold(isPush)
+  write('Build a blog.')
+  await sending.arrived
+  write('Add tags.')
+  const queue = await screen.findByRole('list', { name: 'Queued messages' })
+  expect(queue.textContent).toContain('Add tags.')
+  expect(screen.queryByRole('button', { name: 'Steer with it now' })).toBeNull()
+  expect(screen.getByRole('button', { name: 'Remove' })).toBeTruthy()
+
+  sending.allow()
+  await screen.findByRole('button', { name: 'Steer with it now' })
+})
+
+it('folds a call that waits on a paused turn, and shows edits as a diff', async () => {
+  const { step } = await generating()
+  const { container } = await chat()
+  const call = (name: string) =>
+    container.querySelector(`[data-tool="${name}"]`)
+  step.write(
+    ...toolCall({
+      id: 'b1',
+      name: 'bash',
+      chunks: [JSON.stringify({ command: 'cat > a.txt <<EOF\nhi\nEOF' })],
+    }),
+  )
+  step.finish('tool-calls')
+  await waitFor(() =>
+    expect(call('bash')?.querySelector('.a0-spin')).toBeTruthy(),
+  )
+  expect(call('bash')?.querySelector('details')?.hasAttribute('open')).toBe(
+    true,
+  )
+  fireEvent.click(screen.getByRole('button', { name: 'Pause' }))
+  // Paused, not still pausing: the aborted call has settled in the log,
+  // neither done nor failed, so it runs again on Resume.
+  await waitFor(async () =>
+    expect((await server.session(appId).state()).state.active?.phase).toBe(
+      'paused',
+    ),
+  )
+  const { state } = await server.session(appId).state()
+  expect(
+    state.messages
+      .flatMap((message) => message.parts)
+      .find((part) => 'toolCallId' in part && part.toolCallId === 'b1'),
+  ).toMatchObject({ state: 'input-available' })
+  await paused()
+  expect(call('bash')?.querySelector('details')?.hasAttribute('open')).toBe(
+    false,
+  )
+  expect(call('bash')?.querySelector('.a0-spin')).toBeNull()
+  expect(call('bash')?.querySelector('.a0-step-done')).toBeNull()
+
+  fireEvent.click(screen.getByRole('button', { name: 'Resume' }))
+  bashDone.resolve()
+  const next = await model.step(model.steps.length)
+  // Pi's older edit input: the edit's fields on the call itself.
+  next.write(
+    { type: 'tool-input-start', id: 'e1', toolName: 'edit' },
+    {
+      type: 'tool-input-delta',
+      id: 'e1',
+      delta: '{"path":"a.txt","oldText":"hi","newText":"hello"',
+    },
+    // Nothing written yet: the step's one line, with no empty box.
+    { type: 'tool-input-start', id: 'w1', toolName: 'write' },
+    {
+      type: 'tool-input-delta',
+      id: 'w1',
+      delta: '{"path":"a.txt","content":""',
+    },
+  )
+  await waitFor(() =>
+    expect(call('edit')?.querySelector('.a0-tool-text')?.textContent).toBe(
+      '- hi\n+ hello',
+    ),
+  )
+  await waitFor(() => expect(call('write')).toBeTruthy())
+  expect(call('write')?.classList.contains('a0-step')).toBe(true)
+})
+
+const dropOn = (target: Element, ...files: File[]): void => {
+  const dataTransfer = { files, types: ['Files'], dropEffect: 'none' }
+  fireEvent.dragEnter(target, { dataTransfer })
+  fireEvent.dragOver(target, { dataTransfer })
+  fireEvent.drop(target, { dataTransfer })
+}
+
+const shot = () => new File(['png'], 'shot.png', { type: 'image/png' })
+
+/** jsdom has no DragEvent, so the event carries where the files went itself. */
+const leave = (from: Element, to: Element | null): void => {
+  const event = createEvent.dragLeave(from)
+  Object.defineProperty(event, 'relatedTarget', { value: to })
+  fireEvent(from, event)
+}
+
+/** Every attached file has finished uploading. */
+const uploaded = () =>
+  waitFor(() =>
+    expect(
+      screen
+        .getByRole('list', { name: 'Attachments' })
+        .querySelector('[aria-label="Uploading"]'),
+    ).toBeNull(),
+  )
+
+it('sends images dropped anywhere on the chat with the next message', async () => {
+  await chat()
+  const transcript = screen.getByRole('list', { name: 'Conversation' })
+  fireEvent.dragEnter(transcript, {
+    dataTransfer: { files: [], types: ['Files'] },
+  })
+  expect(screen.getByText('Drop to attach')).toBeTruthy()
+  // Into the composer, still on the chat, then out of the window.
+  leave(transcript, screen.getByRole('textbox', { name: 'Message' }))
+  expect(screen.getByText('Drop to attach')).toBeTruthy()
+  leave(transcript, null)
+  expect(screen.queryByText('Drop to attach')).toBeNull()
+
+  dropOn(transcript, shot())
+  const attached = screen.getByRole('list', { name: 'Attachments' })
+  expect(attached.querySelector('[aria-label="Uploading"]')).toBeTruthy()
+  await uploaded()
+
+  write('Make it look like this.')
+  expect(screen.queryByRole('list', { name: 'Attachments' })).toBeNull()
+  await waitFor(async () =>
+    expect(
+      (await server.session(appId).state()).state.messages[0]?.parts,
+    ).toEqual([
+      {
+        type: 'file',
+        mediaType: 'image/png',
+        filename: 'shot.png',
+        url: blob('shot.png'),
+      },
+      { type: 'text', text: 'Make it look like this.' },
+    ]),
+  )
+})
+
+it("puts a file that isn't an image into the message as its link", async () => {
+  await chat()
+  const box = screen.getByRole('textbox', {
+    name: 'Message',
+  }) as HTMLTextAreaElement
+  fireEvent.change(box, { target: { value: 'Use this copy.' } })
+  dropOn(box, new File(['%PDF'], 'copy.pdf', { type: 'application/pdf' }))
+  await waitFor(() =>
+    expect(box.value).toBe(`Use this copy.\ncopy.pdf: ${blob('copy.pdf')}`),
+  )
+  expect(screen.queryByRole('list', { name: 'Attachments' })).toBeNull()
+})
+
+it('shows a sent link as its file, and a sent image as itself', async () => {
+  await server.session(appId).append(
+    ...inputs.seed({
+      id: 'u1',
+      role: 'user',
+      parts: [
+        {
+          type: 'file',
+          mediaType: 'image/png',
+          filename: 'shot.png',
+          url: blob('shot.png'),
+        },
+        {
+          type: 'text',
+          text: `Use this copy.\ncopy.pdf: ${blob('copy.pdf')}`,
+        },
+      ],
+    }),
+  )
+  await chat()
+  expect(screen.getByRole('img', { name: 'shot.png' })).toBeTruthy()
+  expect(
+    screen.getByRole('link', { name: 'copy.pdf' }).getAttribute('href'),
+  ).toBe(blob('copy.pdf'))
+})
+
+it('says which upload failed, and drops the file', async () => {
+  blobStore.failure = new Error('Vercel Blob: Failed to upload')
+  await chat()
+  dropOn(screen.getByRole('list', { name: 'Conversation' }), shot())
+  expect((await screen.findByRole('alert')).textContent).toBe(
+    "Couldn't upload shot.png: Vercel Blob: Failed to upload",
+  )
+  expect(screen.queryByRole('list', { name: 'Attachments' })).toBeNull()
+})
+
+it('attaches a pasted image instead of pasting it as text', async () => {
+  await chat()
+  const box = screen.getByRole('textbox', {
+    name: 'Message',
+  }) as HTMLTextAreaElement
+  fireEvent.paste(box, { clipboardData: { files: [shot()] } })
+  expect(screen.getByRole('list', { name: 'Attachments' })).toBeTruthy()
+  await uploaded()
+  expect(box.value).toBe('')
+})
+
+it('keeps the attachments of a send that failed, beside any attached since', async () => {
+  await chat()
+  const transcript = screen.getByRole('list', { name: 'Conversation' })
+  dropOn(transcript, shot())
+  await uploaded()
+  const sending = probe.hold(isPush)
+  write('Like this.')
+  dropOn(transcript, new File(['png'], 'more.png', { type: 'image/png' }))
+  await uploaded()
+  await sending.arrived
+  await act(async () => sending.deny())
+  expect((await screen.findByRole('alert')).textContent).toBe(
+    'the operation is not authorized',
+  )
+  expect(
+    screen.getByRole('list', { name: 'Attachments' }).querySelectorAll('li'),
+  ).toHaveLength(2)
+  expect(await controls()).toEqual([])
+})