experimental-a2 0.16.1 → 0.18.0

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

@@ -1,5 +1,6 @@
+import { contract } from 'experimental-a2'
 import type { FinishReason, LanguageModelUsage, UIMessage } from 'ai'
-import { agent, type AIState, type AIStoredState } from 'experimental-a2/ai'
+import { agentContract } from 'experimental-a2/ai'
 import { z } from 'zod'
 
 export type ChatAgentStep = {
@@ -147,28 +148,12 @@ const chatMessageSchema = z.custom<ChatAgentMessage>(
   'invalid chat UI message',
 )
 
-const chatEvents = {
-  'message.sent': z.object({
-    author: z.string().min(1).max(40),
-    text: z.string().min(1).max(2000),
-  }),
-  typed: z.object({ author: z.string().min(1).max(40) }),
-  'agent.ambient.changed': z.object({
-    enabled: z.boolean(),
-    author: z.string().min(1).max(40),
-  }),
-  'agent.message.posted': agentPostEventSchema,
-}
-
-export const chatAgent = agent<ChatAgentMessage, typeof chatEvents>({
-  name: 'chat',
-  events: chatEvents,
+/** The agent's own session: the conversation it sees, one per room. */
+export const chatAgent = agentContract<ChatAgentMessage>({
+  name: 'chat-agent',
   messageSchema: chatMessageSchema,
-  reducerName: 'a2.playground.chat.ai.v3',
 })
 
-export const chat = chatAgent.contract
-
 export type ChatMessage = {
   id: string
   index: number
@@ -187,7 +172,6 @@ export type AgentPost = {
 }
 
 export type ChatState = {
-  ai: AIStoredState
   messages: ChatMessage[]
   posts: AgentPost[]
   lastTypingAt: Record<string, string>
@@ -199,124 +183,94 @@ export type ChatState = {
 }
 
 const initialChatState: ChatState = {
-  ai: chatAgent.reducer.initialState,
   messages: [],
   posts: [],
   lastTypingAt: {},
   ambient: { enabled: false, changedBy: null, changedAt: null },
 }
 
-export const chatReducer = chat
-  .reducer({
-    name: 'a2.playground.chat.v4',
+/**
+ * The room: an ordinary A2 session. The agent runs in its own session, and its
+ * `post` tool appends here, so the room's log orders its messages among
+ * everyone else's.
+ */
+export const chat = contract({
+  name: 'chat-room',
+  events: {
+    'message.sent': z.object({
+      author: z.string().min(1).max(40),
+      text: z.string().min(1).max(2000),
+    }),
+    typed: z.object({ author: z.string().min(1).max(40) }),
+    'agent.ambient.changed': z.object({
+      enabled: z.boolean(),
+      author: z.string().min(1).max(40),
+    }),
+    'agent.message.posted': agentPostEventSchema,
+  },
+  reducer: {
+    name: 'a2.playground.chat-room.v1',
     initialState: initialChatState,
-  })
-  .fold((state, event) => {
-    const previousAI = chatAgent.reducer.view(state.ai)
-    if (
-      event.type === 'ai.tool.called' &&
-      event.payload.toolName === 'post' &&
-      (previousAI.activeGeneration?.generationId !==
-        event.payload.generationId ||
-        previousAI.activeGeneration.responseMessageId !==
-          event.payload.messageId)
-    ) {
-      return state
-    }
-
-    const ai = chatAgent.reducer.fold(state.ai, event)
-    const base = ai === state.ai ? state : { ...state, ai }
-
-    switch (event.type) {
-      case 'message.sent': {
-        const stillTyping = { ...base.lastTypingAt }
-        delete stillTyping[event.payload.author]
-        return {
-          ...base,
-          messages: [
-            ...base.messages,
-            {
-              id: event.id,
-              index: event.index,
-              author: event.payload.author,
-              text: event.payload.text,
-              at: event.createdAt.toISOString(),
-            },
-          ],
-          lastTypingAt: stillTyping,
+    fold: (state, event) => {
+      switch (event.type) {
+        case 'message.sent': {
+          const stillTyping = { ...state.lastTypingAt }
+          delete stillTyping[event.payload.author]
+          return {
+            ...state,
+            messages: [
+              ...state.messages,
+              {
+                id: event.id,
+                index: event.index,
+                author: event.payload.author,
+                text: event.payload.text,
+                at: event.createdAt.toISOString(),
+              },
+            ],
+            lastTypingAt: stillTyping,
+          }
         }
-      }
-      case 'typed':
-        return {
-          ...base,
-          lastTypingAt: {
-            ...base.lastTypingAt,
-            [event.payload.author]: event.createdAt.toISOString(),
-          },
-        }
-      case 'agent.ambient.changed':
-        return {
-          ...base,
-          ambient: {
-            enabled: event.payload.enabled,
-            changedBy: event.payload.author,
-            changedAt: event.createdAt.toISOString(),
-          },
-        }
-      case 'agent.message.posted': {
-        const payload = event.payload
-        const running = previousAI.tools.some(
-          (activity) =>
-            activity.generationId === payload.generationId &&
-            activity.messageId === payload.responseMessageId &&
-            activity.toolCallId === payload.toolCallId &&
-            activity.toolName === 'post' &&
-            activity.status === 'running',
-        )
-        if (!running) return base
-        return {
-          ...base,
-          posts: [
-            ...base.posts,
-            {
-              generationId: payload.generationId,
-              responseMessageId: payload.responseMessageId,
-              toolCallId: payload.toolCallId,
-              index: event.index,
-              text: payload.message,
-              at: event.createdAt.toISOString(),
+        case 'typed':
+          return {
+            ...state,
+            lastTypingAt: {
+              ...state.lastTypingAt,
+              [event.payload.author]: event.createdAt.toISOString(),
             },
-          ],
-        }
-      }
-      case 'ai.message.interrupted': {
-        const cutoff = event.payload.lastSeenIndex
-        if (ai === state.ai || cutoff === undefined) return base
-        return {
-          ...base,
-          posts: base.posts.filter(
-            (post) =>
-              post.responseMessageId !== event.payload.messageId ||
-              (event.payload.generationId !== undefined &&
-                post.generationId !== event.payload.generationId) ||
-              post.index <= cutoff,
-          ),
-        }
+          }
+        case 'agent.ambient.changed':
+          return {
+            ...state,
+            ambient: {
+              enabled: event.payload.enabled,
+              changedBy: event.payload.author,
+              changedAt: event.createdAt.toISOString(),
+            },
+          }
+        case 'agent.message.posted':
+          return {
+            ...state,
+            posts: [
+              ...state.posts,
+              {
+                generationId: event.payload.generationId,
+                responseMessageId: event.payload.responseMessageId,
+                toolCallId: event.payload.toolCallId,
+                index: event.index,
+                text: event.payload.message,
+                at: event.createdAt.toISOString(),
+              },
+            ],
+          }
+        default:
+          return state
       }
-      default:
-        return base
-    }
-  })
+    },
+  },
+})
+export const chatReducer = chat.reducer
 
 const agentMention = /(?:^|\s)@agent(?![\w-])/i
 
 export const mentionsAgent = (text: string): boolean => agentMention.test(text)
-
-export type ChatViewState = Omit<ChatState, 'ai'> & {
-  ai: AIState<ChatAgentMessage>
-}
-
-export const viewChatState = (state: ChatState): ChatViewState => ({
-  ...state,
-  ai: chatAgent.reducer.view(state.ai),
-})

package/examples/playground/app/chat/scheduler/route.ts

@@ -1,5 +1,5 @@
-import { chatScheduler, chatServer } from '../server'
+import { chatAgentServer, chatScheduler, chatServer } from '../server'
 
 export const POST =
-  chatScheduler?.handler(chatServer) ??
+  chatScheduler?.handler(chatServer, chatAgentServer) ??
   (() => new Response('chat scheduler is not configured', { status: 503 }))

package/examples/playground/app/chat/server.test.ts

@@ -0,0 +1,112 @@
+import type { UIMessageChunk } from 'ai'
+import { expect, it, vi } from 'vitest'
+
+vi.mock('server-only', () => ({}))
+vi.mock('@/lib/store', async () => {
+  const { memory } = await import('experimental-a2/store-memory')
+  return { store: memory() }
+})
+
+// The chat's real agent, with a scripted model: two `post` calls, then done.
+const { prompts, step } = vi.hoisted(() => ({
+  prompts: [] as string[],
+  step: (chunks: UIMessageChunk[]) =>
+    new ReadableStream<UIMessageChunk>({
+      start(controller) {
+        for (const chunk of chunks) controller.enqueue(chunk)
+        controller.close()
+      },
+    }),
+}))
+vi.mock('experimental-a2/ai/server', async (importActual) => {
+  const actual =
+    await importActual<typeof import('experimental-a2/ai/server')>()
+  return {
+    ...actual,
+    agent: (options: Parameters<typeof actual.agent>[0]) =>
+      actual.agent({
+        ...options,
+        harness: {
+          ...(options.harness?.type === 'ai-sdk' ? options.harness : {}),
+          type: 'ai-sdk',
+          generate: ({
+            event,
+            messages,
+          }: {
+            event: { payload: { reason: string } }
+            messages: { role: string; id: string }[]
+          }) => {
+            prompts.push(
+              messages.map((message) => `${message.role}:${message.id}`).join(),
+            )
+            return step(
+              event.payload.reason === 'message'
+                ? [
+                    { type: 'start' },
+                    ...['Hello', 'there'].map(
+                      (message, call): UIMessageChunk => ({
+                        type: 'tool-input-available',
+                        toolCallId: `call-${call}`,
+                        toolName: 'post',
+                        input: { message },
+                      }),
+                    ),
+                    { type: 'finish', finishReason: 'tool-calls' },
+                  ]
+                : [{ type: 'start' }, { type: 'finish', finishReason: 'stop' }],
+            )
+          },
+        },
+      } as never),
+  }
+})
+
+it('forwards room messages to the agent, whose posts land in the room', async () => {
+  const { chatAgentServer, chatServer } = await import('./server')
+  const room = chatServer.session('room')
+  const agent = chatAgentServer.session('room')
+  await room.append({
+    type: 'message.sent',
+    payload: { author: 'alice', text: 'just chatting' },
+  })
+  await room.append({
+    type: 'message.sent',
+    payload: { author: 'bob', text: '@agent say hi' },
+  })
+
+  // Without a scheduler, appends run in background drains.
+  await vi.waitFor(
+    async () => {
+      expect((await room.state()).state.posts).toHaveLength(2)
+      expect((await agent.state()).state.status).toBe('idle')
+    },
+    { timeout: 5_000 },
+  )
+
+  const { state } = await room.state()
+  expect(state.messages.map((message) => message.author)).toEqual([
+    'alice',
+    'bob',
+  ])
+  expect(state.posts.map((post) => post.text)).toEqual(['Hello', 'there'])
+  expect(state.posts[0]!.index).toBeGreaterThan(state.messages[1]!.index)
+  // The passive message reached the agent as context without a reply of its
+  // own; only the mention generated, and it saw both.
+  expect(prompts).toHaveLength(2)
+  expect(
+    prompts[0]!.split(',').filter((id) => id.startsWith('user:')),
+  ).toHaveLength(2)
+  const messages = (await agent.state()).state.messages
+  expect(messages.map((message) => message.role)).toEqual([
+    'user',
+    'user',
+    'assistant',
+  ])
+
+  // Replaying the post tool's append converges on the stored event.
+  const [post] = (await room.history()).filter(
+    (event) => event.type === 'agent.message.posted',
+  )
+  await room.append({ type: post!.type, id: post!.id, payload: post!.payload })
+  expect((await room.state()).state.posts).toHaveLength(2)
+})

package/examples/playground/app/chat/server.ts

@@ -1,11 +1,13 @@
 import 'server-only'
 import { tool } from 'ai'
+import { idempotentId } from 'experimental-a2'
 import { handlerContext, inputs } from 'experimental-a2/ai'
-import { createHandlers } from 'experimental-a2/ai/server'
+import { agent } from 'experimental-a2/ai/server'
 import { vercelQueues } from 'experimental-a2/scheduler-vercel'
 import { createServer, type A2Operation } from 'experimental-a2/server'
 import { store } from '@/lib/store'
 import {
+  chat,
   chatAgent,
   chatPostInputSchema,
   mentionsAgent,
@@ -22,12 +24,16 @@ const tools = {
     execute: async ({ message }, { toolCallId }) => {
       const context = handlerContext(chatAgent)
       context.signal.throwIfAborted()
-      await context.session.append('post', {
+      const { generationId, messageId } = context.event.payload.call
+      // The agent's session shares the room's id. The explicit id makes a
+      // retried tool converge on one post.
+      await chatServer.session(context.session.id).append({
         type: 'agent.message.posted',
+        id: await idempotentId('chat.post', generationId, toolCallId),
         payload: {
           message,
-          generationId: context.event.payload.call.generationId,
-          responseMessageId: context.event.payload.call.messageId,
+          generationId,
+          responseMessageId: messageId,
           toolCallId,
         },
       })
@@ -36,139 +42,88 @@ const tools = {
   }),
 }
 
-const aiHandlers = createHandlers({
-  agent: chatAgent,
+export const chatScheduler =
+  process.env.VERCEL || process.env.A2_QUEUES
+    ? vercelQueues({ topic: 'a2-chat' })
+    : undefined
+
+export const chatAgentServer = agent({
+  ...(store ? { store } : {}),
+  ...(chatScheduler ? { scheduler: chatScheduler } : {}),
+  contract: chatAgent,
   model: 'openai/gpt-5.4',
   tools,
+  // Posts become chat bubbles, so one agent's posts must land in call order.
+  toolLane: 'a2.chat.agent-tools',
   instructions: `You are "agent", a participant in a multiplayer chat.
 The latest participant message invited you either by mentioning @agent or by enabling the room's ambient mode. Messages are formatted as "name: message"; earlier passive messages are useful conversation context.
 Sound like a sharp, relaxed person in the chat, not a customer-support bot. Match the room's language and energy. Keep messages specific and conversational; avoid canned acknowledgements, assistant-y disclaimers, and explanations of your own mechanics.
 Only contribute when you have something useful, relevant, or friendly to add. It is fine to stay silent or stop whenever you have nothing else to say.
 To speak, call post once for each visible chat bubble. Never communicate through ordinary assistant text because it is private and is not shown in the room.
 Use as many post calls as the moment naturally calls for. Several short messages are welcome when they create better rhythm, separate distinct thoughts, or fulfill a request for a sequence. A single message is fine too. Do not spam. After post results, continue only when another message genuinely improves the conversation.`,
-  generation: {
-    maxOutputTokens: 500,
-    temperature: 0.4,
+  harness: {
+    type: 'ai-sdk',
+    generation: {
+      maxOutputTokens: 500,
+      temperature: 0.4,
+    },
+    messageMetadata: (context) => {
+      const previous = context.messages.find(
+        (message) => message.id === context.event.payload.responseMessageId,
+      )?.metadata
+      const trigger = context.messages.find(
+        (message) => message.id === context.event.payload.messageId,
+      )?.metadata
+      const steps = previous?.kind === 'agent' ? previous.steps : []
+      const mode: ChatAgentMode | undefined =
+        previous?.kind === 'agent'
+          ? previous.mode
+          : trigger?.kind === 'participant' && trigger.mode !== 'passive'
+            ? trigger.mode
+            : undefined
+      if (mode === undefined) return undefined
+      const metadata = {
+        kind: 'agent' as const,
+        author: 'agent' as const,
+        mode,
+        steps,
+      }
+      if (context.part.type === 'start') return metadata
+      if (context.part.type !== 'finish') return undefined
+      return {
+        ...metadata,
+        steps: [
+          ...steps,
+          {
+            generationId: context.generationId,
+            model: context.model,
+            finishReason: context.part.finishReason,
+            ...(context.part.rawFinishReason === undefined
+              ? {}
+              : { rawFinishReason: context.part.rawFinishReason }),
+            usage: context.part.totalUsage,
+            completedAt: new Date().toISOString(),
+          },
+        ],
+      }
+    },
   },
   progress: {
     maxChunks: 4,
     maxDelayMs: 30,
   },
-  messageMetadata: (context) => {
-    const previous = context.messages.find(
-      (message) => message.id === context.responseMessageId,
-    )?.metadata
-    const trigger = context.messages.find(
-      (message) => message.id === context.request.payload.messageId,
-    )?.metadata
-    const steps = previous?.kind === 'agent' ? previous.steps : []
-    const mode: ChatAgentMode | undefined =
-      previous?.kind === 'agent'
-        ? previous.mode
-        : trigger?.kind === 'participant' && trigger.mode !== 'passive'
-          ? trigger.mode
-          : undefined
-    if (mode === undefined) return undefined
-    const metadata = {
-      kind: 'agent' as const,
-      author: 'agent' as const,
-      mode,
-      steps,
-    }
-    if (context.part.type === 'start') return metadata
-    if (context.part.type !== 'finish') return undefined
-    return {
-      ...metadata,
-      steps: [
-        ...steps,
-        {
-          generationId: context.generationId,
-          model:
-            typeof context.model === 'string'
-              ? context.model
-              : context.model.modelId,
-          finishReason: context.part.finishReason,
-          ...(context.part.rawFinishReason === undefined
-            ? {}
-            : { rawFinishReason: context.part.rawFinishReason }),
-          usage: context.part.totalUsage,
-          completedAt: new Date().toISOString(),
-        },
-      ],
-    }
-  },
 })
 
-export const authorizeChatOperation = (
-  operation: A2Operation<(typeof chatAgent.contract)['events']>,
-): boolean => {
-  if (operation.type !== 'push') return true
-  const serverOnly = operation.events.find(
-    (event) =>
-      event.type === 'ai.message.created' ||
-      event.type === 'agent.message.posted',
-  )
-  if (serverOnly !== undefined) return false
-  const rejected = operation.events.find((event) => {
-    if (
-      event.type === 'message.sent' ||
-      event.type === 'typed' ||
-      event.type === 'agent.ambient.changed'
-    ) {
-      return false
-    }
-    return !(
-      event.type === 'ai.control.requested' &&
-      typeof event.payload === 'object' &&
-      event.payload !== null &&
-      'action' in event.payload &&
-      typeof event.payload.action === 'string' &&
-      [
-        'approval',
-        'input',
-        'interrupt',
-        'stop',
-        'retry',
-        'pause',
-        'resume',
-      ].includes(event.payload.action)
-    )
-  })
-  return rejected === undefined
-}
-
-const toolExecutionHandler = aiHandlers['ai.tool.execution.requested']
-if (toolExecutionHandler === undefined) {
-  throw new Error('A2 AI handlers must include ai.tool.execution.requested')
-}
-const serializedToolExecutionHandler =
-  typeof toolExecutionHandler === 'function'
-    ? { handler: toolExecutionHandler, lane: 'a2.chat.agent-tools' }
-    : { ...toolExecutionHandler, lane: 'a2.chat.agent-tools' }
-
-export const chatScheduler =
-  process.env.VERCEL || process.env.A2_QUEUES
-    ? vercelQueues({ topic: 'a2-chat' })
-    : undefined
-
+/** The room. Each message goes on to the agent's session, marked for whether it should answer. */
 export const chatServer = createServer({
   ...(store ? { store } : {}),
   ...(chatScheduler ? { scheduler: chatScheduler } : {}),
-  contract: chatAgent.contract,
+  contract: chat,
   handlers: {
-    ...aiHandlers,
-    'ai.tool.execution.requested': serializedToolExecutionHandler,
     'message.sent': {
       lane: 'a2.chat.ingress',
       handler: async (context) => {
-        let ambient = false
-        for (const event of await context.session.history()) {
-          if (event.index > context.event.index) break
-          if (event.type === 'agent.ambient.changed') {
-            ambient = event.payload.enabled
-          }
-        }
-
+        const ambient = (await context.session.state()).state.ambient.enabled
         const { author, text } = context.event.payload
         const mode: ChatParticipantMode = ambient
           ? 'ambient'
@@ -176,7 +131,9 @@ export const chatServer = createServer({
             ? 'mention'
             : 'passive'
         const message: ChatAgentMessage = {
-          id: context.event.id,
+          // Event ids are unique across a store's sessions, so the forward
+          // derives its own; a retried forward converges on the same id.
+          id: await idempotentId('chat.message', context.event.id),
           role: 'user',
           metadata: {
             kind: 'participant',
@@ -187,8 +144,47 @@ export const chatServer = createServer({
           },
           parts: [{ type: 'text', text: `${author}: ${text}` }],
         }
-        return inputs.message(message, { generate: mode !== 'passive' })
+        await chatAgentServer
+          .session(context.session.id)
+          .append(...inputs.message(message, { generate: mode !== 'passive' }))
       },
     },
   },
 })
+
+/** Room pushes: people's messages, typing, and ambient mode. Posts are the agent's. */
+export const authorizeChatOperation = (
+  operation: A2Operation<(typeof chat)['events']>,
+): boolean =>
+  operation.type !== 'push' ||
+  operation.events.every(
+    (event) =>
+      event.type === 'message.sent' ||
+      event.type === 'typed' ||
+      event.type === 'agent.ambient.changed',
+  )
+
+const agentControls = new Set([
+  'approval',
+  'input',
+  'interrupt',
+  'stop',
+  'retry',
+  'pause',
+  'resume',
+])
+
+/** Agent pushes: controls only. Messages reach the agent through the room. */
+export const authorizeChatAgentOperation = (
+  operation: A2Operation<(typeof chatAgent)['events']>,
+): boolean =>
+  operation.type !== 'push' ||
+  operation.events.every(
+    (event) =>
+      event.type === 'ai.control.requested' &&
+      typeof event.payload === 'object' &&
+      event.payload !== null &&
+      'action' in event.payload &&
+      typeof event.payload.action === 'string' &&
+      agentControls.has(event.payload.action),
+  )