@opengeni/worker-bundle 2.0.1 → 2.0.3

package/src/activities/retained-screenshots.ts

@@ -6,13 +6,16 @@ import {
   COMPUTER_SCREENSHOT_WORKSPACE_QUOTA_BYTES,
   RetainedArtifactMetadataSchema,
   retainedScreenshotReferenceFromFile,
+  retainedSessionScreenshotKindFromObjectKey,
   type RetainedArtifactMetadata,
   type RetainedArtifactReference,
   type RetainedOutputUnavailableReason,
+  type RetainedSessionScreenshotKind,
 } from "@opengeni/contracts";
 import {
   RetainedScreenshotQuotaExceededError,
   getRetainedScreenshotArtifact,
+  getRetainedScreenshotArtifactForToolCall,
   isDatabasePersistenceFailure,
   isSessionEventPersistenceError,
   prepareRetainedScreenshotArtifact,
@@ -209,14 +212,26 @@ export function toolOutputContainsInlineImage(output: unknown): boolean {
       return (
         (record.type === "input_image" &&
           sdkImageSourceContainsInlineImage(record.image ?? record.image_url ?? record.imageUrl)) ||
-        (record.type === "image" && decodeInlineMcpImage(record.data, record.mimeType) !== null)
+        (record.type === "image" && isDeclaredInlineMcpImage(record))
       );
     });
   }
   if (!output || typeof output !== "object") return false;
   const record = output as Record<string, unknown>;
   if (Array.isArray(record.content)) return toolOutputContainsInlineImage(record.content);
-  return record.type === "image" && sdkImageSourceContainsInlineImage(record.image);
+  return (
+    record.type === "image" &&
+    (sdkImageSourceContainsInlineImage(record.image) || isDeclaredInlineMcpImage(record))
+  );
+}
+
+/** Detect image-bearing MCP blocks even when decoding rejects oversized or malformed bytes. */
+function isDeclaredInlineMcpImage(record: Record<string, unknown>): boolean {
+  return (
+    typeof record.data === "string" &&
+    typeof record.mimeType === "string" &&
+    record.mimeType.toLowerCase().startsWith("image/")
+  );
 }
 
 function decodeInlineMcpImage(
@@ -520,9 +535,14 @@ export function retainedScreenshotIdentity(input: {
   attemptId: string;
   toolCallId: string;
   toolOutputId: string;
+  kind?: RetainedSessionScreenshotKind | undefined;
 }): { artifactId: string; settlementKey: string } {
   const settlementKey = createHash("sha256")
-    .update("opengeni:computer-screenshot:v1\0")
+    .update(
+      input.kind === "browser_screenshot"
+        ? "opengeni:browser-screenshot:v1\0"
+        : "opengeni:computer-screenshot:v1\0",
+    )
     .update(input.sessionId)
     .update("\0")
     .update(input.turnId)
@@ -551,12 +571,13 @@ export function unavailableRetainedSessionImage(input: {
   toolCallId: string;
   toolOutputId: string;
   reason: RetainedOutputUnavailableReason;
+  kind?: RetainedSessionScreenshotKind | undefined;
 }): RetainedArtifactMetadata {
   const identity = retainedScreenshotIdentity(input);
   return unavailable(identity.artifactId, input.reason);
 }
 
-export async function retainComputerScreenshot(input: {
+export async function retainSessionScreenshot(input: {
   db: Database;
   objectStorage: ObjectStorage | null;
   accountId: string;
@@ -565,6 +586,7 @@ export async function retainComputerScreenshot(input: {
   turnId: string;
   attemptId: string;
   output: TypedScreenshotToolOutput;
+  kind?: RetainedSessionScreenshotKind | undefined;
   now?: Date;
   retentionMs?: number;
   workspaceQuotaBytes?: number;
@@ -575,6 +597,7 @@ export async function retainComputerScreenshot(input: {
     attemptId: input.attemptId,
     toolCallId: input.output.callId,
     toolOutputId: input.output.toolOutputId,
+    kind: input.kind,
   });
   let screenshot: ValidatedSessionImage;
   try {
@@ -594,7 +617,8 @@ export async function retainComputerScreenshot(input: {
   const retentionExpiresAt = new Date(
     now.getTime() + (input.retentionMs ?? COMPUTER_SCREENSHOT_RETENTION_MS),
   );
-  const objectKey = `workspaces/${input.workspaceId}/files/${identity.artifactId}/retained/session-image.${screenshot.extension}`;
+  const objectName = input.kind === "browser_screenshot" ? "browser-screenshot" : "session-image";
+  const objectKey = `workspaces/${input.workspaceId}/files/${identity.artifactId}/retained/${objectName}.${screenshot.extension}`;
   let prepared: RetainedScreenshotArtifact;
   try {
     prepared = (
@@ -697,6 +721,9 @@ export async function retainComputerScreenshot(input: {
   }
 }
 
+/** Existing computer callers keep their stable entry point. */
+export const retainComputerScreenshot = retainSessionScreenshot;
+
 /** Replace only new screenshot bytes with a compact receipt before first persistence. */
 export function compactRetainedScreenshotHistory(
   history: Array<Record<string, unknown>>,
@@ -874,6 +901,26 @@ async function materializeRetainedScreenshotHistoryWithCache(
   cache: Map<string, string>,
 ): Promise<Array<Record<string, unknown>>> {
   const now = input.now ?? new Date();
+  const receiptForMarker = async (
+    marker: unknown,
+    callId: string | null,
+  ): Promise<RetainedArtifactMetadata | null> => {
+    if (!isRetainedImageMarker(marker)) return null;
+    const receipt = retainedReceipt(marker.artifact);
+    if (receipt) return receipt;
+    // Older canonical history could truncate the artifact UUID and reason as
+    // ordinary text. Recover only from this exact session and an unambiguous
+    // tool call; otherwise fail before constructing a malformed model image.
+    if (!callId) throw new Error("Retained screenshot receipt has no tool call identity");
+    const artifact = await getRetainedScreenshotArtifactForToolCall(
+      input.db,
+      input.workspaceId,
+      input.sessionId,
+      callId,
+    );
+    if (!artifact) throw new Error("Retained screenshot receipt cannot be recovered");
+    return reference(artifact);
+  };
   const dataUrlForReceipt = async (receipt: RetainedArtifactMetadata): Promise<string> => {
     let dataUrl = cache.get(receipt.artifactId);
     if (!dataUrl) {
@@ -921,8 +968,12 @@ async function materializeRetainedScreenshotHistoryWithCache(
     }
     return dataUrl;
   };
-  const materializeEntry = async (entry: unknown): Promise<unknown> => {
-    const receipt = retainedReceiptFromImageContent(entry);
+  const materializeEntry = async (entry: unknown, callId: string | null): Promise<unknown> => {
+    const image =
+      entry && typeof entry === "object" && !Array.isArray(entry)
+        ? (entry as Record<string, unknown>).image
+        : null;
+    const receipt = await receiptForMarker(image, callId);
     if (!receipt) return entry;
     const dataUrl = await dataUrlForReceipt(receipt);
     return { ...(entry as Record<string, unknown>), image: dataUrl };
@@ -930,7 +981,8 @@ async function materializeRetainedScreenshotHistoryWithCache(
 
   const materialized: Array<Record<string, unknown>> = [];
   for (const item of input.history) {
-    const directReceipt = retainedReceiptFromDirectOutput(item.output);
+    const callId = historyCallId(item);
+    const directReceipt = await receiptForMarker(item.output, callId);
     if (directReceipt) {
       materialized.push({
         ...item,
@@ -954,7 +1006,7 @@ async function materializeRetainedScreenshotHistoryWithCache(
     let changed = false;
     const output: unknown[] = [];
     for (const entry of outputEntries) {
-      const next = await materializeEntry(entry);
+      const next = await materializeEntry(entry, callId);
       changed ||= next !== entry;
       output.push(next);
     }
@@ -1059,12 +1111,15 @@ function reference(artifact: RetainedScreenshotArtifact): RetainedArtifactRefere
   if (!artifact.sessionId) {
     throw new Error(`Ready retained screenshot is detached: ${artifact.artifactId}`);
   }
+  const kind = retainedSessionScreenshotKindFromObjectKey(artifact.file.objectKey);
+  if (!kind) throw new Error(`Ready retained screenshot has unknown key: ${artifact.artifactId}`);
   const value = retainedScreenshotReferenceFromFile({
     ...artifact.file,
     sessionId: artifact.sessionId,
     width: artifact.width,
     height: artifact.height,
     expiresAt: artifact.retentionExpiresAt.toISOString(),
+    kind,
   });
   if (!value)
     throw new Error(`Ready retained screenshot receipt is invalid: ${artifact.artifactId}`);
@@ -1110,6 +1165,18 @@ function retainedReceiptFromImageContent(entry: unknown): RetainedArtifactMetada
   return retainedReceipt(marker.artifact);
 }
 
+function isRetainedImageMarker(value: unknown): value is {
+  type: "retained_artifact";
+  artifact: unknown;
+} {
+  return Boolean(
+    value &&
+    typeof value === "object" &&
+    !Array.isArray(value) &&
+    (value as Record<string, unknown>).type === RETAINED_IMAGE_MARKER,
+  );
+}
+
 function retainedReceiptFromDirectOutput(output: unknown): RetainedArtifactMetadata | null {
   if (!output || typeof output !== "object" || Array.isArray(output)) return null;
   const marker = output as Record<string, unknown>;
@@ -1120,7 +1187,11 @@ function retainedReceiptFromDirectOutput(output: unknown): RetainedArtifactMetad
 function retainedReceipt(value: unknown): RetainedArtifactMetadata | null {
   const parsed = RetainedArtifactMetadataSchema.safeParse(value);
   if (!parsed.success) return null;
-  return parsed.data.available && parsed.data.kind !== "computer_screenshot" ? null : parsed.data;
+  return parsed.data.available &&
+    parsed.data.kind !== "computer_screenshot" &&
+    parsed.data.kind !== "browser_screenshot"
+    ? null
+    : parsed.data;
 }
 
 function isInlineImageContent(entry: unknown): boolean {
@@ -1130,7 +1201,7 @@ function isInlineImageContent(entry: unknown): boolean {
     (record.type === "input_image" &&
       typeof record.image === "string" &&
       record.image.startsWith("data:image/")) ||
-    (record.type === "image" && decodeInlineMcpImage(record.data, record.mimeType) !== null)
+    (record.type === "image" && isDeclaredInlineMcpImage(record))
   );
 }
 

package/src/activities/sandbox-lease.ts

@@ -39,10 +39,10 @@ import {
   countSandboxLeasesByLiveness,
   deferRetainedProcessReconciliation,
   forceDrainOverLimitViewerOnlyBoxes,
-  getBillingBalance,
   listCreditBalancesByAccount,
   listLegacyModalCheckpointSlots,
   listLiveModalSandboxLeaseAttributions,
+  markWarmBillingStopCutoff,
   listMeterableWarmLeases,
   listSandboxViewerForceDrainWorkspaceIds,
   markSandboxCheckpointArtifactDeletePending,
@@ -186,8 +186,7 @@ export type ReapSandboxLeasesResult = {
   skipped: number;
   /** Warm viewer-only leases that accrued warm-seconds this tick (P2.1). */
   metered: number;
-  /** Viewer-only boxes force-drained because their workspace is over a limit
-   *  (0 balance / over the warm cap) — turn-held boxes are never drained (P2.1). */
+  /** Holderless boxes accelerated to draining by the optional warm cap. */
   forceDrained: number;
   /** Provider-side Modal orphan sandboxes terminated by the defensive sweep. */
   modalOrphansTerminated: number;
@@ -256,6 +255,7 @@ export type TerminateBoxFn = (
   diskBackedArchives?: boolean,
   releaseFailedCapture?: () => Promise<void>,
   workspaceId?: string,
+  beforeProviderStop?: () => Promise<void>,
 ) => Promise<boolean | ProviderTerminationOutcome>;
 
 export type SweepModalOrphansFn = (
@@ -482,7 +482,14 @@ export function createSandboxLeaseActivities(
   options: SandboxLeaseActivityOptions = {},
 ) {
   const terminateBox: TerminateBoxFn =
-    options.terminateBox ??
+    (options.terminateBox
+      ? async (...args: Parameters<TerminateBoxFn>) => {
+          // The injected test seam represents the stop itself; the real
+          // implementation invokes this only after capture has succeeded.
+          await args[10]?.();
+          return options.terminateBox!(...args);
+        }
+      : null) ??
     (async (
       settings,
       lease,
@@ -494,6 +501,7 @@ export function createSandboxLeaseActivities(
       diskBackedArchives,
       releaseFailedCapture,
       workspaceId,
+      beforeProviderStop,
     ) =>
       await terminateProviderBox(
         settings,
@@ -508,6 +516,7 @@ export function createSandboxLeaseActivities(
         diskBackedArchives,
         releaseFailedCapture,
         workspaceId,
+        beforeProviderStop,
       ));
   const sweepModalOrphans: SweepModalOrphansFn =
     options.sweepModalOrphans ?? sweepModalOrphansForConfiguredBackend;
@@ -735,8 +744,7 @@ export function createSandboxLeaseActivities(
       metered = meterResult.accrued;
       const forceDrainWorkspaceIds = new Set<string>();
       if (
-        settings.billingMode === "stripe" ||
-        settings.usageLimitsMode === "managed" ||
+        settings.sandboxWarmBillingMode === "credits" ||
         settings.sandboxMaxWarmSecondsPerWorkspace > 0
       ) {
         for (const workspaceId of meterResult.workspaceIds) {
@@ -1507,6 +1515,25 @@ async function reconcileTerminalRetainedProcesses(
         }
       }
 
+      if (
+        observation?.status === "deferred" &&
+        process.providerBackend === "modal" &&
+        !process.providerBindingKey &&
+        !process.deadlineCancellationRequestedAt
+      ) {
+        // An unbound historical command cannot be signalled safely. Still
+        // record deadline intent when its lease matches: a failed observer
+        // must not block capture until the provider expires.
+        try {
+          await requestRetainedProcessDeadlineCancellation(db, processScope);
+        } catch (error) {
+          observability.warn("sandbox reaper: legacy deadline cancellation intent failed", {
+            processId: process.id,
+            error: error instanceof Error ? error.message : String(error),
+          });
+        }
+      }
+
       if (observation === null) {
         const lease = await readLease(db, process.workspaceId, process.sandboxGroupId).catch(
           () => null,
@@ -1583,11 +1610,19 @@ async function reconcileTerminalRetainedProcesses(
               );
               supervisionMetric("cancellation_intent");
             }
-            observation = await probe(
+            const legacyDeadlineStop =
+              lease!.rotationReason === "provider_deadline" &&
+              !supervised &&
+              !process.deadlineCancellationRequestedAt;
+            const probeOutcome = await probe(
               settings,
               lease!,
               process,
-              claim.ownerState === "background_stopping" ? "cancel" : "observe",
+              (claim.ownerState === "background_stopping" &&
+                !process.deadlineCancellationRequestedAt) ||
+                legacyDeadlineStop
+                ? "cancel"
+                : "observe",
               async (result, chunkId, stream, streamFidelity) => {
                 if (
                   typeof result !== "string" ||
@@ -1610,7 +1645,20 @@ async function reconcileTerminalRetainedProcesses(
                     .catch(() => undefined);
               },
               commandPersistence,
+            ).then(
+              (value) => ({ ok: true as const, value }),
+              (error: unknown) => ({ ok: false as const, error }),
             );
+            // Legacy PTY input reserves Ctrl-C before the cancellation fence.
+            // Non-PTY has no command signal in the Modal SDK: probe it without
+            // injecting Ctrl-C as data, then fence and save after the grace.
+            if (
+              legacyDeadlineStop &&
+              !(await requestRetainedProcessDeadlineCancellation(db, processScope))
+            )
+              throw new Error("Deadline cancellation no longer owns its rotating lease");
+            if (!probeOutcome.ok) throw probeOutcome.error;
+            observation = probeOutcome.value;
             if (supervised) {
               supervisionMetric(
                 (await commandPersistence.loadSupervisionReceipt())
@@ -1899,7 +1947,13 @@ export async function probeRetainedProcessAtProvider(
   if (pending) {
     if (!captureOutput) throw new Error("Pending command output requires its persistence callback");
     await captureRetainedProbeOutput(process.id, pending.result, captureOutput);
-    return classifyRetainedProcessPollResult(pending.result, process.providerSessionId);
+    const pendingObservation = classifyRetainedProcessPollResult(
+      pending.result,
+      process.providerSessionId,
+    );
+    // Capture of a previous page must not consume the only deadline stop
+    // attempt. Terminal proof needs no further provider call.
+    if (mode !== "cancel" || pendingObservation.status === "proved") return pendingObservation;
   }
   const envelopeBackend = (lease.resumeState as { backendId?: unknown }).backendId;
   if (
@@ -1996,6 +2050,7 @@ export async function probeRetainedProcessAtProvider(
   if (typeof session.writeStdin !== "function") {
     return { status: "deferred", reason: "backend_unsupported" };
   }
+  let legacyPty = false;
   if (process.providerBackend === "modal") {
     let liveBinding: Awaited<ReturnType<typeof resolveModalCheckpointProviderBindingForSession>>;
     try {
@@ -2027,6 +2082,7 @@ export async function probeRetainedProcessAtProvider(
     if (!command || !providerPersistence || !session.bindProviderCommand)
       return { status: "deferred", reason: "process_observation_unavailable" };
     session.bindProviderCommand(process.providerSessionId, command, providerPersistence);
+    legacyPty = command.pty === true;
   }
   const capturePage = async (value: unknown): Promise<void> => {
     if (!captureOutput) {
@@ -2062,7 +2118,14 @@ export async function probeRetainedProcessAtProvider(
     result = await withRetainedProcessProbeTimeout(
       session.writeStdin({
         sessionId: process.providerSessionId,
-        chars: mode === "cancel" && !supervisedCancelled ? "\u0003" : "",
+        // Ctrl-C is a signal only through a PTY. For a non-PTY command it is
+        // ordinary stdin data and can corrupt the command's input.
+        chars:
+          mode === "cancel" &&
+          !supervisedCancelled &&
+          (process.providerBackend !== "modal" || legacyPty)
+            ? "\u0003"
+            : "",
         yieldTimeMs: 1_000,
         maxOutputTokens: 2_000,
       }),
@@ -2092,6 +2155,9 @@ export async function probeRetainedProcessAtProvider(
     observation.status === "deferred" &&
     observation.reason === "provider_running" &&
     lease.rotationRequestedAt !== null &&
+    mode !== "cancel" &&
+    !process.cancellationRequestedAt &&
+    legacyPty &&
     !(await providerPersistence
       ?.load()
       .then((command) => command?.kind === "modal-router-v1" && command.supervision))
@@ -2197,19 +2263,25 @@ async function accrueWarmTick(
   }
   await forEachWithConcurrency(leases, SANDBOX_MAINTENANCE_ITEM_CONCURRENCY, async (lease) => {
     try {
-      const rate = sandboxWarmRateMicrosPerSecond(settings, lease.backend);
+      const rate =
+        settings.sandboxWarmBillingMode === "usage_only"
+          ? 0
+          : sandboxWarmRateMicrosPerSecond(settings, lease.backend);
       const result = await accrueWarmSeconds(db, {
         accountId: lease.accountId,
         workspaceId: lease.workspaceId,
         sandboxGroupId: lease.sandboxGroupId,
         expectedEpoch: lease.leaseEpoch,
         warmRateMicrosPerSecond: rate,
+        billingMode: settings.sandboxWarmBillingMode,
         subjectId: lease.sandboxGroupId,
       });
       if (result.accrued) {
         accrued += 1;
       }
-      recordCreditMicros(observability, "usage", result.costMicros);
+      if (result.costMicros > 0 && settings.sandboxWarmBillingMode === "credits") {
+        recordCreditMicros(observability, "usage", result.costMicros);
+      }
     } catch (error) {
       observability.warn("sandbox reaper: warm-seconds accrual failed for lease", {
         workspaceId: lease.workspaceId,
@@ -2325,8 +2397,8 @@ async function refreshSandboxInventoryGauge(
 
 /**
  * The per-workspace warm-cap + force-drain pass (P2.1). For each workspace with a
- * live warm box, under the usage lock: if it is at 0 balance (when a billing /
- * managed mode is on) or over its warm cap, force-drain its VIEWER-ONLY boxes
+ * live warm box, under the usage lock: if it is over its optional warm cap,
+ * accelerate only holderless boxes to draining
  * (guarded turn_holders=0 — a paying turn is never killed). Returns the count of
  * viewer-only boxes force-drained. Per-workspace best-effort.
  */
@@ -2336,8 +2408,7 @@ async function forceDrainOverLimitWorkspaces(
   workspaceIds: Set<string>,
   observability: ActivityServices["observability"],
 ): Promise<number> {
-  const enforceBalance =
-    settings.billingMode === "stripe" || settings.usageLimitsMode === "managed";
+  const enforceBalance = false;
   const cap = settings.sandboxMaxWarmSecondsPerWorkspace;
   let forceDrained = 0;
   await forEachWithConcurrency(
@@ -2345,13 +2416,8 @@ async function forceDrainOverLimitWorkspaces(
     SANDBOX_MAINTENANCE_ITEM_CONCURRENCY,
     async (workspaceId) => {
       try {
-        const { accountId } = await rlsContextForWorkspace(db, workspaceId);
-        const balance = enforceBalance
-          ? await getBillingBalance(db, accountId)
-          : ({ balanceMicros: 1 } as { balanceMicros: number });
         const result = await forceDrainOverLimitViewerOnlyBoxes(db, {
           workspaceId,
-          balanceMicros: balance.balanceMicros,
           enforceBalance,
           maxWarmSecondsPerWorkspace: cap,
           idleGraceMs: settings.sandboxIdleGraceMs,
@@ -2939,6 +3005,20 @@ async function terminateDrainableBox(
   ) {
     return false;
   }
+  // Stamp the charge horizon after verified capture and immediately before
+  // provider stop. Capture can fail and release its claim, allowing re-arm.
+  // A retry after stop success cannot charge post-termination wall time.
+  let stopCutoffMarked = false;
+  const beforeProviderStop = async (): Promise<void> => {
+    const cutoff = await markWarmBillingStopCutoff(db, {
+      accountId,
+      workspaceId: row.workspaceId,
+      sandboxGroupId: row.sandboxGroupId,
+      expectedEpoch: row.leaseEpoch,
+    });
+    if (!cutoff) throw new Error("sandbox stop cutoff was fenced before provider termination");
+    stopCutoffMarked = true;
+  };
   const termination: ProviderTerminationOutcome | boolean = providerMissingBeforeCapture
     ? { terminated: true, providerMissingBeforeCapture: true }
     : await terminateBox(
@@ -2962,12 +3042,49 @@ async function terminateDrainableBox(
           });
         },
         row.workspaceId,
+        beforeProviderStop,
       );
   const terminated = typeof termination === "boolean" ? termination : termination.terminated;
   if (!terminated) {
     return false;
   }
 
+  // Self-hosted/none and an unassigned lease have no provider stop callback;
+  // finish their usage at logical detach without touching customer compute.
+  // A managed instance that was stopped without its callback must fail closed.
+  const providerMissingDuringStop =
+    providerMissingBeforeCapture ||
+    (typeof termination !== "boolean" && termination.providerMissingBeforeCapture);
+  if (!stopCutoffMarked && !providerMissingDuringStop) {
+    if (managedProvider && lease.instanceId) {
+      throw new Error("managed sandbox stopped without a warm usage cutoff");
+    }
+    await beforeProviderStop();
+  }
+
+  // Stop was confirmed but the lease is not cold yet. This final epoch-fenced
+  // tick settles the interval since the last heartbeat/reaper sweep, including
+  // a box drained before its first periodic tick. If settlement fails, keep
+  // the draining receipt for retry instead of silently losing the charge.
+  const finalRate =
+    settings.sandboxWarmBillingMode === "usage_only"
+      ? 0
+      : sandboxWarmRateMicrosPerSecond(settings, backend);
+  if (!providerMissingDuringStop) {
+    const finalTick = await accrueWarmSeconds(db, {
+      accountId,
+      workspaceId: row.workspaceId,
+      sandboxGroupId: row.sandboxGroupId,
+      expectedEpoch: row.leaseEpoch,
+      warmRateMicrosPerSecond: finalRate,
+      billingMode: settings.sandboxWarmBillingMode,
+      finalDrain: true,
+    });
+    if (finalTick.costMicros > 0 && settings.sandboxWarmBillingMode === "credits") {
+      recordCreditMicros(observability, "usage", finalTick.costMicros);
+    }
+  }
+
   const providerMissing =
     providerMissingBeforeCapture ||
     (typeof termination === "boolean" ? false : termination.providerMissingBeforeCapture);
@@ -3097,6 +3214,7 @@ export async function terminateProviderBox(
   diskBackedArchives = false,
   releaseFailedCapture?: () => Promise<void>,
   workspaceId?: string,
+  beforeProviderStop?: () => Promise<void>,
 ): Promise<ProviderTerminationOutcome> {
   const durableBackendId = (lease.resumeBackendId ?? lease.backend) as string;
   const backend = sandboxBackendForSdkBackendId(durableBackendId) ?? durableBackendId;
@@ -3196,6 +3314,7 @@ export async function terminateProviderBox(
       return { terminated: false, providerMissingBeforeCapture: false };
     }
     try {
+      await beforeProviderStop?.();
       await terminateModalById(settings, lease.instanceId);
     } catch (error) {
       if (!isProviderSandboxNotFoundError("modal", error)) {
@@ -3439,6 +3558,7 @@ export async function terminateProviderBox(
       throw new Error(`sandbox backend ${backend} returned no terminable session state`);
     }
     prepareProviderForTeardownAfterCapture(backend, session);
+    await beforeProviderStop?.();
     await terminateManagedSandboxSession(client, sessionState, session);
     return { terminated: true, providerMissingBeforeCapture: false };
   } catch (error) {

package/src/sandbox-resume.ts

@@ -24,6 +24,7 @@ import {
   effectiveSandboxLifecycle,
   sandboxArchiveCaptureTimeoutMs,
   sandboxLifecycleTransitionWaitMs,
+  sandboxWarmRateMicrosPerSecond,
   type Settings,
 } from "@opengeni/config";
 import { randomUUID } from "node:crypto";
@@ -1153,6 +1154,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 +1288,7 @@ export async function resumeBoxForTurn(
         holderId,
         expectedEpoch: heartbeat.expectedEpoch,
         leaseTtlMs: heartbeat.leaseTtlMs,
+        billingMode: services.settings.sandboxWarmBillingMode,
       })
         .then(async (status) => {
           if (heartbeat !== holderLeaseHeartbeat) return;

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;
   }