@policyengine/calibration-tree 0.1.0 → 0.1.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@policyengine/calibration-tree",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "Builds, validates, reads and publishes Microcosm calibration tree artifacts (tree schema 6), with the release and staging readers they are built from.",
   "license": "MIT",
   "repository": {
@@ -30,6 +30,7 @@
     "./calibration-tree-blob": "./src/calibration-tree-blob.ts",
     "./calibration-tree-bundle": "./src/calibration-tree-bundle.ts",
     "./calibration-tree-manifest": "./src/calibration-tree-manifest.ts",
+    "./calibration-tree-publication": "./src/calibration-tree-publication.ts",
     "./calibration-tree-worker": "./src/calibration-tree-worker.ts",
     "./chronicle-metadata": "./src/chronicle-metadata.ts",
     "./countries": "./src/countries.ts",
@@ -60,7 +61,8 @@
     "publish:calibration-tree": "bun run src/publish-calibration-tree.ts",
     "test": "bun test",
     "typecheck": "tsc --noEmit",
-    "worker": "bun run src/calibration-tree-worker.ts"
+    "worker": "bun run src/calibration-tree-worker.ts",
+    "publication": "bun run src/calibration-tree-publication.ts"
   },
   "dependencies": {
     "@vercel/blob": "^2.6.1"

package/src/calibration-tree-publication.ts

@@ -0,0 +1,399 @@
+// Publishes a tree the calibration tree worker wrote to disk: uploads each of
+// its gzip parts to the part's immutable path, then adds the build's entry to
+// the manifest with a compare-and-swap. The collector runs it as its own Bun
+// process, after the worker and under its one manifest lock:
+//
+//   bun run src/calibration-tree-publication.ts PUBLICATION.json
+//
+// with the store's read-write token in CALIBRATION_TREE_BLOB_READ_WRITE_TOKEN
+// and PUBLICATION.json holding:
+//
+//   {
+//     "bundleDirectory": "<the worker's output directory>",
+//     "worker": <the worker's printed result>,
+//     "makeLatest": false
+//   }
+//
+// Each part is read from disk and checked against the worker's record of its
+// raw bytes before anything is uploaded. A part already stored with the same
+// bytes is kept; one stored with other bytes stops the publication. The index
+// goes last. The manifest is created only if it is still absent, and an
+// existing one is replaced only if its ETag is still the one read, so a
+// manifest another writer changed is never overwritten. A build asked to
+// become the country's latest does so only when it is at least as new
+// (createdAt) as the current latest, so a retried or reprocessed older build
+// never takes the latest back. On success it prints one JSON line and exits
+// 0. The other exit codes:
+//
+//   1  the request could not be followed: no token, a malformed request, or a
+//      part on disk that differs from the worker's record;
+//   2  the build is refused: another build is already recorded for the same
+//      release or staged run, or the same build with other content (another
+//      submission of the same files, whose index records another time or
+//      label), or a staging build was to become the latest release; running
+//      again on the same tree cannot help;
+//   3  anything else, including a manifest another writer changed meanwhile
+//      and any error from the store; running again may succeed.
+import { createHash } from "node:crypto";
+import { readFileSync } from "node:fs";
+import { join } from "node:path";
+import { gunzipSync } from "node:zlib";
+import { get, put } from "@vercel/blob";
+
+import { calibrationTreePartPath, type CalibrationTreePart } from "./calibration-tree-artifact";
+import {
+  calibrationTreeBlobText,
+  getCalibrationTreeBlob,
+  readCalibrationTreeManifest,
+  type CalibrationTreeBlobClient,
+} from "./calibration-tree-blob";
+import {
+  CALIBRATION_TREE_MANIFEST_PATH,
+  CALIBRATION_TREE_MANIFEST_SCHEMA_VERSION,
+  parseCalibrationTreeManifest,
+  serializeCalibrationTreeManifest,
+  withCalibrationTreeManifestEntry,
+  type CalibrationTreeManifest,
+  type CalibrationTreeManifestEntry,
+} from "./calibration-tree-manifest";
+import { isCountry, type MicrocosmCountry } from "./countries";
+
+export const CALIBRATION_TREE_PUBLICATION_EXIT_REQUEST_ERROR = 1;
+export const CALIBRATION_TREE_PUBLICATION_EXIT_ENTRY_REFUSED = 2;
+export const CALIBRATION_TREE_PUBLICATION_EXIT_PUBLICATION_ERROR = 3;
+export const CALIBRATION_TREE_PUBLICATION_TOKEN_VARIABLE = "CALIBRATION_TREE_BLOB_READ_WRITE_TOKEN";
+
+// As the package's own uploads write them.
+const IMMUTABLE_CACHE_SECONDS = 31_536_000;
+const MANIFEST_CACHE_SECONDS = 60;
+const SHA256_RE = /^[0-9a-f]{64}$/;
+
+/** The request could not be followed; nothing was published. */
+export class CalibrationTreePublicationRequestError extends Error {}
+
+/** The build is refused; nothing in the manifest changed. */
+export class CalibrationTreeEntryRefusedError extends Error {}
+
+/** Another writer changed the manifest after it was read; nothing was overwritten. */
+export class CalibrationTreeManifestConflictError extends Error {}
+
+export interface CalibrationTreePublicationPart {
+  part: CalibrationTreePart;
+  path: string;
+  sha256: string;
+  rawBytes: number;
+}
+
+export interface CalibrationTreePublicationRequest {
+  bundleDirectory: string;
+  country: MicrocosmCountry;
+  buildArtifactId: string;
+  files: CalibrationTreePublicationPart[];
+  manifestEntry: CalibrationTreeManifestEntry;
+  makeLatest: boolean;
+}
+
+export interface CalibrationTreePublicationResult {
+  buildArtifactId: string;
+  kind: CalibrationTreeManifestEntry["kind"];
+  uploadedParts: number;
+  keptParts: number;
+  madeLatest: boolean;
+}
+
+type JsonObject = Record<string, unknown>;
+
+function requestObject(value: unknown, label: string): JsonObject {
+  if (!value || typeof value !== "object" || Array.isArray(value)) {
+    throw new CalibrationTreePublicationRequestError(`${label} must be a JSON object.`);
+  }
+  return value as JsonObject;
+}
+
+function requestString(value: JsonObject, field: string, label: string): string {
+  const text = value[field];
+  if (typeof text !== "string" || text.length === 0) {
+    throw new CalibrationTreePublicationRequestError(`${label} must be a non-empty string.`);
+  }
+  return text;
+}
+
+function publicationPart(value: unknown, country: MicrocosmCountry, buildArtifactId: string): CalibrationTreePublicationPart {
+  const file = requestObject(value, "worker.files[]");
+  const part = requestString(file, "part", "worker.files[].part") as CalibrationTreePart;
+  const path = requestString(file, "path", "worker.files[].path");
+  const sha256 = requestString(file, "sha256", "worker.files[].sha256");
+  const rawBytes = file.rawBytes;
+  let expectedPath: string;
+  try {
+    expectedPath = calibrationTreePartPath(country, buildArtifactId, part);
+  } catch {
+    throw new CalibrationTreePublicationRequestError(`worker.files[] names an unknown part ${part}.`);
+  }
+  // A part is only ever written to its own immutable path.
+  if (path !== expectedPath) {
+    throw new CalibrationTreePublicationRequestError(`Part ${part} must be stored at ${expectedPath}.`);
+  }
+  if (!SHA256_RE.test(sha256) || !Number.isSafeInteger(rawBytes) || (rawBytes as number) < 0) {
+    throw new CalibrationTreePublicationRequestError(`Part ${part} needs a SHA-256 and a raw size.`);
+  }
+  return { part, path, sha256, rawBytes: rawBytes as number };
+}
+
+/** Checks a publication request read from JSON. */
+export function parseCalibrationTreePublicationRequest(value: unknown): CalibrationTreePublicationRequest {
+  const request = requestObject(value, "The request");
+  const unknown = Object.keys(request).filter((field) => !["bundleDirectory", "worker", "makeLatest"].includes(field));
+  if (unknown.length > 0) {
+    throw new CalibrationTreePublicationRequestError(`The request has unknown fields: ${unknown.join(", ")}.`);
+  }
+  if (typeof request.makeLatest !== "boolean") {
+    throw new CalibrationTreePublicationRequestError("makeLatest must be true or false.");
+  }
+  const bundleDirectory = requestString(request, "bundleDirectory", "bundleDirectory");
+  const worker = requestObject(request.worker, "worker");
+  const country = requestString(worker, "country", "worker.country");
+  if (!isCountry(country)) {
+    throw new CalibrationTreePublicationRequestError(`worker.country ${country} is not a registered country.`);
+  }
+  const buildArtifactId = requestString(worker, "buildArtifactId", "worker.buildArtifactId");
+  if (!SHA256_RE.test(buildArtifactId)) {
+    throw new CalibrationTreePublicationRequestError("worker.buildArtifactId must be a build ID.");
+  }
+  if (!Array.isArray(worker.files) || worker.files.length === 0) {
+    throw new CalibrationTreePublicationRequestError("worker.files must list the tree's parts.");
+  }
+  const files = worker.files.map((file) => publicationPart(file, country, buildArtifactId));
+  if (files.filter((file) => file.part === "index").length !== 1) {
+    throw new CalibrationTreePublicationRequestError("worker.files must hold exactly one index.");
+  }
+  const manifestEntry = requestObject(worker.manifestEntry, "worker.manifestEntry") as unknown as CalibrationTreeManifestEntry;
+  // The manifest's own rules, checked before anything is uploaded: a
+  // malformed entry is a defect in the request, not a refusal of the build.
+  try {
+    parseCalibrationTreeManifest({
+      schemaVersion: CALIBRATION_TREE_MANIFEST_SCHEMA_VERSION,
+      countries: { [country]: { latestReleaseBuildArtifactId: null, builds: [manifestEntry] } },
+    });
+  } catch (error) {
+    throw new CalibrationTreePublicationRequestError(
+      `worker.manifestEntry is not a manifest entry: ${error instanceof Error ? error.message : String(error)}`,
+    );
+  }
+  const index = files.find((file) => file.part === "index")!;
+  if (
+    manifestEntry.buildArtifactId !== buildArtifactId ||
+    manifestEntry.indexSha256 !== index.sha256 ||
+    manifestEntry.indexBytes !== index.rawBytes
+  ) {
+    throw new CalibrationTreePublicationRequestError("worker.manifestEntry does not describe the tree's index.");
+  }
+  return { bundleDirectory, country, buildArtifactId, files, manifestEntry, makeLatest: request.makeLatest };
+}
+
+interface PartOnDisk extends CalibrationTreePublicationPart {
+  compressed: Uint8Array;
+  serialized: string;
+}
+
+// Every part is read and checked before anything is uploaded, so a damaged
+// output directory publishes nothing.
+function partsOnDisk(request: CalibrationTreePublicationRequest): PartOnDisk[] {
+  return request.files.map((file) => {
+    let compressed: Uint8Array;
+    let raw: Uint8Array;
+    try {
+      compressed = new Uint8Array(readFileSync(join(request.bundleDirectory, file.path)));
+      raw = new Uint8Array(gunzipSync(compressed));
+    } catch (error) {
+      throw new CalibrationTreePublicationRequestError(
+        `Part ${file.part} could not be read from ${request.bundleDirectory}: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+    if (raw.byteLength !== file.rawBytes || createHash("sha256").update(raw).digest("hex") !== file.sha256) {
+      throw new CalibrationTreePublicationRequestError(`Part ${file.part} on disk differs from the worker's record.`);
+    }
+    return { ...file, compressed, serialized: new TextDecoder("utf-8").decode(raw) };
+  });
+}
+
+// As the package's uploadCalibrationTreeFile: an existing part must hold the
+// same raw JSON, and an uploaded one is read back.
+async function uploadPart(
+  part: PartOnDisk,
+  request: CalibrationTreePublicationRequest,
+  token: string,
+  client: CalibrationTreeBlobClient,
+): Promise<boolean> {
+  const stored = () =>
+    getCalibrationTreeBlob({
+      country: request.country,
+      buildArtifactId: request.buildArtifactId,
+      part: part.part,
+      token,
+      consistent: true,
+      client,
+    });
+  const existing = await stored();
+  if (existing) {
+    if (await calibrationTreeBlobText(existing.stream) !== part.serialized) {
+      // The build ID covers the files, not the time or label the index
+      // records, so another submission of the same files lands here, as does
+      // a rebuild whose output changed.
+      throw new CalibrationTreeEntryRefusedError(
+        `Immutable calibration tree part ${part.path} already exists with different content: ` +
+          "a tree for these files is already stored with another index.",
+      );
+    }
+    return false;
+  }
+  await client.put(part.path, Buffer.from(part.compressed), {
+    access: "private",
+    token,
+    addRandomSuffix: false,
+    multipart: true,
+    contentType: "application/gzip",
+    cacheControlMaxAge: IMMUTABLE_CACHE_SECONDS,
+  });
+  const uploaded = await stored();
+  if (!uploaded || await calibrationTreeBlobText(uploaded.stream) !== part.serialized) {
+    throw new Error(`Uploaded calibration tree part ${part.path} failed verification.`);
+  }
+  return true;
+}
+
+// When a build was published: createdAt, else updatedAt; an undated build is
+// older than any dated one.
+function publishedInstant(entry: CalibrationTreeManifestEntry): number {
+  const instant = Date.parse(entry.createdAt ?? entry.updatedAt ?? "");
+  return Number.isFinite(instant) ? instant : Number.NEGATIVE_INFINITY;
+}
+
+// Whether the build is at least as new as the country's current latest.
+function isAtLeastAsNewAsLatest(
+  manifest: CalibrationTreeManifest,
+  country: MicrocosmCountry,
+  entry: CalibrationTreeManifestEntry,
+): boolean {
+  const current = manifest.countries[country];
+  const latest = current?.builds.find((build) => build.buildArtifactId === current.latestReleaseBuildArtifactId);
+  return !latest || latest.buildArtifactId === entry.buildArtifactId ||
+    publishedInstant(entry) >= publishedInstant(latest);
+}
+
+async function addManifestEntry(
+  request: CalibrationTreePublicationRequest,
+  token: string,
+  client: CalibrationTreeBlobClient,
+): Promise<boolean> {
+  const current = await readCalibrationTreeManifest({ token, consistent: true, client });
+  // An existing manifest is replaced only on its ETag; without one, the write
+  // would carry no condition.
+  if (current.etag !== null && !current.etag) {
+    throw new CalibrationTreeManifestConflictError(
+      "The stored calibration tree manifest was read without an ETag, so it cannot be replaced safely.",
+    );
+  }
+  const makeLatest = request.makeLatest &&
+    (request.manifestEntry.kind !== "release" ||
+      isAtLeastAsNewAsLatest(current.manifest, request.country, request.manifestEntry));
+  let updated;
+  try {
+    updated = withCalibrationTreeManifestEntry(current.manifest, request.country, request.manifestEntry, makeLatest);
+  } catch (error) {
+    throw new CalibrationTreeEntryRefusedError(error instanceof Error ? error.message : String(error));
+  }
+  try {
+    await client.put(CALIBRATION_TREE_MANIFEST_PATH, serializeCalibrationTreeManifest(updated), {
+      access: "private",
+      token,
+      addRandomSuffix: false,
+      // Create only while absent; otherwise replace only the version read.
+      allowOverwrite: current.etag != null,
+      ifMatch: current.etag ?? undefined,
+      contentType: "application/json; charset=utf-8",
+      cacheControlMaxAge: MANIFEST_CACHE_SECONDS,
+    });
+  } catch (error) {
+    throw new CalibrationTreeManifestConflictError(
+      `The calibration tree manifest was not written, because another writer changed it or the store refused the write: ${error instanceof Error ? error.message : String(error)}`,
+    );
+  }
+  const verified = await readCalibrationTreeManifest({ token, consistent: true, client });
+  const recorded = verified.manifest.countries[request.country];
+  const entry = recorded?.builds.find((build) => build.buildArtifactId === request.buildArtifactId);
+  if (
+    !entry ||
+    entry.indexSha256 !== request.manifestEntry.indexSha256 ||
+    (makeLatest && recorded?.latestReleaseBuildArtifactId !== request.buildArtifactId)
+  ) {
+    throw new Error("Calibration tree manifest verification failed.");
+  }
+  return makeLatest;
+}
+
+const DEFAULT_CLIENT: CalibrationTreeBlobClient = { get, put };
+
+/** Uploads a built tree and records it in the manifest. */
+export async function publishBuiltCalibrationTree(options: {
+  request: CalibrationTreePublicationRequest;
+  token: string;
+  client?: CalibrationTreeBlobClient;
+}): Promise<CalibrationTreePublicationResult> {
+  const { request, token, client = DEFAULT_CLIENT } = options;
+  const parts = partsOnDisk(request);
+  const ordered = [...parts.filter((part) => part.part !== "index"), ...parts.filter((part) => part.part === "index")];
+  let uploadedParts = 0;
+  for (const part of ordered) {
+    if (await uploadPart(part, request, token, client)) uploadedParts += 1;
+  }
+  const madeLatest = await addManifestEntry(request, token, client);
+  return {
+    buildArtifactId: request.buildArtifactId,
+    kind: request.manifestEntry.kind,
+    uploadedParts,
+    keptParts: parts.length - uploadedParts,
+    madeLatest,
+  };
+}
+
+/** The exit code for an error that stopped the publication. */
+export function calibrationTreePublicationExitCode(error: unknown): number {
+  if (error instanceof CalibrationTreePublicationRequestError) return CALIBRATION_TREE_PUBLICATION_EXIT_REQUEST_ERROR;
+  if (error instanceof CalibrationTreeEntryRefusedError) return CALIBRATION_TREE_PUBLICATION_EXIT_ENTRY_REFUSED;
+  return CALIBRATION_TREE_PUBLICATION_EXIT_PUBLICATION_ERROR;
+}
+
+/** The command line: returns the process exit code. */
+export async function calibrationTreePublicationMain(
+  argv: string[],
+  environment: Record<string, string | undefined> = process.env,
+): Promise<number> {
+  try {
+    if (argv.length !== 1) {
+      throw new CalibrationTreePublicationRequestError("Usage: bun run src/calibration-tree-publication.ts PUBLICATION.json");
+    }
+    const token = environment[CALIBRATION_TREE_PUBLICATION_TOKEN_VARIABLE]?.trim();
+    if (!token) {
+      throw new CalibrationTreePublicationRequestError(`${CALIBRATION_TREE_PUBLICATION_TOKEN_VARIABLE} is required.`);
+    }
+    let parsed: unknown;
+    try {
+      parsed = JSON.parse(readFileSync(argv[0], "utf8"));
+    } catch (error) {
+      throw new CalibrationTreePublicationRequestError(
+        `The request ${argv[0]} could not be read as JSON: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+    const result = await publishBuiltCalibrationTree({ request: parseCalibrationTreePublicationRequest(parsed), token });
+    console.log(JSON.stringify(result));
+    return 0;
+  } catch (error) {
+    console.error(error instanceof Error ? error.message : String(error));
+    return calibrationTreePublicationExitCode(error);
+  }
+}
+
+if (import.meta.main) {
+  process.exitCode = await calibrationTreePublicationMain(Bun.argv.slice(2));
+}

package/src/calibration-tree-worker.ts

@@ -1,7 +1,7 @@
-// The calibration tree worker: builds one release's tree bundle from local
-// copies of the files the publisher reads from Hugging Face, and writes each
-// part's gzip bytes under an output directory, at the path the publisher
-// uploads it to. The collector runs it as its own Bun process:
+// The calibration tree worker: builds one release's or staged run's tree
+// bundle from local copies of a build's files, and writes each part's gzip
+// bytes under an output directory, at the path the publisher uploads it to.
+// The collector runs it as its own Bun process:
 //
 //   bun run src/calibration-tree-worker.ts REQUEST.json OUTPUT_DIRECTORY
 //
@@ -24,12 +24,25 @@
 //     }
 //   }
 //
-// It builds the bundle exactly as the publisher's `publishRelease` does from
-// the same files, commit and publication time, so both write the same bytes:
-// the golden tests hold each to the recorded trees. Like the publisher, it
-// records the publication time as Date.prototype.toISOString() writes it. On
-// success it prints one JSON line describing the written parts and exits 0.
-// It writes nothing unless the whole bundle builds. The other exit codes:
+// or a staged run's, built from a staged bundle's submitted files:
+//
+//   {
+//     "kind": "staging",
+//     "country": "uk",
+//     "runId": "<the staged run>",
+//     "label": "<the candidate release, or the run id when it names none>",
+//     "hfRepo": ..., "hfCommitSha": ..., "createdAt": ..., "artifacts": {...}
+//   }
+//
+// A release bundle is built exactly as the publisher's `publishRelease` builds
+// it from the same files, commit and publication time, so both write the same
+// bytes: the golden tests hold each to the recorded trees. A staged run's is
+// built as the publisher's `publishStagingSource` builds a staging tree, but
+// from the files given here rather than the staging repo's run files. Like the
+// publisher, it records the publication time as Date.prototype.toISOString()
+// writes it. On success it prints one JSON line describing the written parts,
+// with the manifest entry the publisher would record for the build, and exits
+// 0. It writes nothing unless the whole bundle builds. The other exit codes:
 //
 //   1  the request could not be followed: a malformed request, a missing
 //      file, an output directory that is not empty;
@@ -49,10 +62,11 @@ import { mkdirSync, readdirSync, readFileSync, writeFileSync } from "node:fs";
 import { dirname, join } from "node:path";
 import { createHash } from "node:crypto";
 
-import type { CalibrationTreeSourceArtifact } from "./calibration-tree-artifact";
+import { CALIBRATION_TREE_SCHEMA_VERSION, type CalibrationTreeSourceArtifact } from "./calibration-tree-artifact";
 import { buildCalibrationTreeBundle, type CalibrationTreeBundle } from "./calibration-tree-bundle";
+import type { CalibrationTreeManifestEntry } from "./calibration-tree-manifest";
 import { isCountry, type MicrocosmCountry } from "./countries";
-import { buildCalibration, InvalidReleaseIdError } from "./latest-artifact";
+import { assertSafeReleaseId, buildCalibration, InvalidReleaseIdError } from "./latest-artifact";
 import { CalibrationTreeArtifactSizeError, enforceConfiguredSizeLimits } from "./publish-calibration-tree";
 import { isBuildFolderId, isReleaseDirectory } from "./release-record";
 import { IncompatibleStagingDataError } from "./staging-contract";
@@ -76,6 +90,16 @@ const REQUEST_FIELDS = [
   "createdAt",
   "artifacts",
 ] as const;
+const STAGING_REQUEST_FIELDS = [
+  "kind",
+  "country",
+  "runId",
+  "label",
+  "hfRepo",
+  "hfCommitSha",
+  "createdAt",
+  "artifacts",
+] as const;
 const HF_COMMIT_RE = /^[0-9a-f]{40,64}$/;
 // A date and time with seconds and an explicit offset, so the instant never
 // depends on the machine's time zone.
@@ -91,7 +115,14 @@ export interface CalibrationTreeWorkerArtifact {
   file: string;
 }
 
-export interface CalibrationTreeWorkerRequest {
+export interface CalibrationTreeWorkerArtifacts {
+  calibrationDiagnostics: CalibrationTreeWorkerArtifact;
+  buildManifest: CalibrationTreeWorkerArtifact | null;
+  releaseManifest: CalibrationTreeWorkerArtifact | null;
+  demographics: CalibrationTreeWorkerArtifact | null;
+}
+
+export interface CalibrationTreeWorkerReleaseRequest {
   kind: "release";
   country: MicrocosmCountry;
   releaseId: string;
@@ -99,14 +130,26 @@ export interface CalibrationTreeWorkerRequest {
   hfRepo: string;
   hfCommitSha: string;
   createdAt: string | null;
-  artifacts: {
-    calibrationDiagnostics: CalibrationTreeWorkerArtifact;
-    buildManifest: CalibrationTreeWorkerArtifact | null;
-    releaseManifest: CalibrationTreeWorkerArtifact | null;
-    demographics: CalibrationTreeWorkerArtifact | null;
-  };
+  artifacts: CalibrationTreeWorkerArtifacts;
+}
+
+/** A staged run, built from its staged bundle's files. */
+export interface CalibrationTreeWorkerStagingRequest {
+  kind: "staging";
+  country: MicrocosmCountry;
+  runId: string;
+  /** The candidate release the run was staged for, or the run id. */
+  label: string;
+  hfRepo: string;
+  hfCommitSha: string;
+  createdAt: string | null;
+  artifacts: CalibrationTreeWorkerArtifacts;
 }
 
+export type CalibrationTreeWorkerRequest =
+  | CalibrationTreeWorkerReleaseRequest
+  | CalibrationTreeWorkerStagingRequest;
+
 export interface CalibrationTreeWorkerPart {
   part: string;
   path: string;
@@ -117,6 +160,7 @@ export interface CalibrationTreeWorkerPart {
 
 export interface CalibrationTreeWorkerResult {
   country: MicrocosmCountry;
+  /** The release id for a release; the label for a staged run. */
   releaseId: string;
   hfRepo: string;
   hfCommitSha: string;
@@ -124,6 +168,8 @@ export interface CalibrationTreeWorkerResult {
   buildArtifactId: string;
   index: Omit<CalibrationTreeWorkerPart, "part">;
   files: CalibrationTreeWorkerPart[];
+  /** The entry the publisher would add to the manifest for this build. */
+  manifestEntry: CalibrationTreeManifestEntry;
 }
 
 /** The request could not be followed; nothing about the files is known. */
@@ -224,13 +270,23 @@ function workerArtifact(value: unknown, name: ArtifactName): CalibrationTreeWork
   };
 }
 
+function stagedRunId(value: JsonObject): string {
+  const runId = requestString(value, "runId");
+  try {
+    if (runId === "latest") throw new InvalidReleaseIdError("Invalid run id");
+    return assertSafeReleaseId(runId, "run");
+  } catch {
+    throw new CalibrationTreeWorkerRequestError("runId must be a run id: letters, digits, dots, underscores and hyphens.");
+  }
+}
+
 /** Checks a request read from JSON, refusing anything the worker would guess at. */
 export function parseCalibrationTreeWorkerRequest(value: unknown): CalibrationTreeWorkerRequest {
   const request = requestObject(value, "The request");
-  onlyFields(request, REQUEST_FIELDS, "The request");
-  if (request.kind !== "release") {
-    throw new CalibrationTreeWorkerRequestError('kind must be "release".');
+  if (request.kind !== "release" && request.kind !== "staging") {
+    throw new CalibrationTreeWorkerRequestError('kind must be "release" or "staging".');
   }
+  onlyFields(request, request.kind === "release" ? REQUEST_FIELDS : STAGING_REQUEST_FIELDS, "The request");
   const country = requestString(request, "country");
   if (!isCountry(country)) {
     throw new CalibrationTreeWorkerRequestError(`country ${country} is not a registered country.`);
@@ -244,13 +300,31 @@ export function parseCalibrationTreeWorkerRequest(value: unknown): CalibrationTr
     throw new CalibrationTreeWorkerRequestError("hfCommitSha must be a full lowercase commit SHA.");
   }
   const createdAt = publicationTime(request.createdAt);
-  const artifacts = requestObject(request.artifacts, "artifacts");
-  onlyFields(artifacts, ARTIFACT_NAMES, "artifacts");
+  const artifactFields = requestObject(request.artifacts, "artifacts");
+  onlyFields(artifactFields, ARTIFACT_NAMES, "artifacts");
   for (const name of ARTIFACT_NAMES) {
-    if (!(name in artifacts)) {
+    if (!(name in artifactFields)) {
       throw new CalibrationTreeWorkerRequestError(`artifacts.${name} is required (null when absent).`);
     }
   }
+  const artifacts = {
+    calibrationDiagnostics: workerArtifact(artifactFields.calibrationDiagnostics, "calibrationDiagnostics")!,
+    buildManifest: workerArtifact(artifactFields.buildManifest, "buildManifest"),
+    releaseManifest: workerArtifact(artifactFields.releaseManifest, "releaseManifest"),
+    demographics: workerArtifact(artifactFields.demographics, "demographics"),
+  };
+  if (request.kind === "staging") {
+    return {
+      kind: "staging",
+      country,
+      runId: stagedRunId(request),
+      label: requestString(request, "label"),
+      hfRepo,
+      hfCommitSha,
+      createdAt,
+      artifacts,
+    };
+  }
   return {
     kind: "release",
     country,
@@ -259,12 +333,7 @@ export function parseCalibrationTreeWorkerRequest(value: unknown): CalibrationTr
     hfRepo,
     hfCommitSha,
     createdAt,
-    artifacts: {
-      calibrationDiagnostics: workerArtifact(artifacts.calibrationDiagnostics, "calibrationDiagnostics")!,
-      buildManifest: workerArtifact(artifacts.buildManifest, "buildManifest"),
-      releaseManifest: workerArtifact(artifacts.releaseManifest, "releaseManifest"),
-      demographics: workerArtifact(artifacts.demographics, "demographics"),
-    },
+    artifacts,
   };
 }
 
@@ -302,61 +371,183 @@ function readArtifact(artifact: CalibrationTreeWorkerArtifact | null): ReadArtif
   };
 }
 
-/** Builds the release's bundle in memory, as the publisher's publishRelease does. */
+/**
+ * Builds the requested bundle in memory: a release as the publisher's
+ * publishRelease does, a staged run as its publishStagingSource does.
+ */
 export function buildCalibrationTreeWorkerBundle(request: CalibrationTreeWorkerRequest): CalibrationTreeBundle {
   const [diagnostics, buildManifest, releaseManifest, demographics] = ARTIFACT_NAMES.map((name) =>
     readArtifact(request.artifacts[name])
   );
   if (!diagnostics) throw new CalibrationTreeWorkerRequestError("Calibration diagnostics are missing.");
-  const { country, releaseId: id, releaseDirectory, hfCommitSha } = request;
   try {
-    const calibration = buildCalibration(
-      diagnostics.json,
-      id,
-      null,
-      buildManifest?.json ?? {},
-      releaseManifest?.json ?? {},
-      demographics?.json ?? {},
-      country,
-      "huggingface_immutable",
-      diagnostics.source.sha256,
-      hfCommitSha,
-    );
-    const bundle = buildCalibrationTreeBundle({
-      country,
-      createdAt: request.createdAt,
-      releaseId: id,
-      hfRepo: request.hfRepo,
-      hfCommitSha,
-      sourceArtifacts: {
-        calibrationDiagnostics: diagnostics.source,
-        buildManifest: buildManifest?.source ?? null,
-        releaseManifest: releaseManifest?.source ?? null,
-        demographics: demographics?.source ?? null,
-      },
-      rows: calibration.rows,
-      calibrationProvenance: calibration.calibration_provenance,
-      lossAttributionAvailable: calibration.target_loss_attribution.status !== "unavailable",
-      comparison: {
-        // As in publishRelease: a tag or a build folder compares under its
-        // own id, a release folder under the release its diagnostics name.
-        releaseId: releaseDirectory === id && !isBuildFolderId(id) ? calibration.release_id : id,
-        calibrationProvenance: calibration.calibration_provenance,
-        status: calibration.target_loss_attribution.status,
-        aggregate: calibration.target_loss_attribution.aggregate,
-        cap: calibration.target_loss_attribution.cap,
-        basisIdentifier: calibration.target_loss_attribution.basis_identifier,
-        targetRepresentation: calibration.target_schema.target_representation,
-      },
+    const sourceArtifacts = {
+      calibrationDiagnostics: diagnostics.source,
+      buildManifest: buildManifest?.source ?? null,
+      releaseManifest: releaseManifest?.source ?? null,
+      demographics: demographics?.source ?? null,
+    };
+    if (request.kind === "staging") {
+      return buildStagingBundle(request, sourceArtifacts, {
+        diagnostics: diagnostics.json,
+        buildManifest: buildManifest?.json ?? {},
+        releaseManifest: releaseManifest?.json ?? {},
+        demographics: demographics?.json ?? {},
+      });
+    }
+    return buildReleaseBundle(request, sourceArtifacts, {
+      diagnostics: diagnostics.json,
+      buildManifest: buildManifest?.json ?? {},
+      releaseManifest: releaseManifest?.json ?? {},
+      demographics: demographics?.json ?? {},
     });
-    for (const file of bundle.files) enforceConfiguredSizeLimits(file);
-    return bundle;
   } catch (error) {
     if (!isCalibrationTreeContentRefusal(error)) throw error;
     throw new CalibrationTreeUnbuildableInputError((error as Error).message);
   }
 }
 
+interface WorkerDocuments {
+  diagnostics: JsonObject;
+  buildManifest: JsonObject;
+  releaseManifest: JsonObject;
+  demographics: JsonObject;
+}
+
+type WorkerSourceArtifacts = Parameters<typeof buildCalibrationTreeBundle>[0]["sourceArtifacts"];
+
+// As publishStagingSource builds a new staged run's tree: a staging build of
+// the run, labelled and released under its candidate release, comparing
+// under the release its diagnostics name.
+function buildStagingBundle(
+  request: CalibrationTreeWorkerStagingRequest,
+  sourceArtifacts: WorkerSourceArtifacts,
+  documents: WorkerDocuments,
+): CalibrationTreeBundle {
+  const { country, runId, label, hfRepo, hfCommitSha, createdAt } = request;
+  const calibration = buildCalibration(
+    documents.diagnostics,
+    label,
+    createdAt,
+    documents.buildManifest,
+    documents.releaseManifest,
+    documents.demographics,
+    country,
+    "huggingface_staged_bundle",
+    sourceArtifacts.calibrationDiagnostics?.sha256 ?? null,
+    hfCommitSha,
+  );
+  const bundle = buildCalibrationTreeBundle({
+    country,
+    buildKind: "staging",
+    sourceId: runId,
+    label,
+    createdAt,
+    releaseId: label,
+    hfRepo,
+    hfCommitSha,
+    sourceArtifacts,
+    rows: calibration.rows,
+    calibrationProvenance: calibration.calibration_provenance,
+    lossAttributionAvailable: calibration.target_loss_attribution.status !== "unavailable",
+    comparison: {
+      releaseId: calibration.release_id,
+      calibrationProvenance: calibration.calibration_provenance,
+      status: calibration.target_loss_attribution.status,
+      aggregate: calibration.target_loss_attribution.aggregate,
+      cap: calibration.target_loss_attribution.cap,
+      basisIdentifier: calibration.target_loss_attribution.basis_identifier,
+      targetRepresentation: calibration.target_schema.target_representation,
+    },
+  });
+  for (const file of bundle.files) enforceConfiguredSizeLimits(file);
+  return bundle;
+}
+
+function buildReleaseBundle(
+  request: CalibrationTreeWorkerReleaseRequest,
+  sourceArtifacts: WorkerSourceArtifacts,
+  documents: WorkerDocuments,
+): CalibrationTreeBundle {
+  const { country, releaseId: id, releaseDirectory, hfCommitSha } = request;
+  const calibration = buildCalibration(
+    documents.diagnostics,
+    id,
+    null,
+    documents.buildManifest,
+    documents.releaseManifest,
+    documents.demographics,
+    country,
+    "huggingface_immutable",
+    sourceArtifacts.calibrationDiagnostics?.sha256 ?? null,
+    hfCommitSha,
+  );
+  const bundle = buildCalibrationTreeBundle({
+    country,
+    createdAt: request.createdAt,
+    releaseId: id,
+    hfRepo: request.hfRepo,
+    hfCommitSha,
+    sourceArtifacts,
+    rows: calibration.rows,
+    calibrationProvenance: calibration.calibration_provenance,
+    lossAttributionAvailable: calibration.target_loss_attribution.status !== "unavailable",
+    comparison: {
+      // As in publishRelease: a tag or a build folder compares under its
+      // own id, a release folder under the release its diagnostics name.
+      releaseId: releaseDirectory === id && !isBuildFolderId(id) ? calibration.release_id : id,
+      calibrationProvenance: calibration.calibration_provenance,
+      status: calibration.target_loss_attribution.status,
+      aggregate: calibration.target_loss_attribution.aggregate,
+      cap: calibration.target_loss_attribution.cap,
+      basisIdentifier: calibration.target_loss_attribution.basis_identifier,
+      targetRepresentation: calibration.target_schema.target_representation,
+    },
+  });
+  for (const file of bundle.files) enforceConfiguredSizeLimits(file);
+  return bundle;
+}
+
+// The entry the publisher adds to the manifest for the build: releaseManifestEntry
+// for a release, publishStagingSource's new entry for a staged run.
+function manifestEntry(
+  request: CalibrationTreeWorkerRequest,
+  bundle: CalibrationTreeBundle,
+  index: Omit<CalibrationTreeWorkerPart, "part">,
+): CalibrationTreeManifestEntry {
+  const build = bundle.index.build;
+  const shared = {
+    buildArtifactId: build.buildArtifactId,
+    hfRepo: request.hfRepo,
+    hfCommitSha: request.hfCommitSha,
+    treeSchemaVersion: CALIBRATION_TREE_SCHEMA_VERSION,
+    indexSha256: index.sha256,
+    indexBytes: index.rawBytes,
+    createdAt: build.createdAt,
+    updatedAt: build.createdAt,
+  };
+  if (request.kind === "staging") {
+    return {
+      ...shared,
+      kind: "staging",
+      sourceId: request.runId,
+      label: request.label,
+      releaseId: null,
+      stagingRunId: request.runId,
+    };
+  }
+  return {
+    ...shared,
+    kind: "release",
+    sourceId: build.sourceId,
+    label: build.label,
+    releaseId: request.releaseId,
+    // Recorded only for a tag read from its release's folder.
+    ...(request.releaseDirectory !== request.releaseId ? { releaseDirectory: request.releaseDirectory } : {}),
+    stagingRunId: null,
+  };
+}
+
 function emptyOutputDirectory(outputDirectory: string): void {
   try {
     mkdirSync(outputDirectory, { recursive: true });
@@ -394,13 +585,14 @@ export function runCalibrationTreeWorker(
   const { part: _index, ...index } = files.find((file) => file.part === "index")!;
   return {
     country: request.country,
-    releaseId: request.releaseId,
+    releaseId: request.kind === "staging" ? request.label : request.releaseId,
     hfRepo: request.hfRepo,
     hfCommitSha: request.hfCommitSha,
     createdAt: request.createdAt,
     buildArtifactId: bundle.index.buildArtifactId,
     index,
     files,
+    manifestEntry: manifestEntry(request, bundle, index),
   };
 }