Drive live run and resource recovery from scoped durable deadlines

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

server/src/services/environment-cleanup-notifications.ts

@@ -5,13 +5,15 @@ import { notifyNativeEnvironmentSettlement, type NativeEnvironmentSettlementScop
 import type { Db } from "@paperclipai/db";
 import { DELIVERY_QUEUES, notifyDeliveryWork } from "./delivery-work-notifications.js";
 
-export async function notifyEnvironmentCleanup(transaction: object, scope?: NativeEnvironmentSettlementScope): Promise<void> {
+export type EnvironmentCleanupWorkScope = NativeEnvironmentSettlementScope & { leaseId?: string };
+
+export async function notifyEnvironmentCleanup(transaction: object, scope?: EnvironmentCleanupWorkScope): Promise<void> {
   await notifyLiveResourceRecovery(transaction, scope);
-  if (scope?.runId && typeof (transaction as Db).select === "function") {
+  if (scope && (scope.runId || scope.leaseId) && typeof (transaction as Db).select === "function") {
     const peers = await (transaction as Db).select({ id: environmentLeases.id, companyId: environmentLeases.companyId,
       runId: environmentLeases.heartbeatRunId }).from(environmentLeases).where(and(eq(environmentLeases.status, "active"),
       sql`exists (select 1 from environment_leases changed where changed.company_id = ${scope.companyId}
-        and changed.heartbeat_run_id = ${scope.runId} and changed.provider_lease_id is not null
+        and ${scope.leaseId ? sql`changed.id = ${scope.leaseId}` : sql`changed.heartbeat_run_id = ${scope.runId}`} and changed.provider_lease_id is not null
         and changed.provider = ${environmentLeases.provider} and changed.provider_lease_id = ${environmentLeases.providerLeaseId})`));
     for (const peer of peers) await notifyDeliveryWork(transaction, DELIVERY_QUEUES.orphanActiveLeases,
       JSON.stringify({ companyId: peer.companyId, leaseId: peer.id }));
@@ -21,7 +23,7 @@ export async function notifyEnvironmentCleanup(transaction: object, scope?: Nati
   await notifyDeliveryWork(transaction, DELIVERY_QUEUES.historicalCleanup);
 }
 
-export function writeEnvironmentCleanup<T>(db: Db, write: (tx: Db) => PromiseLike<T>, scope: NativeEnvironmentSettlementScope): Promise<T> {
+export function writeEnvironmentCleanup<T>(db: Db, write: (tx: Db) => PromiseLike<T>, scope: EnvironmentCleanupWorkScope): Promise<T> {
   return db.transaction(async tx => {
     await notifyEnvironmentCleanup(tx, scope);
     return write(tx as unknown as Db);

server/src/services/heartbeat/live-recovery-work.ts

@@ -66,10 +66,12 @@ export function registerLiveRecoveryWork(db: Db, coordinator: ReturnType<typeof
               scope?.runId ? eq(heartbeatRuns.id, scope.runId) : undefined,
               cursor ? gt(heartbeatRuns.id, cursor) : undefined,
               kind === "native" ? and(eq(heartbeatRuns.runtimeMode, "native"), inArray(heartbeatRuns.status, ["running", "failed"]),
-                eq(nativeRunFinalizations.phase, "retryable_failure"), isNull(nativeRunFinalizations.resultId))
+                eq(nativeRunFinalizations.phase, "retryable_failure"), isNull(nativeRunFinalizations.resultId),
+                sql`${heartbeatRuns.contextSnapshot}->'nativeQuestionCancellation' is null`)
                 : kind === "question" ? and(inArray(heartbeatRuns.status, ["queued", "running", "scheduled_retry"]),
                   sql`${heartbeatRuns.contextSnapshot}->'nativeQuestionCancellation' is not null`)
-                  : kind === "active" ? eq(heartbeatRuns.status, "running")
+                  : kind === "active" ? and(eq(heartbeatRuns.status, "running"),
+                    sql`${heartbeatRuns.contextSnapshot}->'nativeQuestionCancellation' is null`)
                     : and(eq(heartbeatRuns.runtimeMode, "native"), inArray(heartbeatRuns.status, ["failed", "cancelled", "timed_out", "interrupted"]),
                       isNotNull(heartbeatRuns.finishedAt), or(isNotNull(heartbeatRuns.processPid),
                         isNotNull(heartbeatRuns.processGroupId), isNotNull(heartbeatRuns.processStartedAt))),

server/src/services/heartbeat/recovery.ts

@@ -916,8 +916,8 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD
     };
   }
 
-  async function recoverNativeRunsAfterRestart(now = new Date(), runIds?: string[]) {
-    if (!runIds) {
+  async function recoverNativeRunsAfterRestart(now = new Date(), runIds?: string[], options?: { skipDelegatedWork?: boolean }) {
+    if (!runIds && !options?.skipDelegatedWork) {
       await reconcileRetainedWorkspaces(db, environmentRuntime, resumeRemoteStopComments);
       await reconcileNativeFinalizations(db, undefined, {
         environmentRuntime,
@@ -1529,11 +1529,11 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD
         if (lease.heartbeatRunId && !activeRunExecutions.has(lease.heartbeatRunId)
           && !adapterExecutionControls.has(lease.heartbeatRunId)
           && await preserveOrphanedNativeWorkspace(db, lease)) recovered += 1;
-        await deferOrphanedActiveLease(lease.id);
+        if (!opts.leaseId) await deferOrphanedActiveLease(lease.id);
         continue;
       }
       if (!(await allowLegacyShutdownWorkspaceCleanup(db, lease))) {
-        await deferOrphanedActiveLease(lease.id);
+        if (!opts.leaseId) await deferOrphanedActiveLease(lease.id);
         continue;
       }
       // A provider resource id names one physical sandbox. A different lease
@@ -1560,7 +1560,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD
           // Defer this row so the fixed-size page reaches the rows behind
           // it next tick, instead of re-selecting the same guarded rows
           // forever.
-          await deferOrphanedActiveLease(lease.id);
+          if (!opts.leaseId) await deferOrphanedActiveLease(lease.id);
           continue;
         }
       }
@@ -1582,7 +1582,7 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD
             eq(environmentLeases.status, "active"),
           ),
         )
-        .returning({ id: environmentLeases.id }), { companyId: lease.companyId, runId: lease.heartbeatRunId ?? undefined });
+        .returning({ id: environmentLeases.id }), { companyId: lease.companyId, runId: lease.heartbeatRunId ?? undefined, leaseId: lease.id });
       if (flipped.length > 0) recovered += 1;
     }
 
@@ -2292,6 +2292,10 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD
         runningProcesses.has(run.id) ||
         activeRunExecutions.has(run.id) ||
         coordinatorOwnedByCurrentController;
+      if (nativeRun && !locallyTracked && nativeControllerLeaseExpiresAt && nativeControllerLeaseExpiresAt > now) {
+        scope.onDeferredDeadline?.(nativeControllerLeaseExpiresAt.getTime());
+        continue;
+      }
       if (
         nativeRun &&
         ([
@@ -2302,6 +2306,8 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD
         ].includes(nativeRecoveryState ?? "") ||
           nativeCoordinatorPhase === "retryable_failure")
       ) {
+        if (nativeControllerLeaseExpiresAt && nativeControllerLeaseExpiresAt > now)
+          scope.onDeferredDeadline?.(nativeControllerLeaseExpiresAt.getTime());
         continue;
       }
       const observedOwnerUnverified =
@@ -2691,7 +2697,10 @@ export function createHeartbeatRecovery(db: Db, dependencies: HeartbeatRecoveryD
       active: scope => recoverActiveExecutions(scope),
       settled: async scope => {
         if (!scope.runId || activeRunExecutions.has(scope.runId) || runningProcesses.has(scope.runId)
-          || adapterExecutionControls.has(scope.runId)) return;
+          || adapterExecutionControls.has(scope.runId)) {
+          scope.onDeferredDeadline(Date.now() + 30_000);
+          return;
+        }
         await releaseSettledNativeRunProcess(db, scope.runId, scope);
       },
       orphanLease: scope => sweepOrphanedActiveLeases({ ...scope, backoffMs: 5 * 60_000 }),

server/src/services/instance-settings.ts

@@ -1,4 +1,4 @@
-import { notifyActiveRunSilence } from "./live-run-recovery-notifications.js";
+import { notifyLiveRunRecovery, notifyLiveResourceRecovery } from "./live-run-recovery-notifications.js";
 import { notifyDeferredExecutionEligibility } from "./deferred-execution-notifications.js";
 import { notifyRunDispatch } from "./run-dispatch-notifications.js";
 import { notifySkillVersionRetentionReconstruction } from "./skill-version-retention-notifications.js";
@@ -635,7 +635,8 @@ export function instanceSettingsService(db: Db, options: InstanceSettingsService
         if (patch.enableWorktreeRunExecution !== undefined) {
           await notifyRunDispatch(tx);
           await notifyDeferredExecutionEligibility(tx);
-          await notifyActiveRunSilence(tx);
+          await notifyLiveRunRecovery(tx);
+          await notifyLiveResourceRecovery(tx);
           await notifyDeliveryWork(tx, DELIVERY_QUEUES.agentTimers);
           await notifyExternalAgentSettlement(tx);
           await notifyInstructionCommit(tx);

server/src/services/live-run-recovery-notifications.ts

@@ -8,10 +8,12 @@ export async function notifyActiveRunSilence(tx: object, scope?: LiveRunRecovery
 }
 
 export async function notifyLiveRunRecovery(tx: object, scope?: LiveRunRecoveryScope) {
-  const hint = scope ? JSON.stringify({ companyId: scope.companyId, runId: scope.runId, agentId: scope.agentId }) : "*";
+  const hint = scope ? JSON.stringify({ companyId: scope.companyId, runId: scope.runId }) : "*";
   for (const queue of [DELIVERY_QUEUES.nativeProcessSessions, DELIVERY_QUEUES.nativeQuestionCancellations,
-    DELIVERY_QUEUES.activeExecutionRecovery, DELIVERY_QUEUES.settledNativeProcesses, DELIVERY_QUEUES.strandedAgentStatus])
+    DELIVERY_QUEUES.activeExecutionRecovery, DELIVERY_QUEUES.settledNativeProcesses])
     await notifyDeliveryWork(tx, queue, hint);
+  await notifyDeliveryWork(tx, DELIVERY_QUEUES.strandedAgentStatus, scope
+    ? JSON.stringify({ companyId: scope.companyId, agentId: scope.agentId }) : "*");
   await notifyActiveRunSilence(tx, scope);
 }
 
@@ -19,9 +21,12 @@ export type LiveResourceRecoveryScope = { companyId: string; runId?: string; lea
 
 export async function notifyLiveResourceRecovery(tx: object, scope?: LiveResourceRecoveryScope) {
   const hint = scope ? JSON.stringify({ companyId: scope.companyId, runId: scope.runId, leaseId: scope.leaseId }) : "*";
-  for (const queue of [DELIVERY_QUEUES.orphanActiveLeases, DELIVERY_QUEUES.legacyShutdownSettlement,
-    DELIVERY_QUEUES.strandedAgentStatus, DELIVERY_QUEUES.settledNativeProcesses])
+  for (const queue of [DELIVERY_QUEUES.orphanActiveLeases, DELIVERY_QUEUES.legacyShutdownSettlement])
     await notifyDeliveryWork(tx, queue, hint);
+  await notifyDeliveryWork(tx, DELIVERY_QUEUES.settledNativeProcesses, scope
+    ? JSON.stringify({ companyId: scope.companyId, runId: scope.runId }) : "*");
+  await notifyDeliveryWork(tx, DELIVERY_QUEUES.strandedAgentStatus, scope
+    ? JSON.stringify({ companyId: scope.companyId }) : "*");
 }
 
 export async function notifyStrandedAgentStatus(tx: object, scope: { companyId: string; agentId: string }) {

server/src/services/native-local-process-stop.ts

@@ -147,7 +147,7 @@ export async function readNativeLocalProcessStop(db: Db, companyId: string, runI
 type SettledProcessRelease = "released" | "pending" | "skipped";
 
 export async function releaseSettledNativeRunProcess(db: Db, runId: string, options: { onDeferredDeadline?: (at: number) => void } = {}): Promise<SettledProcessRelease> {
-  return db.transaction(async transaction => {
+  const outcome = await db.transaction(async transaction => {
     const tx = transaction as unknown as Db;
     const [run] = await tx.select({ id: heartbeatRuns.id, companyId: heartbeatRuns.companyId, agentId: heartbeatRuns.agentId,
       runtimeMode: heartbeatRuns.runtimeMode, status: heartbeatRuns.status, finishedAt: heartbeatRuns.finishedAt,
@@ -173,7 +173,8 @@ export async function releaseSettledNativeRunProcess(db: Db, runId: string, opti
         eq(heartbeatRunEvents.eventType, PROCESS_IDENTITY_RECORDED), isNull(heartbeatRunEvents.sourceEventId),
       )).orderBy(desc(heartbeatRunEvents.seq)).limit(1);
       if (identity?.payload?.targetKind !== "local" || identity.payload.hostname !== os.hostname() ||
-          identity.payload.processPid !== run.processPid || identity.payload.processGroupId !== run.processGroupId) return "skipped";
+          identity.payload.processPid !== run.processPid || identity.payload.processGroupId !== run.processGroupId ||
+          !run.processStartedAt || identity.payload.processStartedAt !== run.processStartedAt.toISOString()) return "skipped";
       if (!(await recordNativeLocalProcessStop(tx, run))) {
         options.onDeferredDeadline?.(Date.now() + 30_000);
         return "pending";
@@ -184,8 +185,10 @@ export async function releaseSettledNativeRunProcess(db: Db, runId: string, opti
     await tx.update(heartbeatRuns).set({
       processPid: null, processGroupId: null, processStartedAt: null, updatedAt: new Date(),
     }).where(eq(heartbeatRuns.id, run.id));
-    return "released";
+    return "released" as const;
   });
+  if (outcome !== "pending") pendingSettledProcessReleases.delete(runId);
+  return outcome;
 }
 
 const MAX_SETTLED_PROCESS_RELEASE_ATTEMPTS = 20;