UK source stages on the shared manifest: source_stages.json + StagePlan assembly

[juaristi22]

Workstream E1 of the UK migration (#145, master epic #665). Closes #677.

What

Migrates the UK country package off its bespoke source-stage machinery onto the shared, country-agnostic path — the structural prerequisite for every other WS-E ingest/imputation stage. Behavior does not change: this is a dispatch migration, verified by a payload-equivalence test and the full suite.

• uk/source_stages.json — a real SourceManifest, carrying the two runtime stages. The frozen hmrc_income_source_stages.json declared one stage holding all 10 operations; the copy splits at operation granularity into frs_hmrc_retained_leaves (2 ops, the 5 retained-leaf outputs) and hmrc_spi_income (8 ops, the original 14 outputs), order preserved. The frozen original stays byte-identical as the June evidence anchor, behind a sha-pinned lockstep test that allows differences only at three enumerated citation paths — rephrased because FORBIDDEN_SOURCE_DEPENDENCIES substring-matches the incumbent repo's name inside them (every DOI, UKDS study number, sha, and size is unchanged). It retires at the FRS 2024-25 refresh (E11).
• source_manifest.py — the 10 UK operation kinds plus a new generic materialize_rules_engine_predictors join ALLOWED_SOURCE_OPERATION_KINDS. The validator itself is untouched.
• uk_runtime/national_build.py — build_uk_national_dataset accepts shared plan.Stage entries or a StagePlan; execution goes through StagePlan.run, so every stage now yields a StageRecord (surfaced on UKNationalBuildResult, including through checkpoint resume). validate_uk_national_frame still runs after every stage. UKNationalStage remains as a one-release deprecated shim, converted internally.
• tools/build_uk_national_dataset.py — assembles stages via country_stage_plan(load_country_spec("uk"), uk_stage_implementations(...)), the BE whole-stage-implementations pattern. Stage names are unchanged, so checkpoint identities remain valid.
• uk_runtime/source_runtime.py (new) — the implementations-map builder plus the per-operation dispatch seam later WS-E children extend.
• frame/rules.py — shared assert_rules_engine_country and materialize_rules_engine_predictors (refuses missing/mismatched country declarations, empty predictor lists, already-present columns, and shape drift). frame/adapters/policyengine_uk.py (new) — a thin deferred-import PolicyEngine-UK adapter (country = "uk"); the US adapter now declares country = "us"; microcosm-frame gains an opt-in [uk] extra. This is the ratified WS-E rules-engine seam (Port UK compact source stages and 2024-25 refresh from uk-data #145 U1): stages that need formula-computed predictors declare them in the manifest and materialize through the adapter, with the engine's country required to match the dataset's.

UK-specificity note

Per the country-agnostic doctrine, the UK-specific surface of this PR is: the stage declarations in uk/source_stages.json (HMRC/SPI/FRS/UKDS identities), the whole-stage wiring in uk_runtime/source_runtime.py, the UK CLI driver boundary, and the policyengine_uk adapter + [uk] extra. Everything else — the allowlist entries, the materialize_rules_engine_predictors operation, the country assertion, the StagePlan execution path — is shared and usable by any country.

Implementation choices (no forced deviations from the reviewed plan)

• Checkpoint/resume executes one-stage StagePlans per boundary after full-plan assembly validation, so stage names survive resume.
• The shared country check and predictor materializer live in frame/rules.py (no existing US helper to extend).
• Adapter tests use a stub engine so ordinary test runs stay import-light; the protocol test proves construction defers the policyengine-uk import.
• frs_hmrc_retained_leaves stage metadata derives from the runtime's FRS_HMRC_RETAINED_LEAF_SOURCE_EVIDENCE (DonorSpec refuses an unnamed donor).

Verification

• Full suite locally: 6,279 tests, zero failures (the [uk] extra installed — note CI's us-only lane skips requires_uk tests, the known WS-A gap, so the local run is the UK evidence).
• New tests: manifest/lockstep/plan-assembly (test_uk_source_stages.py), runtime seam (test_uk_source_runtime.py), adapter (test_policyengine_uk_adapter.py), and a shim-vs-StagePlan payload-equivalence test on the national-build fixture.
• ruff check clean; changed files format-clean (26 pre-existing repo-wide format --check failures are outside this surface and untouched).
• Frozen evidence files byte-identical (git diff clean on all uk/*.json except the two declared edits: country_package.json resource line, new source_stages.json).
• The tree-scan incumbent-reference test passes; the new manifest contains zero forbidden substrings.

The licensed-data acceptance from #677 is complete (see the acceptance comment below): full-rung control (origin/main) and candidate builds from identical licensed inputs are byte-identical at both stage-boundary checkpoints, with all 13 terminal-gate statuses and detail payloads identical. Both sides block at uk_weight_ratio (#630 Finding 2, reproducing bit-for-bit on the new carrier) — the pre-existing release blocker, independent of this PR. The run surfaced and fixed one manifest defect only reachable with licensed data: the frozen state_pension output is a SPI concept, not a persisted column (corrected to hmrc_spi_state_pension_income, declared in the lockstep test).

🤖 Generated with Claude Code

[juaristi22]

Codex adversarial review (branch diff vs origin/main): verdict needs-attention — one high finding, fixed in f30e643.

• [high] Country-assertion bypass in the UK predictor handler (uk_runtime/source_runtime.py): the handler read country as an optional runtime extra; when absent, the shared materialize_rules_engine_predictors skipped assert_rules_engine_country, so the UK-specific handler could silently materialize from a wrong-country engine. Fix: the handler now fails closed — the dataset country defaults to "uk" (this is the UK handler map, by construction) and a runtime context declaring any other country is refused before materialization. Regression tests cover both the wrong-country engine with no country extra and the non-UK context.

Review focus covered: lockstep fidelity to the frozen manifest, StagePlan/checkpoint semantics, the Logbook-preamble rebase reconciliation, the shared frame/rules.py helpers, and UKDS/incumbent-name discipline — no other findings survived.

[juaristi22]

Licensed-data acceptance complete (E1 exit criterion): full-rung control-vs-candidate builds are byte-identical at every stage boundary, with identical gate verdicts. All licensed inputs stayed local; this comment carries only hashes, statuses, and already-public observables.

Setup: two full-rung staging builds from identical inputs (certified populace_uk_2023.h5, FRS 2023-24 tabs, SPI 2022-23 donor, HMRC collated ODS; seed 42, 100 estimators) — control at origin/main (f362f63, post-Logbook, pre-E1) and candidate at this branch. Both with --checkpoint-dir (checkpoints persist before the terminal battery, track_times=False makes equivalent frames byte-identical).

Results:

• Stage-boundary frames byte-identical on both carriers: 000_frs_hmrc_retained_leaves.frame.h5 sha256 6eb79b04…ceb3c52 and 001_hmrc_spi_income.frame.h5 sha256 de977dda…8ecbd87 (1.35 GB / 1.53 GB) — same hashes from the old UKNationalStage path and the new country_stage_plan/StagePlan path.
• Terminal battery: 13/13 gate statuses identical, all gate detail payloads identical, evidence_sha256 identical. 7 passed; 5 evidence_absent (target/export/weighted-integrity trio — absent by staging design); 1 failed on both sides: uk_weight_ratio 1590.5346779161957 vs reviewed max 1151.2542195939373 — National staging rebuild fails the current terminal gates: constant source_year release column and SPI weight tail above the reviewed maximum #630 Finding 2, reproducing bit-for-bit on the new carrier too. Neither side stages an H5 (write-then-block), so National staging rebuild fails the current terminal gates: constant source_year release column and SPI weight tail above the reviewed maximum #630 remains the release blocker independent of this PR, exactly as tracked.

One defect found and fixed by this run (c4cd761): the frozen manifest's declared output state_pension is a SPI concept, not a persisted frame column — the stage writes hmrc_spi_state_pension_income (SPI_HMRC_STATE_PENSION_INCOME_COLUMN), and the model-input state_pension is formula-owned. This was harmless while outputs were purely declarative, but E1 made them load-bearing (country_stage_plan compiles them into StagePlan.produces), so the full-rung build refused to stage — a class of defect only reachable with the licensed data, which is exactly why this acceptance run exists. The canonical copy now declares the persisted column; operation payloads keep the concept name; the lockstep test records the correction as a declared exception. A smoke-rung run plus the full-rung build above confirm the corrected produces set empirically.

Both failed attempts and all four builds landed hash-chained Logbook rows (#666 wiring exercised end to end, including failure dispositions).

[vahid-ahmadi]

Byte-identical checkpoints from identical licensed inputs on both sides, with all 13 terminal-gate statuses and detail payloads matching, is about as good as a dispatch-migration acceptance gets. Both sides blocking at uk_weight_ratio and reproducing #630 Finding 2 bit-for-bit on the new carrier is a nice bonus signal — it shows the migration preserved even the failure.

I went in expecting the lockstep test to be the soft spot, since "allows differences only at three enumerated paths" is usually where drift hides. It isn't:

assert stage1["operations"] + stage2["operations"] == expected_operations

Full structural equality of the concatenated operations against the frozen original, plus artifacts, grain and outputs compared the same way, with each deviation constructed explicitly rather than skipped by a wildcard — and test_frozen_original_bytes_are_pinned making the baseline itself immutable. That's the right shape and it earns the split.

The thing I'd act on before the WS-E children land: none of this PR's UK verification runs in CI.

.github/workflows/test.yml does:

run: uv sync --all-packages --extra us

This PR introduces the [uk] extra and the policyengine_uk adapter, so test_policyengine_uk_adapter.py, test_uk_source_runtime.py, test_uk_source_stages.py and the shim-vs-StagePlan payload-equivalence test are all requires_uk — skipped on every CI run. You say as much ("the local run is the UK evidence"), and for this PR the licensed acceptance more than compensates.

The problem is forward-looking. This is the structural prerequisite for eleven more WS-E children (#678–#688), each of which extends the per-operation dispatch seam introduced here. Every one of them will land with its UK surface verified only by whatever local run its author remembered to do, and no regression in this seam will be caught by CI on any subsequent PR — including PRs that aren't about the UK at all.

Epic #665 lists "UK CI lane (to be filed)" under workstream A. I checked, and it's still unfiled — there's no open issue for it. This PR is the one that makes it acute rather than theoretical, so it's worth filing now with this PR as the motivating case, and worth deciding whether it blocks the WS-E children or merely races them.

One observation worth carrying into the workstream

The state_pension → hmrc_spi_state_pension_income correction is more interesting than a typo fix. Your own note says outputs "became load-bearing when country_stage_plan compiled them into StagePlan.produces" — so the frozen manifest carried a wrong output name harmlessly for as long as nothing read it, and the migration is what made it detectable. That implies the other frozen manifests may carry the same class of latent error, dormant until they too become load-bearing.

If that's cheap to check ahead of time — validating declared outputs against persisted columns for the manifests not yet migrated — it would turn a series of licensed-run surprises into one upfront sweep. Given each discovery currently costs a credentialed build to find, that seems worth an hour.

Smaller

• UKNationalStage kept as a one-release deprecated shim with the conversion internal is the right call, and having the payload-equivalence test compare shim against StagePlan rather than asserting the shim away is what makes the deprecation safe.
• Executing one-stage StagePlans per checkpoint boundary after full-plan assembly validation is the detail that keeps stage names — and therefore checkpoint identities — valid across resume. Easy to have gotten wrong in the other order.
• Requiring the adapter's declared country to match the dataset's, and refusing already-present columns in materialize_rules_engine_predictors, are both the fail-closed choices. Good that the protocol test proves the policyengine-uk import stays deferred, so ordinary runs don't pay for it.