@opengeni/worker-bundle 2.0.1 → 2.0.3-canary.36199476632001

package/src/sandbox-resume.ts

@@ -24,6 +24,7 @@ import {
   effectiveSandboxLifecycle,
   sandboxArchiveCaptureTimeoutMs,
   sandboxLifecycleTransitionWaitMs,
+  sandboxWarmRateMicrosPerSecond,
   type Settings,
 } from "@opengeni/config";
 import { randomUUID } from "node:crypto";
@@ -80,6 +81,7 @@ import {
   withoutSandboxProviderIdentity,
   type EstablishedSandboxSession,
   type RuntimeMetricsHooks,
+  type SandboxReadinessReplacementOutcome,
   type WorkspaceArchiveDescriptor,
 } from "@opengeni/runtime";
 import {
@@ -149,6 +151,24 @@ export type SandboxResumeServices = {
   /** Test seam for the bounded command-readiness proof performed before an
    * attached provider box is handed to the agent. */
   verifyAttachedSandboxReadiness?: (established: EstablishedSandboxSession) => Promise<void>;
+  /** Test seam for the bounded command-readiness proof the elected spawner runs
+   * on its freshly created box before publishing the lease warm. Production
+   * uses {@link waitForSandboxExecReadiness} with the shared 60 s budget. */
+  verifySpawnedSandboxReadiness?: (
+    established: EstablishedSandboxSession,
+    identity: { sandboxGroupId: string },
+  ) => Promise<void>;
+  /** Test seam for the jittered pause before the single fresh-box readiness
+   * replacement. Production uses {@link freshSandboxReadinessReplacementDelayMs}.
+   * It may be async so a test can observe the rolled-back lease in between. */
+  freshSandboxReadinessReplacementDelayMs?: () => number | Promise<number>;
+  /**
+   * The turn attempt's fresh-box readiness replacement budget. The lazy
+   * provisioner may call resumeBoxForTurn again after a typed lease
+   * supersession, so the turn creates the budget once and shares it with every
+   * call. Absent, the call gets its own single replacement.
+   */
+  freshSandboxReadinessReplacementBudget?: FreshSandboxReadinessReplacementBudget;
   /** Called only by the observer that wins the exact warm->cold loss CAS. */
   onSandboxLost?: (input: {
     sandboxGroupId: string;
@@ -317,6 +337,80 @@ export class SandboxSiblingWarmingTimeoutError extends SandboxWarmingTimeoutErro
   }
 }
 
+/**
+ * Out-of-band proof attached to one exact {@link SandboxExecReadinessTimeoutError}
+ * thrown by the elected spawner: the box that missed its readiness budget was
+ * freshly created by this caller (not an attached, resumed, or provider-continuity
+ * box), it was never published warm, the provider terminated it, and the holder
+ * was released. Only such an error may be replaced by re-entering admission.
+ */
+const disposedUnreadyFreshSandboxErrors = new WeakSet<object>();
+
+/** True only for a spawner readiness timeout whose unready fresh box is proven gone. */
+export function isReplaceableFreshSandboxReadinessTimeout(error: unknown): boolean {
+  return (
+    error instanceof SandboxExecReadinessTimeoutError &&
+    disposedUnreadyFreshSandboxErrors.has(error)
+  );
+}
+
+/**
+ * A burst of simultaneous turn starts can leave many freshly created provider
+ * boxes command-unready at once. The single replacement re-enters admission
+ * after a jittered pause so the replacement creates do not re-synchronize into
+ * the same burst. Bounded to [2 s, 10 s).
+ */
+export const FRESH_SANDBOX_READINESS_REPLACEMENT_BASE_DELAY_MS = 2_000;
+export const FRESH_SANDBOX_READINESS_REPLACEMENT_JITTER_MS = 8_000;
+
+export function freshSandboxReadinessReplacementDelayMs(
+  random: () => number = Math.random,
+): number {
+  const sample = random();
+  const unit = Number.isFinite(sample) ? Math.min(Math.max(sample, 0), 1 - Number.EPSILON) : 0;
+  return (
+    FRESH_SANDBOX_READINESS_REPLACEMENT_BASE_DELAY_MS +
+    Math.floor(unit * FRESH_SANDBOX_READINESS_REPLACEMENT_JITTER_MS)
+  );
+}
+
+/** Replacements allowed per turn attempt, across every resumeBoxForTurn call. */
+export const FRESH_SANDBOX_READINESS_REPLACEMENTS_PER_TURN_ATTEMPT = 1;
+
+export type FreshSandboxReadinessReplacementBudget = { remaining: number };
+
+export function createFreshSandboxReadinessReplacementBudget(): FreshSandboxReadinessReplacementBudget {
+  return { remaining: FRESH_SANDBOX_READINESS_REPLACEMENTS_PER_TURN_ATTEMPT };
+}
+
+function recordFreshSandboxReadinessReplacement(
+  metrics: RuntimeMetricsHooks | undefined,
+  backend: string,
+  outcome: SandboxReadinessReplacementOutcome,
+): void {
+  try {
+    metrics?.onSandboxReadinessReplacement?.({ backend, outcome });
+  } catch {
+    // Metrics emission must never affect sandbox recovery or error propagation.
+  }
+}
+
+async function sleepUnlessCancelled(ms: number, signal: AbortSignal | undefined): Promise<boolean> {
+  if (signal?.aborted) return false;
+  if (ms <= 0) return true;
+  return await new Promise<boolean>((resolve) => {
+    const onAbort = () => {
+      clearTimeout(timer);
+      resolve(false);
+    };
+    const timer = setTimeout(() => {
+      signal?.removeEventListener("abort", onAbort);
+      resolve(true);
+    }, ms);
+    signal?.addEventListener("abort", onAbort, { once: true });
+  });
+}
+
 /** The exact attached caller that won warm->cold after proving the provider
  * instance gone. It is a lease supersession (the same logical turn recovers),
  * plus the lost id needed for one durable observability event. */
@@ -1056,12 +1150,103 @@ async function persistWarmWorkspaceSnapshot(
  * holderId is the globally unique durable turn-attempt id. It must not be a
  * Temporal activity id, because activity ids are only workflow-local and
  * collide when sibling sessions share one sandbox group.
+ *
+ * A freshly created box that misses its bounded command-readiness budget is
+ * replaced at most once per turn attempt (the budget is shared by every call
+ * the turn's provisioner makes). The first establish has already terminated that exact
+ * unpublished box, rolled its exact warming epoch back to cold (advancing the
+ * epoch), and released its holder, so the replacement is an ordinary new
+ * admission: it re-runs the cold->warming CAS, persists its own provider
+ * instance before readiness, and may instead attach to a sibling that won the
+ * CAS first. Nothing model- or tool-visible ran on the discarded box, so no
+ * side effect is replayed. A second readiness timeout fails the turn exactly
+ * as before; attached/resumed boxes and unconfirmed terminations never
+ * replace.
  */
 export async function resumeBoxForTurn(
   services: SandboxResumeServices,
   ids: ResumeBoxIds,
   kind: "turn",
   holderId: TurnSandboxLeaseHolderId,
+): Promise<ResumedTurnSandbox> {
+  try {
+    return await resumeBoxForTurnOnce(services, ids, kind, holderId);
+  } catch (error) {
+    if (!isReplaceableFreshSandboxReadinessTimeout(error)) throw error;
+    const readiness = error as SandboxExecReadinessTimeoutError;
+    const warn = (message: string, fields: Parameters<Observability["warn"]>[1]) => {
+      if (services.observability) services.observability.warn(message, fields);
+      else console.warn(message, fields);
+    };
+    const identity = {
+      workspaceId: ids.workspaceId,
+      sessionId: ids.sessionId,
+      sandboxGroupId: ids.sandboxGroupId,
+      backend: readiness.backend,
+      instanceId: readiness.instanceId,
+      readinessTimeoutMs: readiness.timeoutMs,
+    };
+    const budget =
+      services.freshSandboxReadinessReplacementBudget ??
+      createFreshSandboxReadinessReplacementBudget();
+    if (budget.remaining <= 0) {
+      recordFreshSandboxReadinessReplacement(
+        services.sandboxMetrics,
+        readiness.backend,
+        "budget_spent",
+      );
+      warn(
+        "sandbox command-readiness timed out on a fresh box; this turn attempt already used its replacement",
+        identity,
+      );
+      throw error;
+    }
+    budget.remaining -= 1;
+    const delayMs = await (
+      services.freshSandboxReadinessReplacementDelayMs ?? freshSandboxReadinessReplacementDelayMs
+    )();
+    warn("sandbox command-readiness timed out on a fresh box; replacing it once after jitter", {
+      ...identity,
+      replacementDelayMs: delayMs,
+    });
+    if (!(await sleepUnlessCancelled(delayMs, services.cancellationSignal))) {
+      // Cancellation owns the turn boundary; surface the original typed
+      // readiness failure rather than starting a replacement.
+      recordFreshSandboxReadinessReplacement(
+        services.sandboxMetrics,
+        readiness.backend,
+        "cancelled",
+      );
+      throw error;
+    }
+    try {
+      const replaced = await resumeBoxForTurnOnce(services, ids, kind, holderId);
+      recordFreshSandboxReadinessReplacement(
+        services.sandboxMetrics,
+        readiness.backend,
+        "replaced",
+      );
+      return replaced;
+    } catch (replacementError) {
+      recordFreshSandboxReadinessReplacement(
+        services.sandboxMetrics,
+        readiness.backend,
+        services.cancellationSignal?.aborted
+          ? "cancelled"
+          : replacementError instanceof SandboxExecReadinessTimeoutError
+            ? "failed_again"
+            : "replacement_failed",
+      );
+      throw replacementError;
+    }
+  }
+}
+
+async function resumeBoxForTurnOnce(
+  services: SandboxResumeServices,
+  ids: ResumeBoxIds,
+  kind: "turn",
+  holderId: TurnSandboxLeaseHolderId,
 ): Promise<ResumedTurnSandbox> {
   const { db, settings } = services;
   const os = ids.os ?? "linux";
@@ -1153,6 +1338,10 @@ export async function resumeBoxForTurn(
     holderId,
     subjectId: ids.sessionId,
     backend: ids.backend,
+    warmBilling: {
+      mode: settings.sandboxWarmBillingMode,
+      rateMicrosPerSecond: sandboxWarmRateMicrosPerSecond(settings, ids.backend),
+    },
     os,
     // IMAGE IS SHARED STATE (B3): thread the resolved image so the lease stamps it +
     // conflicts on a live box already running a different image. A
@@ -1283,6 +1472,7 @@ export async function resumeBoxForTurn(
         holderId,
         expectedEpoch: heartbeat.expectedEpoch,
         leaseTtlMs: heartbeat.leaseTtlMs,
+        billingMode: services.settings.sandboxWarmBillingMode,
       })
         .then(async (status) => {
           if (heartbeat !== holderLeaseHeartbeat) return;
@@ -1650,11 +1840,17 @@ export async function resumeBoxForTurn(
       // A sandbox handle is not sufficient evidence that an asynchronous
       // provider's command router is live. Do not publish a warm lease until
       // one bounded no-op exec works.
-      // On timeout the catch below terminates the box and rolls warming -> cold,
-      // so the next turn cold-creates instead of hanging forever on first use.
-      await waitForSandboxExecReadiness(established, MODAL_EXEC_READINESS_TIMEOUT_MS, {
-        sandboxGroupId: ids.sandboxGroupId,
-      });
+      // On timeout the catch below terminates the box and rolls warming -> cold;
+      // resumeBoxForTurn may then replace a proven-terminated fresh box once.
+      if (services.verifySpawnedSandboxReadiness) {
+        await services.verifySpawnedSandboxReadiness(established, {
+          sandboxGroupId: ids.sandboxGroupId,
+        });
+      } else {
+        await waitForSandboxExecReadiness(established, MODAL_EXEC_READINESS_TIMEOUT_MS, {
+          sandboxGroupId: ids.sandboxGroupId,
+        });
+      }
       await maybeRenewProviderExpiration(true);
       throwIfReleasedOrCancelled();
       // Fold the LIVE box into a re-resumable envelope and persist it as the
@@ -1790,6 +1986,16 @@ export async function resumeBoxForTurn(
       }
       await release();
       recordSandboxWarmingTimeout(services.sandboxMetrics, error);
+      if (
+        terminated &&
+        createdEstablished &&
+        error instanceof SandboxExecReadinessTimeoutError &&
+        error.instanceId === createdEstablished.instanceId &&
+        (createdEstablished.origin === "created" || createdEstablished.origin === "restored") &&
+        !createdEstablished.providerContinuity
+      ) {
+        disposedUnreadyFreshSandboxErrors.add(error);
+      }
       throw sandboxProvisionStageError(rematerialization ? "archive_recovery" : "create", error);
     }
   }

package/src/workflows/activities.ts

@@ -1,4 +1,9 @@
-import { ActivityCancellationType, proxyActivities } from "@temporalio/workflow";
+import {
+  ActivityCancellationType,
+  ActivityFailure,
+  ApplicationFailure,
+  proxyActivities,
+} from "@temporalio/workflow";
 import type * as activities from "../activities";
 import {
   KNOWLEDGE_SOURCE_SYNC_ACTIVITY_HEARTBEAT_TIMEOUT_MS,
@@ -160,6 +165,14 @@ export const knowledgeSourceSyncActivity = proxyActivities<
 });
 
 export function workflowFailureMessage(error: unknown): string {
+  if (
+    error instanceof ActivityFailure &&
+    error.activityType === "runAgentTurn" &&
+    error.cause instanceof ApplicationFailure &&
+    error.cause.type === "TurnExecutionPolicyDefinitionMismatchError"
+  ) {
+    return error.cause.message;
+  }
   if (error instanceof Error) {
     return error.message;
   }

package/src/workflows/session.ts

@@ -58,6 +58,50 @@ const HUMAN_INPUT_EXPIRY_STALE_RETRY_MS = 1_000;
  */
 const ROTATION_IDLE_FLOOR_MS = 60_000; // 60s
 
+/**
+ * Provider capacity waits are shared: every waiter of one exhausted pool learns
+ * the same authoritative reset time, and one capacity mutation (for example the
+ * bounded refresh that verifies a quota reset) wakes every waiter at once, both
+ * as the typed capacity signal and as the generic durable workflow wake. Without
+ * spread, a reset resumes the whole backlog in the same few seconds and its
+ * turns stampede sandbox creation and the provider. Each workflow therefore
+ * delays its reconciliation by a bounded, replay-deterministic (Temporal-seeded
+ * `Math.random`) jitter: up to one minute past a scheduled reset timer and up to
+ * 30 seconds after a capacity or queue wake or an already-due waiter. A wake
+ * that lands inside a waiter's timer spread does not shorten it. The Postgres
+ * waiter stays authoritative; jitter only delays the same reconciliation
+ * activity and never creates queue rows, input, or inference. Interruptions are
+ * never delayed.
+ *
+ * The patch marker is not understood by workers built before it: rolling the
+ * worker image back past this change while a session has recorded the marker
+ * (it waited on capacity in its current run) fails that workflow's tasks as
+ * nondeterministic until a patched worker returns.
+ */
+export const CAPACITY_WAKE_JITTER_PATCH = "session-capacity-wake-jitter-v1";
+export const CAPACITY_TIMER_WAKE_JITTER_MAX_MS = 60_000;
+export const CAPACITY_WAKE_JITTER_MAX_MS = 30_000;
+
+/**
+ * Pure + exported so the bound is unit-testable without a workflow environment.
+ * `overrideMaxMs` is the test-only SessionWorkflowInput ceiling; it can only
+ * narrow the production bound.
+ */
+export function capacityWakeJitterMs(
+  kind: "timer" | "wake",
+  sample: number,
+  overrideMaxMs?: number,
+): number {
+  const productionMax =
+    kind === "timer" ? CAPACITY_TIMER_WAKE_JITTER_MAX_MS : CAPACITY_WAKE_JITTER_MAX_MS;
+  const max =
+    overrideMaxMs !== undefined && Number.isFinite(overrideMaxMs)
+      ? Math.min(productionMax, Math.max(0, Math.trunc(overrideMaxMs)))
+      : productionMax;
+  const unit = Number.isFinite(sample) ? Math.min(Math.max(sample, 0), 1 - Number.EPSILON) : 0;
+  return Math.floor(unit * max);
+}
+
 /**
  * How long the continuation loop must hold before re-admitting the next turn. 0 ⇒ no
  * hold (re-dispatch immediately — a rotation candidate is ready, or no idle delay was
@@ -369,6 +413,9 @@ export type SessionWorkflowInput = {
   // Test-only override for the durable capacity-wait continue-as-new
   // backstop. Production uses CODEX_CAPACITY_CHECKS_PER_RUN_BACKSTOP.
   maxCapacityChecksPerRun?: number;
+  // Test-only ceiling for capacity-wake jitter (0 disables it) so real-server
+  // workflow tests stay fast. Production omits it and uses the bounds above.
+  capacityWakeJitterMaxMs?: number;
 };
 
 export async function sessionWorkflow(input: SessionWorkflowInput): Promise<void> {
@@ -476,7 +523,14 @@ export async function sessionWorkflow(input: SessionWorkflowInput): Promise<void
       const timerMs = Number.isFinite(parsedDeadline)
         ? Math.max(0, parsedDeadline - Date.now())
         : 0;
+      const observedCause = (): activities.ReconcileCodexCapacityWaitInput["cause"] =>
+        wakeups !== seenWakeups
+          ? "queue"
+          : capacityWakeups !== seenCapacityWakeups
+            ? "signal"
+            : "timer";
       let cause: activities.ReconcileCodexCapacityWaitInput["cause"] = "timer";
+      let spread = false;
       if (wakeups !== seenWakeups) {
         cause = "queue";
       } else if (capacityWakeups !== seenCapacityWakeups) {
@@ -492,12 +546,48 @@ export async function sessionWorkflow(input: SessionWorkflowInput): Promise<void
         if (interruptionWakeups !== seenInterruptionWakeups) {
           return;
         }
-        cause =
-          wakeups !== seenWakeups
-            ? "queue"
-            : capacityWakeups !== seenCapacityWakeups
-              ? "signal"
-              : "timer";
+        cause = observedCause();
+        // Histories recorded before the jitter replay without the marker and
+        // reconcile right at the deadline, exactly as they did.
+        if (cause === "timer" && patched(CAPACITY_WAKE_JITTER_PATCH)) {
+          // The shared reset deadline itself fired. The first waiter to
+          // reconcile usually refreshes usage and wakes every other waiter; a
+          // waiter already inside its own spread keeps it instead of starting a
+          // second, shorter one, so the backlog resumes across the whole window.
+          spread = true;
+          const timerJitterMs = capacityWakeJitterMs(
+            "timer",
+            Math.random(),
+            input.capacityWakeJitterMaxMs,
+          );
+          if (timerJitterMs > 0) {
+            await condition(() => interruptionWakeups !== seenInterruptionWakeups, timerJitterMs);
+            if (interruptionWakeups !== seenInterruptionWakeups) {
+              return;
+            }
+          }
+          cause = observedCause();
+        }
+      }
+      if (!spread && patched(CAPACITY_WAKE_JITTER_PATCH)) {
+        // One capacity mutation wakes every waiter of the pool at once (the typed
+        // capacity signal and the generic durable wake), and a fresh run after
+        // continue-as-new, restart, or Resume sees every unobserved wake as an
+        // already-due waiter. Spread those reconciliations too. A later wake
+        // does not cut the pause short (that would re-synchronize the herd); a
+        // queued prompt stays behind the blocked turn either way, so only an
+        // interruption (Pause/Steer/Cancel) wins immediately.
+        const wakeJitterMs = capacityWakeJitterMs(
+          "wake",
+          Math.random(),
+          input.capacityWakeJitterMaxMs,
+        );
+        if (wakeJitterMs > 0) {
+          await condition(() => interruptionWakeups !== seenInterruptionWakeups, wakeJitterMs);
+          if (interruptionWakeups !== seenInterruptionWakeups) {
+            return;
+          }
+        }
       }
       const reconcileInput = {
         accountId: input.accountId,
@@ -536,6 +626,9 @@ export async function sessionWorkflow(input: SessionWorkflowInput): Promise<void
           ...(input.maxCapacityChecksPerRun !== undefined
             ? { maxCapacityChecksPerRun: input.maxCapacityChecksPerRun }
             : {}),
+          ...(input.capacityWakeJitterMaxMs !== undefined
+            ? { capacityWakeJitterMaxMs: input.capacityWakeJitterMaxMs }
+            : {}),
         });
       }
       current = result;
@@ -582,6 +675,9 @@ export async function sessionWorkflow(input: SessionWorkflowInput): Promise<void
           ...(input.maxCapacityChecksPerRun !== undefined
             ? { maxCapacityChecksPerRun: input.maxCapacityChecksPerRun }
             : {}),
+          ...(input.capacityWakeJitterMaxMs !== undefined
+            ? { capacityWakeJitterMaxMs: input.capacityWakeJitterMaxMs }
+            : {}),
         });
       }
     }