experimental-a2 0.19.0 → 0.20.0

package/examples/playground/test/model.ts

@@ -0,0 +1,149 @@
+/**
+ * The model behind an agent under test. Agents take AI Gateway model IDs;
+ * the AI SDK's global provider serves every one of them from this model,
+ * whose steps the test writes part by part.
+ */
+
+import { MockLanguageModelV4 } from 'ai/test'
+import { vi } from 'vitest'
+
+export type StreamPart =
+  Awaited<
+    ReturnType<MockLanguageModelV4['doStream']>
+  >['stream'] extends ReadableStream<infer Part>
+    ? Part
+    : never
+
+export type CallOptions = MockLanguageModelV4['doStreamCalls'][number]
+
+/** One model step, open until the test ends it or the agent aborts it. */
+export type ModelStep = {
+  /** What the agent sent the model for this step. */
+  options: CallOptions
+  write(...parts: StreamPart[]): void
+  /** Ends the step the way a model does: with a finish part. */
+  finish(reason?: 'stop' | 'tool-calls'): void
+}
+
+const usage = {
+  inputTokens: { total: 1, noCache: 1, cacheRead: 0, cacheWrite: 0 },
+  outputTokens: { total: 1, text: 1, reasoning: 0 },
+}
+
+/** A text part, whole. */
+export const say = ({ id, text }: { id: string; text: string }) =>
+  [
+    { type: 'text-start', id },
+    { type: 'text-delta', id, delta: text },
+    { type: 'text-end', id },
+  ] satisfies StreamPart[]
+
+/** A tool call whose input streams in as the given chunks of its JSON. */
+export const toolCall = ({
+  id,
+  name,
+  chunks,
+}: {
+  id: string
+  name: string
+  chunks: string[]
+}) =>
+  [
+    { type: 'tool-input-start', id, toolName: name },
+    ...chunks.map((delta) => ({
+      type: 'tool-input-delta' as const,
+      id,
+      delta,
+    })),
+    { type: 'tool-input-end', id },
+    {
+      type: 'tool-call',
+      toolCallId: id,
+      toolName: name,
+      input: chunks.join(''),
+    },
+  ] satisfies StreamPart[]
+
+let models = 0
+
+/**
+ * Installs a model the test drives step by step, served for `id` alone, so
+ * an agent left over from an earlier test can't take its steps. Remove it
+ * with `vi.unstubAllGlobals()`.
+ */
+export function testModel({
+  id = `test-model-${++models}`,
+}: {
+  /** The model ID the agent asks for. */
+  id?: string
+} = {}): {
+  /** The agent's `model`. */
+  id: string
+  /** Every step the agent asked for, in order. */
+  steps: ModelStep[]
+  /** The `n`th step (from 0), once the agent asks for it. */
+  step(n: number): Promise<ModelStep>
+  /** Ends every step still open, and each one asked for from now on. */
+  close(): void
+} {
+  const steps: ModelStep[] = []
+  const arrivals: PromiseWithResolvers<ModelStep>[] = []
+  const arrival = (n: number): PromiseWithResolvers<ModelStep> =>
+    (arrivals[n] ??= Promise.withResolvers<ModelStep>())
+  const ends: (() => void)[] = []
+  let closed = false
+  const model = new MockLanguageModelV4({
+    modelId: id,
+    doStream: async (options) => {
+      let controller!: ReadableStreamDefaultController<StreamPart>
+      let open = true
+      const end = (): void => {
+        if (!open) return
+        open = false
+        controller.close()
+      }
+      const stream = new ReadableStream<StreamPart>({
+        start(value) {
+          controller = value
+          controller.enqueue({ type: 'stream-start', warnings: [] })
+        },
+      })
+      options.abortSignal?.addEventListener('abort', end, { once: true })
+      const step: ModelStep = {
+        options,
+        write(...parts) {
+          for (const part of parts) if (open) controller.enqueue(part)
+        },
+        finish(reason = 'stop') {
+          step.write({
+            type: 'finish',
+            finishReason: { unified: reason, raw: reason },
+            usage,
+          })
+          end()
+        },
+      }
+      ends.push(end)
+      arrival(steps.length).resolve(step)
+      steps.push(step)
+      if (closed) step.finish()
+      return { stream }
+    },
+  })
+  vi.stubGlobal('AI_SDK_DEFAULT_PROVIDER', {
+    specificationVersion: 'v4',
+    languageModel: (requested: string) => {
+      if (requested !== id) throw new Error(`no test model '${requested}'`)
+      return model
+    },
+  })
+  return {
+    id,
+    steps,
+    step: (n) => arrival(n).promise,
+    close: () => {
+      closed = true
+      for (const end of ends) end()
+    },
+  }
+}

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.19.0",
+  "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,

package/src/ai-fx.ts

@@ -20,6 +20,7 @@ import {
   Divergence,
   gatewayCredential,
   type GatewayCredential,
+  gatewaySessionAffinity,
   type Harness,
   type HarnessGeneration,
   harnessImages,
@@ -244,19 +245,27 @@ async function runStep(step: HarnessStep): Promise<HarnessGeneration> {
       return pending(init.signal)
     }
     state.live = true
-    return live(state, url, init, compaction)
+    return live(state, url, { ...init, body }, compaction)
   }
 
+  const affinity = gatewaySessionAffinity({
+    agentName: step.agentName,
+    sessionId: step.sessionId,
+  })
   let compactions = 0
   const live = async (
     state: TurnState,
     url: string,
-    init: RequestInit,
+    init: RequestInit & { body: string },
     compaction: boolean,
   ): Promise<Response> => {
+    const headers = withAuth(init.headers, credential)
+    // The kernel's own affinity is new with every rebuild; A2's lasts the session.
+    headers.set('x-session-affinity', affinity)
     const response = await fetch(url, {
       ...init,
-      headers: withAuth(init.headers, credential),
+      body: withGatewayCaching(init.body),
+      headers,
       signal: network,
     })
     // Abandoned before its headers came: the kernel has moved on without it.
@@ -774,6 +783,31 @@ async function loadLibfx(): Promise<typeof import('libfx')> {
   }
 }
 
+/**
+ * A model request with Gateway's automatic caching, as the AI SDK harness
+ * asks for it; options the request already has take precedence. Every other
+ * byte stays as the kernel wrote it, a tool input's spacing included.
+ */
+function withGatewayCaching(body: string): string {
+  const request = JSON.parse(body) as {
+    providerOptions?: Record<string, Record<string, unknown> | undefined>
+  }
+  const options = request.providerOptions
+  if (options === undefined)
+    return body.replace(
+      /^\{/,
+      '{"providerOptions":{"gateway":{"caching":"auto"}},',
+    )
+  // The kernel writes its own options last; any other shape goes unchanged.
+  const own = `"providerOptions":${JSON.stringify(options)}}`
+  if (!body.endsWith(own)) return body
+  const merged = {
+    ...options,
+    gateway: { caching: 'auto', ...options['gateway'] },
+  }
+  return `${body.slice(0, -own.length)}"providerOptions":${JSON.stringify(merged)}}`
+}
+
 function withAuth(
   headers: HeadersInit | undefined,
   credential: GatewayCredential,

package/src/ai-harness.ts

@@ -14,6 +14,7 @@ import type {
   ReasoningLevel,
   ToolResultPayload,
 } from './ai.ts'
+import { createHash } from 'node:crypto'
 import { STEERED_REASON } from './ai-control.ts'
 import { boundedCache, type BoundedCache } from './bounded-cache.ts'
 import type { ContractEvent, EventDefs } from './contract.ts'
@@ -598,6 +599,23 @@ export function gatewayCredential(harness: string): GatewayCredential {
   throw new Error(`${harness}() needs AI_GATEWAY_API_KEY or VERCEL_OIDC_TOKEN`)
 }
 
+/**
+ * The `x-session-affinity` an agent session's Gateway calls carry: stable for
+ * the session, so supporting providers can reuse their cache, and free of its
+ * raw ID.
+ */
+export function gatewaySessionAffinity({
+  agentName,
+  sessionId,
+}: {
+  agentName: string
+  sessionId: string
+}): string {
+  return createHash('sha256')
+    .update(JSON.stringify([agentName, sessionId]))
+    .digest('hex')
+}
+
 // Where Vercel Functions keep the current request, as @vercel/oidc reads it.
 const REQUEST_CONTEXT = Symbol.for('@vercel/request-context')
 

package/src/errors.ts

@@ -7,7 +7,7 @@ export type A2ServerErrorCode =
   | 'SESSION_CLEARED' // a read or call reached across the session's latest clear
   | 'CLAIM_EXPIRED' // ctx.signal abort reason: the attempt's lease lapsed without renewal
   | 'STORE_UNAVAILABLE' // store backend failure — retryable, nothing was written
-  | 'STORE_NOT_CONFIGURED' // production boot with no `store` passed
+  | 'STORE_NOT_CONFIGURED' // no `store` in production, a missing peer dependency, or Postgres lacks A2's tables
   | 'UNKNOWN_PRESENCE_FIELD' // presence field not in the contract's `presence` map
   | 'PRESENCE_NOT_SUPPORTED' // contract declares presence but the store lacks the capability