Add country-specific staging repositories

README.md

@@ -118,6 +118,9 @@ Optional env: `POPULACE_HF_REPO`, `POPULACE_HF_REVISION` to point at a different
 US dataset/revision; `POPULACE_UK_HF_REPO`, `POPULACE_UK_HF_REVISION` for the UK;
 and `POPULACE_BE_HF_REPO`, `POPULACE_BE_HF_REVISION` for Belgium. The Belgium
 repository defaults in code to `policyengine/populace-be-private`. Set `HF_TOKEN`
-or `HUGGINGFACE_TOKEN` to read private datasets. US staging defaults to
-`policyengine/populace-us-staging`; override it with `POPULACE_STAGING_HF_REPO`
-and `POPULACE_STAGING_HF_REVISION`.
+or `HUGGINGFACE_TOKEN` to read private datasets. Each country with the
+`staging` capability declares its own staging repository in the country
+registry; it never falls back to another country's repository. US staging
+defaults to `policyengine/populace-us-staging`; override it with the
+backward-compatible `POPULACE_STAGING_HF_REPO` and
+`POPULACE_STAGING_HF_REVISION` variables.

frontend/lib/microcosm/countries.test.ts

@@ -35,6 +35,18 @@ test("every registration carries the full shape", () => {
       expect(registration.repo_env).toMatch(/^[A-Z][A-Z0-9_]*$/);
       expect(registration.revision_env).toMatch(/^[A-Z][A-Z0-9_]*$/);
     }
+    if (registration.staging) {
+      expect(registration.staging.repo).toMatch(/^[a-z0-9-]+\/[a-z0-9-]+$/);
+      expect(registration.staging.revision.length).toBeGreaterThan(0);
+      if (
+        registration.staging.repo_env != null ||
+        registration.staging.revision_env != null
+      ) {
+        expect(registration.staging.repo_env).toMatch(/^[A-Z][A-Z0-9_]*$/);
+        expect(registration.staging.revision_env).toMatch(/^[A-Z][A-Z0-9_]*$/);
+      }
+    }
+    expect(hasCapability(country, "staging")).toBe(registration.staging != null);
   }
 });
 
@@ -49,6 +61,11 @@ test("keeps the live registrations on their published repositories and labels",
     geography_id: "0100000US",
     visibility: "public",
     capabilities: COUNTRY_CAPABILITIES,
+    staging: {
+      repo: "policyengine/populace-us-staging",
+      repo_env: "POPULACE_STAGING_HF_REPO",
+      revision_env: "POPULACE_STAGING_HF_REVISION",
+    },
   });
   expect(countryRegistration("uk")).toMatchObject({
     repo: "policyengine/populace-uk-private",
@@ -71,12 +88,15 @@ test("keeps the live registrations on their published repositories and labels",
 test("selectable countries follow registry order and exclude fixtures", () => {
   expect(selectableCountries()).toEqual(["us", "uk", "be"]);
   expect(countryRegistration("zz").fixture).toBe(true);
+  expect(countryRegistration("am").fixture).toBe(true);
   expect(selectableCountries()).not.toContain("zz");
 });
 
 test("fixture registrations are valid countries without being selectable", () => {
   expect(isCountry("zz")).toBe(true);
+  expect(isCountry("am")).toBe(true);
   expect(parseCountry("zz")).toBe("zz");
+  expect(parseCountry("am")).toBe("am");
 });
 
 test("country parsing is exact and defaults to the registry default", () => {
@@ -96,6 +116,7 @@ test("country parsing is exact and defaults to the registry default", () => {
 
 test("capability gates read the registration", () => {
   expect(hasCapability("us", "staging")).toBe(true);
+  expect(hasCapability("am", "staging")).toBe(true);
   expect(hasCapability("us", "model_coverage")).toBe(true);
   expect(hasCapability("uk", "staging")).toBe(false);
   expect(hasCapability("be", "pipeline")).toBe(false);

frontend/lib/microcosm/countries.ts

@@ -38,6 +38,14 @@ export interface CountryRegistration {
   geography_id: string | null;
   visibility: RepositoryVisibility;
   capabilities: readonly CountryCapability[];
+  // Optional staging telemetry repository. A staging-capable country must
+  // declare this instead of inheriting another country's repository.
+  staging?: {
+    repo: string;
+    revision: string;
+    repo_env?: string;
+    revision_env?: string;
+  };
   // Other jurisdiction codes a cross-dataset bundle may use for this country.
   jurisdiction_aliases?: readonly string[];
   // Conformance-only registration: a valid country that is never listed in
@@ -68,6 +76,14 @@ export const COUNTRY_REGISTRY = {
     geography_id: "0100000US",
     visibility: "public",
     capabilities: ALL_CAPABILITIES,
+    staging: {
+      // Deprecated upstream identifier: the US staging publisher still uses
+      // the former Populace repository and deployment-variable names.
+      repo: "policyengine/populace-us-staging",
+      revision: "main",
+      repo_env: "POPULACE_STAGING_HF_REPO",
+      revision_env: "POPULACE_STAGING_HF_REVISION",
+    },
   },
   uk: {
     repo: "policyengine/populace-uk-private",
@@ -106,6 +122,24 @@ export const COUNTRY_REGISTRY = {
     capabilities: ["calibration", "targets", "compare"],
     fixture: true,
   },
+  // Synthetic non-US staging registration. It verifies that staging artifact
+  // resolution is country-scoped without exposing an unfinished country in
+  // selectors or treating its fixture repository as a published dataset.
+  am: {
+    repo: "policyengine/microcosm-am-fixture",
+    revision: "main",
+    label: "Armenia",
+    dataset_label: "Microcosm Armenia",
+    geography: "Armenia",
+    geography_id: null,
+    visibility: "private",
+    capabilities: ["calibration", "targets", "compare", "staging"],
+    staging: {
+      repo: "policyengine/microcosm-am-staging-fixture",
+      revision: "main",
+    },
+    fixture: true,
+  },
 } satisfies Record<string, CountryRegistration>;
 
 export type MicrocosmCountry = keyof typeof COUNTRY_REGISTRY;

frontend/lib/microcosm/staging-artifact.test.ts

@@ -5,10 +5,90 @@ import {
   loadStagingRun,
   loadStagingRuns,
   loadStagingTargetDiagnostics,
+  MICROCOSM_STAGING_HF_REPO,
+  MICROCOSM_STAGING_HF_REVISION,
+  stagingRepository,
+  stagingResolveUrl,
   stagingTargetChangeCacheTtlSeconds,
   stagingUnavailableReason,
 } from "./staging-artifact";
 
+test("resolves staging repositories from each country registration", () => {
+  expect(stagingRepository("us")).toEqual({
+    repo: "policyengine/populace-us-staging",
+    revision: "main",
+  });
+  expect(MICROCOSM_STAGING_HF_REPO).toBe("policyengine/populace-us-staging");
+  expect(MICROCOSM_STAGING_HF_REVISION).toBe("main");
+
+  // Armenia is a fixture-only registration: this asserts country-specific
+  // repository selection without publishing a country or telemetry artifact.
+  expect(stagingRepository("am")).toEqual({
+    repo: "policyengine/microcosm-am-staging-fixture",
+    revision: "main",
+  });
+  expect(stagingResolveUrl("runs/am-fixture/progress.json", "am")).toBe(
+    "https://huggingface.co/datasets/policyengine/microcosm-am-staging-fixture/resolve/main/runs/am-fixture/progress.json",
+  );
+
+  const originalRepo = process.env.POPULACE_STAGING_HF_REPO;
+  const originalRevision = process.env.POPULACE_STAGING_HF_REVISION;
+  try {
+    process.env.POPULACE_STAGING_HF_REPO = "policyengine/us-staging-override";
+    process.env.POPULACE_STAGING_HF_REVISION = "test-revision";
+    expect(stagingRepository("us")).toEqual({
+      repo: "policyengine/us-staging-override",
+      revision: "test-revision",
+    });
+    expect(stagingRepository("am")).toEqual({
+      repo: "policyengine/microcosm-am-staging-fixture",
+      revision: "main",
+    });
+  } finally {
+    if (originalRepo === undefined) delete process.env.POPULACE_STAGING_HF_REPO;
+    else process.env.POPULACE_STAGING_HF_REPO = originalRepo;
+    if (originalRevision === undefined) delete process.env.POPULACE_STAGING_HF_REVISION;
+    else process.env.POPULACE_STAGING_HF_REVISION = originalRevision;
+  }
+});
+
+test("loads Armenia staging telemetry from Armenia's registered repository", async () => {
+  const originalFetch = globalThis.fetch;
+  const urls: string[] = [];
+  globalThis.fetch = (async (input: RequestInfo | URL) => {
+    const url = String(input);
+    urls.push(url);
+    if (url.endsWith("/runs.json")) {
+      return Response.json({
+        runs: [
+          {
+            run_id: "am-fixture",
+            candidate_release_id: "am-candidate",
+            status: "running",
+          },
+        ],
+      });
+    }
+    if (url.includes("/tree/main/runs?recursive=true")) return Response.json([]);
+    return new Response(null, { status: 404 });
+  }) as typeof fetch;
+
+  try {
+    await expect(loadStagingRuns(0, "am")).resolves.toMatchObject({
+      available: true,
+      source_repo: "policyengine/microcosm-am-staging-fixture",
+      revision: "main",
+      runs: [{ run_id: "am-fixture", candidate_release_id: "am-candidate" }],
+    });
+    expect(urls).toEqual([
+      "https://huggingface.co/datasets/policyengine/microcosm-am-staging-fixture/resolve/main/runs.json",
+      "https://huggingface.co/api/datasets/policyengine/microcosm-am-staging-fixture/tree/main/runs?recursive=true",
+    ]);
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
 test("names the country when staging is unavailable", () => {
   expect(stagingUnavailableReason("us")).toBeNull();
   expect(stagingUnavailableReason("uk")).toBe(