perf(heartbeat): drive retry promotion from durable deadlines (#18387)

Reviewed and verified at f1c48c924d.

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/modules/agent-lifecycle/adapters/records.ts

@@ -80,22 +80,20 @@ export function agentRecords(db: Db, effects: LifecycleEffects) {
         throw conflict("Only agents in error status can have their error cleared");
       }
 
-      const updated = await db
-        .update(agents)
-        .set({
-          status: "idle",
-          pauseReason: null,
-          pausedAt: null,
-          errorReason: null,
-          updatedAt: new Date(),
-        })
-        .where(and(eq(agents.id, id), eq(agents.status, "error"), eq(agents.lifecycleState, "ready")))
-        .returning()
-        .then((rows) => rows[0] ?? null);
+      const updated = await effects.transaction(db, existing.companyId, async tx => {
+        const [current] = await tx.select({ status: agents.status, lifecycleState: agents.lifecycleState })
+          .from(agents).where(and(eq(agents.id, id), eq(agents.companyId, existing.companyId))).for("no key update");
+        if (current?.status !== "error" || current.lifecycleState !== "ready") {
+          throw conflict("Only agents in error status can have their error cleared");
+        }
+        await effects.beforeStatusChange(tx, existing.companyId, id, "error", "idle");
+        const [saved] = await tx.update(agents).set({
+          status: "idle", pauseReason: null, pausedAt: null, errorReason: null, updatedAt: new Date(),
+        }).where(and(eq(agents.id, id), eq(agents.status, "error"), eq(agents.lifecycleState, "ready"))).returning();
+        if (!saved) throw conflict("Only agents in error status can have their error cleared");
+        return saved;
+      });
 
-      if (!updated) {
-        throw conflict("Only agents in error status can have their error cleared");
-      }
       return getById(updated.id);
     },
 

server/src/modules/run-dispatch/adapters/postgres.ts

@@ -511,6 +511,15 @@ export function createPostgresRunDispatchAdapter(
     return decideScheduledRetryGate(factsResult.facts, input.now);
   }
 
+  async function nextScheduledRetryAt(input: { cutoff: Date | null }): Promise<number | null> {
+    const [row] = await db.select({ at: heartbeatRuns.scheduledRetryAt }).from(heartbeatRuns)
+      .where(and(eq(heartbeatRuns.status, "scheduled_retry"),
+        sql`jsonb_typeof(${heartbeatRuns.resultJson}->'startupCancellation') is distinct from 'object'`,
+        input.cutoff ? gte(heartbeatRuns.createdAt, input.cutoff) : undefined))
+      .orderBy(asc(heartbeatRuns.scheduledRetryAt)).limit(1);
+    return row?.at?.getTime() ?? null;
+  }
+
   async function listDueRetries(input: ListDueRetriesInput): Promise<DueRetryRun[]> {
     const rows = await db
       .select()
@@ -518,6 +527,7 @@ export function createPostgresRunDispatchAdapter(
       .where(
         and(
           eq(heartbeatRuns.status, "scheduled_retry"),
+          sql`jsonb_typeof(${heartbeatRuns.resultJson}->'startupCancellation') is distinct from 'object'`,
           lte(heartbeatRuns.scheduledRetryAt, input.now),
           input.cutoff ? gte(heartbeatRuns.createdAt, input.cutoff) : undefined,
         ),
@@ -1414,6 +1424,7 @@ export function createPostgresRunDispatchAdapter(
 
   return {
     evaluateScheduledRetryGate,
+    nextScheduledRetryAt,
     listDueRetries,
     cancelStaleQueuedRun,
     dispatchResolvedInteractionIfCurrent,

server/src/modules/run-dispatch/application/ports.ts

@@ -24,6 +24,7 @@ export type EvaluateScheduledRetryGateInput = {
 
 /** Read-only scheduled-retry operations exposed to application use cases. */
 export interface ScheduledRetryReader {
+  nextScheduledRetryAt(input: { cutoff: Date | null }): Promise<number | null>;
   evaluateScheduledRetryGate(input: EvaluateScheduledRetryGateInput): Promise<GateDecision>;
   listDueRetries(input: ListDueRetriesInput): Promise<DueRetryRun[]>;
 }

server/src/modules/run-dispatch/application/use-cases.test.ts

@@ -12,6 +12,7 @@ function fakeReader(dueRuns: DueRetryRun[] = []): ScheduledRetryReader & { evalu
   const evaluateCalls: unknown[] = [];
   return {
     evaluateCalls,
+    async nextScheduledRetryAt() { return null; },
     async evaluateScheduledRetryGate(input) {
       evaluateCalls.push(input);
       return { allowed: true };

server/src/modules/run-dispatch/index.ts

@@ -67,6 +67,7 @@ export function createRunDispatch(db: Db, deps: RunDispatchDeps = {}) {
   });
 
   return {
+    nextScheduledRetryAt: (input: { cutoff: Date | null }) => adapter.nextScheduledRetryAt(input),
     evaluateScheduledRetryGate: createEvaluateScheduledRetryGate({ reader: adapter }),
     promoteScheduledRetry,
     promoteDueScheduledRetries: createPromoteDueScheduledRetries({

server/src/modules/wake-queue/adapters/postgres.ts

@@ -1,3 +1,4 @@
+import { notifyRunDispatch } from "../../../services/run-dispatch-notifications.js";
 import { notifyExternalConversationSettlement } from "../../../services/external-conversation-settlement-notifications.js";
 import { notifyNativeStatusReassessment } from "../../../services/native-status-reassessment-notifications.js";
 import { notifyUnsafeWorkspaceExportRecovery } from "../../../services/unsafe-workspace-export-notifications.js";
@@ -637,6 +638,7 @@ function buildTransaction(tx: Db, deps: WakeQueuePostgresAdapterDeps, db: Db, ru
         .returning()
         .then((rows) => rows[0]);
 
+      await notifyRunDispatch(tx, { companyId });
       const queuedRun = await tx
         .insert(heartbeatRuns)
         .values({
@@ -715,6 +717,7 @@ function buildTransaction(tx: Db, deps: WakeQueuePostgresAdapterDeps, db: Db, ru
         .returning()
         .then((rows) => rows[0]);
 
+      await notifyRunDispatch(tx, { companyId });
       const queuedRun = await tx
         .insert(heartbeatRuns)
         .values({