experimental-a2 0.16.1 → 0.18.0

package/src/ai-progress-batches.ts

@@ -1,4 +1,30 @@
 import type { UIMessageChunk } from 'ai'
+import type { GenerationProgressPayload } from './ai.ts'
+
+/** Progress already in a log may name its generation by ids instead of `startIndex`. */
+export type StoredProgressPayload =
+  | GenerationProgressPayload
+  | {
+      requestId: string
+      generationId: string
+      responseMessageId: string
+      chunks: UIMessageChunk[]
+    }
+
+export const progressOwnedBy = ({
+  payload,
+  owner,
+  startIndex,
+}: {
+  payload: StoredProgressPayload
+  owner: { requestId: string; generationId: string; responseMessageId: string }
+  startIndex: number
+}): boolean =>
+  'startIndex' in payload
+    ? payload.startIndex === startIndex
+    : payload.requestId === owner.requestId &&
+      payload.generationId === owner.generationId &&
+      payload.responseMessageId === owner.responseMessageId
 
 export type ProgressBatch = { index: number; chunks: UIMessageChunk[] }
 type BatchBlock = ProgressBatch[] | { left: BatchBlock; right: BatchBlock }

package/src/ai-projector.ts

@@ -1,6 +1,11 @@
+import { sameValue } from './equal-json.ts'
+import type { ToolCalledPayload } from './ai.ts'
 import type { UIMessage, UIMessageChunk } from 'ai'
 import { parsePartialJSON } from './parse-partial-json.ts'
 
+export type ProjectionChunk =
+  UIMessageChunk | { type: 'a2-tool-called'; call: ToolCalledPayload }
+
 type MutablePart = Record<string, unknown>
 
 type PartialToolCall = {
@@ -22,9 +27,22 @@ const cloneMessage = <M extends UIMessage>(messageId: string, base?: M): M =>
     : {
         ...base,
         id: messageId,
-        parts: base.parts.map((part) => ({ ...part })),
+        parts: [...base.parts],
       }) as M
 
+const editPart = ({
+  parts,
+  part,
+}: {
+  parts: MutablePart[]
+  part: MutablePart
+}): MutablePart => {
+  const index = parts.indexOf(part)
+  const next = { ...part }
+  parts[index] = next
+  return next
+}
+
 const currentStepParts = (parts: MutablePart[]): MutablePart[] => {
   let start = parts.length - 1
   while (start >= 0 && parts[start]?.['type'] !== 'step-start') start -= 1
@@ -49,6 +67,30 @@ const findToolPart = (
   )
 }
 
+export const hasToolCall = ({
+  message,
+  call,
+}: {
+  message: UIMessage
+  call: ToolCalledPayload
+}): boolean => {
+  const part = findToolPart(
+    message.parts as unknown as MutablePart[],
+    call.toolCallId,
+  )
+  return (
+    part !== undefined &&
+    part['state'] !== 'input-streaming' &&
+    sameValue(part['input'], call.input) &&
+    (part['callProviderMetadata'] !== undefined ||
+      call.providerMetadata === undefined) &&
+    (part['toolMetadata'] !== undefined || call.toolMetadata === undefined) &&
+    (call.providerExecuted === undefined ||
+      part['providerExecuted'] === call.providerExecuted) &&
+    (call.title === undefined || part['title'] === call.title)
+  )
+}
+
 const findApprovalPart = (
   parts: MutablePart[],
   approvalId: string,
@@ -92,7 +134,7 @@ const setToolPart = (options: {
       toolCallId: options.toolCallId,
     }
     options.parts.push(part)
-  }
+  } else part = editPart({ parts: options.parts, part })
 
   part['state'] = options.state
   part['input'] = options.input
@@ -134,7 +176,7 @@ export type UIMessageCursor = {
 export function projectUIMessage<M extends UIMessage>(options: {
   messageId: string
   base?: M
-  chunks: readonly UIMessageChunk[]
+  chunks: readonly ProjectionChunk[]
 }): M {
   return projectUIMessageStream(options).message
 }
@@ -143,7 +185,7 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
   messageId: string
   base?: M
   cursor?: UIMessageCursor
-  chunks: readonly UIMessageChunk[]
+  chunks: readonly ProjectionChunk[]
 }): { message: M; cursor: UIMessageCursor } {
   const message = cloneMessage(options.messageId, options.base)
   const record = message as unknown as Record<string, unknown>
@@ -189,14 +231,17 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
           activeText.set(chunk.id, part)
           parts.push(part)
         }
+        part = editPart({ parts, part })
+        activeText.set(chunk.id, part)
         part['text'] =
           `${typeof part['text'] === 'string' ? part['text'] : ''}${chunk.delta}`
         setOptional(part, 'providerMetadata', chunk.providerMetadata)
         break
       }
       case 'text-end': {
-        const part = activeText.get(chunk.id)
-        if (part) {
+        const current = activeText.get(chunk.id)
+        if (current) {
+          const part = editPart({ parts, part: current })
           part['state'] = 'done'
           setOptional(part, 'providerMetadata', chunk.providerMetadata)
           activeText.delete(chunk.id)
@@ -223,14 +268,17 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
           activeReasoning.set(chunk.id, part)
           parts.push(part)
         }
+        part = editPart({ parts, part })
+        activeReasoning.set(chunk.id, part)
         part['text'] =
           `${typeof part['text'] === 'string' ? part['text'] : ''}${chunk.delta}`
         setOptional(part, 'providerMetadata', chunk.providerMetadata)
         break
       }
       case 'reasoning-end': {
-        const part = activeReasoning.get(chunk.id)
-        if (part) {
+        const current = activeReasoning.get(chunk.id)
+        if (current) {
+          const part = editPart({ parts, part: current })
           part['state'] = 'done'
           setOptional(part, 'providerMetadata', chunk.providerMetadata)
           activeReasoning.delete(chunk.id)
@@ -327,6 +375,37 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
         })
         break
       }
+      case 'a2-tool-called': {
+        const call = chunk.call
+        const current = findToolPart(parts, call.toolCallId)
+        if (current) {
+          const existing = editPart({ parts, part: current })
+          existing['input'] = call.input
+          if (existing['state'] === 'input-streaming')
+            existing['state'] = 'input-available'
+          if (existing['callProviderMetadata'] === undefined)
+            setOptional(existing, 'callProviderMetadata', call.providerMetadata)
+          if (existing['toolMetadata'] === undefined)
+            setOptional(existing, 'toolMetadata', call.toolMetadata)
+          setOptional(existing, 'providerExecuted', call.providerExecuted)
+          setOptional(existing, 'title', call.title)
+        } else {
+          setToolPart({
+            parts,
+            toolCallId: call.toolCallId,
+            toolName: call.toolName,
+            dynamic: call.dynamic === true,
+            state: 'input-available',
+            input: call.input,
+            providerExecuted: call.providerExecuted,
+            providerMetadata: call.providerMetadata,
+            toolMetadata: call.toolMetadata,
+            title: call.title,
+          })
+        }
+        partialTools.delete(call.toolCallId)
+        break
+      }
       case 'tool-input-available':
         setToolPart({
           parts,
@@ -363,8 +442,9 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
         break
       }
       case 'tool-approval-request': {
-        const part = findToolPart(parts, chunk.toolCallId)
-        if (!part) break
+        const current = findToolPart(parts, chunk.toolCallId)
+        if (!current) break
+        const part = editPart({ parts, part: current })
         part['state'] = 'approval-requested'
         part['approval'] = {
           id: chunk.approvalId,
@@ -376,8 +456,9 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
         break
       }
       case 'tool-approval-response': {
-        const part = findApprovalPart(parts, chunk.approvalId)
-        if (!part) break
+        const current = findApprovalPart(parts, chunk.approvalId)
+        if (!current) break
+        const part = editPart({ parts, part: current })
         const approval = isRecord(part['approval']) ? part['approval'] : {}
         part['state'] = 'approval-responded'
         part['approval'] = {
@@ -441,7 +522,7 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
       }
       case 'tool-output-denied': {
         const part = findToolPart(parts, chunk.toolCallId)
-        if (part) part['state'] = 'output-denied'
+        if (part) editPart({ parts, part })['state'] = 'output-denied'
         break
       }
       case 'start-step':
@@ -496,7 +577,7 @@ export function projectUIMessageStream<M extends UIMessage>(options: {
                 (part) =>
                   part['type'] === chunk.type && part['id'] === chunk.id,
               )
-        if (existing) existing['data'] = chunk.data
+        if (existing) editPart({ parts, part: existing })['data'] = chunk.data
         else
           parts.push({
             type: chunk.type,

package/src/ai-react.ts

@@ -0,0 +1,146 @@
+'use client'
+/**
+ * experimental-a2/ai/react — one hook for an agent session: its live
+ * optimistic `AIState`, and actions that fill in protocol identities.
+ *
+ * ```tsx
+ * import type { assistant } from '@/server/assistant' // erased at build
+ * const { messages, indicator, sendMessage, stop } =
+ *   useAgent<typeof assistant>({ api: '/api/agent', id, initial })
+ * ```
+ *
+ * With a shared `agentContract()`, pass it as `contract` to fold its
+ * application state in the browser too.
+ */
+
+import type { AIEventDefs, AIState } from './ai.ts'
+import { withAgentActions, type UseAgentResult } from './ai-agent-actions.ts'
+import {
+  createAgentClient,
+  type AgentClient,
+  type AgentMessage,
+  type AnyAgent,
+} from './ai-client.ts'
+import type { Contract, EventDefs } from './contract.ts'
+import type { InitialSnapshot } from './initial-snapshot.ts'
+import { createClient, type A2Client } from './client.ts'
+import type { Reducer } from './reducer.ts'
+import { useSession, type UseSessionResult } from './react.ts'
+
+export type UseAgentOptions<A extends AnyAgent> = {
+  /** The session id. */
+  id: string
+  /** The server-rendered `assistant.session(id).state()`. */
+  initial?: InitialSnapshot<AIState<AgentMessage<A>>>
+  participant?: string
+} & (
+  | {
+      /** The route serving the agent's `fetch`; its client is shared. */
+      api: string
+      contract?: never
+    }
+  | { client: AgentClient<A> }
+)
+
+/**
+ * Use a session of an agent served from a shared `agentContract()`, folding
+ * its application state in the browser. The contract is a real import here.
+ */
+export type UseContractAgentOptions<D extends EventDefs, S extends AIState> = {
+  /** The session id. */
+  id: string
+  /** The server-rendered `agent.session(id).state()`. */
+  initial?: InitialSnapshot<S>
+  participant?: string
+} & (
+  | {
+      /** The route serving the agent's `fetch`; its client is shared. */
+      api: string
+      contract: Contract<D, Record<never, never>, Reducer<D, S>>
+    }
+  | {
+      /** `createClient({ contract, api })`, shared with pushes outside React. */
+      client: A2Client<D, S>
+    }
+)
+
+const clients = new Map<string, AgentClient>()
+const contractClients = new WeakMap<object, Map<string, object>>()
+
+/**
+ * One client per route for the page's lifetime; its idle sessions expire
+ * after the default `gcTime`. It reads `globalThis.fetch` per request, so
+ * instrumentation installed later still applies.
+ */
+const clientFor = (api: string): AgentClient => {
+  let client = clients.get(api)
+  if (client === undefined) {
+    client = createAgentClient({
+      api,
+      fetch: (input, init) => globalThis.fetch(input, init),
+    })
+    clients.set(api, client)
+  }
+  return client
+}
+
+/** One client per contract and route, like `clientFor`. */
+const contractClientFor = <D extends EventDefs, S>(
+  contract: Contract<D, Record<never, never>, Reducer<D, S>>,
+  api: string,
+): A2Client<D, S> => {
+  let byApi = contractClients.get(contract)
+  if (byApi === undefined) {
+    byApi = new Map()
+    contractClients.set(contract, byApi)
+  }
+  let client = byApi.get(api) as A2Client<D, S> | undefined
+  if (client === undefined) {
+    client = createClient({
+      contract,
+      api,
+      fetch: (input, init) => globalThis.fetch(input, init),
+    })
+    byApi.set(api, client)
+  }
+  return client
+}
+
+/** Use a session of an agent served from a shared `agentContract()`. */
+export function useAgent<D extends EventDefs, S extends AIState>(
+  options: UseContractAgentOptions<D, S>,
+): UseAgentResult<D, S>
+/**
+ * Use one agent session. Pass `typeof` the server's `agent()` for its
+ * message types; the import is type-only and erased.
+ */
+export function useAgent<A extends AnyAgent = AnyAgent>(
+  options: UseAgentOptions<A>,
+): UseAgentResult<AIEventDefs<AgentMessage<A>>, AIState<AgentMessage<A>>>
+export function useAgent(
+  options:
+    UseAgentOptions<AnyAgent> | UseContractAgentOptions<AIEventDefs, AIState>,
+): object {
+  const session: UseSessionResult<AIEventDefs, AIState> = useSession(
+    'client' in options
+      ? options.client
+      : options.contract !== undefined
+        ? contractClientFor(options.contract, options.api)
+        : clientFor(options.api),
+    options.id,
+    {
+      ...(options.initial === undefined ? {} : { initial: options.initial }),
+      ...(options.participant === undefined
+        ? {}
+        : { participant: options.participant }),
+    },
+  )
+  return withAgentActions(session)
+}
+
+export type { AgentClient, AgentMessage, AnyAgent }
+export type {
+  AgentActions,
+  AgentIndicator,
+  UseAgentResult,
+} from './ai-agent-actions.ts'

package/src/ai-sdk-step.ts

@@ -8,7 +8,6 @@ import type {
   FinishReason,
   InferUIMessageChunk,
   Instructions,
-  LanguageModel,
   LanguageModelUsage,
   ModelMessage,
   TextStreamPart,
@@ -125,7 +124,8 @@ export type AISDKStepSettings<T extends ToolSet> = Partial<
 }
 
 export type AISDKStepInput<M extends UIMessage, T extends ToolSet> = {
-  model: LanguageModel
+  /** An AI Gateway model ID, resolved through the AI SDK's global provider. */
+  model: string
   tools: T
   messages: M[]
   modelMessages?: ModelMessage[]

package/src/ai-server.browser.ts

@@ -1,5 +1,5 @@
-/** experimental-a2/ai/server is server-only; agent definitions live in experimental-a2/ai. */
+/** experimental-a2/ai/server is server-only; browsers use experimental-a2/ai/react or /ai/client. */
 throw new Error(
-  'experimental-a2/ai/server is server-only — a client bundle imported it. Import agent from experimental-a2/ai in browser code instead.',
+  "experimental-a2/ai/server is server-only — a client bundle imported it. Import its agent's type with `import type`, and use experimental-a2/ai/react or experimental-a2/ai/client in browser code.",
 )
 export {}