experimental-a2 0.19.0 → 0.20.0

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

@@ -1,136 +1,218 @@
 // @vitest-environment jsdom
 import {
+  act,
   cleanup,
   createEvent,
   fireEvent,
   render,
   screen,
   waitFor,
+  within,
 } from '@testing-library/react'
-import type { ReactElement } from '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 type { AppMessage } from './model'
-import { defaultModelChoice } from './model'
-
-const push = vi.fn<(...events: unknown[]) => Promise<unknown[]>>()
-let state: Record<string, unknown>
-const appId = '6f1c2a54-9a3e-4c51-8f0e-2b7d8c9e1a00'
-const upload =
-  vi.fn<
-    (
-      pathname: string,
-      file: File,
-      options: object,
-    ) => Promise<Record<string, string>>
-  >()
+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: vi.fn<(href: string) => void>() }),
+  useRouter: () => ({ replace: () => {} }),
   useSearchParams: () => new URLSearchParams(),
 }))
-vi.mock('./session', () => ({ useSession: () => ({ state, push }) }))
-vi.mock('@vercel/blob/client', () => ({ upload }))
+// 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,
+    }
+  },
+}))
 
-const { Chat } = await import('./chat')
+/** 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: [] }
+    },
+  }),
+}
 
-const user = (id: string, text: string): AppMessage => ({
-  id,
-  role: 'user',
-  parts: [{ type: 'text', text }],
+/** 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 {}
+    },
+  )
 })
 
-const idle = () => ({
-  modelChoice: defaultModelChoice,
-  messages: [] as AppMessage[],
-  starting: false,
-  status: 'idle',
-  active: null,
-  inbox: { paused: false, items: [] as unknown[] },
-  error: null,
+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 turn = { turnId: 't1', phase: 'generating', input: {} }
-
-const generating = () => ({
-  ...idle(),
-  messages: [user('u1', 'Build a blog.')],
-  status: 'generating',
-  active: turn,
+const user = (id: string, text: string): AppMessage => ({
+  id,
+  role: 'user',
+  parts: [{ type: 'text', text }],
 })
 
-const chat = () =>
+const chat = async () =>
   render(
-    <Chat appId={appId} harness="pi" onPrompt={() => {}} autoFocus={false} />,
+    <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' })
 }
 
-const pushed = (): unknown[] => push.mock.calls.flat()
+/** 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}`
 
-beforeEach(() => {
-  push.mockReset().mockResolvedValue([])
-  // Blob stores each file under its path.
-  upload.mockReset().mockImplementation(async (pathname, file) => ({
-    url: blob(pathname.split('/').at(-1) ?? ''),
-    contentType: file.type,
-  }))
-  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 {}
-    },
+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')!,
   )
-})
 
-afterEach(() => {
-  cleanup()
-  vi.unstubAllGlobals()
-})
+const chevron = () => screen.getByRole('button', { name: 'Send mode' })
 
-it('sends a message between turns', () => {
-  state = idle()
-  chat()
+it('sends a message between turns', async () => {
+  await chat()
   expect(screen.queryByRole('button', { name: 'Send mode' })).toBeNull()
   write('Build a blog.')
-  expect(pushed()).toContainEqual(
-    expect.objectContaining({
-      payload: expect.objectContaining({ action: 'send' }),
-    }),
+  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', () => {
-  state = generating()
-  chat()
+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.')
-  expect(pushed()).toContainEqual(
-    expect.objectContaining({
-      payload: expect.objectContaining({
+  await waitFor(async () =>
+    expect(await controls()).toContainEqual(
+      expect.objectContaining({
         action: 'steer',
-        turnId: 't1',
+        turnId,
         message: expect.objectContaining({
           parts: [{ type: 'text', text: 'Use a dark theme.' }],
         }),
       }),
-    }),
+    ),
   )
 
-  push.mockClear()
-  fireEvent.click(screen.getByRole('button', { name: 'Send mode' }))
+  fireEvent.click(chevron())
   expect(
     screen
       .getByRole('menuitemradio', { name: 'Steer' })
@@ -139,208 +221,236 @@ it('steers the active turn by default, and queues behind it from the chevron', (
   fireEvent.click(screen.getByRole('menuitemradio', { name: 'Queue' }))
   expect(screen.getByRole('button', { name: 'Queue' })).toBeTruthy()
   write('Then add an about page.')
-  expect(pushed()).toContainEqual(
-    expect.objectContaining({
-      payload: expect.objectContaining({ action: 'send' }),
-    }),
+  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', () => {
-  state = generating()
-  const { rerender } = chat()
+it('pauses, resumes, and stops the turn like a player', async () => {
+  const { turnId } = await generating()
+  await chat()
   fireEvent.click(screen.getByRole('button', { name: 'Pause' }))
-  fireEvent.click(screen.getByRole('button', { name: 'Stop' }))
-  expect(pushed()).toEqual([
-    expect.objectContaining({ payload: { action: 'pause', when: 'now' } }),
-    expect.objectContaining({ payload: { action: 'stop', turnId: 't1' } }),
-  ])
-
-  push.mockClear()
-  state = {
-    ...generating(),
-    status: 'paused',
-    active: { ...turn, phase: 'paused' },
-    inbox: { paused: true, items: [] },
-  }
-  rerender(
-    <Chat appId={appId} harness="pi" onPrompt={() => {}} autoFocus={false} />,
-  )
-  expect(screen.getByText('Paused')).toBeTruthy()
+  await paused()
   expect(screen.queryByText('Thinking')).toBeNull()
   fireEvent.click(screen.getByRole('button', { name: 'Resume' }))
-  expect(pushed()).toEqual([
-    expect.objectContaining({ payload: { action: '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', () => {
-  state = {
-    ...generating(),
-    inbox: {
-      paused: false,
-      items: [{ id: 'q1', message: user('q1', 'Add tags.'), generate: true }],
-    },
-  }
-  chat()
-  const queued = screen.getByRole('list', { name: 'Queued messages' })
-  expect(queued.textContent).toContain('Add tags.')
-  fireEvent.click(screen.getByRole('button', { name: 'Steer with it now' }))
-  fireEvent.click(screen.getByRole('button', { name: 'Remove' }))
-  expect(pushed()).toEqual([
-    expect.objectContaining({
-      payload: { action: 'send-now', inputId: 'q1', turnId: 't1' },
-    }),
-    expect.objectContaining({ payload: { action: 'remove', inputId: 'q1' } }),
-  ])
+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', () => {
-  const writing = (part: Record<string, unknown>) => ({
-    ...generating(),
-    messages: [
-      user('u1', 'Build a blog.'),
-      { id: 'a1', role: 'assistant', parts: [part] },
-    ],
-  })
-  state = writing({
-    type: 'tool-write',
-    toolCallId: 'w1',
-    state: 'input-streaming',
-    input: { path: 'app/page.tsx', content: 'export default function' },
-  })
-  const { container, rerender } = chat()
-  const details = () => container.querySelector('.a0-tool details')
-  expect(details()?.hasAttribute('open')).toBe(true)
-  expect(container.querySelector('.a0-tool-text')?.textContent).toBe(
-    'export default function',
-  )
-
-  state = writing({
-    type: 'tool-write',
-    toolCallId: 'w1',
-    state: 'output-available',
-    input: {
-      path: 'app/page.tsx',
-      content: 'export default function Page() {}',
+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',
     },
-    output: { content: [] },
-  })
-  rerender(
-    <Chat appId={appId} harness="pi" onPrompt={() => {}} autoFocus={false} />,
   )
-  expect(details()?.hasAttribute('open')).toBe(false)
+  await waitFor(() =>
+    expect(call('write')?.querySelector('.a0-tool-text')?.textContent).toBe(
+      'export default function',
+    ),
+  )
+  expect(call('write')?.querySelector('details')?.hasAttribute('open')).toBe(
+    true,
+  )
 
-  // A one-line command fits its step's line.
-  state = writing({
-    type: 'tool-bash',
-    toolCallId: 'b1',
-    state: 'input-available',
-    input: { command: 'ls' },
-  })
-  rerender(
-    <Chat appId={appId} harness="pi" onPrompt={() => {}} autoFocus={false} />,
+  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,
+    ),
   )
-  expect(container.querySelector('.a0-tool')).toBeNull()
-  expect(screen.getByText('ls')).toBeTruthy()
-})
 
-const again = (rerender: (ui: ReactElement) => void): void =>
-  rerender(
-    <Chat appId={appId} harness="pi" onPrompt={() => {}} autoFocus={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', () => {
-  state = generating()
-  const { rerender } = chat()
-  const chevron = screen.getByRole('button', { name: 'Send mode' })
-  fireEvent.click(chevron)
+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)
+  fireEvent.pointerDown(chevron())
   expect(screen.getAllByRole('menuitemradio')).toHaveLength(2)
-  fireEvent.click(chevron)
+  fireEvent.click(chevron())
   expect(screen.queryByRole('menuitemradio')).toBeNull()
 
-  fireEvent.click(chevron)
-  state = idle()
-  again(rerender)
-  state = generating()
-  again(rerender)
+  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', () => {
-  state = {
-    ...generating(),
-    status: 'paused',
-    active: { ...turn, phase: 'paused' },
-    inbox: { paused: true, items: [] },
-  }
-  chat()
+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' }))
-  expect(pushed()).toEqual([
-    expect.objectContaining({ payload: { action: 'stop', turnId: 't1' } }),
-    expect.objectContaining({ payload: { action: 'resume' } }),
-  ])
+  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', () => {
-  state = {
-    ...idle(),
-    starting: true,
-    inbox: {
-      paused: false,
-      items: [{ id: 'q1', message: user('q1', 'Add tags.'), generate: true }],
-    },
-  }
-  chat()
+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', () => {
-  const step = (part: Record<string, unknown>, status = 'generating') => ({
-    ...generating(),
-    status,
-    messages: [
-      user('u1', 'Build a blog.'),
-      { id: 'a1', role: 'assistant', parts: [part] },
-    ],
-  })
-  const bash = {
-    type: 'tool-bash',
-    toolCallId: 'b1',
-    state: 'input-available',
-    input: { command: 'cat > a.txt <<EOF\nhi\nEOF' },
-  }
-  state = step(bash, 'paused')
-  const { container, rerender } = chat()
+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(
-    container.querySelector('.a0-tool details')?.hasAttribute('open'),
-  ).toBe(false)
-  expect(container.querySelector('.a0-tool .a0-spin')).toBeNull()
+    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.
-  state = step({
-    type: 'tool-edit',
-    toolCallId: 'e1',
-    state: 'input-streaming',
-    input: { path: 'a.txt', oldText: 'hi', newText: 'hello' },
-  })
-  again(rerender)
-  expect(container.querySelector('.a0-tool-text')?.textContent).toBe(
-    '- hi\n+ hello',
+  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":""',
+    },
   )
-
-  // Nothing written yet: the step's one line, with no empty box.
-  state = step({
-    type: 'tool-write',
-    toolCallId: 'w1',
-    state: 'input-streaming',
-    input: { path: 'a.txt', content: '' },
-  })
-  again(rerender)
-  expect(container.querySelector('.a0-tool')).toBeNull()
+  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 => {
@@ -359,9 +469,18 @@ const leave = (from: Element, to: Element | null): void => {
   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 () => {
-  state = idle()
-  chat()
+  await chat()
   const transcript = screen.getByRole('list', { name: 'Conversation' })
   fireEvent.dragEnter(transcript, {
     dataTransfer: { files: [], types: ['Files'] },
@@ -376,44 +495,27 @@ it('sends images dropped anywhere on the chat with the next message', async () =
   dropOn(transcript, shot())
   const attached = screen.getByRole('list', { name: 'Attachments' })
   expect(attached.querySelector('[aria-label="Uploading"]')).toBeTruthy()
-  await waitFor(() =>
-    expect(attached.querySelector('[aria-label="Uploading"]')).toBeNull(),
-  )
-  expect(upload).toHaveBeenCalledWith(
-    `a0/${appId}/shot.png`,
-    expect.any(File),
-    expect.objectContaining({
-      access: 'public',
-      handleUploadUrl: '/a0/uploads',
-      clientPayload: appId,
-    }),
-  )
+  await uploaded()
 
   write('Make it look like this.')
-  expect(pushed()).toContainEqual(
-    expect.objectContaining({
-      payload: expect.objectContaining({
-        action: 'send',
-        message: expect.objectContaining({
-          parts: [
-            {
-              type: 'file',
-              mediaType: 'image/png',
-              filename: 'shot.png',
-              url: blob('shot.png'),
-            },
-            { type: 'text', text: '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 () => {
-  state = idle()
-  chat()
+  await chat()
   const box = screen.getByRole('textbox', {
     name: 'Message',
   }) as HTMLTextAreaElement
@@ -425,29 +527,26 @@ it("puts a file that isn't an image into the message as its link", async () => {
   expect(screen.queryByRole('list', { name: 'Attachments' })).toBeNull()
 })
 
-it('shows a sent link as its file, and a sent image as itself', () => {
-  state = {
-    ...idle(),
-    messages: [
-      {
-        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')}`,
-          },
-        ],
-      },
-    ],
-  }
-  chat()
+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'),
@@ -455,9 +554,8 @@ it('shows a sent link as its file, and a sent image as itself', () => {
 })
 
 it('says which upload failed, and drops the file', async () => {
-  state = idle()
-  upload.mockRejectedValue(new Error('Vercel Blob: Failed to upload'))
-  chat()
+  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",
@@ -466,39 +564,32 @@ it('says which upload failed, and drops the file', async () => {
 })
 
 it('attaches a pasted image instead of pasting it as text', async () => {
-  state = idle()
-  chat()
-  fireEvent.paste(screen.getByRole('textbox', { name: 'Message' }), {
-    clipboardData: { files: [shot()] },
-  })
+  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 waitFor(() => expect(upload).toHaveBeenCalled())
+  await uploaded()
+  expect(box.value).toBe('')
 })
 
-/** Every attached file has finished uploading. */
-const uploaded = () =>
-  waitFor(() =>
-    expect(
-      screen
-        .getByRole('list', { name: 'Attachments' })
-        .querySelector('[aria-label="Uploading"]'),
-    ).toBeNull(),
-  )
-
 it('keeps the attachments of a send that failed, beside any attached since', async () => {
-  state = idle()
-  chat()
+  await chat()
   const transcript = screen.getByRole('list', { name: 'Conversation' })
   dropOn(transcript, shot())
   await uploaded()
-  let fail!: (cause: Error) => void
-  push.mockReturnValueOnce(new Promise((_, reject) => (fail = reject)))
+  const sending = probe.hold(isPush)
   write('Like this.')
   dropOn(transcript, new File(['png'], 'more.png', { type: 'image/png' }))
   await uploaded()
-  fail(new Error('offline'))
-  expect((await screen.findByRole('alert')).textContent).toBe('offline')
+  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([])
 })