@opengeni/worker-bundle 0.28.1-canary.0 → 0.28.3-canary.0

package/src/sandbox-archive-storage.ts

@@ -1,4 +1,6 @@
-import type { ObjectStorage } from "@opengeni/storage";
+import { uploadWorkspaceArchiveSpool, type ObjectStorage } from "@opengeni/storage";
+import { randomUUID } from "node:crypto";
+import type { VerifiedHostWorkspaceArchive } from "@opengeni/runtime/sandbox";
 import {
   parseWorkspaceArchiveObjectRef,
   workspaceArchiveObjectKey,
@@ -35,43 +37,102 @@ export async function putTarWorkspaceArchiveObject(input: {
   accountId: string;
   workspaceId: string;
   sandboxGroupId: string;
-  archive: { bytes: Uint8Array; descriptor: WorkspaceArchiveDescriptor };
+  archive:
+    | { bytes: Uint8Array; descriptor: WorkspaceArchiveDescriptor }
+    | VerifiedHostWorkspaceArchive;
 }): Promise<WorkspaceArchiveObjectRef> {
-  if (input.archive.descriptor.version !== 1) {
+  // Snapshot caller-owned metadata before any provider callback can mutate it.
+  const descriptor = structuredClone(input.archive.descriptor);
+  if (descriptor.version !== 1) {
     throw new Error("Object-storage workspace archives are portable tar only");
   }
   const key = workspaceArchiveObjectKey({
     accountId: input.accountId,
     workspaceId: input.workspaceId,
     sandboxGroupId: input.sandboxGroupId,
-    revision: input.archive.descriptor.revision,
-  });
-  await input.objectStorage.putObject({
-    key,
-    contentType: "application/x-tar",
-    body: input.archive.bytes,
-    sha256: input.archive.descriptor.archiveSha256,
+    revision: descriptor.revision,
+    uploadId: randomUUID(),
   });
+  const archive = input.archive;
+  if ("spool" in archive) {
+    if (
+      archive.spool.byteSize !== descriptor.archiveBytes ||
+      archive.spool.sha256 !== descriptor.archiveSha256
+    ) {
+      throw new Error("Workspace archive spool does not match its verified descriptor");
+    }
+  }
+  const source =
+    "spool" in archive
+      ? archive.spool
+      : {
+          path: "",
+          byteSize: archive.bytes.byteLength,
+          sha256: descriptor.archiveSha256,
+          async *open() {
+            for (let offset = 0; offset < archive.bytes.length; offset += 1024 * 1024)
+              yield archive.bytes.subarray(offset, offset + 1024 * 1024);
+          },
+          async dispose() {},
+        };
+  try {
+    await uploadWorkspaceArchiveSpool(input.objectStorage, key, {
+      ...source,
+      byteSize: descriptor.archiveBytes,
+      sha256: descriptor.archiveSha256,
+      open: source.open.bind(source),
+    });
+  } catch (error) {
+    // This fresh locator has never been offered to database publication. An
+    // ambiguous provider upload can leave an orphan, never overwrite a retained
+    // archive or justify deleting another attempt's object.
+    await input.objectStorage.deleteObject(key).catch(() => undefined);
+    throw error;
+  }
   return {
     schema: "sandbox_archive_object_v1",
     key,
-    sha256: input.archive.descriptor.archiveSha256,
-    bytes: input.archive.descriptor.archiveBytes,
+    sha256: descriptor.archiveSha256,
+    bytes: descriptor.archiveBytes,
     backend: input.objectStorage.backend,
   };
 }
 
+/** The caller owns a fresh candidate returned by this module. A thrown or lost
+ * database response is UNKNOWN: it may have committed, so never delete on catch.
+ * Only the transaction's exact candidate disposition can authorize cleanup. */
+export async function persistWorkspaceArchiveCandidate<
+  T extends {
+    wrote: boolean;
+    candidateDisposition?: "adopted" | "already_referenced" | "unused";
+  },
+>(input: {
+  objectStorage?: ObjectStorage | null;
+  ref?: WorkspaceArchiveObjectRef;
+  persist: () => Promise<T>;
+}): Promise<T> {
+  const objectStorage = input.objectStorage;
+  const ref = input.ref ? { ...input.ref } : undefined;
+  const result = await input.persist();
+  if (objectStorage && ref && result.candidateDisposition === "unused") {
+    await deleteUnpublishedWorkspaceArchiveObject(objectStorage, ref);
+  }
+  return result;
+}
+
 export async function putVersion1TarArchiveOrInline(input: {
   backend: string;
   objectStorage?: ObjectStorage | null | undefined;
   accountId: string;
   workspaceId: string;
   sandboxGroupId: string;
-  archive: {
-    bytes: Uint8Array;
-    descriptor: WorkspaceArchiveDescriptor;
-    base64: string;
-  };
+  archive:
+    | VerifiedHostWorkspaceArchive
+    | {
+        bytes: Uint8Array;
+        descriptor: WorkspaceArchiveDescriptor;
+        base64: string;
+      };
   metrics?: {
     onWorkspaceArchiveObject?: (input: {
       outcome: "put" | "put_failed" | "deleted_unpublished";
@@ -83,12 +144,15 @@ export async function putVersion1TarArchiveOrInline(input: {
   workspaceArchiveRef?: WorkspaceArchiveObjectRef;
 }> {
   if (input.archive.descriptor.version !== 1) {
+    if ("spool" in input.archive) throw new Error("Native snapshots cannot use a portable spool");
     return { workspaceArchive: input.archive.base64 };
   }
   if (input.backend === "opensandbox" && !input.objectStorage) {
     throw new WorkspaceArchiveObjectStorageRequiredError("opensandbox");
   }
   if (!input.objectStorage) {
+    if ("spool" in input.archive)
+      throw new WorkspaceArchiveObjectStorageRequiredError(input.backend);
     return { workspaceArchive: input.archive.base64 };
   }
   try {
@@ -97,7 +161,7 @@ export async function putVersion1TarArchiveOrInline(input: {
       accountId: input.accountId,
       workspaceId: input.workspaceId,
       sandboxGroupId: input.sandboxGroupId,
-      archive: { bytes: input.archive.bytes, descriptor: input.archive.descriptor },
+      archive: input.archive,
     });
     try {
       input.metrics?.onWorkspaceArchiveObject?.({

package/src/sandbox-resume.ts

@@ -55,7 +55,9 @@ import {
   type LeaseHolderKind,
 } from "@opengeni/db";
 import {
-  captureVerifiedWorkspaceArchive,
+  captureWorkspaceArchiveForStorage,
+  disposeWorkspaceArchive,
+  type VerifiedWorkspaceArchivePayload,
   describeLegacyNativeSnapshotArchive,
   inlineWorkspaceArchiveForRestore,
   MODAL_EXEC_READINESS_TIMEOUT_MS,
@@ -79,13 +81,15 @@ import {
   type RuntimeMetricsHooks,
   type WorkspaceArchiveDescriptor,
 } from "@opengeni/runtime";
-import type { ObjectStorage } from "@opengeni/storage";
+import {
+  downloadWorkspaceArchiveSpool,
+  WorkspaceArchiveStorageError,
+  type ObjectStorage,
+} from "@opengeni/storage";
 import type { Observability } from "@opengeni/observability";
 import { parseWorkspaceArchiveObjectRef } from "@opengeni/contracts";
 import {
-  collectWorkspaceArchiveObjectKeys,
-  deleteUnpublishedWorkspaceArchiveObject,
-  deleteWorkspaceArchiveObjectKeys,
+  persistWorkspaceArchiveCandidate,
   putVersion1TarArchiveOrInline,
 } from "./sandbox-archive-storage";
 
@@ -582,6 +586,14 @@ async function materializeSpawnEnvelopeArchive(
       "workspace archive object storage is not configured",
     );
   }
+  const descriptor = parseWorkspaceArchiveDescriptor(sessionState.workspaceArchiveMeta);
+  if (
+    process.platform === "linux" &&
+    descriptor?.version === 1 &&
+    descriptor.workspace.projection === "sdk_local_archive_v1"
+  ) {
+    return envelope;
+  }
   return {
     ...record,
     sessionState: await inlineWorkspaceArchiveForRestore(sessionState, (key) =>
@@ -898,9 +910,10 @@ async function persistWarmWorkspaceSnapshot(
     // first (the bounded-wait race Bugbot flagged).
     const capturedAtMs = Date.now();
     const registerCandidate = async (
-      archive: Awaited<ReturnType<typeof captureVerifiedWorkspaceArchive>>,
+      archive: VerifiedWorkspaceArchivePayload,
     ): Promise<{ id: string } | null> => {
       if (
+        archive.kind === "host_spool" ||
         archive.descriptor.version !== 2 ||
         (archive.descriptor.provider !== "modal_snapshot_filesystem" &&
           archive.descriptor.provider !== "modal_snapshot_directory")
@@ -938,82 +951,65 @@ async function persistWarmWorkspaceSnapshot(
     // turn signal resolves first. Its finally block is the only normal release
     // of the exact admission gate; a late callback cannot release a successor.
     const captureAndPublish = (async (): Promise<boolean> => {
+      let archive: VerifiedWorkspaceArchivePayload | undefined;
       let candidate: { id: string } | null = null;
-      let workspaceArchiveRef: Awaited<
-        ReturnType<typeof putVersion1TarArchiveOrInline>
-      >["workspaceArchiveRef"];
+      let publicationAttempted = false;
       try {
-        const archive = await captureVerifiedWorkspaceArchive(session, capturedAtMs, {
-          requestId: claimed.claim.providerRequestId,
-          strategy: capturePolicy.strategy,
-        });
-        candidate = await registerCandidate(archive);
-        const priorLease = await readLease(db, ids.workspaceId, ids.sandboxGroupId);
-        const priorKeys = collectWorkspaceArchiveObjectKeys(
-          (priorLease?.resumeState as Record<string, unknown> | null | undefined) ?? null,
+        archive = await captureWorkspaceArchiveForStorage(
+          session,
+          capturedAtMs,
+          {
+            requestId: claimed.claim.providerRequestId,
+            strategy: capturePolicy.strategy,
+          },
+          Boolean(services.objectStorage),
         );
+        candidate = await registerCandidate(archive);
+        const archiveDescriptor = archive.descriptor;
         const published = await putVersion1TarArchiveOrInline({
           backend: lease.backend,
           objectStorage: services.objectStorage,
           accountId: ids.accountId,
           workspaceId: ids.workspaceId,
           sandboxGroupId: ids.sandboxGroupId,
-          archive: {
-            bytes: archive.bytes,
-            descriptor: archive.descriptor,
-            base64: archive.base64,
-          },
+          archive,
           ...(services.sandboxMetrics ? { metrics: services.sandboxMetrics } : {}),
         });
-        workspaceArchiveRef = published.workspaceArchiveRef;
-        const { wrote } = await persistWarmSnapshot(db, {
-          accountId: ids.accountId,
-          workspaceId: ids.workspaceId,
-          sessionId: ids.sessionId,
-          turnId: ids.turnId,
-          attemptId: ids.attemptId,
-          sandboxGroupId: ids.sandboxGroupId,
-          expectedEpoch: leaseEpoch,
-          expectedInstanceId: instanceId,
-          expectedWorkspaceGeneration: claimed.claim.workspaceGeneration,
-          captureId,
-          workspaceArchiveMeta: archive.descriptor,
-          ...published,
-          checkpointArtifactId: candidate?.id ?? null,
-          minIntervalMs: force ? 0 : intervalMs,
-          capturedAtMs,
+        publicationAttempted = true;
+        const { wrote } = await persistWorkspaceArchiveCandidate({
+          ...(services.objectStorage ? { objectStorage: services.objectStorage } : {}),
+          ...(published.workspaceArchiveRef ? { ref: published.workspaceArchiveRef } : {}),
+          persist: () =>
+            persistWarmSnapshot(db, {
+              accountId: ids.accountId,
+              workspaceId: ids.workspaceId,
+              sessionId: ids.sessionId,
+              turnId: ids.turnId,
+              attemptId: ids.attemptId,
+              sandboxGroupId: ids.sandboxGroupId,
+              expectedEpoch: leaseEpoch,
+              expectedInstanceId: instanceId,
+              expectedWorkspaceGeneration: claimed.claim.workspaceGeneration,
+              captureId,
+              workspaceArchiveMeta: archiveDescriptor,
+              ...published,
+              checkpointArtifactId: candidate?.id ?? null,
+              minIntervalMs: force ? 0 : intervalMs,
+              capturedAtMs,
+            }),
         });
-        if (published.workspaceArchiveRef && services.objectStorage) {
-          if (!wrote) {
-            await deleteUnpublishedWorkspaceArchiveObject(
-              services.objectStorage,
-              published.workspaceArchiveRef,
-              services.sandboxMetrics,
-            );
-          } else {
-            const afterLease = await readLease(db, ids.workspaceId, ids.sandboxGroupId);
-            const afterKeys = collectWorkspaceArchiveObjectKeys(
-              (afterLease?.resumeState as Record<string, unknown> | null | undefined) ?? null,
-            );
-            await deleteWorkspaceArchiveObjectKeys(
-              services.objectStorage,
-              [...priorKeys].filter((key) => !afterKeys.has(key)),
-            ).catch(() => undefined);
-          }
-        }
         if (!wrote && candidate) {
           await abandonCandidate(candidate.id, "snapshot_publication_fenced");
         }
         return wrote;
       } catch (error) {
-        await deleteUnpublishedWorkspaceArchiveObject(
-          services.objectStorage,
-          workspaceArchiveRef,
-          services.sandboxMetrics,
-        );
-        if (candidate) await abandonCandidate(candidate.id, "snapshot_capture_failed");
+        if (candidate && !publicationAttempted)
+          await abandonCandidate(candidate.id, "snapshot_capture_failed");
         throw error;
       } finally {
+        await disposeWorkspaceArchive(archive).catch(() => {
+          console.warn("workspace archive spool cleanup failed");
+        });
         await releaseWorkspaceArchiveCapture(db, {
           accountId: ids.accountId,
           workspaceId: ids.workspaceId,
@@ -1504,6 +1500,26 @@ export async function resumeBoxForTurn(
       const established = await establishSandboxSessionFromEnvelope(settings, hydrateEnvelope, {
         sessionId: ids.sessionId,
         recovery: "create-or-restore",
+        ...(services.objectStorage
+          ? {
+              loadHostWorkspaceArchive: async (ref) => {
+                try {
+                  return await downloadWorkspaceArchiveSpool(services.objectStorage!, ref.key, {
+                    bytes: ref.bytes,
+                    sha256: ref.sha256,
+                  });
+                } catch (error) {
+                  if (error instanceof WorkspaceArchiveStorageError) {
+                    throw new WorkspaceArchiveIntegrityError(error.code, error.message, {
+                      retryable: error.retryable,
+                      cause: error,
+                    });
+                  }
+                  throw error;
+                }
+              },
+            }
+          : {}),
         backendOverride: ids.backend as never,
         ...(ids.environment ? { environment: ids.environment } : {}),
         ...(services.sandboxMetrics ? { metrics: services.sandboxMetrics } : {}),