fix(agents-server): dropped parent wake on parallel sub-agent spawn

.changeset/wake-registry-parallel-spawn-clobber.md

@@ -0,0 +1,5 @@
+---
+"@electric-ax/agents-server": patch
+---
+
+Fix a dropped parent wake when a parent spawns sub-agents in parallel. Each child's `runFinished` wake is registered from two paths (spawn + manifest-sync) keyed by the same `manifestKey`; the second insert hits `uq_wake_registration` and takes the conflict branch. That branch called `loadRegistrations()`, a full clear-and-rebuild of the in-memory registration cache from a snapshot read across an `await`. Under parallel spawn several such reloads interleave, and a stale snapshot landing last evicts a sibling's newer registration from the cache — so when that sibling finishes, `evaluate()` finds no match and the wake is silently dropped (no error). Sequential spawn never overlaps the reloads, which is why only the parallel fan-out reproduced it. The conflict branch now re-reads only the single conflicting row and caches just that entry, leaving sibling registrations untouched.

packages/agents-server/src/entity-manager.ts

@@ -2977,23 +2977,26 @@ export class EntityManager {
     const spec = resolveCronScheduleSpec(req.expression, req.timezone)
 
     const manifestKey = `schedule:${req.id}`
-    await this.wakeRegistry.unregisterByManifestKey(
+    // Gap-free reconcile (register-first, then prune stale). The old
+    // unregister-then-register briefly left this manifest key's wake absent from
+    // the cache; a cron tick landing in that window would be missed. See
+    // WakeRegistry.reconcileManifestRegistration.
+    await this.wakeRegistry.reconcileManifestRegistration(
       entityUrl,
       manifestKey,
+      {
+        subscriberUrl: entityUrl,
+        sourceUrl: getCronStreamPath(spec.expression, spec.timezone),
+        condition: {
+          on: `change`,
+        },
+        debounceMs: req.debounceMs,
+        timeoutMs: req.timeoutMs,
+        oneShot: false,
+        manifestKey,
+      },
       this.tenantId
     )
-    await this.wakeRegistry.register({
-      tenantId: this.tenantId,
-      subscriberUrl: entityUrl,
-      sourceUrl: getCronStreamPath(spec.expression, spec.timezone),
-      condition: {
-        on: `change`,
-      },
-      debounceMs: req.debounceMs,
-      timeoutMs: req.timeoutMs,
-      oneShot: false,
-      manifestKey,
-    })
     await this.getOrCreateCronStream(spec.expression, spec.timezone)
 
     const txid = randomUUID()
@@ -3162,24 +3165,26 @@ export class EntityManager {
     )
 
     // The manifest is the durable source of truth. Register side effects after
-    // it is appended so failures can be repaired by manifest replay.
-    await this.wakeRegistry.unregisterByManifestKey(
+    // it is appended so failures can be repaired by manifest replay. Gap-free
+    // reconcile (register-first, then prune stale) so a webhook event landing
+    // during a re-subscribe isn't missed — see
+    // WakeRegistry.reconcileManifestRegistration.
+    await this.wakeRegistry.reconcileManifestRegistration(
       entityUrl,
       manifestKey,
+      {
+        subscriberUrl: entityUrl,
+        sourceUrl: req.subscription.sourceUrl,
+        condition: {
+          on: `change`,
+          collections: [`webhook_event`],
+          ops: [`insert`],
+        },
+        oneShot: false,
+        manifestKey,
+      },
       this.tenantId
     )
-    await this.wakeRegistry.register({
-      tenantId: this.tenantId,
-      subscriberUrl: entityUrl,
-      sourceUrl: req.subscription.sourceUrl,
-      condition: {
-        on: `change`,
-        collections: [`webhook_event`],
-        ops: [`insert`],
-      },
-      oneShot: false,
-      manifestKey,
-    })
 
     return { txid, subscription: req.subscription }
   }

packages/agents-server/src/runtime.ts

@@ -242,31 +242,34 @@ export class ElectricAgentsTenantRuntime {
       if (!manifestKey) continue
 
       if (operation === `delete`) {
-        await this.manager.wakeRegistry.unregisterByManifestKey(
+        await this.manager.wakeRegistry.reconcileManifestRegistration(
           subscriberUrl,
           manifestKey,
+          null,
           this.serviceId
         )
         continue
       }
 
-      await this.manager.wakeRegistry.unregisterByManifestKey(
+      // Reconcile idempotently and WITHOUT a delivery gap. The old
+      // unregister-then-register sequence briefly removed the registration
+      // from the cache; a source (e.g. a sibling sub-agent) that finished in
+      // that window had its wake dropped. reconcileManifestRegistration
+      // registers the desired reg first, then prunes only stale rows.
+      const reg = value
+        ? buildManifestWakeRegistration(subscriberUrl, value, manifestKey)
+        : null
+      if (reg) {
+        reg.tenantId = this.serviceId
+      }
+      await this.manager.wakeRegistry.reconcileManifestRegistration(
         subscriberUrl,
         manifestKey,
+        reg,
         this.serviceId
       )
 
       if (value) {
-        const reg = buildManifestWakeRegistration(
-          subscriberUrl,
-          value,
-          manifestKey
-        )
-        if (reg) {
-          reg.tenantId = this.serviceId
-          await this.manager.wakeRegistry.register(reg)
-        }
-
         const cronSpec = extractManifestCronSpec(value)
         if (cronSpec) {
           void this.manager

packages/agents-server/src/wake-registry.ts

@@ -3,7 +3,7 @@ import {
   isChangeMessage,
   isControlMessage,
 } from '@electric-sql/client'
-import { and, eq } from 'drizzle-orm'
+import { and, eq, isNull, ne } from 'drizzle-orm'
 import { wakeRegistrations } from './db/schema.js'
 import { serverLog } from './utils/log.js'
 import { electricUrlWithPath } from './utils/electric-url.js'
@@ -325,7 +325,15 @@ export class WakeRegistry {
     return this.syncRecoveryPromise
   }
 
-  async register(reg: WakeRegistration): Promise<void> {
+  /**
+   * Register a wake subscription. Returns the dbId of the row now backing the
+   * cached registration — the freshly-inserted row, or (on unique-constraint
+   * conflict) the pre-existing row that was re-read and re-cached. Returns -1
+   * only in the pathological case where the conflicting row could not be
+   * re-read. Callers reconciling a set of registrations use the returned id to
+   * prune siblings unambiguously (see reconcileManifestRegistration).
+   */
+  async register(reg: WakeRegistration): Promise<number> {
     const tenantId = this.resolveTenantId(reg.tenantId)
     const result = await this.db
       .insert(wakeRegistrations)
@@ -344,10 +352,50 @@ export class WakeRegistry {
       .returning({ id: wakeRegistrations.id })
 
     if (result.length === 0) {
-      // Another path (e.g. manifest-sync) may have created the row first.
-      // Refresh the cache so this process still sees the effective registration.
-      await this.loadRegistrations()
-      return
+      // Another path (e.g. manifest-sync) created the row first. Re-read only
+      // the conflicting row and cache just that entry. A full loadRegistrations()
+      // here clears and rebuilds the entire cache from a snapshot taken across an
+      // await; when several register() calls conflict concurrently (parallel
+      // sub-agent spawn), a stale snapshot landing last can evict a sibling's
+      // newer registration, silently dropping its wake. See uq_wake_registration.
+      const existing = await this.db
+        .select()
+        .from(wakeRegistrations)
+        .where(
+          and(
+            eq(wakeRegistrations.tenantId, tenantId),
+            eq(wakeRegistrations.subscriberUrl, reg.subscriberUrl),
+            eq(wakeRegistrations.sourceUrl, reg.sourceUrl),
+            eq(wakeRegistrations.oneShot, reg.oneShot),
+            eq(wakeRegistrations.debounceMs, reg.debounceMs ?? 0),
+            eq(wakeRegistrations.timeoutMs, reg.timeoutMs ?? 0),
+            eq(wakeRegistrations.condition, reg.condition),
+            reg.manifestKey == null
+              ? isNull(wakeRegistrations.manifestKey)
+              : eq(wakeRegistrations.manifestKey, reg.manifestKey)
+          )
+        )
+        .limit(1)
+
+      const row = existing[0]
+      if (row) {
+        this.upsertCachedRegistration({
+          tenantId: row.tenantId,
+          subscriberUrl: row.subscriberUrl,
+          sourceUrl: row.sourceUrl,
+          condition: row.condition as WakeRegistration[`condition`],
+          debounceMs: row.debounceMs || undefined,
+          timeoutMs: row.timeoutMs || undefined,
+          oneShot: row.oneShot,
+          includeResponse: row.includeResponse === false ? false : undefined,
+          manifestKey: row.manifestKey ?? undefined,
+          dbId: row.id,
+          createdAt: row.createdAt,
+          timeoutConsumed: row.timeoutConsumed,
+        })
+        return row.id
+      }
+      return -1
     }
 
     const dbId = result[0]!.id
@@ -358,6 +406,7 @@ export class WakeRegistry {
       createdAt: new Date(),
       timeoutConsumed: false,
     })
+    return dbId
   }
 
   private startTimeoutTimer(reg: CachedWakeRegistration, dbId: number): void {
@@ -413,6 +462,76 @@ export class WakeRegistry {
     }
   }
 
+  /**
+   * Idempotently reconcile the single wake registration anchored to a manifest
+   * entry, without ever leaving a delivery gap.
+   *
+   * The obvious sequence — `unregisterByManifestKey()` then `register()` — drops
+   * the registration from the cache and only re-adds it after an async DB
+   * round-trip. A source whose run finishes inside that window evaluates against
+   * an empty cache and its wake is silently lost. Parallel sub-agent spawn hits
+   * this constantly: every child's manifest entry re-syncs a registration that
+   * is *identical* to the one the spawn already created, so each child gets a
+   * remove→re-add churn, and any sibling that finishes mid-churn is dropped.
+   *
+   * Instead we register the desired registration FIRST — register() returns the
+   * id of the row it left in the cache (a fresh insert, or the pre-existing row
+   * re-read on conflict), so an equivalent registration is continuously present
+   * — THEN delete only the rows for this manifest key that differ from it.
+   * Pruning by that exact id (never by a re-derived field key, which can diverge
+   * from what register() actually cached) guarantees we never delete the row we
+   * just kept. `desired == null` (a manifest delete, or an entry carrying no
+   * wake) prunes them all, matching the previous unregister-only behaviour.
+   */
+  async reconcileManifestRegistration(
+    subscriberUrl: string,
+    manifestKey: string,
+    desired: WakeRegistration | null,
+    tenantId?: string
+  ): Promise<void> {
+    const resolvedTenantId = this.resolveTenantId(tenantId)
+
+    let keptDbId = -1
+    if (desired) {
+      keptDbId = await this.register({
+        ...desired,
+        tenantId: resolvedTenantId,
+        manifestKey,
+      })
+    }
+
+    // Delete every other row anchored to this manifest key in one predicate —
+    // covering rows that are not currently cached, so a later loadRegistrations()
+    // can't resurrect them. `keptDbId === -1` (a delete/no-wake reconcile)
+    // matches all rows for the key, preserving the old unregister-only delete.
+    await this.db
+      .delete(wakeRegistrations)
+      .where(
+        and(
+          eq(wakeRegistrations.tenantId, resolvedTenantId),
+          eq(wakeRegistrations.subscriberUrl, subscriberUrl),
+          eq(wakeRegistrations.manifestKey, manifestKey),
+          ne(wakeRegistrations.id, keptDbId)
+        )
+      )
+
+    const staleDbIds = Array.from(this.registrationCache.values()).flatMap(
+      (regs) =>
+        regs
+          .filter(
+            (r) =>
+              r.tenantId === resolvedTenantId &&
+              r.subscriberUrl === subscriberUrl &&
+              r.manifestKey === manifestKey &&
+              r.dbId !== keptDbId
+          )
+          .map((r) => r.dbId)
+    )
+    for (const dbId of staleDbIds) {
+      this.removeCachedRegistrationByDbId(dbId)
+    }
+  }
+
   async unregisterBySubscriber(
     subscriberUrl: string,
     tenantId?: string