experimental-a2 0.17.0 → 0.18.1

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.ts

@@ -6,6 +6,7 @@
 
 export { agent, runWithHandlerContext } from './ai-agent.ts'
 export type {
+  AISDKHarness,
   Agent,
   AgentFromContractOptions,
   AgentGenerate,
@@ -19,8 +20,9 @@ export type {
   CompactionOptions,
   CompactionPolicy,
   Instructions,
-  LanguageModel,
   ModelSettings,
   Resolvable,
+  ToolRecovery,
   ToolSet,
 } from './ai-agent.ts'
+export type { Harness } from './ai-harness.ts'

package/src/ai-tool-input.ts

@@ -0,0 +1,45 @@
+import {
+  asSchema,
+  InvalidToolInputError,
+  type ToolSet,
+  TypeValidationError,
+} from 'ai'
+
+export type SchemaInput =
+  { success: true; value: unknown } | { success: false; errorText: string }
+
+/**
+ * A call's input through its tool's schema, as the AI SDK step reads it: the
+ * schema's output with its defaults, or the AI SDK's error text. A name the
+ * tool set does not own has no schema.
+ */
+export async function schemaInput({
+  tools,
+  toolName,
+  value,
+}: {
+  tools: ToolSet
+  toolName: string
+  value: unknown
+}): Promise<SchemaInput> {
+  const schema = Object.hasOwn(tools, toolName)
+    ? asSchema(tools[toolName]!.inputSchema)
+    : undefined
+  if (schema?.validate === undefined) return { success: true, value }
+  let cause: unknown
+  try {
+    const result = await schema.validate(value)
+    if (result.success) return { success: true, value: result.value }
+    cause = result.error
+  } catch (error) {
+    cause = error
+  }
+  return {
+    success: false,
+    errorText: new InvalidToolInputError({
+      toolName,
+      toolInput: JSON.stringify(value),
+      cause: TypeValidationError.wrap({ value, cause }),
+    }).message,
+  }
+}

package/src/internal.ts

@@ -73,6 +73,13 @@ export type ServerInternals = {
 
 export const serverInternals: WeakMap<object, ServerInternals> = new WeakMap()
 
+/**
+ * Handler contexts whose previous attempt ended in a failure the handler
+ * threw, as opposed to a claim that lapsed (a crash, a timeout, a frozen
+ * process). Only the attempt's own outcome is known, not how far it got.
+ */
+export const failedPreviousAttempt: WeakSet<object> = new WeakSet()
+
 const objectLike = (value: unknown): value is object =>
   (typeof value === 'object' && value !== null) || typeof value === 'function'
 
@@ -267,12 +274,18 @@ export const DRAIN_TIMINGS: {
   watchdogGraceMs: number
   /** Maximum time a top-level append waits for its initial scheduler arm. */
   schedulerArmTimeoutMs: number
+  /** First retry delay after a caught failure, doubling per failure. */
+  failureRetryBaseMs: number
+  /** Ceiling of the caught-failure retry delay. */
+  failureRetryMaxMs: number
 } = {
   claimTtlMs: 5_000,
   claimHeartbeatMs: 2_000,
   renewRetryMs: 400,
   watchdogGraceMs: 1_000,
   schedulerArmTimeoutMs: 2_000,
+  failureRetryBaseMs: 1_000,
+  failureRetryMaxMs: 30_000,
 }
 
 /**

package/src/libfx.d.ts

@@ -0,0 +1,34 @@
+// The subset of libfx's agent API that experimental-a2/ai/fx uses; libfx
+// ships no type declarations.
+declare module 'libfx' {
+  export type FxHostTool = {
+    name: string
+    description: string
+    inputSchema: object
+    execute(input: unknown, context: { signal: AbortSignal }): unknown
+  }
+
+  export type FxTurn = AsyncIterable<{ type: string }> & {
+    result: Promise<{ stopReason: string }>
+    steer(input: string): Promise<void>
+    cancel(): void
+  }
+
+  export type FxAgent = {
+    prompt(input: string, options?: { signal?: AbortSignal }): FxTurn
+    checkpoint(): Promise<Uint8Array>
+    close(): Promise<void>
+  }
+
+  export function createFxAgent(options: {
+    apiKey: string
+    model?: string | { id: string; effort?: string; fast?: boolean }
+    instructions?: string | string[]
+    tools?: FxHostTool[]
+    checkpoint?: Uint8Array
+    fetch?: (
+      input: string | URL | Request,
+      init?: RequestInit,
+    ) => Promise<Response>
+  }): Promise<FxAgent>
+}

package/src/reducer.ts

@@ -62,7 +62,7 @@ export function foldEvent<D extends EventDefs, S>(
  * the browser cache) are keyed by it, which makes it the invalidation
  * knob too: changed the fold's logic? Change the name. Snapshots under
  * the old name are ignored and the next read refolds from raw events.
- * See docs/concepts/04-state.mdx.
+ * See docs/contract/02-state.mdx.
  */
 export type Reducer<
   D extends EventDefs = EventDefs,

package/src/scheduler-qstash.ts

@@ -912,7 +912,7 @@ export function qstash(options: QStashSchedulerOptions = {}): A2Scheduler {
         }
         if (result.settled) return response.ok()
 
-        if (result.outcome === 'busy') {
+        if (result.outcome === 'busy' || result.outcome === 'stalled') {
           if (result.schedulerArm) await result.schedulerArm
           else {
             await sendTask({
@@ -926,10 +926,10 @@ export function qstash(options: QStashSchedulerOptions = {}): A2Scheduler {
           return response.ok()
         }
 
-        // QStash's per-message retry default can be lower than A2's failure
-        // budget. Hand stalled work to a fresh durable message so recovery
-        // continues until the store reports success or dead-letter. If this
-        // publish fails, the throw keeps the current delivery retryable.
+        // A drain whose outcome is unknown hands its work to a fresh durable
+        // message: QStash's per-message retry default can be lower than A2's
+        // failure budget. If this publish fails, the throw keeps the current
+        // delivery retryable.
         await sendTask({
           version: 1,
           kind: 'drain',

package/src/scheduler-vercel.ts

@@ -14,8 +14,7 @@
  *   the configured deployment scope for sending and receiving.
  * - `handler(...servers)` is the delivery route: look up the server for
  *   the contract named in the message and drain it. Settled drains ack;
- *   contention arms a later watchdog before acking; handler failure keeps
- *   the current message and redelivers with backoff.
+ *   contention and handler failure arm a later watchdog before acking.
  *
  * `@vercel/queue` is an optional peer dependency, imported lazily. The
  * `transport` option is the injection seam — tests run a real little
@@ -44,12 +43,12 @@ import {
 } from './server.ts'
 
 /**
- * The "not settled yet" signal for a stalled handler — an expected state,
- * not a transport failure. The default transport classifies it in
- * `@vercel/queue`'s `retry` option and reschedules the delivery quietly
+ * The "not settled yet" signal for a drain whose outcome is unknown — an
+ * expected state, not a transport failure. The default transport classifies
+ * it in `@vercel/queue`'s `retry` option and reschedules the delivery quietly
  * (`{ afterSeconds }` → visibility change → 200 response, no error log)
- * instead of letting it surface as a thrown 500. Busy deliveries use a
- * fresh watchdog instead. Custom transports should treat this class the
+ * instead of letting it surface as a thrown 500. Busy and stalled deliveries
+ * use a fresh watchdog instead. Custom transports should treat this class the
  * same way when their protocol has a polite reschedule.
  */
 export class DrainUnsettled extends Error {}
@@ -354,11 +353,12 @@ export function vercelQueues(options: VercelQueuesOptions = {}): A2Scheduler {
               }
           if (result.settled) return
 
-          if (result.outcome === 'busy') {
-            // The live owner normally arms its own claim windows.
-            // Ensure a later slot here as well before relinquishing this
-            // durable delivery. If the queue is down, the throw leaves the
-            // current message unacked for visibility-timeout redelivery.
+          if (result.outcome === 'busy' || result.outcome === 'stalled') {
+            // The live owner arms its own claim windows, and the store holds
+            // a failed event until its retry time. Ensure a later slot here
+            // as well before relinquishing this durable delivery. If the queue is down, the
+            // throw leaves the current message unacked for visibility-timeout
+            // redelivery.
             if (result.schedulerArm) await result.schedulerArm
             else {
               await sendTask({
@@ -372,9 +372,8 @@ export function vercelQueues(options: VercelQueuesOptions = {}): A2Scheduler {
             return
           }
 
-          // Handler failure keeps the current durable message and uses the
-          // transport retry path. Dead-lettered sessions count as settled
-          // above, so retries stop there.
+          // A drain whose outcome is unknown keeps the current durable message
+          // and uses the transport retry path.
           throw new DrainUnsettled(
             `a2 scheduler: session '${task.sessionId}' of '${task.contract}' not settled yet (${result.outcome}) — rescheduled for redelivery`,
           )

package/src/server.ts

@@ -25,6 +25,7 @@ import {
   DRAIN_TIMINGS,
   POLL_TIMINGS,
   MAX_DATE_MS,
+  failedPreviousAttempt,
   RESERVED_PARTICIPANT_IDS,
   STATE_CACHE,
   markSchedulerSendFailure,
@@ -55,6 +56,7 @@ import type {
   StoreClaimAvailableResult,
   StoreClearResult,
   StoreReclaimResult,
+  StoreRedriveResult,
   StoreSnapshotWrite,
   StoreStateRead,
   PresenceRow,
@@ -188,6 +190,48 @@ export type SessionClearOptions = {
 
 export type SessionClearResult = { cleared: boolean; index: number }
 
+export type SessionRedriveOptions = {
+  /** `retry` runs the handler again with a fresh failure budget; `skip` settles the event without running it. */
+  action: 'retry' | 'skip'
+}
+
+export type SessionRedriveResult = {
+  /** Whether the event was dead-lettered and is now resolved. */
+  redriven: boolean
+}
+
+/** A dead-lettered event, with what an operator needs to resolve it. */
+export type DeadLetter<D extends EventDefs> = ContractEvent<D> & {
+  lane: string | null
+  /** Durable dispatch claims, including claims abandoned by hard kills. */
+  attemptCount: number
+  /**
+   * Failed attempts since the event was appended or last retried: caught
+   * handler failures, and attempts that ended before they finished.
+   */
+  failureCount: number
+  /**
+   * The last error a handler threw, stringified; null when every failed
+   * attempt ended before it finished.
+   */
+  lastError: string | null
+  failedAt: Date
+}
+
+/** Dead-letter resolution; root sessions only. */
+export type SessionDeadLetters<D extends EventDefs> = {
+  /** The session's dead-lettered events, oldest first. */
+  deadLetters(): Promise<DeadLetter<D>[]>
+  /**
+   * Atomically resolve one dead-lettered event, then wake the session like
+   * an append. Any other event is left alone and reports `redriven: false`.
+   */
+  redrive(
+    index: number,
+    options: SessionRedriveOptions,
+  ): Promise<SessionRedriveResult>
+}
+
 /**
  * The presence members of a session — intersected in via
  * `WithPresence`, so they exist exactly when the contract declares
@@ -293,7 +337,9 @@ export type A2Server<
     (request: Request): Promise<Response>
     (request: Request, options: ServerFetchOptions<D, P>): Promise<Response>
   }
-  session(id: string): Session<D, SessionAppend<D>, P, PublicState<R>>
+  session(
+    id: string,
+  ): Session<D, SessionAppend<D>, P, PublicState<R>> & SessionDeadLetters<D>
   /**
    * Process every currently eligible event. `settled` means nothing
    * actionable remains, including work blocked behind a dead letter.
@@ -459,6 +505,13 @@ export type ServerOptions<
 
 const MAX_FAILURES = 10
 
+/** How long an event waits after its `failureCount`th failed attempt. */
+const failureRetryDelayMs = (failureCount: number): number =>
+  Math.min(
+    DRAIN_TIMINGS.failureRetryBaseMs * 2 ** Math.max(failureCount - 1, 0),
+    DRAIN_TIMINGS.failureRetryMaxMs,
+  )
+
 /** Abort a lost attempt's handler and release its drain slot. */
 function endSuperseded(
   entry: {
@@ -487,6 +540,7 @@ type ClaimWindow = {
   expiresAtMs: number
   watchdogAtMs: number
   deadlineCapped: boolean
+  pastDeadline: boolean
 }
 
 type DrainSignal = {
@@ -561,9 +615,11 @@ function pruneAbortHub(hub: AbortHub): void {
 }
 
 /**
- * Deadlines tighten the claim/watchdog window; they never decide whether a
- * handler starts. A platform timeout therefore follows the same path as any
- * other process death: the in-flight event remains pending and is retried.
+ * Deadlines tighten the claim/watchdog window, and a passed deadline ends
+ * claiming: a claim then could not run, only lapse and be claimed again until
+ * the platform ends the process. A platform timeout otherwise follows the
+ * same path as any other process death: the in-flight event remains pending
+ * and a fresh invocation retries it.
  */
 function nextClaimWindow(): ClaimWindow {
   const nowMs = Date.now()
@@ -579,6 +635,7 @@ function nextClaimWindow(): ClaimWindow {
     expiresAtMs,
     watchdogAtMs: expiresAtMs + DRAIN_TIMINGS.watchdogGraceMs,
     deadlineCapped,
+    pastDeadline: deadlineMs !== null && deadlineTtlMs <= 0,
   }
 }
 
@@ -603,6 +660,27 @@ const devDefaultStore = retryableLazy(() =>
   import('./store-sqlite.ts').then((m) => m.sqlite()),
 )
 
+/** Waits for a scheduler arm for at most `schedulerArmTimeoutMs`. */
+async function joinSchedulerArm(arm: Promise<void>): Promise<void> {
+  let timeout: ReturnType<typeof setTimeout> | null = null
+  try {
+    await Promise.race([
+      arm,
+      new Promise<never>((_, reject) => {
+        timeout = setTimeout(() => {
+          reject(
+            new Error(
+              `a2 scheduler arm timed out after ${DRAIN_TIMINGS.schedulerArmTimeoutMs}ms`,
+            ),
+          )
+        }, DRAIN_TIMINGS.schedulerArmTimeoutMs)
+      }),
+    ])
+  } finally {
+    if (timeout) clearTimeout(timeout)
+  }
+}
+
 function environment(): 'development' | 'test' | 'production' {
   const env =
     typeof process === 'undefined' ? undefined : process.env?.['NODE_ENV']
@@ -1081,26 +1159,12 @@ export function createServer<
           await dispatchArm
         }
         if (initialArm) {
-          let timeout: ReturnType<typeof setTimeout> | null = null
           try {
-            await Promise.race([
-              initialArm,
-              new Promise<never>((_, reject) => {
-                timeout = setTimeout(() => {
-                  reject(
-                    new Error(
-                      `a2 scheduler arm timed out after ${DRAIN_TIMINGS.schedulerArmTimeoutMs}ms`,
-                    ),
-                  )
-                }, DRAIN_TIMINGS.schedulerArmTimeoutMs)
-              }),
-            ])
+            await joinSchedulerArm(initialArm)
           } catch (err) {
             span.setAttribute('a2.append.armed', false)
             span.recordError(err)
             if (requireDurableInitialArm) throw err
-          } finally {
-            if (timeout) clearTimeout(timeout)
           }
         }
         return appended
@@ -2483,6 +2547,102 @@ export function createServer<
     return { cleared: result.cleared, index: result.index }
   }
 
+  // ── dead letters ─────────────────────────────────────────────────
+
+  const redrivingStore = async (
+    method: 'deadLetters' | 'redrive',
+  ): Promise<A2Store & Required<Pick<A2Store, 'deadLetters' | 'redrive'>>> => {
+    const store = await resolveStore()
+    if (!store.deadLetters || !store.redrive) {
+      throw new TypeError(
+        `session.${method} requires a store that supports redriving dead letters`,
+      )
+    }
+    return store as A2Store & Required<Pick<A2Store, 'deadLetters' | 'redrive'>>
+  }
+
+  const listDeadLetters = async (
+    sessionId: string,
+  ): Promise<DeadLetter<D>[]> => {
+    const store = await redrivingStore('deadLetters')
+    let rows: StoredEvent[]
+    try {
+      rows = await store.deadLetters(nsId(sessionId))
+    } catch (err) {
+      throw asStoreUnavailable(err)
+    }
+    return rows.flatMap((row) =>
+      row.failedAt === null
+        ? []
+        : [
+            {
+              ...toPublic(row),
+              lane: row.lane,
+              attemptCount: row.attemptCount,
+              failureCount: row.failureCount,
+              lastError: row.lastError,
+              failedAt: row.failedAt,
+            },
+          ],
+    )
+  }
+
+  const redriveEvent = async (
+    sessionId: string,
+    index: number,
+    redriveOptions: SessionRedriveOptions,
+  ): Promise<SessionRedriveResult> => {
+    assertRedriveIndex(index)
+    const action = redriveOptions?.action
+    if (action !== 'retry' && action !== 'skip') {
+      throw new TypeError("redrive requires action 'retry' or 'skip'")
+    }
+    return telemetry.span(
+      'a2.redrive',
+      {
+        'a2.contract': name,
+        'a2.session_id': sessionId,
+        'a2.redrive.index': index,
+        'a2.redrive.action': action,
+      },
+      async (span) => {
+        const store = await redrivingStore('redrive')
+        let result: StoreRedriveResult
+        try {
+          result = await store.redrive({
+            sessionId: nsId(sessionId),
+            index,
+            action,
+          })
+        } catch (err) {
+          throw asStoreUnavailable(err)
+        }
+        span.setAttribute('a2.redrive.redriven', result.redriven)
+        // Healing like a top-level append, even when nothing was redriven, so
+        // repeating a redrive whose wakeup was lost resumes the session.
+        if (result.hasPending) {
+          const watchdogAt = scheduler
+            ? schedulerSlot(nextClaimWindow().watchdogAtMs)
+            : undefined
+          const arm =
+            watchdogAt === undefined
+              ? null
+              : startSchedulerArm(sessionId, watchdogAt)
+          scheduleDrain(sessionId, watchdogAt)
+          if (arm) {
+            try {
+              await joinSchedulerArm(arm)
+            } catch (err) {
+              span.setAttribute('a2.redrive.armed', false)
+              span.recordError(err)
+            }
+          }
+        }
+        return { redriven: result.redriven }
+      },
+    )
+  }
+
   const makeSession = <Append>(
     id: string,
     append: Append,
@@ -2584,7 +2744,7 @@ export function createServer<
         generatedIds,
       )
     }
-    return {
+    const ctx: HandlerContext<D> = {
       event: toPublic(trigger),
       attempt,
       // The handler-scoped `setPresence` stamps `seen` with the
@@ -2597,6 +2757,9 @@ export function createServer<
       }),
       signal,
     }
+    if (attempt > 1 && trigger.lastFailedAttempt === attempt - 1)
+      failedPreviousAttempt.add(ctx)
+    return ctx
   }
 
   const abortHubs = new Map<string, AbortHub>()
@@ -2922,6 +3085,25 @@ export function createServer<
           return tracked
         }
 
+        /** Records an event the store dead-lettered on takeover; it never runs. */
+        const reportExhausted = (event: StoredEvent): Promise<void> =>
+          telemetry.span(
+            'a2.event',
+            {
+              'a2.contract': name,
+              'a2.session_id': sessionId,
+              'a2.event.type': event.type,
+              'a2.event.index': event.index,
+              'a2.event.id': event.id,
+              'a2.event.attempt': event.attemptCount,
+              'a2.event.handled': handlers.has(event.type),
+              ...(event.lane === null ? {} : { 'a2.event.lane': event.lane }),
+            },
+            async (eventSpan) => {
+              eventSpan.setAttribute('a2.event.outcome', 'dead_lettered')
+            },
+          )
+
         const runEvent = async (
           event: StoredEvent,
           supersedeSignal: AbortSignal,
@@ -2981,9 +3163,8 @@ export function createServer<
                 if (supersedeSignal.aborted && err === supersedeSignal.reason) {
                   // The handler surrendered to its own abort — identity, not
                   // code, so it cannot arrive from another context. A lapse
-                  // is not a handler failure: no completion, no failure
-                  // budget spent; the claim simply expires and recovery
-                  // retries.
+                  // is not a handler failure: no completion; the claim simply
+                  // expires, and the claim that takes it over counts it.
                   eventSpan.setAttribute('a2.event.outcome', 'surrendered')
                   return {
                     status: 'surrendered',
@@ -3016,6 +3197,7 @@ export function createServer<
                       // budget before dead-lettering.
                       maxFailures:
                         err instanceof NonRetriableError ? 1 : MAX_FAILURES,
+                      retryAfterMs: failureRetryDelayMs(event.failureCount + 1),
                     })
                     commitSpan.setAttribute('a2.commit.outcome', result.outcome)
                     return result
@@ -3147,6 +3329,10 @@ export function createServer<
             let claim: StoreClaimAvailableResult | undefined
             if (shouldClaim) {
               const window = nextClaimWindow()
+              if (window.pastDeadline) {
+                outcome = 'busy'
+                break
+              }
               claim = await telemetry.span(
                 'a2.claim',
                 { 'a2.contract': name, 'a2.session_id': sessionId },
@@ -3155,6 +3341,7 @@ export function createServer<
                     sessionId: ns,
                     holder,
                     ttlMs: window.ttlMs,
+                    maxFailures: MAX_FAILURES,
                     ...(window.deadlineCapped
                       ? { expiresAtMs: window.expiresAtMs }
                       : {}),
@@ -3171,11 +3358,16 @@ export function createServer<
                 },
               )
               if (claim.outcome === 'claimed') {
-                requestWatchdog(window.watchdogAtMs)
+                // A takeover that spent the last failed attempt comes back
+                // dead-lettered instead of claimed.
+                const runnable = claim.events.filter(
+                  (event) => event.failedAt === null,
+                )
+                for (const event of claim.events)
+                  if (event.failedAt !== null) await reportExhausted(event)
+                if (runnable.length > 0) requestWatchdog(window.watchdogAtMs)
                 if (window.deadlineCapped) stopRenewing = true
-                for (const event of claim.events) {
-                  start(event, window.expiresAtMs)
-                }
+                for (const event of runnable) start(event, window.expiresAtMs)
                 shouldClaim = false
               }
             }
@@ -3276,7 +3468,9 @@ export function createServer<
     )
   }
 
-  const session = (id: string): Session<D, SessionAppend<D>, P> => {
+  const session = (
+    id: string,
+  ): Session<D, SessionAppend<D>, P> & SessionDeadLetters<D> => {
     assertSessionId(id)
     const appendInline = async (...events: AppendInput<D>[]) =>
       appendExternal(id, 'inline', events)
@@ -3285,7 +3479,11 @@ export function createServer<
     const append = Object.assign(appendInline, {
       dispatch: appendDispatch,
     }) as SessionAppend<D>
-    return makeSession(id, append, null)
+    return Object.assign(makeSession(id, append, null), {
+      deadLetters: () => listDeadLetters(id),
+      redrive: (index: number, redriveOptions: SessionRedriveOptions) =>
+        redriveEvent(id, index, redriveOptions),
+    })
   }
 
   const historyPage = async (
@@ -3393,6 +3591,12 @@ function assertStateIndex(value: number): void {
   }
 }
 
+function assertRedriveIndex(value: number): void {
+  if (!Number.isSafeInteger(value) || value < 1) {
+    throw new TypeError('redrive requires a positive safe-integer index')
+  }
+}
+
 function assertClearIndex(value: number): void {
   if (!Number.isSafeInteger(value) || value < 0) {
     throw new TypeError('ifLastIndex must be a non-negative safe integer')