experimental-a2 0.18.1 → 0.20.0

package/examples/playground/test/serve.ts

@@ -0,0 +1,184 @@
+/**
+ * A real HTTP server for the app's route handlers, so the browser's clients
+ * reach A2 servers over the actual wire format.
+ */
+
+import { createServer, type ServerResponse } from 'node:http'
+import type { AddressInfo, Socket } from 'node:net'
+import type { A2Operation } from 'experimental-a2/server'
+import { setGlobalOrigin } from 'undici'
+
+type Handler = (request: Request) => Response | Promise<Response>
+
+export type Served = {
+  /** Where the server listens; relative URLs resolve against it until it closes. */
+  origin: string
+  close(): Promise<void>
+}
+
+const writeResponse = async (
+  res: ServerResponse,
+  response: Response,
+): Promise<void> => {
+  res.writeHead(response.status, Object.fromEntries(response.headers.entries()))
+  res.flushHeaders()
+  if (!response.body) {
+    res.end()
+    return
+  }
+  const reader = response.body.getReader()
+  // A disconnecting client must cancel the web stream, which closes the
+  // store subscription behind it.
+  res.on('close', () => {
+    reader.cancel().catch(() => {})
+  })
+  try {
+    for (;;) {
+      // oxlint-disable-next-line no-await-in-loop -- stream pump
+      const { done, value } = await reader.read()
+      if (done) break
+      res.write(value)
+    }
+    res.end()
+  } catch {
+    res.destroy()
+  }
+}
+
+/**
+ * Serves each route's handler at its path. The page and `fetch` resolve
+ * relative URLs, like a module-level client's `api`, against this server
+ * until it closes.
+ */
+export async function serve({
+  routes,
+}: {
+  routes: Readonly<Record<string, Handler>>
+}): Promise<Served> {
+  let origin = ''
+  const server = createServer((req, res) => {
+    void (async () => {
+      try {
+        const url = new URL(req.url ?? '/', origin)
+        const handler = routes[url.pathname]
+        if (handler === undefined) {
+          res.writeHead(404).end()
+          return
+        }
+        const chunks: Buffer[] = []
+        for await (const chunk of req) chunks.push(chunk as Buffer)
+        const request = new Request(url, {
+          method: req.method ?? 'GET',
+          headers: Object.entries(req.headers).flatMap(([name, value]) =>
+            value === undefined
+              ? []
+              : (Array.isArray(value) ? value : [value]).map(
+                  (item): [string, string] => [name, item],
+                ),
+          ),
+          ...(chunks.length === 0 ? {} : { body: Buffer.concat(chunks) }),
+        })
+        await writeResponse(res, await handler(request))
+      } catch (error) {
+        // A handler's bug answers as one, not as a dropped connection the
+        // client would retry.
+        if (res.headersSent) res.destroy()
+        else res.writeHead(500).end(String(error))
+      }
+    })()
+  })
+  // close() must never wait on a request a test parked.
+  const sockets = new Set<Socket>()
+  server.on('connection', (socket) => {
+    sockets.add(socket)
+    socket.on('close', () => sockets.delete(socket))
+  })
+  // The address clients use: on macOS, another process can bind 127.0.0.1 on
+  // a port this server holds on every address, and take its connections.
+  await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
+  const { port } = server.address() as AddressInfo
+  origin = `http://127.0.0.1:${port}`
+  setGlobalOrigin(origin)
+  // Vitest's jsdom environment exposes its JSDOM, which sets the page's URL.
+  ;(
+    globalThis as { jsdom?: { reconfigure(settings: { url: string }): void } }
+  ).jsdom?.reconfigure({ url: `${origin}/` })
+  return {
+    origin,
+    close: () =>
+      new Promise((resolve) => {
+        setGlobalOrigin(undefined)
+        for (const socket of sockets) socket.destroy()
+        server.close(() => resolve())
+      }),
+  }
+}
+
+type Match = (operation: A2Operation) => boolean
+
+export type Held = {
+  /** The parked operation, once it arrives. */
+  arrived: Promise<A2Operation>
+  allow(): void
+  deny(): void
+}
+
+/**
+ * An `authorize` for `server.fetch` that records every operation it sees,
+ * and can park or deny the next one a test names.
+ */
+export function authorizer(): {
+  authorize: (operation: A2Operation) => Promise<boolean>
+  /** Every operation so far, in arrival order. */
+  operations: A2Operation[]
+  /** Parks the next matching operation until the test decides it. */
+  hold(match: Match): Held
+  /** Denies the next matching operation. */
+  deny(match: Match): void
+  /** Forgets the rules not yet used, and denies what is still parked. */
+  release(): void
+} {
+  const operations: A2Operation[] = []
+  const rules: {
+    match: Match
+    decide: (operation: A2Operation) => Promise<boolean>
+  }[] = []
+  const parked = new Set<Held>()
+  return {
+    operations,
+    authorize: async (operation) => {
+      operations.push(operation)
+      const at = rules.findIndex((rule) => rule.match(operation))
+      return at === -1 || rules.splice(at, 1)[0]!.decide(operation)
+    },
+    hold(match) {
+      const arrival = Promise.withResolvers<A2Operation>()
+      const decision = Promise.withResolvers<boolean>()
+      const decide = (allowed: boolean): void => {
+        parked.delete(held)
+        decision.resolve(allowed)
+      }
+      const held: Held = {
+        arrived: arrival.promise,
+        allow: () => decide(true),
+        deny: () => decide(false),
+      }
+      parked.add(held)
+      rules.push({
+        match,
+        decide: (operation) => {
+          arrival.resolve(operation)
+          return decision.promise
+        },
+      })
+      return held
+    },
+    deny(match) {
+      rules.push({ match, decide: async () => false })
+    },
+    release() {
+      rules.length = 0
+      for (const held of parked) held.deny()
+    },
+  }
+}

package/examples/playground/test/server-only.ts

@@ -0,0 +1 @@
+export {}

package/examples/playground/vitest.config.ts

@@ -0,0 +1,27 @@
+import { fileURLToPath } from 'node:url'
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  resolve: {
+    tsconfigPaths: true,
+    alias: {
+      'server-only': fileURLToPath(
+        new URL('./test/server-only.ts', import.meta.url),
+      ),
+    },
+  },
+  test: {
+    // The app's servers run in-process: with no store URL they use A2's
+    // memory store, never a database a dev shell points at, and with no
+    // queues their work drains in the background of the append. Without
+    // credentials, nothing can reach the AI Gateway.
+    env: {
+      REDIS_URL: '',
+      DATABASE_URL: '',
+      VERCEL: '',
+      A2_QUEUES: '',
+      AI_GATEWAY_API_KEY: '',
+      VERCEL_OIDC_TOKEN: '',
+    },
+  },
+})

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "experimental-a2",
-  "version": "0.18.1",
+  "version": "0.20.0",
   "description": "Syncs your immutable event log, from server to browser: one contract, handlers, derived state, and a live client per session.",
   "license": "MIT",
   "type": "module",
@@ -21,6 +21,7 @@
     "dist",
     "docs",
     "examples",
+    "sql",
     "src",
     "AGENTS.md",
     "CHANGELOG.md",
@@ -120,6 +121,7 @@
     "@types/ws": "^8.18.1",
     "@vercel/queue": "^0.2.0",
     "pg": "^8.16.0",
+    "postgres": "^3.4.9",
     "ws": "^8.21.3",
     "@testing-library/react": "^16.1.0",
     "@types/react": "^19.0.0",

package/sql/postgres/notifications.sql

@@ -0,0 +1,24 @@
+-- Notification triggers for experimental-a2/store-postgres, needed only with
+-- `listenConnectionString`. Apply after schema.sql. They wake live streams
+-- for every write to A2's event and presence tables, including writes from
+-- stores without the option. Safe to re-run.
+
+create or replace function a2_notify_events() returns trigger language plpgsql as $$
+begin
+  perform pg_notify('a2_events_' || md5(new.session_id), '');
+  return null;
+end $$;
+create or replace function a2_notify_presence() returns trigger language plpgsql as $$
+begin
+  perform pg_notify('a2_presence_' || md5(case when TG_OP = 'DELETE' then old.ns else new.ns end),
+    coalesce(current_setting('a2.presence_source', true), ''));
+  return null;
+end $$;
+do $$ begin
+  if not exists (select 1 from pg_trigger where tgname = 'a2_notify_events' and tgrelid = 'a2_events'::regclass) then
+    create trigger a2_notify_events after insert on a2_events for each row execute function a2_notify_events();
+  end if;
+  if not exists (select 1 from pg_trigger where tgname = 'a2_notify_presence' and tgrelid = 'a2_presence'::regclass) then
+    create trigger a2_notify_presence after insert or update or delete on a2_presence for each row execute function a2_notify_presence();
+  end if;
+end $$;

package/sql/postgres/schema.sql

@@ -0,0 +1,137 @@
+-- A2's tables for experimental-a2/store-postgres.
+--
+-- Apply this file once before the first deploy, and again whenever a release
+-- raises the schema version. Every statement creates only what is missing and
+-- leaves existing data alone, so the file is safe to re-run on any A2
+-- database, whichever release created it.
+--
+-- Each index sits behind a catalog check: `create index if not exists` takes
+-- a share lock on its table, and requires owning it, even when the index
+-- exists. The check makes a re-run on a live database wait for no writer.
+
+create table if not exists a2_events (
+  session_id   text    not null,
+  idx          bigint  not null,
+  event_type   text    not null,
+  payload      jsonb   not null,
+  event_id     text    not null,
+  created_at   bigint  not null,
+  cause        jsonb,
+  lane         text,
+  lane_ready   boolean not null default false,
+  processed_at bigint,
+  processed_by_attempt integer,
+  returned_event_ids jsonb,
+  first_claimed_at bigint,
+  last_claimed_at bigint,
+  attempt_count integer not null default 0,
+  failure_count integer not null default 0,
+  last_failed_at bigint,
+  last_failed_attempt integer,
+  last_error   text,
+  failed_at    bigint,
+  claim_holder text,
+  claim_expires_at bigint,
+  primary key (session_id, idx)
+);
+
+do $$ begin
+  if to_regclass('a2_events_event_id') is null then
+    create unique index if not exists a2_events_event_id on a2_events (event_id);
+  end if;
+  if to_regclass('a2_events_unprocessed') is null then
+    create index if not exists a2_events_unprocessed
+      on a2_events (session_id, idx) where processed_at is null;
+  end if;
+  if to_regclass('a2_events_lane_pending') is null then
+    create index if not exists a2_events_lane_pending
+      on a2_events (session_id, lane, idx) where processed_at is null;
+  end if;
+  if to_regclass('a2_events_dispatch_ready') is null then
+    create index if not exists a2_events_dispatch_ready
+      on a2_events (session_id, idx)
+      where lane_ready and processed_at is null and failed_at is null;
+  end if;
+  if to_regclass('a2_events_claim_expiry') is null then
+    create index if not exists a2_events_claim_expiry
+      on a2_events (session_id, claim_expires_at)
+      where lane_ready and processed_at is null and failed_at is null;
+  end if;
+end $$;
+
+create table if not exists a2_snapshots (
+  session_id   text   not null,
+  reducer_name text   not null,
+  up_to_index  bigint not null,
+  state        jsonb  not null,
+  updated_at   bigint not null,
+  primary key (session_id, reducer_name)
+);
+
+create table if not exists a2_snapshot_history (
+  session_id   text   not null,
+  reducer_name text   not null,
+  up_to_index  bigint not null,
+  state        jsonb  not null,
+  updated_at   bigint not null,
+  primary key (session_id, reducer_name, up_to_index)
+);
+
+create table if not exists a2_snapshot_pins (
+  session_id   text   not null,
+  event_index  bigint not null,
+  reducer_name text   not null,
+  up_to_index  bigint not null,
+  primary key (session_id, event_index, reducer_name, up_to_index)
+);
+
+do $$ begin
+  if to_regclass('a2_snapshot_pins_checkpoint') is null then
+    create index if not exists a2_snapshot_pins_checkpoint
+      on a2_snapshot_pins (session_id, reducer_name, up_to_index);
+  end if;
+end $$;
+
+create table if not exists a2_clears (
+  session_id  text   primary key,
+  clear_index bigint not null
+);
+
+create table if not exists a2_presence (
+  ns          text   not null,
+  participant text   not null,
+  field       text   not null,
+  value       jsonb  not null,
+  seen        bigint not null,
+  at          bigint not null,
+  expires_at  bigint not null,
+  primary key (ns, participant, field)
+);
+
+create table if not exists a2_schema (
+  only_row boolean primary key default true check (only_row),
+  version  integer not null
+);
+
+-- The version is recorded last, and only once everything above exists, so a
+-- run that skipped a failed statement records nothing, whichever tool ran it.
+do $$ begin
+  if to_regclass('a2_events') is null
+    or to_regclass('a2_events_event_id') is null
+    or to_regclass('a2_events_unprocessed') is null
+    or to_regclass('a2_events_lane_pending') is null
+    or to_regclass('a2_events_dispatch_ready') is null
+    or to_regclass('a2_events_claim_expiry') is null
+    or to_regclass('a2_snapshots') is null
+    or to_regclass('a2_snapshot_history') is null
+    or to_regclass('a2_snapshot_pins') is null
+    or to_regclass('a2_snapshot_pins_checkpoint') is null
+    or to_regclass('a2_clears') is null
+    or to_regclass('a2_presence') is null
+  then
+    raise exception 'A2 schema incomplete: an earlier statement in schema.sql failed. Fix it and apply the file again.';
+  end if;
+  insert into a2_schema (version) values (1)
+  on conflict (only_row)
+  do update set version = greatest(a2_schema.version, excluded.version);
+end $$;

package/src/ai-agent.ts

@@ -20,7 +20,6 @@ import type { Reducer } from './reducer.ts'
 import { aiEventId, AI_EVENT_ID_PREFIX } from './ai-id.ts'
 import type { ControlCommit, ControlCommand } from './ai-control.ts'
 import { AsyncLocalStorage } from 'node:async_hooks'
-import { createHash } from 'node:crypto'
 import { convertToModelMessages, isToolUIPart } from 'ai'
 import type {
   Experimental_SandboxSession,
@@ -35,6 +34,7 @@ import type {
 } from 'ai'
 import { generateAISDKStep, type AISDKStepSettings } from './ai-sdk-step.ts'
 import {
+  gatewaySessionAffinity,
   harnessDriver,
   harnessLog,
   harnessStates,
@@ -627,9 +627,10 @@ async function generateWithAISDK<
             ? headers
             : {
                 ...headers,
-                'x-session-affinity': createHash('sha256')
-                  .update(JSON.stringify([agentName, context.event.sessionId]))
-                  .digest('hex'),
+                'x-session-affinity': gatewaySessionAffinity({
+                  agentName,
+                  sessionId: context.event.sessionId,
+                }),
               },
           providerOptions: {
             ...context.generation.providerOptions,
@@ -1282,6 +1283,10 @@ export function createHandlers<
         throw new TypeError(
           `a harness runs tools without approval; '${name}' sets needsApproval`,
         )
+      if (tool.type === 'provider')
+        throw new TypeError(
+          `a harness runs only tools A2 executes; '${name}' is provider-defined`,
+        )
     }
   }
 

package/src/ai-client-state.ts

@@ -1,6 +1,11 @@
 import type { UIMessage } from 'ai'
 import type { AIState } from './ai.ts'
-import type { ControlCommand, InboxItem } from './ai-control.ts'
+import {
+  steeringEnd,
+  steersOn,
+  type ControlCommand,
+  type InboxItem,
+} from './ai-control.ts'
 import {
   interruptGenerationProjection,
   settleAbandonedCompaction,
@@ -92,10 +97,10 @@ export function projectAIClientState<M extends UIMessage>(
   const admit = (): void => {
     if (turnId !== null || paused || items.length === 0) return
     let count = 0
-    for (const item of items) {
+    for (const [index, item] of items.entries()) {
       count += 1
       show(item.message)
-      if (item.generate) {
+      if (item.generate && !steersOn(item, items.slice(index + 1))) {
         turnId = item.id
         starting = true
         suspended = false
@@ -118,18 +123,21 @@ export function projectAIClientState<M extends UIMessage>(
             ? { afterStepOf: command.turnId }
             : {}),
           id: command.message.id,
+          role: command.message.role,
           message: command.message,
           generate:
             command.message.role === 'user' &&
             (command.action === 'steer' || command.generate !== false),
         }
-        items = command.action === 'steer' ? [item, ...items] : [...items, item]
+        if (command.action === 'steer') {
+          const at = steeringEnd(items, command.turnId)
+          items = [...items.slice(0, at), item, ...items.slice(at)]
+        } else items = [...items, item]
         break
       }
       case 'edit':
         items = items.map((item) =>
-          item.id === command.message.id &&
-          item.message.role === command.message.role
+          item.id === command.message.id && item.role === command.message.role
             ? { ...item, message: command.message }
             : item,
         )
@@ -147,20 +155,24 @@ export function projectAIClientState<M extends UIMessage>(
       case 'send-now': {
         if (turnId !== command.turnId) break
         const selected = items.find((item) => item.id === command.inputId)
-        if (!selected || selected.message.role !== 'user') break
+        if (!selected || selected.role !== 'user') break
+        const { afterStepOf: _, ...promoted } = selected
         items = items.map((item) =>
           item.id === selected.id
             ? {
-                ...item,
+                ...promoted,
                 generate: true,
                 ...(turnId === null ? {} : { afterStepOf: turnId }),
               }
             : item,
         )
+        const others = items.filter((item) => item.id !== selected.id)
         items = moveItem({
           items,
           inputId: command.inputId,
-          beforeId: items[0]?.id ?? null,
+          beforeId:
+            others[turnId === null ? 0 : steeringEnd(others, turnId)]?.id ??
+            null,
         })
         break
       }

package/src/ai-context-schema.ts

@@ -1,6 +1,7 @@
 import { assertMessageBase, restoreMessageBase } from './ai-message-base.ts'
 import type { UIMessage } from 'ai'
 import { events } from './ai.ts'
+import { isUIMessage } from './ai-message-shape.ts'
 import type { AIContextState } from './ai-context.ts'
 import type { StandardSchemaV1 } from './standard-schema.ts'
 import { validateSync } from './validate.ts'
@@ -14,13 +15,34 @@ const id = (value: unknown): value is string =>
 const nullableId = (value: unknown): boolean => value === null || id(value)
 const matches = (schema: StandardSchemaV1<unknown>, value: unknown): boolean =>
   validateSync(schema, value, 'AI context checkpoint').issues === undefined
-const isMessage = (value: unknown): value is UIMessage =>
-  matches(events['ai.message.created'], { message: value })
 const indexes = (value: unknown, messages: readonly UIMessage[]): boolean =>
   record(value) &&
   Object.values(value).every(index) &&
   messages.every((message) => Object.hasOwn(value, message.id))
 
+const inputs = (value: unknown): boolean =>
+  record(value) &&
+  Object.keys(value).length === 2 &&
+  Array.isArray(value['commands']) &&
+  value['commands'].every(
+    (held) =>
+      record(held) &&
+      Object.keys(held).every((key) =>
+        ['commandId', 'inputId', 'message'].includes(key),
+      ) &&
+      id(held['commandId']) &&
+      id(held['inputId']) &&
+      (held['message'] === undefined || isUIMessage(held['message'])),
+  ) &&
+  Array.isArray(value['waiting']) &&
+  value['waiting'].every(
+    (input) =>
+      record(input) &&
+      Object.keys(input).length === 2 &&
+      id(input['inputId']) &&
+      isUIMessage(input['message']),
+  )
+
 const compaction = (value: unknown): boolean =>
   value === null ||
   (record(value) &&
@@ -109,7 +131,7 @@ const restorableBase = (options: {
   current: UIMessage | undefined
 }): boolean => {
   assertMessageBase(options.base)
-  return isMessage(
+  return isUIMessage(
     restoreMessageBase({ current: options.current, base: options.base }),
   )
 }
@@ -230,6 +252,7 @@ export const contextSchema = <M extends UIMessage>(
         'compaction',
         'modelMetadata',
         'closed',
+        'inputs',
         'authority',
         'generation',
         'activeProjection',
@@ -249,9 +272,10 @@ export const contextSchema = <M extends UIMessage>(
         const current = value['activeProjection']
         if (
           !Array.isArray(messages) ||
-          !messages.every(isMessage) ||
+          !messages.every(isUIMessage) ||
           !compaction(value['compaction']) ||
           typeof value['closed'] !== 'boolean' ||
+          !inputs(value['inputs']) ||
           !authority(value['authority']) ||
           !record(metadata) ||
           !Object.entries(metadata).every(