experimental-a2 0.18.0 → 0.18.1

package/src/ai-pi.ts

@@ -59,6 +59,7 @@ import {
   type RecordedGeneration,
   saveHarness,
 } from './ai-harness.ts'
+import { type SchemaInput, schemaInput } from './ai-tool-input.ts'
 import type { CompactionPartData, ToolResultPayload } from './ai.ts'
 import { boundedCache } from './bounded-cache.ts'
 import { ambientToolScopeStorage } from './internal.ts'
@@ -110,11 +111,14 @@ type PiModules = {
   ai: typeof import('@earendil-works/pi-ai')
 }
 
-/** A content block without the text the transcript already holds. */
+/**
+ * A content block without the text the transcript already holds. A tool call
+ * keeps its arguments only where the transcript's input is not the model's.
+ */
 type StoredBlock =
   | (Omit<TextContent, 'text'> & { ref: string })
   | (Omit<ThinkingContent, 'thinking'> & { ref: string })
-  | Omit<ToolCall, 'arguments'>
+  | (Omit<ToolCall, 'arguments'> & Partial<Pick<ToolCall, 'arguments'>>)
 
 /** What a step's transcript parts do not hold of Pi's response. */
 type StoredResponse = Omit<AssistantMessage, 'content'> & {
@@ -251,7 +255,7 @@ async function runStep(
     })
     const live = kind === 'step' ? liveStep() : undefined
     for await (const event of upstream) {
-      live?.(event)
+      await live?.(event)
       if (event.type === 'done' || event.type === 'error') {
         const response = event.type === 'done' ? event.message : event.error
         if (kind === 'compaction') state.compactions.push(response)
@@ -267,11 +271,39 @@ async function runStep(
     }
   }
 
+  /**
+   * A live call's input as Pi's loop, then the AI SDK step, read it. Pi
+   * answers a call to no tool, or whose arguments it cannot prepare or
+   * validate, itself. The agent's tool reads the arguments Pi accepts through
+   * its own schema; A2 runs Pi's own tools from the model's arguments, as
+   * Pi's loop does.
+   */
+  const callInput = async (block: ToolCall): Promise<SchemaInput> => {
+    const definition = tools.find((tool) => tool.name === block.name)
+    if (definition === undefined)
+      return { success: false, errorText: `Tool ${block.name} not found` }
+    let args: unknown
+    try {
+      const prepared = definition.prepareArguments?.(block.arguments)
+      args = modules.ai.validateToolArguments(
+        definition,
+        prepared === undefined
+          ? block
+          : { ...block, arguments: prepared as JsonObject },
+      )
+    } catch (error) {
+      return { success: false, errorText: asError(error).message }
+    }
+    if (builtinOf(step.tools, block.name) !== undefined)
+      return { success: true, value: block.arguments }
+    return schemaInput({ tools: step.tools, toolName: block.name, value: args })
+  }
+
   /** Streams one live step response as UI chunks and records it. */
   const liveStep = () => {
     const ref = (index: number) => `pi:${step.requestIndex}:${index}`
     const started = new Set<string>()
-    return (event: AssistantMessageEvent): void => {
+    return async (event: AssistantMessageEvent): Promise<void> => {
       switch (event.type) {
         case 'start':
           emit({ type: 'start-step' })
@@ -325,6 +357,8 @@ async function runStep(
         case 'done':
         case 'error': {
           const response = event.type === 'done' ? event.message : event.error
+          // Calls whose transcript input is not the model's arguments.
+          const changed = new Set<string>()
           if (event.type === 'done') {
             for (const block of response.content) {
               if (block.type !== 'toolCall') continue
@@ -335,26 +369,41 @@ async function runStep(
                   toolName: block.name,
                 })
               // Pi does not run the calls of a cut-off response.
-              emit(
+              const input: SchemaInput =
                 response.stopReason === 'length'
                   ? {
-                      type: 'tool-input-error',
-                      toolCallId: block.id,
-                      toolName: block.name,
-                      input: block.arguments,
+                      success: false,
                       errorText:
                         'The response hit the output token limit before this call was complete.',
                     }
-                  : {
+                  : await callInput(block)
+              if (
+                input.success &&
+                JSON.stringify(input.value) !== JSON.stringify(block.arguments)
+              )
+                changed.add(block.id)
+              emit(
+                input.success
+                  ? {
                       type: 'tool-input-available',
                       toolCallId: block.id,
                       toolName: block.name,
+                      input: input.value,
+                    }
+                  : {
+                      type: 'tool-input-error',
+                      toolCallId: block.id,
+                      toolName: block.name,
                       input: block.arguments,
+                      errorText: input.errorText,
                     },
               )
             }
           }
-          record({ call: 'step', response: storedResponse(response, ref) })
+          record({
+            call: 'step',
+            response: storedResponse(response, ref, changed),
+          })
           completion = {
             finishReason: finishReason(response.stopReason),
             usage: modelUsage(response.usage),
@@ -620,6 +669,7 @@ function recordedCalls(generation: RecordedGeneration): RecordedCall[] {
 function storedResponse(
   response: AssistantMessage,
   ref: (index: number) => string,
+  changed: ReadonlySet<string>,
 ): StoredResponse {
   return {
     ...response,
@@ -634,6 +684,7 @@ function storedResponse(
           return { ...rest, ref: ref(index) }
         }
         case 'toolCall': {
+          if (changed.has(block.id)) return block
           const { arguments: _, ...rest } = block
           return rest
         }
@@ -653,7 +704,8 @@ function restoredResponse(
       if (block.type === 'toolCall')
         return {
           ...block,
-          arguments: (inputs.get(block.id) ?? {}) as JsonObject,
+          arguments:
+            block.arguments ?? ((inputs.get(block.id) ?? {}) as JsonObject),
         }
       const { ref, ...rest } = block
       const text = texts.get(ref) ?? ''
@@ -826,6 +878,15 @@ async function recordedResult(
   return { ...kept, details: { ...details, truncation } }
 }
 
+/** The Pi tool a name in the agent's tool set runs, unless the agent replaced it. */
+const builtinOf = (tools: ToolSet, name: string): Builtin | undefined =>
+  BUILTINS.find(
+    (candidate) =>
+      candidate === name &&
+      Object.hasOwn(tools, name) &&
+      tools[name] === builtinTools[name],
+  )
+
 // Each definition is typed by its parameters; a session takes any of them.
 const untyped = (definition: object) => definition as ToolDefinition
 
@@ -879,9 +940,7 @@ async function sessionTools(
     Object.entries(tools)
       .filter(([, tool]) => tool.type !== 'provider')
       .map(async ([name, tool]): Promise<ToolDefinition> => {
-        const builtin = BUILTINS.find(
-          (candidate) => candidate === name && builtinTools[name] === tool,
-        )
+        const builtin = builtinOf(tools, name)
         if (builtin !== undefined)
           // Pi's own definition, for its prompt; A2 already ran the call.
           return Object.assign(builtinDefinition(modules, builtin, cwd), {

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

@@ -274,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/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

@@ -205,9 +205,15 @@ export type DeadLetter<D extends EventDefs> = ContractEvent<D> & {
   lane: string | null
   /** Durable dispatch claims, including claims abandoned by hard kills. */
   attemptCount: number
-  /** Caught handler failures since the event was appended or last retried. */
+  /**
+   * Failed attempts since the event was appended or last retried: caught
+   * handler failures, and attempts that ended before they finished.
+   */
   failureCount: number
-  /** The last handler failure, stringified. */
+  /**
+   * The last error a handler threw, stringified; null when every failed
+   * attempt ended before it finished.
+   */
   lastError: string | null
   failedAt: Date
 }
@@ -499,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: {
@@ -527,6 +540,7 @@ type ClaimWindow = {
   expiresAtMs: number
   watchdogAtMs: number
   deadlineCapped: boolean
+  pastDeadline: boolean
 }
 
 type DrainSignal = {
@@ -601,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()
@@ -619,6 +635,7 @@ function nextClaimWindow(): ClaimWindow {
     expiresAtMs,
     watchdogAtMs: expiresAtMs + DRAIN_TIMINGS.watchdogGraceMs,
     deadlineCapped,
+    pastDeadline: deadlineMs !== null && deadlineTtlMs <= 0,
   }
 }
 
@@ -3068,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,
@@ -3127,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',
@@ -3162,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
@@ -3293,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 },
@@ -3301,6 +3341,7 @@ export function createServer<
                     sessionId: ns,
                     holder,
                     ttlMs: window.ttlMs,
+                    maxFailures: MAX_FAILURES,
                     ...(window.deadlineCapped
                       ? { expiresAtMs: window.expiresAtMs }
                       : {}),
@@ -3317,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
               }
             }

package/src/store-memory.ts

@@ -282,12 +282,35 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
     }
   }
 
-  const claim = (
-    row: Row,
-    holder: string,
-    now: Date,
-    expiresAt: Date,
-  ): StoredEvent => {
+  const claim = ({
+    sessionId,
+    row,
+    holder,
+    now,
+    expiresAt,
+    maxFailures,
+  }: {
+    sessionId: string
+    row: Row
+    holder: string
+    now: Date
+    expiresAt: Date
+    maxFailures: number | undefined
+  }): StoredEvent => {
+    // A claim still held here lapsed: its attempt never finished.
+    if (row.claimHolder !== null) {
+      row.failureCount += 1
+      if (maxFailures !== undefined && row.failureCount >= maxFailures) {
+        // Advancing the attempt fences the lapsed one; nothing runs.
+        row.attemptCount += 1
+        row.failedAt = new Date(now)
+        row.claimHolder = null
+        row.claimExpiresAt = null
+        dispatchOf(sessionId).ready.delete(row.index)
+        releaseSnapshotPins(sessionId, row.index)
+        return toStored(row)
+      }
+    }
     row.attemptCount += 1
     row.firstClaimedAt ??= new Date(now)
     row.lastClaimedAt = new Date(now)
@@ -547,6 +570,7 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
       ttlMs,
       expiresAtMs,
       excludeIndexes,
+      maxFailures,
     }) {
       const claimedAt = clock.now()
       const now = claimedAt.getTime()
@@ -574,7 +598,14 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
       return {
         outcome: 'claimed',
         events: eligible.map((row) => {
-          const claimed = claim(row, holder, claimedAt, expiresAt)
+          const claimed = claim({
+            sessionId,
+            row,
+            holder,
+            now: claimedAt,
+            expiresAt,
+            maxFailures,
+          })
           if (clearIndex > 0) claimed.clearIndex = clearIndex
           return claimed
         }),
@@ -681,7 +712,14 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
       return { outcome: 'completed', events: inserted.map(toStored) }
     },
 
-    async failAttempt({ sessionId, index, attempt, error, maxFailures }) {
+    async failAttempt({
+      sessionId,
+      index,
+      attempt,
+      error,
+      maxFailures,
+      retryAfterMs,
+    }) {
       if (index <= clearIndexOf(sessionId)) {
         return { outcome: 'superseded', failureCount: 0 }
       }
@@ -708,7 +746,10 @@ export function memory(options: MemoryStoreOptions = {}): A2Store {
       row.lastFailedAt = new Date(failedAt)
       row.lastFailedAttempt = attempt
       row.claimHolder = null
-      row.claimExpiresAt = null
+      row.claimExpiresAt =
+        retryAfterMs === undefined || row.failureCount >= maxFailures
+          ? null
+          : new Date(failedAt.getTime() + retryAfterMs)
       if (row.failureCount >= maxFailures) {
         row.failedAt = new Date(failedAt)
         dispatchOf(sessionId).ready.delete(row.index)