From 251c4290926dc925fa23e3cc22cfd70ec76d461d Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:13:19 +0400
Subject: [PATCH 01/15] Add a typed operation scope next to the span scope

scopes.operation(runtime, name, attributes=..., remote_context=...)
returns runtime.operation(...) unchanged, typed like scopes.span as a
context manager that doesn't suppress exceptions, so Pyright doesn't
take a Modal function's operation scope to swallow the exception raised
in its body. It takes the two keyword arguments the executor's Modal
functions pass.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../scopes.py                                 | 25 +++++++++--
 .../tests/test_scopes.py                      | 44 ++++++++++++++++++-
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py
index c79bc552c..d1a31400b 100644
--- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py
+++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/scopes.py
@@ -1,4 +1,4 @@
-"""The observability runtime's span scope, typed as never suppressing an exception.
+"""The observability runtime's scopes, typed as never suppressing an exception.
 
 policyengine-observability annotates its scope managers' ``__exit__`` as
 returning ``bool``, although it always returns ``False``. A type checker reads a
@@ -7,9 +7,10 @@
 the block assigned as possibly unbound, and a function that returns from inside
 the block as missing a return.
 
-``span`` returns the runtime's own scope unchanged, typed as a context manager
-whose ``__exit__`` returns ``bool | None``: the type ``contextlib.contextmanager``
-scopes have, which type checkers don't treat as suppressing exceptions.
+``span`` and ``operation`` return the runtime's own scope unchanged, typed as a
+context manager whose ``__exit__`` returns ``bool | None``: the type
+``contextlib.contextmanager`` scopes have, which type checkers don't treat as
+suppressing exceptions.
 """
 
 from __future__ import annotations
@@ -35,3 +36,19 @@ def span(
     the block, and lets that exception propagate.
     """
     return runtime.span(name, attributes=attributes)
+
+
+def operation(
+    runtime: ObservabilityRuntime,
+    name: str,
+    *,
+    attributes: Mapping[str, Any] | None = None,
+    remote_context: Mapping[str, Any] | None = None,
+) -> AbstractContextManager[object]:
+    """Return ``runtime.operation(name, ...)`` with these attributes and context.
+
+    The scope times its block as an operation, continuing the trace in
+    ``remote_context`` when there is one, records an exception raised in the
+    block, and lets that exception propagate.
+    """
+    return runtime.operation(name, attributes=attributes, remote_context=remote_context)
diff --git a/libs/policyengine-simulation-observability/tests/test_scopes.py b/libs/policyengine-simulation-observability/tests/test_scopes.py
index c86cad316..7b524bc45 100644
--- a/libs/policyengine-simulation-observability/tests/test_scopes.py
+++ b/libs/policyengine-simulation-observability/tests/test_scopes.py
@@ -1,4 +1,4 @@
-"""Tests for the typed span scope of the observability runtime."""
+"""Tests for the typed span and operation scopes of the observability runtime."""
 
 from contextlib import nullcontext
 
@@ -8,15 +8,20 @@
 
 
 class RecordingRuntime:
-    """Records the spans it is asked to open, and enters each with a marker."""
+    """Records the scopes it is asked to open, and enters each with a marker."""
 
     def __init__(self) -> None:
         self.opened_spans: list[tuple[str, object]] = []
+        self.opened_operations: list[tuple[str, object, object]] = []
 
     def span(self, name, *, attributes=None):
         self.opened_spans.append((name, attributes))
         return nullcontext("entered span state")
 
+    def operation(self, name, *, attributes=None, remote_context=None):
+        self.opened_operations.append((name, attributes, remote_context))
+        return nullcontext("entered operation state")
+
 
 @pytest.fixture
 def runtime(monkeypatch):
@@ -61,3 +66,38 @@ def test_span_lets_an_exception_raised_inside_it_propagate(runtime):
         scopes.span(runtime, "simulation.parse"),
     ):
         raise ValueError(message)
+
+
+def test_operation_opens_the_runtime_operation_with_its_context():
+    recording_runtime = RecordingRuntime()
+
+    with scopes.operation(
+        recording_runtime,
+        "simulation.execution",
+        attributes={"country": "us"},
+        remote_context={"traceparent": "00-trace"},
+    ) as entered_state:
+        assert recording_runtime.opened_operations == [
+            ("simulation.execution", {"country": "us"}, {"traceparent": "00-trace"})
+        ]
+
+    assert entered_state == "entered operation state"
+
+
+def test_operation_defaults_to_no_attributes_or_remote_context():
+    recording_runtime = RecordingRuntime()
+
+    with scopes.operation(recording_runtime, "simulation.execution"):
+        pass
+
+    assert recording_runtime.opened_operations == [("simulation.execution", None, None)]
+
+
+def test_operation_lets_an_exception_raised_inside_it_propagate(runtime):
+    message = "simulation failed"
+
+    with (
+        pytest.raises(RuntimeError, match=message),
+        scopes.operation(runtime, "simulation.execution"),
+    ):
+        raise RuntimeError(message)

From e5907a63c8faf39d008d1198104982d6d2696a79 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:15:44 +0400
Subject: [PATCH 02/15] Open the executor's remaining spans and operations with
 typed scopes

The Modal functions' operation scopes (app.py, v2_app.py) now use
scopes.operation, and the remaining spans in the budget-window batch and
scheduler, the segmented national reduce, the Stage 12 fan-out and
comparison, and the output builder use scopes.span. Pyright no longer
reports the 13 variables assigned inside them as possibly unbound, or the
4 functions that return from inside one as missing a return.

The Stage 12 failure-cause scope already declared its type
(AbstractContextManager[object]) and is unchanged. The executor's
baseline drops from 72 entries to 55.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                       |  2 +-
 .../src/modal/app.py                          | 10 ++--
 .../src/modal/budget_window_batch.py          |  3 +-
 .../src/modal/budget_window_scheduler.py      | 46 +++++++++++--------
 .../src/modal/segmented_national.py           |  5 +-
 .../src/modal/v2_app.py                       | 13 ++++--
 .../simulation_output_builder.py              |  7 ++-
 .../stage12_runtime/comparison.py             |  3 +-
 .../stage12_runtime/fanout.py                 | 12 +++--
 rules/baselines/pyright-executor.txt          | 17 -------
 rules/registry.toml                           |  2 +-
 11 files changed, 65 insertions(+), 55 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 8e880c3c8..f60044a18 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 72 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 55 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py
index 3ee231c03..0e5fcd1f3 100644
--- a/projects/policyengine-simulation-executor/src/modal/app.py
+++ b/projects/policyengine-simulation-executor/src/modal/app.py
@@ -24,6 +24,7 @@
     get_bundled_country_model_version,
     get_bundled_package_version,
 )
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.observability import (
     init_process_observability,
     modal_image_environment,
@@ -366,7 +367,8 @@ def run_simulation(
     }
     try:
         propagated = normalize_observability_context(observability_context)
-        with runtime.operation(
+        with scopes.operation(
+            runtime,
             ANNUAL_IMPACT_STAGES.name(Stage.ANNUAL_EXECUTION),
             attributes=redacted_params,
             remote_context=propagated,
@@ -438,7 +440,8 @@ def run_simulation_segment(
     }
     try:
         propagated = normalize_observability_context(observability_context)
-        with runtime.operation(
+        with scopes.operation(
+            runtime,
             SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_EXECUTION),
             attributes=redacted_params,
             remote_context=propagated,
@@ -478,7 +481,8 @@ def run_budget_window_batch(
     }
     try:
         propagated = normalize_observability_context(observability_context)
-        with runtime.operation(
+        with scopes.operation(
+            runtime,
             BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_EXECUTION),
             attributes=redacted_params,
             remote_context=propagated,
diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py
index 90ec90208..d03a4c065 100644
--- a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py
+++ b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py
@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING, Any
 
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.stages import BUDGET_WINDOW_STAGES, Stage
 from src.modal.budget_window_context import build_batch_context
 from src.modal.budget_window_scheduler import BudgetWindowBatchRunner
@@ -36,7 +37,7 @@ def run_budget_window_batch_impl(
         msg = "run_budget_window_batch must run inside a Modal function call"
         raise RuntimeError(msg)
     runtime.set_context(batch_job_id=batch_job_id)
-    with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CONTEXT)):
+    with scopes.span(runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CONTEXT)):
         context = build_batch_context(
             params,
             batch_job_id=batch_job_id,
diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
index 3478c46b2..d90144a3e 100644
--- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
+++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
@@ -20,6 +20,7 @@
     put_batch_job_state,
 )
 from policyengine_simulation_contract.spm import spm_error_detail
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.errors import log_and_redact_exception
 from policyengine_simulation_observability.stages import BUDGET_WINDOW_STAGES, Stage
 from src.modal.budget_window_context import (
@@ -71,8 +72,8 @@ def serialize_batch_status(
     runtime: ObservabilityRuntime,
 ) -> dict[str, Any]:
     """Render ``state`` as the JSON-ready batch status response."""
-    with runtime.span(
-        BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION)
+    with scopes.span(
+        runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATUS_SERIALIZATION)
     ):
         return build_batch_status_response(state).model_dump(mode="json")
 
@@ -83,7 +84,9 @@ def load_or_create_batch_state(
     runtime: ObservabilityRuntime,
 ) -> BudgetWindowBatchState:
     """Return the batch's submit-time seed state, creating and storing it if absent."""
-    with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)):
+    with scopes.span(
+        runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)
+    ):
         state = get_batch_job_seed(context.batch_job_id)
     if state is None:
         state = create_initial_batch_state(
@@ -94,7 +97,9 @@ def load_or_create_batch_state(
             ),
             bundle=context.bundle,
         )
-        with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)):
+        with scopes.span(
+            runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
+        ):
             put_batch_job_seed(state)
     return state
 
@@ -155,8 +160,8 @@ def run(self) -> dict[str, Any]:
         failure fails the whole batch.
         """
         mark_batch_running(self.state)
-        with self.runtime.span(
-            BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
+        with scopes.span(
+            self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
         ):
             put_batch_job_state(self.state)
 
@@ -196,7 +201,8 @@ def spawn_until_capacity(self) -> None:
             and self.state.queued_years
         ):
             simulation_year = self.state.queued_years[0]
-            with self.runtime.span(
+            with scopes.span(
+                self.runtime,
                 BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CHILD_REQUEST_BUILD),
                 attributes={"simulation_year": simulation_year},
             ):
@@ -205,7 +211,8 @@ def spawn_until_capacity(self) -> None:
                     simulation_year=simulation_year,
                 )
             observability_context = self.runtime.capture_context()
-            with self.runtime.span(
+            with scopes.span(
+                self.runtime,
                 BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_CHILD_SPAWN),
                 attributes={"simulation_year": simulation_year},
             ):
@@ -223,8 +230,8 @@ def spawn_until_capacity(self) -> None:
                 year=simulation_year,
                 child_job_id=call.object_id,
             )
-            with self.runtime.span(
-                BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
+            with scopes.span(
+                self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
             ):
                 put_batch_job_state(self.state)
 
@@ -274,7 +281,8 @@ def poll_running_children_once(self) -> bool:
                 self._publish_poll_stats()
 
             try:
-                with self.runtime.span(
+                with scopes.span(
+                    self.runtime,
                     BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_RESULT_PARSE),
                     attributes={"simulation_year": simulation_year},
                 ):
@@ -313,8 +321,8 @@ def poll_running_children_once(self) -> bool:
                 year=simulation_year,
                 annual_impact=annual_impact,
             )
-            with self.runtime.span(
-                BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
+            with scopes.span(
+                self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
             ):
                 put_batch_job_state(self.state)
             progress_made = True
@@ -357,15 +365,15 @@ def fail_batch_for_child_error(
         """Mark the year's child and the whole batch failed and store the state."""
         mark_child_failed(self.state, year=simulation_year, error=error)
         mark_batch_failed(self.state, error=error)
-        with self.runtime.span(
-            BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
+        with scopes.span(
+            self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
         ):
             put_batch_job_state(self.state)
 
     def complete_batch(self) -> dict[str, Any]:
         """Aggregate annual impacts, mark the batch complete, and return its status."""
-        with self.runtime.span(
-            BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_AGGREGATION)
+        with scopes.span(
+            self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_AGGREGATION)
         ):
             annual_impacts = [
                 self.state.partial_annual_impacts[simulation_year]
@@ -378,8 +386,8 @@ def complete_batch(self) -> dict[str, Any]:
                 annual_impacts=annual_impacts,
             )
         mark_batch_complete(self.state, result=result)
-        with self.runtime.span(
-            BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
+        with scopes.span(
+            self.runtime, BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_WRITE)
         ):
             put_batch_job_state(self.state)
         return serialize_batch_status(self.state, runtime=self.runtime)
diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py
index 0a291da68..955b42d20 100644
--- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py
+++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py
@@ -587,8 +587,9 @@ def _reduce(
         # The canonical internal-key stripper, plus the opt-out knob.
         simulation_params, _, _ = split_internal_payload(self.params)
         simulation_params.pop("segmented", None)
-        with self.runtime.span(
-            SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_REDUCE)
+        with scopes.span(
+            self.runtime,
+            SEGMENTED_NATIONAL_STAGES.name(Stage.SEGMENTED_NATIONAL_REDUCE),
         ):
             output = build_national_output(
                 child_results,
diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py
index 92642d673..cee879432 100644
--- a/projects/policyengine-simulation-executor/src/modal/v2_app.py
+++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py
@@ -30,6 +30,7 @@
     UK_SIMULATION_TIMEOUT_SECONDS,
     US_SIMULATION_TIMEOUT_SECONDS,
 )
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.observability import (
     init_process_observability,
     modal_image_environment,
@@ -280,7 +281,8 @@ def run_single_simulation_segment_us(
     try:
         propagated = normalize_observability_context(observability_context)
         result = None
-        with runtime.operation(
+        with scopes.operation(
+            runtime,
             STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION),
             attributes={
                 "runner_name": "stage12",
@@ -337,7 +339,8 @@ def segmented_calculator(
     try:
         propagated = normalize_observability_context(observability_context)
         result = None
-        with runtime.operation(
+        with scopes.operation(
+            runtime,
             STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION),
             attributes={
                 "runner_name": "stage12",
@@ -387,7 +390,8 @@ def run_single_simulation_uk(
     try:
         propagated = normalize_observability_context(observability_context)
         result = None
-        with runtime.operation(
+        with scopes.operation(
+            runtime,
             STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION),
             attributes={
                 "runner_name": "stage12",
@@ -453,7 +457,8 @@ def coordinate_report(
     try:
         propagated = normalize_observability_context(observability_context)
         result = None
-        with runtime.operation(
+        with scopes.operation(
+            runtime,
             stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION),
             attributes={
                 "runner_name": "stage12",
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py
index 5a86a0482..9e851922c 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py
@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING, Any
 
 from policyengine_simulation_contract.job_progress import JobStage
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.stages import (
     ANNUAL_IMPACT_STAGES,
     Stage,
@@ -74,7 +75,11 @@ def __post_init__(self) -> None:
         self.country = self.country.lower()
 
     def _span(self, name: str) -> AbstractContextManager[object]:
-        return self.runtime.span(name) if self.runtime is not None else nullcontext()
+        return (
+            scopes.span(self.runtime, name)
+            if self.runtime is not None
+            else nullcontext()
+        )
 
     @property
     def analysis(self) -> object:
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py
index fefc8281b..7524c2525 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py
@@ -13,6 +13,7 @@
     ResultComparisonStatus,
     Stage12InvocationContext,
 )
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.stages import Stage, StagePlan
 
 from policyengine_simulation_executor.stage12_result_comparison import compare_results
@@ -59,7 +60,7 @@ def compare_completed_report(
         return
     try:
         comparison_span = (
-            runtime.span(stage_plan.name(Stage.STAGE12_RESULT_COMPARISON))
+            scopes.span(runtime, stage_plan.name(Stage.STAGE12_RESULT_COMPARISON))
             if runtime is not None and stage_plan is not None
             else nullcontext()
         )
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py
index 83194b9f5..147c5549a 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py
@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Any
 
 import modal
+from policyengine_simulation_observability import scopes
 from policyengine_simulation_observability.stages import (
     STAGE12_SIMULATION_STAGES,
     Stage,
@@ -131,7 +132,8 @@ def run(self) -> SimulationCalculation:
         try:
             for segment in segments:
                 call = None
-                with self.runtime.span(
+                with scopes.span(
+                    self.runtime,
                     STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH),
                     attributes={"segment_index": segment.segment_index},
                 ):
@@ -145,16 +147,16 @@ def run(self) -> SimulationCalculation:
             raise
 
         payloads: list[object] | None = None
-        with self.runtime.span(
-            STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT)
+        with scopes.span(
+            self.runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT)
         ):
             payloads = self._collect(handles)
         if payloads is None:
             msg = "Stage 12 segment wait returned no payloads"
             raise RuntimeError(msg)
         merged: SimulationCalculation | None = None
-        with self.runtime.span(
-            STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE)
+        with scopes.span(
+            self.runtime, STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE)
         ):
             merged = merge_segment_results(
                 self.simulation,
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index ec72836f6..45560f7b5 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -7,25 +7,8 @@
 # occur with
 #   python -m scripts.rules pyright projects/policyengine-simulation-executor --write-baseline
 # then set the rule's [rule.baseline] count in rules/registry.toml to match.
-projects/policyengine-simulation-executor/src/modal/app.py :: reportReturnType :: Function with declared return type "dict[Unknown, Unknown]" must return value on all code paths
-projects/policyengine-simulation-executor/src/modal/app.py :: reportReturnType :: Function with declared return type "dict[Unknown, Unknown]" must return value on all code paths
-projects/policyengine-simulation-executor/src/modal/app.py :: reportReturnType :: Function with declared return type "dict[Unknown, Unknown]" must return value on all code paths
-projects/policyengine-simulation-executor/src/modal/budget_window_batch.py :: reportPossiblyUnboundVariable :: "context" is possibly unbound
 projects/policyengine-simulation-executor/src/modal/budget_window_results.py :: reportArgumentType :: Argument of type "float" cannot be assigned to parameter "year" of type "Literal['Total']" in function "__init__"
 projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "annual_impact" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "call" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "call" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "call" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "child_request" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "result" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "state" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "state" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportReturnType :: Function with declared return type "dict[str, Any]" must return value on all code paths
-projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound
-projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportPossiblyUnboundVariable :: "output" is possibly unbound
 projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py :: reportOptionalMemberAccess :: "split" is not a known attribute of "None"
 projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportArgumentType :: Argument of type "dict[Unknown, Unknown]" cannot be assigned to parameter "value" of type "BundleManifestMetadata" in function "__setitem__"
 projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportGeneralTypeIssues :: Could not assign item in TypedDict
diff --git a/rules/registry.toml b/rules/registry.toml
index 0c792addf..46e4723d6 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 72
+count = 55
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From 4f9ce986a2142b5ef0f56f436038c4e4bec3d9f4 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:20:00 +0400
Subject: [PATCH 03/15] Type the version-registry utility's routing state as it
 is stored

update_version_registry's 8 Pyright errors, without changing what it
writes or prints:

- LatestVersions is no longer total. The empty routing state, and any
  state before a release of a kind is routed, has no latest entry for
  it; readers already used .get. The few direct reads now use .get
  too: the plan's advances_latest flag and the post-publish and
  post-restore reports, which run after validate_routing_state has
  required those entries, and the legacy-seed merge, which copied a
  kind's latest only when present.
- RoutingState.bundles holds Mapping[str, Any]. A stored state is read
  back unvalidated, and the coercion and legacy seeding copy any
  mapping in; validate_routing_state checks the manifests against
  their routes.
- _json_copy is generic: a JSON round trip of JSON-safe data keeps its
  shape, so the planned state stays a RoutingState.
- finalize_routing_publication returns "no-pending" as soon as there is
  no routing-state Dict, the outcome it reached anyway; the finalize
  command treats a missing Dict like no pending publication, as it did.

The executor's baseline drops from 55 entries to 47.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                       |  2 +-
 .../modal/utils/update_version_registry.py    | 46 ++++++++++++-------
 rules/baselines/pyright-executor.txt          |  8 ----
 rules/registry.toml                           |  2 +-
 4 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index f60044a18..c0890dd9b 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 55 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 47 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py
index b11acff9b..0c2b5fd9b 100644
--- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py
+++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py
@@ -11,11 +11,11 @@
 import json
 import os
 import uuid
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Mapping
 from copy import deepcopy
 from dataclasses import dataclass
 from pathlib import Path
-from typing import NotRequired, TypedDict
+from typing import Any, NotRequired, TypedDict
 
 import modal
 from packaging.version import InvalidVersion, Version
@@ -66,8 +66,11 @@ class BundleManifestMetadata(TypedDict):
     job_progress: NotRequired[int]
 
 
-class LatestVersions(TypedDict):
-    """The versions that unpinned requests route to, per route kind."""
+class LatestVersions(TypedDict, total=False):
+    """The versions that unpinned requests route to, per route kind.
+
+    A kind has no entry until a release of that kind is routed.
+    """
 
     policyengine: str
     us: str
@@ -89,7 +92,11 @@ class RoutingState(TypedDict):
     generation: str
     latest: LatestVersions
     routes: RouteMaps
-    bundles: dict[str, BundleManifestMetadata]
+    # Each policyengine version's bundle manifest, as publishing writes it
+    # (BundleManifestMetadata). A stored state is read back unvalidated, so
+    # its manifests are only known to be mappings until validate_routing_state
+    # checks them against their routes.
+    bundles: dict[str, Mapping[str, Any]]
 
 
 def _is_newer_version(candidate: str, current: str | None) -> bool:
@@ -456,7 +463,8 @@ def routing_state_digest(state: object) -> str:
     return hashlib.sha256(text.encode("utf-8")).hexdigest()
 
 
-def _json_copy(value: object) -> object:
+def _json_copy[JsonValue](value: JsonValue) -> JsonValue:
+    """Return a deep copy of JSON-safe ``value``; the round trip keeps its shape."""
     return json.loads(json.dumps(value))
 
 
@@ -496,7 +504,7 @@ def plan_routing_publication(
         "after": after,
         "before_sha256": routing_state_digest(before),
         "after_sha256": routing_state_digest(after),
-        "advances_latest": after["latest"]["policyengine"] == policyengine_version,
+        "advances_latest": after["latest"].get("policyengine") == policyengine_version,
     }
 
 
@@ -810,9 +818,12 @@ def finalize_routing_publication(
         msg = "operator recovery requires gates_passed=false"
         raise ValueError(msg)
     store = _routing_store(environment, create_if_missing=False)
+    if store is None:
+        # No routing-state Dict, so nothing was ever published or recorded.
+        return "no-pending"
     pending = _unfinished_publication(store)
     if pending is None:
-        record = store.get(ROUTING_PUBLICATION_KEY) if store is not None else None
+        record = store.get(ROUTING_PUBLICATION_KEY)
         if (
             record is not None
             and publication_run is not None
@@ -1048,8 +1059,9 @@ def _merge_legacy_seed_into_current(
 
     for kind in ("policyengine", "us", "uk"):
         merged["routes"][kind].update(current["routes"][kind])
-        if kind in current["latest"]:
-            merged["latest"][kind] = current["latest"][kind]
+        current_latest = current["latest"].get(kind)
+        if current_latest is not None:
+            merged["latest"][kind] = current_latest
 
     merged["bundles"].update(current["bundles"])
     validate_routing_state(merged)
@@ -1372,7 +1384,7 @@ def _finalize_command(args: argparse.Namespace) -> None:
         raise _UsageError(msg)
     store = _routing_store(args.environment, create_if_missing=False)
     pending = _unfinished_publication(store)
-    if pending is None:
+    if store is None or pending is None:
         _report("No pending routing publication; routing is unchanged.")
     else:
         _report("Pending routing publication:")
@@ -1437,9 +1449,9 @@ def _restore_latest_command(args: argparse.Namespace) -> None:
         expected_latest=args.expected_latest,
     )
     _report(f"Restored routing state in Modal environment: {args.environment}")
-    _report(f"  latest.policyengine: {state['latest']['policyengine']}")
-    _report(f"  latest.us: {state['latest']['us']}")
-    _report(f"  latest.uk: {state['latest']['uk']}")
+    _report(f"  latest.policyengine: {state['latest'].get('policyengine')}")
+    _report(f"  latest.us: {state['latest'].get('us')}")
+    _report(f"  latest.uk: {state['latest'].get('uk')}")
 
 
 def _expected_latest_without_restore(_args: argparse.Namespace) -> None:
@@ -1518,9 +1530,9 @@ def _publish_command(args: argparse.Namespace) -> None:
     )
     _report()
     _report(f"  {ROUTING_STATE_DICT_NAME}[{ROUTING_STATE_ACTIVE_KEY}]: updated")
-    _report(f"  latest.policyengine: {state['latest']['policyengine']}")
-    _report(f"  latest.us: {state['latest']['us']}")
-    _report(f"  latest.uk: {state['latest']['uk']}")
+    _report(f"  latest.policyengine: {state['latest'].get('policyengine')}")
+    _report(f"  latest.us: {state['latest'].get('us')}")
+    _report(f"  latest.uk: {state['latest'].get('uk')}")
     _report()
     _report("Routing state published successfully.")
 
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index 45560f7b5..ccdc9ef77 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -10,14 +10,6 @@
 projects/policyengine-simulation-executor/src/modal/budget_window_results.py :: reportArgumentType :: Argument of type "float" cannot be assigned to parameter "year" of type "Literal['Total']" in function "__init__"
 projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
 projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py :: reportOptionalMemberAccess :: "split" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportArgumentType :: Argument of type "dict[Unknown, Unknown]" cannot be assigned to parameter "value" of type "BundleManifestMetadata" in function "__setitem__"
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportGeneralTypeIssues :: Could not assign item in TypedDict
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportIndexIssue :: "__getitem__" method not defined on type "object"
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportIndexIssue :: "__setitem__" method not defined on type "object"
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py :: reportReturnType :: Type "dict[str, int | str | RouteMaps | dict[str, BundleManifestMetadata]]" is not assignable to return type "RoutingState"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py :: reportAttributeAccessIssue :: Cannot access attribute "resolve_entity_variables" for class "TaxBenefitModelVersion"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportArgumentType :: Argument of type "str | None" cannot be assigned to parameter "outcome" of type "ArtifactOutcome | None" in function "__init__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation"
diff --git a/rules/registry.toml b/rules/registry.toml
index 46e4723d6..3b5594d09 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 55
+count = 47
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From b0899b509b229d7de3449796f1afbc46924cba94 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:25:52 +0400
Subject: [PATCH 04/15] Check the shape of a simulation's output data before
 reading it

policyengine types Dataset.data as an optional BaseModel, so Pyright
rejected every read of a simulation's output frames: entity_data in the
microdata export and the precompute determinism check, household in the
UK local-authority and constituency breakdowns, and the output dataset
after ensure() in _entity_data (16 errors).

output_frames adds two runtime-checkable protocols, EntityFrames and
HouseholdFrame. Each read site now checks
`isinstance(<simulation>.output_dataset.data, ...)` first and raises
TypeError if the frames are missing (before, AttributeError on None).
Pyright narrows that same expression, so each read is unchanged. That
matters for the microdata weight checks, which recognise weighted frames
by expressions such as <simulation>.output_dataset.data.household, and
whose stripping baseline lists one of these reads verbatim. The checks
need only the attribute, so the tests' stand-in datasets still pass. The
protocols type frames as pandas DataFrames, which a MicroDataFrame is,
since microdf types a column read as MicroSeries | MicroDataFrame.

The executor's baseline drops from 47 entries to 31.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                       |  2 +-
 .../output_frames.py                          | 46 +++++++++++++++++++
 .../precompute.py                             | 10 ++++
 .../simulation_microdata.py                   |  7 +++
 .../simulation_output_common.py               | 10 ++++
 .../simulation_output_geographic.py           |  6 +++
 .../tests/test_output_frames.py               | 38 +++++++++++++++
 rules/baselines/pyright-executor.txt          | 16 -------
 rules/registry.toml                           |  2 +-
 9 files changed, 119 insertions(+), 18 deletions(-)
 create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py
 create mode 100644 projects/policyengine-simulation-executor/tests/test_output_frames.py

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index c0890dd9b..3145fe74f 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 47 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 31 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py
new file mode 100644
index 000000000..81d63116c
--- /dev/null
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py
@@ -0,0 +1,46 @@
+"""The shapes of a simulation's output data, as types code can check.
+
+policyengine types ``Dataset.data`` as an optional ``BaseModel``, but each
+country's year data holds weighted entity frames. Code that reads them checks
+the shape first, for example::
+
+    if not isinstance(simulation.output_dataset.data, EntityFrames):
+        raise TypeError(...)
+    frames = simulation.output_dataset.data.entity_data
+
+Pyright then knows the type of that same expression. Reading through the
+expression itself, rather than a helper's return value, keeps it one the
+microdata weight checks recognise as a weighted frame.
+
+The protocols are runtime-checkable, so the check only needs the attribute:
+real year data and the tests' stand-ins both pass.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    import pandas as pd
+
+
+@runtime_checkable
+class EntityFrames(Protocol):
+    """Year data that holds each entity's weighted frame, by entity name."""
+
+    @property
+    def entity_data(self) -> Mapping[str, pd.DataFrame]:
+        """Return each entity's weighted frame (a microdf MicroDataFrame)."""
+        ...
+
+
+@runtime_checkable
+class HouseholdFrame(Protocol):
+    """Year data with a weighted household frame."""
+
+    @property
+    def household(self) -> pd.DataFrame:
+        """Return the weighted household frame (a microdf MicroDataFrame)."""
+        ...
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
index 0e87f9dc8..75b237a74 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
@@ -44,6 +44,7 @@
     BaselineArtifactIdentity,
     baseline_artifact_filename,
 )
+from policyengine_simulation_executor.output_frames import EntityFrames
 from policyengine_simulation_executor.precompute_models import (
     ArtifactManifest,
     BaselineComputeResult,
@@ -778,6 +779,15 @@ def verify_determinism_impl(
     )
     fresh.run()
 
+    missing_frames = "Determinism gate found a simulation with no output entity frames"
+    if baseline.output_dataset is None or not isinstance(
+        baseline.output_dataset.data, EntityFrames
+    ):
+        raise TypeError(missing_frames)
+    if fresh.output_dataset is None or not isinstance(
+        fresh.output_dataset.data, EntityFrames
+    ):
+        raise TypeError(missing_frames)
     differences = []
     loaded = baseline.output_dataset.data.entity_data
     recomputed = fresh.output_dataset.data.entity_data
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py
index 063eeed1b..d2acb0ba6 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py
@@ -20,6 +20,8 @@
 
 import pandas as pd
 
+from policyengine_simulation_executor.output_frames import EntityFrames
+
 if TYPE_CHECKING:
     from policyengine.core import Simulation
 
@@ -45,6 +47,11 @@ def dump(
     ) -> tuple[dict[str, dict[str, list]], dict[str, dict[str, str]]]:
         columns: dict[str, dict[str, list]] = {}
         dtypes: dict[str, dict[str, str]] = {}
+        if simulation.output_dataset is None or not isinstance(
+            simulation.output_dataset.data, EntityFrames
+        ):
+            msg = "the simulation has no output entity frames; run it first"
+            raise TypeError(msg)
         for entity, frame in simulation.output_dataset.data.entity_data.items():
             df = pd.DataFrame(frame)
             columns[entity] = df.to_dict("list")
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
index 4ce0d2feb..e0b01db01 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
@@ -11,6 +11,8 @@
 import pandas as pd
 from policyengine_simulation_contract.spm import spm_error_detail
 
+from policyengine_simulation_executor.output_frames import HouseholdFrame
+
 if TYPE_CHECKING:
     from _typeshed import ConvertibleToFloat
     from policyengine.core import Simulation
@@ -60,6 +62,9 @@ def _output_model_dump(value: object) -> object:
 def _entity_data(simulation: Simulation, entity: str) -> pd.DataFrame:
     if simulation.output_dataset is None or simulation.output_dataset.data is None:
         simulation.ensure()
+    if simulation.output_dataset is None:
+        msg = "ensure() left the simulation without an output dataset"
+        raise RuntimeError(msg)
     return getattr(simulation.output_dataset.data, entity)
 
 
@@ -91,6 +96,11 @@ def _unweighted_household_frame(simulation: Simulation) -> pd.DataFrame:
     shrink. It goes once the breakdown hands policyengine the simulations, as
     the constituency and US congressional district breakdowns already do.
     """
+    if simulation.output_dataset is None or not isinstance(
+        simulation.output_dataset.data, HouseholdFrame
+    ):
+        msg = "the simulation has no household output frame; run it first"
+        raise TypeError(msg)
     return pd.DataFrame(simulation.output_dataset.data.household)
 
 
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py
index 849c6c7a1..b19a555b2 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py
@@ -10,6 +10,7 @@
     CONSTITUENCY_ASSET_SPEC,
 )
 
+from policyengine_simulation_executor.output_frames import HouseholdFrame
 from policyengine_simulation_executor.simulation_macro_output import (
     CongressionalDistrictImpactOutput,
     CongressionalDistrictImpactRecord,
@@ -264,6 +265,11 @@ def build_uk_constituency_impact(
     """
     if country != "uk":
         return None
+    if baseline.output_dataset is None or not isinstance(
+        baseline.output_dataset.data, HouseholdFrame
+    ):
+        msg = "the baseline has no household output frame; run it first"
+        raise TypeError(msg)
     _require_uk_area_column(
         baseline.output_dataset.data.household, "constituency_code_oa", "constituency"
     )
diff --git a/projects/policyengine-simulation-executor/tests/test_output_frames.py b/projects/policyengine-simulation-executor/tests/test_output_frames.py
new file mode 100644
index 000000000..16e6bed9d
--- /dev/null
+++ b/projects/policyengine-simulation-executor/tests/test_output_frames.py
@@ -0,0 +1,38 @@
+"""The checkable shapes of a simulation's output data."""
+
+from types import SimpleNamespace
+
+import pandas as pd
+from microdf import MicroDataFrame
+from policyengine.tax_benefit_models.uk.datasets import UKYearData
+from policyengine_simulation_executor.output_frames import EntityFrames, HouseholdFrame
+
+
+def _uk_year_data() -> UKYearData:
+    def frame(entity: str) -> MicroDataFrame:
+        return MicroDataFrame(
+            pd.DataFrame({f"{entity}_id": [1], f"{entity}_weight": [2.0]}),
+            weights=f"{entity}_weight",
+        )
+
+    return UKYearData(
+        person=frame("person"), benunit=frame("benunit"), household=frame("household")
+    )
+
+
+def test_country_year_data_has_entity_frames_and_a_household_frame():
+    year_data = _uk_year_data()
+
+    assert isinstance(year_data, EntityFrames)
+    assert isinstance(year_data, HouseholdFrame)
+
+
+def test_data_with_frames_by_attribute_matches_the_shapes():
+    assert isinstance(SimpleNamespace(entity_data={}), EntityFrames)
+    assert isinstance(SimpleNamespace(household=None), HouseholdFrame)
+
+
+def test_missing_output_data_matches_neither_shape():
+    assert not isinstance(None, EntityFrames)
+    assert not isinstance(None, HouseholdFrame)
+    assert not isinstance(SimpleNamespace(person=None), HouseholdFrame)
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index ccdc9ef77..03ffc022c 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -14,31 +14,15 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/b
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportArgumentType :: Argument of type "str | None" cannot be assigned to parameter "outcome" of type "ArtifactOutcome | None" in function "__init__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: Argument missing for parameter "schema"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: No parameter named "manifest_schema"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_microdata.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "model_dump" for class "object"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "actual_values" of type "Series" in function "_compare_column"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "expected_values" of type "Series" in function "_compare_column"
diff --git a/rules/registry.toml b/rules/registry.toml
index 3b5594d09..0e1a8c9bf 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 47
+count = 31
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From 17848b0cf77e44f0729c1eb475b5d064b3f038ce Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:28:20 +0400
Subject: [PATCH 05/15] Narrow the output helpers' column sums and aggregate
 results

simulation_output_common's 7 Pyright errors, without changing a total:

- _sum_output_variable and _change_output_variable sum a column chosen
  by name, which pandas types as a Series or, for a repeated name, a
  DataFrame. They now require a Series (a MicroSeries is one, so the
  sum stays weighted) and raise TypeError for a repeated name, where
  float() of a frame's per-column sums raised TypeError before.
- After Aggregate.run() and ChangeAggregate.run(), a result that is
  still None raises RuntimeError, where float(None) raised TypeError.
- _output_model_dump reads model_dump with getattr instead of testing
  hasattr, which Pyright doesn't narrow on.

The executor's baseline drops from 31 entries to 24.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                       |  2 +-
 .../simulation_output_common.py               | 26 ++++++++++++++++---
 rules/baselines/pyright-executor.txt          |  7 -----
 rules/registry.toml                           |  2 +-
 4 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 3145fe74f..cf1c4d87f 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 31 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 24 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
index e0b01db01..bcffc4e91 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
@@ -52,8 +52,9 @@ def _collection_records(collection: object) -> list[dict[str, Any]]:
 def _output_model_dump(value: object) -> object:
     if value is None:
         return None
-    if hasattr(value, "model_dump"):
-        return value.model_dump(mode="json")
+    model_dump = getattr(value, "model_dump", None)
+    if model_dump is not None:
+        return model_dump(mode="json")
     if isinstance(value, Mapping):
         return dict(value)
     return None
@@ -142,7 +143,11 @@ def _warn_on_mismatched_weights(
 def _sum_output_variable(simulation: Simulation, variable: str, entity: str) -> float:
     data = _entity_data(simulation, entity)
     if variable in data.columns:
-        return float(data[variable].sum())
+        column = data[variable]
+        if not isinstance(column, pd.Series):
+            msg = f"{entity} output has more than one {variable!r} column"
+            raise TypeError(msg)
+        return float(column.sum())
 
     from policyengine.outputs import (  # noqa: PLC0415  # policyengine stays out of this module's import-time imports
         Aggregate,
@@ -156,6 +161,9 @@ def _sum_output_variable(simulation: Simulation, variable: str, entity: str) ->
         aggregate_type=AggregateType.SUM,
     )
     output.run()
+    if output.result is None:
+        msg = f"Aggregate of {entity}.{variable} produced no result"
+        raise RuntimeError(msg)
     return float(output.result)
 
 
@@ -166,7 +174,14 @@ def _change_output_variable(
     reform_data = _entity_data(reform, entity)
     _warn_on_mismatched_weights(baseline_data, reform_data, entity)
     if variable in baseline_data.columns and variable in reform_data.columns:
-        return float((reform_data[variable] - baseline_data[variable]).sum())
+        reform_column = reform_data[variable]
+        baseline_column = baseline_data[variable]
+        if not isinstance(reform_column, pd.Series) or not isinstance(
+            baseline_column, pd.Series
+        ):
+            msg = f"{entity} output has more than one {variable!r} column"
+            raise TypeError(msg)
+        return float((reform_column - baseline_column).sum())
 
     from policyengine.outputs import (  # noqa: PLC0415  # policyengine stays out of this module's import-time imports
         ChangeAggregate,
@@ -181,6 +196,9 @@ def _change_output_variable(
         aggregate_type=ChangeAggregateType.SUM,
     )
     output.run()
+    if output.result is None:
+        msg = f"ChangeAggregate of {entity}.{variable} produced no result"
+        raise RuntimeError(msg)
     return float(output.result)
 
 
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index 03ffc022c..933e249fe 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -16,13 +16,6 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/p
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: Argument missing for parameter "schema"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: No parameter named "manifest_schema"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Any | None" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportArgumentType :: Argument of type "Unknown | Any | Series" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py :: reportAttributeAccessIssue :: Cannot access attribute "model_dump" for class "object"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "actual_values" of type "Series" in function "_compare_column"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "expected_values" of type "Series" in function "_compare_column"
diff --git a/rules/registry.toml b/rules/registry.toml
index 0e1a8c9bf..2a68787f2 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 31
+count = 24
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From b6a587f68949cc21de90a8b7472a37501314c561 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:31:20 +0400
Subject: [PATCH 06/15] Type the baseline artifact outcome and the precompute
 manifest

The 6 Pyright errors in precompute.py and baseline_artifacts.py,
without changing behaviour:

- build_manifest passes the manifest schema by its alias, `schema`.
  ArtifactManifest also accepts `manifest_schema` (populate_by_name),
  but Pyright's Pydantic support knows only the alias as a constructor
  keyword.
- ArtifactBaselineSimulation records its outcome as ArtifactOutcome,
  the Literal the precompute result model declares (and a test already
  ties to the OUTCOME_* constants), instead of str.
- The determinism check reads the outcome with getattr, as the module
  already reads storage_id: the cohort baseline is declared a
  Simulation, which has no artifact_outcome.
- _missing_output_columns checks the model version against a
  runtime-checkable EntityVariableResolver protocol before calling
  resolve_entity_variables, which policyengine defines on
  MicrosimulationModelVersion rather than the declared
  TaxBenefitModelVersion; the tests' stand-in model versions pass.

The executor's baseline drops from 24 entries to 18.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                       |  2 +-
 .../baseline_artifacts.py                     | 36 ++++++++++++++-----
 .../precompute.py                             | 11 ++++--
 rules/baselines/pyright-executor.txt          |  6 ----
 rules/registry.toml                           |  2 +-
 5 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index cf1c4d87f..253da21b6 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 24 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 18 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py
index 7f0b82cfb..c0e54e47d 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py
@@ -22,18 +22,20 @@
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable
 
 from policyengine.core.scoping_strategy import RegionGroupStrategy
 from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail
 from pydantic import PrivateAttr
 
 from policyengine_simulation_executor import artifact_keys
+from policyengine_simulation_executor.precompute_models import ArtifactOutcome
 from policyengine_simulation_executor.stage_reporting_simulation import (
     StageReportingSimulation,
 )
 
 if TYPE_CHECKING:
+    from policyengine.core import Simulation
     from policyengine.core.scoping_strategy import ScopingStrategy
 
 logger = logging.getLogger(__name__)
@@ -43,9 +45,23 @@
 # "incomplete": loaded but missing requested columns or a valid SPM receipt
 # -> recomputed;
 # "miss": no artifact -> computed (today's behavior).
-OUTCOME_HIT = "hit"
-OUTCOME_INCOMPLETE = "incomplete"
-OUTCOME_MISS = "miss"
+OUTCOME_HIT: Final = "hit"
+OUTCOME_INCOMPLETE: Final = "incomplete"
+OUTCOME_MISS: Final = "miss"
+
+
+@runtime_checkable
+class EntityVariableResolver(Protocol):
+    """A model version that lists the output variables a simulation must carry.
+
+    policyengine defines ``resolve_entity_variables`` on
+    MicrosimulationModelVersion, not on the TaxBenefitModelVersion a
+    Simulation declares.
+    """
+
+    def resolve_entity_variables(self, simulation: Simulation) -> dict[str, list[str]]:
+        """Return each entity's required output variables."""
+        ...
 
 
 @dataclass(frozen=True)
@@ -154,10 +170,10 @@ class ArtifactBaselineSimulation(StageReportingSimulation):
     """A Simulation whose ``ensure()`` validates what it loads."""
 
     _computed_this_process: bool = PrivateAttr(default=False)
-    _artifact_outcome: str | None = PrivateAttr(default=None)
+    _artifact_outcome: ArtifactOutcome | None = PrivateAttr(default=None)
 
     @property
-    def artifact_outcome(self) -> str | None:
+    def artifact_outcome(self) -> ArtifactOutcome | None:
         """Set by ``ensure()``; None until then."""
         return self._artifact_outcome
 
@@ -166,7 +182,7 @@ def run(self) -> None:
         self._computed_this_process = True
         super().run()
 
-    def _record_outcome(self, outcome: str) -> None:
+    def _record_outcome(self, outcome: ArtifactOutcome) -> None:
         """First outcome wins for the life of this instance.
 
         Every request ensures its baseline more than once (the economic
@@ -247,7 +263,11 @@ def ensure(self) -> None:
     def _missing_output_columns(self) -> list[tuple[str, str]]:
         data = getattr(self.output_dataset, "data", None)
         frames = getattr(data, "entity_data", None)
-        resolved = self.tax_benefit_model_version.resolve_entity_variables(self)
+        model_version = self.tax_benefit_model_version
+        if not isinstance(model_version, EntityVariableResolver):
+            msg = "the baseline's model version can't list its output variables"
+            raise TypeError(msg)
+        resolved = model_version.resolve_entity_variables(self)
         if not frames:
             return [(entity, "<no output data>") for entity in resolved]
         missing: list[tuple[str, str]] = []
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
index 75b237a74..50d0561f6 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
@@ -285,7 +285,9 @@ def build_manifest(plan: PrecomputePlan) -> ArtifactManifest:
         for entry in plan.state_datasets
     ]
     return ArtifactManifest(
-        manifest_schema=MANIFEST_SCHEMA,
+        # The field's alias, which Pydantic accepts (and which its type
+        # checker support knows) at construction.
+        schema=MANIFEST_SCHEMA,
         country=PRECOMPUTE_COUNTRY,
         receipt=plan.receipt,
         artifacts=artifacts,
@@ -761,10 +763,13 @@ def verify_determinism_impl(
     if not artifact_file.exists():
         store.download_file(expected.path, artifact_file)
     baseline.ensure()
-    if baseline.artifact_outcome != OUTCOME_HIT:
+    # A cohort baseline is an ArtifactBaselineSimulation, which records the
+    # outcome; the attribute isn't on the Simulation type it is declared as.
+    outcome = getattr(baseline, "artifact_outcome", None)
+    if outcome != OUTCOME_HIT:
         msg = (
             "Determinism gate could not load the artifact it verifies "
-            f"(outcome={baseline.artifact_outcome})"
+            f"(outcome={outcome})"
         )
         raise RuntimeError(msg)
 
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index 933e249fe..c6ee247e3 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -10,12 +10,6 @@
 projects/policyengine-simulation-executor/src/modal/budget_window_results.py :: reportArgumentType :: Argument of type "float" cannot be assigned to parameter "year" of type "Literal['Total']" in function "__init__"
 projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
 projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py :: reportOptionalMemberAccess :: "split" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py :: reportAttributeAccessIssue :: Cannot access attribute "resolve_entity_variables" for class "TaxBenefitModelVersion"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportArgumentType :: Argument of type "str | None" cannot be assigned to parameter "outcome" of type "ArtifactOutcome | None" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportAttributeAccessIssue :: Cannot access attribute "artifact_outcome" for class "Simulation"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: Argument missing for parameter "schema"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py :: reportCallIssue :: No parameter named "manifest_schema"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "actual_values" of type "Series" in function "_compare_column"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "expected_values" of type "Series" in function "_compare_column"
diff --git a/rules/registry.toml b/rules/registry.toml
index 2a68787f2..68834ec10 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 24
+count = 18
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From 9dd6d4d4c96cec4238600bb42beea9ae824e29fa Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:33:23 +0400
Subject: [PATCH 07/15] Build the Stage 12 simulation's two types in separate
 branches

build_stage12_simulation chose its class, then passed the cache id by
unpacking `{"id": ...}`, a dict[str, str], which Pyright matches against
every keyword parameter (7 errors: created_at, updated_at, dynamic and
the rest). It now constructs a Simulation when there is no cache
identity and a Stage12CachedSimulation with id= and its cache
configuration when there is one, as the annual runtime's
_build_simulation does. The optional spm keyword is still passed only
when there is a selection.

The executor's baseline drops from 18 entries to 11.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                       |  2 +-
 .../stage12_runtime/simulation.py             | 19 +++++++++++++------
 rules/baselines/pyright-executor.txt          |  7 -------
 rules/registry.toml                           |  2 +-
 4 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 253da21b6..5ead97bae 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 18 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 11 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py
index 3b2bb5b56..43fad970e 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py
@@ -115,17 +115,24 @@ def build_stage12_simulation(
         scoping_strategy=resolved.region.scoping_strategy,
         spm=selection,
     )
-    simulation_type = Stage12CachedSimulation if identity is not None else Simulation
-    model = simulation_type(
-        **({"spm": selection} if selection is not None else {}),
-        **({"id": identity.simulation_id} if identity is not None else {}),
+    spm_kwargs: dict[str, Any] = {"spm": selection} if selection is not None else {}
+    if identity is None:
+        return Simulation(
+            **spm_kwargs,
+            dataset=resolved.dataset,
+            tax_benefit_model_version=resolved.country_module.model,
+            policy=policy,
+            scoping_strategy=resolved.region.scoping_strategy,
+        )
+    model = Stage12CachedSimulation(
+        **spm_kwargs,
+        id=identity.simulation_id,
         dataset=resolved.dataset,
         tax_benefit_model_version=resolved.country_module.model,
         policy=policy,
         scoping_strategy=resolved.region.scoping_strategy,
     )
-    if isinstance(model, Stage12CachedSimulation):
-        model.configure_stage12_cache(execution.output_plan)
+    model.configure_stage12_cache(execution.output_plan)
     return model
 
 
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index c6ee247e3..ff4cc4319 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -18,10 +18,3 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/s
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportIndexIssue :: "__getitem__" method not defined on type "object"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "created_at" of type "datetime" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "dynamic" of type "Dynamic | dict[str, Any] | None" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "extra_variables" of type "dict[str, list[str]]" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "output_dataset" of type "Dataset | None" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "spm" of type "SPMSelection | None" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "spm_receipt" of type "SPMProvenance | None" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "updated_at" of type "datetime" in function "__init__"
diff --git a/rules/registry.toml b/rules/registry.toml
index 68834ec10..c3e498971 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 18
+count = 11
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From f8e351d4a11b3317a18097656784fd9098cddbfe Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:35:44 +0400
Subject: [PATCH 08/15] Narrow the Stage 12 parity check's report values and
 columns

stage12_parity's 7 Pyright errors:

- _is_report_list and _is_report_number now return TypeGuards
  (Sequence[object], and int | float), so a value that passed one is
  typed as that in the comparison that follows. TypeGuard, unlike
  TypeIs, says nothing about a value that fails, which matters because
  a str is a Sequence the list check rejects.
- A frame column selected by name is a Series, or a DataFrame when the
  name repeats. _compare_entity_frames now requires a Series on both
  sides and raises TypeError naming the column for a repeated one; the
  column comparison would have failed with AttributeError on .dtype.

The executor's baseline drops from 11 entries to 4.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                          |  2 +-
 .../stage12_parity.py                            | 16 ++++++++++++----
 rules/baselines/pyright-executor.txt             |  7 -------
 rules/registry.toml                              |  2 +-
 4 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 5ead97bae..20718ca95 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 11 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 4 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py
index 2e68cebb2..a4da63563 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py
@@ -5,6 +5,7 @@
 import math
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
+from typing import TypeGuard
 
 import numpy as np
 import pandas as pd
@@ -100,10 +101,17 @@ def _compare_entity_frames(
         raise ParityMismatchError(error_code, entity)
     for column in expected_frame.columns:
         path = f"{entity}.{column}"
+        expected_values = expected_frame[column]
+        actual_values = actual_frame[column]
+        if not isinstance(expected_values, pd.Series) or not isinstance(
+            actual_values, pd.Series
+        ):
+            msg = f"{entity} has more than one {column!r} column"
+            raise TypeError(msg)
         _compare_column(
             path,
-            expected_frame[column],
-            actual_frame[column],
+            expected_values,
+            actual_values,
             tolerances.get(path, NumericalTolerance()),
         )
 
@@ -148,14 +156,14 @@ def compare_aggregate_reports(
     _compare_report_values(incumbent, candidate, "", tolerances or {})
 
 
-def _is_report_list(value: object) -> bool:
+def _is_report_list(value: object) -> TypeGuard[Sequence[object]]:
     """Whether a report value is a list-like sequence rather than text or bytes."""
     return isinstance(value, Sequence) and not isinstance(
         value, (str, bytes, bytearray)
     )
 
 
-def _is_report_number(value: object) -> bool:
+def _is_report_number(value: object) -> TypeGuard[int | float]:
     """Whether a report value is an int or float; booleans are not numbers here."""
     return isinstance(value, (int, float)) and not isinstance(value, bool)
 
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index ff4cc4319..5e285d1ab 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -11,10 +11,3 @@ projects/policyengine-simulation-executor/src/modal/budget_window_results.py ::
 projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
 projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py :: reportOptionalMemberAccess :: "split" is not a known attribute of "None"
 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "actual_values" of type "Series" in function "_compare_column"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "Series | Unknown | Self@DataFrame" cannot be assigned to parameter "expected_values" of type "Series" in function "_compare_column"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "expected" of type "Sequence[Unknown]" in function "_compare_report_lists"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "obj" of type "Sized" in function "len"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "x" of type "ConvertibleToFloat" in function "__new__"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_parity.py :: reportIndexIssue :: "__getitem__" method not defined on type "object"
diff --git a/rules/registry.toml b/rules/registry.toml
index c3e498971..de273a847 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 11
+count = 4
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From 07e815f04d2aa2c892cc7b65008727a9bc1cfb58 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 19:38:05 +0400
Subject: [PATCH 09/15] Fix the executor's last four Pyright errors

- The budget-window totals are passed to BudgetWindowTotals by name
  instead of by unpacking a dict[str, float], which Pyright matched
  against the model's year field too.
- The scheduler's poll reads the child call handle, which
  ChildSimulationHandle types as optional, and raises if it is missing;
  resolve_child_handle always returns a handle with its call.
- release_sentinel's argument parser description reads __doc__, which
  is None when Python strips docstrings (-OO), through `or ""`.
- The US SPM capability validates the bundle's configured defaults with
  SPMCapability.model_validate, Pydantic's entry point for untyped
  input, instead of passing an object to the constructor; the
  validation is the same.

The executor now has no Pyright errors, and its baseline is empty.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                                   | 2 +-
 .../src/modal/budget_window_results.py                    | 8 +++++++-
 .../src/modal/budget_window_scheduler.py                  | 8 ++++++--
 .../src/modal/utils/release_sentinel.py                   | 3 ++-
 .../src/policyengine_simulation_executor/spm.py           | 2 +-
 rules/baselines/pyright-executor.txt                      | 4 ----
 rules/registry.toml                                       | 2 +-
 7 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 20718ca95..2d205e006 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 4 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 0 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates
diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py
index 4c1c06ae1..c2b082ae3 100644
--- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py
+++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py
@@ -115,7 +115,13 @@ def sum_annual_impacts(
         )
         totals["budgetaryImpact"] += _as_decimal(annual_impact.budgetaryImpact)
 
-    return BudgetWindowTotals(**{key: float(value) for key, value in totals.items()})
+    return BudgetWindowTotals(
+        taxRevenueImpact=float(totals["taxRevenueImpact"]),
+        federalTaxRevenueImpact=float(totals["federalTaxRevenueImpact"]),
+        stateTaxRevenueImpact=float(totals["stateTaxRevenueImpact"]),
+        benefitSpendingImpact=float(totals["benefitSpendingImpact"]),
+        budgetaryImpact=float(totals["budgetaryImpact"]),
+    )
 
 
 def build_budget_window_result(
diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
index d90144a3e..e85e5ab54 100644
--- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
+++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
@@ -244,11 +244,15 @@ def poll_running_children_once(self) -> bool:
         progress_made = False
 
         for simulation_year in list(self.state.running_years):
-            handle = self.resolve_child_handle(simulation_year)
+            call = self.resolve_child_handle(simulation_year).call
+            if call is None:
+                # resolve_child_handle always returns a handle with its call.
+                msg = f"the {simulation_year} child has no call handle"
+                raise RuntimeError(msg)
 
             poll_started = time.perf_counter()
             try:
-                child_result = handle.call.get(timeout=0)
+                child_result = call.get(timeout=0)
             except TimeoutError:
                 continue
             except Exception as exc:  # noqa: BLE001  # call.get re-raises whatever the child worker raised; every failure fails the batch, logged with its traceback by log_and_redact_exception
diff --git a/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py b/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py
index 6b4eb35fd..55e93f3f1 100644
--- a/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py
+++ b/projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py
@@ -411,7 +411,8 @@ def _report(line: str, *, error: bool = False) -> None:
 
 def main(argv: list[str] | None = None) -> int:
     """Run the ``expected`` or ``check`` command and return its exit status."""
-    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    # __doc__ is None when Python runs with -OO, which strips docstrings.
+    parser = argparse.ArgumentParser(description=(__doc__ or "").split("\n\n")[0])
     commands = parser.add_subparsers(dest="command", required=True)
     expected_parser = commands.add_parser("expected", help="compute in-process values")
     expected_parser.add_argument("--output", required=True)
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py
index b5c2490a7..b4b4ade8e 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py
@@ -103,7 +103,7 @@ def _certified_us_spm_capability(configured: object) -> SPMCapability:
     ):
         msg = "The installed US model does not support canonical SPM"
         raise ValueError(msg)
-    capability = SPMCapability(defaults=configured)
+    capability = SPMCapability.model_validate({"defaults": configured})
     forecast = _forecast(capability.defaults.forecast_content_sha256)
     forecast.entry(
         forecast.years[0],
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index 5e285d1ab..aafb93a73 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -7,7 +7,3 @@
 # occur with
 #   python -m scripts.rules pyright projects/policyengine-simulation-executor --write-baseline
 # then set the rule's [rule.baseline] count in rules/registry.toml to match.
-projects/policyengine-simulation-executor/src/modal/budget_window_results.py :: reportArgumentType :: Argument of type "float" cannot be assigned to parameter "year" of type "Literal['Total']" in function "__init__"
-projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportOptionalMemberAccess :: "get" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/modal/utils/release_sentinel.py :: reportOptionalMemberAccess :: "split" is not a known attribute of "None"
-projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__"
diff --git a/rules/registry.toml b/rules/registry.toml
index de273a847..bd7c05e0a 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -1143,7 +1143,7 @@ job = "pyright"
 invocation = "python -m scripts.rules pyright"
 [rule.baseline]
 file = "rules/baselines/pyright-executor.txt"
-count = 4
+count = 0
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From 69f4cbdf5c1845d4abdc2eb6a3c633a25eab7da3 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 20:24:23 +0400
Subject: [PATCH 10/15] Fail the batch, as before, when a child's call handle
 is missing

The previous fix narrowed the child's optional call handle with a check
before the poll's try, so a missing handle would have raised out of the
scheduler and left the stored batch running. Before that fix the same
case failed inside the try (None.get) and the handler failed the batch.
The check now lives in _child_call, called inside the try, so a missing
handle fails the batch again. resolve_child_handle always returns a
handle with its call, so this guards an unreachable case. A new test
covers it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../src/modal/budget_window_scheduler.py      | 21 +++++++++----
 .../tests/test_budget_window_batch.py         | 30 +++++++++++++++++++
 2 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
index e85e5ab54..3213ab1c5 100644
--- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
+++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
@@ -104,6 +104,19 @@ def load_or_create_batch_state(
     return state
 
 
+def _child_call(handle: ChildSimulationHandle) -> modal.FunctionCall[Any]:
+    """Return the child's Modal call, raising if the handle has none.
+
+    resolve_child_handle always returns a handle with its call. The poll calls
+    this inside its try, so a missing call would fail the batch like any other
+    poll error.
+    """
+    if handle.call is None:
+        msg = f"the {handle.simulation_year} child has no call handle"
+        raise RuntimeError(msg)
+    return handle.call
+
+
 class BudgetWindowBatchRunner:
     """Runs a parent budget-window batch job to completion."""
 
@@ -244,15 +257,11 @@ def poll_running_children_once(self) -> bool:
         progress_made = False
 
         for simulation_year in list(self.state.running_years):
-            call = self.resolve_child_handle(simulation_year).call
-            if call is None:
-                # resolve_child_handle always returns a handle with its call.
-                msg = f"the {simulation_year} child has no call handle"
-                raise RuntimeError(msg)
+            handle = self.resolve_child_handle(simulation_year)
 
             poll_started = time.perf_counter()
             try:
-                child_result = call.get(timeout=0)
+                child_result = _child_call(handle).get(timeout=0)
             except TimeoutError:
                 continue
             except Exception as exc:  # noqa: BLE001  # call.get re-raises whatever the child worker raised; every failure fails the batch, logged with its traceback by log_and_redact_exception
diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py
index 72ede3566..dba29ed02 100644
--- a/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py
+++ b/projects/policyengine-simulation-executor/tests/test_budget_window_batch.py
@@ -13,6 +13,7 @@
     PolicyEngineBundle,
 )
 from src.modal.budget_window_batch import run_budget_window_batch_impl
+from src.modal.budget_window_context import ChildSimulationHandle
 from src.modal.fanout import PollBackoff
 
 
@@ -315,6 +316,35 @@ def test_run_budget_window_batch_impl_marks_failure(
     assert state.child_jobs["2027"].status == "cancelled"
 
 
+def test_a_child_without_a_call_handle_fails_the_batch(
+    monkeypatch, mock_batch_modal, observability_runtime
+):
+    """A missing call handle fails the batch like any other child poll error."""
+    request, payload = _build_parent_payload(window_size=1)
+    _seed_parent_batch(request, mock_batch_modal["parent_call_id"])
+    run_simulation = MockRunSimulationFunction(
+        tracker=SpawnTracker(),
+        results_by_year={"2026": [TimeoutError()]},
+        call_registry=mock_batch_modal["call_registry"],
+    )
+    mock_batch_modal["functions"][
+        ("policyengine-simulation-py4-10-0", "run_simulation")
+    ] = run_simulation
+    monkeypatch.setattr(
+        scheduler_module.BudgetWindowBatchRunner,
+        "resolve_child_handle",
+        lambda _runner, simulation_year: ChildSimulationHandle(
+            simulation_year=simulation_year, job_id="child-2026", call=None
+        ),
+    )
+
+    result = run_budget_window_batch_impl(payload, runtime=observability_runtime)
+
+    assert result["status"] == "failed"
+    assert result["failed_years"] == ["2026"]
+    assert result["error"].startswith("Simulation failed")
+
+
 def test_scheduler_sleep_exponentially_backs_off_then_resets_on_progress(
     monkeypatch, mock_batch_modal, observability_runtime
 ):

From 923ff85a446e266dcf4ed3e00e888322137c0c2f Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 20:26:28 +0400
Subject: [PATCH 11/15] Raise TypeError wherever a simulation's output data is
 missing

The output-frame reads raised TypeError for missing frames, but
_entity_data raised RuntimeError when ensure() left no output dataset
and AttributeError when it left no data. It now raises TypeError in
both cases, keeping its getattr(<simulation>.output_dataset.data, ...)
read for the microdata checks. The UK weight-matrix check in the
simulation runtime now narrows the input dataset with the same
EntityFrames protocol instead of its own getattr and dict check.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../simulation_output_common.py               |  6 ++---
 .../simulation_runtime.py                     |  6 ++---
 .../tests/test_output_frames.py               | 24 +++++++++++++++++++
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
index bcffc4e91..8bc282fab 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py
@@ -63,9 +63,9 @@ def _output_model_dump(value: object) -> object:
 def _entity_data(simulation: Simulation, entity: str) -> pd.DataFrame:
     if simulation.output_dataset is None or simulation.output_dataset.data is None:
         simulation.ensure()
-    if simulation.output_dataset is None:
-        msg = "ensure() left the simulation without an output dataset"
-        raise RuntimeError(msg)
+    if simulation.output_dataset is None or simulation.output_dataset.data is None:
+        msg = "the simulation has no output data after ensure()"
+        raise TypeError(msg)
     return getattr(simulation.output_dataset.data, entity)
 
 
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py
index 2336e48f1..ab5395ad3 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py
@@ -35,6 +35,7 @@
     NULL_JOB_PROGRESS,
     JobProgressReporter,
 )
+from policyengine_simulation_executor.output_frames import EntityFrames
 from policyengine_simulation_executor.release_bundle import (
     get_country_release_bundle,
     resolve_bundle_dataset_name,
@@ -514,11 +515,10 @@ def release(reference: ArtifactPathReference) -> tuple[str, str]:
             msg = f"{matrix_name} entry {year} is not a weight matrix"
             raise TypeError(msg)
         matrix_households = weights.shape[-1]
-    entity_data = getattr(dataset.data, "entity_data", None)
-    if not isinstance(entity_data, dict):
+    if not isinstance(dataset.data, EntityFrames):
         msg = f"the {dataset_selection.name!r} dataset has no entity data loaded"
         raise TypeError(msg)
-    households = len(pd.DataFrame(entity_data["household"]))
+    households = len(pd.DataFrame(dataset.data.entity_data["household"]))
     if households != matrix_households:
         msg = (
             f"UK region {region!r} reweights households with {matrix_name}, "
diff --git a/projects/policyengine-simulation-executor/tests/test_output_frames.py b/projects/policyengine-simulation-executor/tests/test_output_frames.py
index 16e6bed9d..b3dbf408b 100644
--- a/projects/policyengine-simulation-executor/tests/test_output_frames.py
+++ b/projects/policyengine-simulation-executor/tests/test_output_frames.py
@@ -3,9 +3,11 @@
 from types import SimpleNamespace
 
 import pandas as pd
+import pytest
 from microdf import MicroDataFrame
 from policyengine.tax_benefit_models.uk.datasets import UKYearData
 from policyengine_simulation_executor.output_frames import EntityFrames, HouseholdFrame
+from policyengine_simulation_executor.simulation_output_common import _entity_data
 
 
 def _uk_year_data() -> UKYearData:
@@ -36,3 +38,25 @@ def test_missing_output_data_matches_neither_shape():
     assert not isinstance(None, EntityFrames)
     assert not isinstance(None, HouseholdFrame)
     assert not isinstance(SimpleNamespace(person=None), HouseholdFrame)
+
+
+class _SimulationWithoutOutput:
+    """A simulation whose ensure() leaves the output it was given."""
+
+    def __init__(self, output_dataset):
+        self.output_dataset = output_dataset
+
+    def ensure(self):
+        pass
+
+
+@pytest.mark.parametrize(
+    "output_dataset",
+    [None, SimpleNamespace(data=None)],
+    ids=["no output dataset", "no output data"],
+)
+def test_reading_an_entity_frame_without_output_data_raises_type_error(
+    output_dataset,
+):
+    with pytest.raises(TypeError, match="no output data"):
+        _entity_data(_SimulationWithoutOutput(output_dataset), "household")

From 09a737b8388b61deeac20df11a5ff4932f9c4d96 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 20:27:31 +0400
Subject: [PATCH 12/15] Test the output-frame and model-version protocols on
 real US types

The protocol tests covered UK year data only. They now also check that
policyengine's US year data passes EntityFrames and HouseholdFrame, and
that MicrosimulationModelVersion, which the US and UK model versions
derive from, satisfies EntityVariableResolver. Both pass on the
existing code; they guard the narrowing against a policyengine change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../tests/test_output_frames.py               | 52 +++++++++++++++----
 1 file changed, 43 insertions(+), 9 deletions(-)

diff --git a/projects/policyengine-simulation-executor/tests/test_output_frames.py b/projects/policyengine-simulation-executor/tests/test_output_frames.py
index b3dbf408b..84fdfc3fb 100644
--- a/projects/policyengine-simulation-executor/tests/test_output_frames.py
+++ b/projects/policyengine-simulation-executor/tests/test_output_frames.py
@@ -5,30 +5,64 @@
 import pandas as pd
 import pytest
 from microdf import MicroDataFrame
+from policyengine.tax_benefit_models.common.model_version import (
+    MicrosimulationModelVersion,
+)
 from policyengine.tax_benefit_models.uk.datasets import UKYearData
+from policyengine.tax_benefit_models.us.datasets import USYearData
+from policyengine_simulation_executor.baseline_artifacts import EntityVariableResolver
 from policyengine_simulation_executor.output_frames import EntityFrames, HouseholdFrame
 from policyengine_simulation_executor.simulation_output_common import _entity_data
 
 
-def _uk_year_data() -> UKYearData:
-    def frame(entity: str) -> MicroDataFrame:
-        return MicroDataFrame(
-            pd.DataFrame({f"{entity}_id": [1], f"{entity}_weight": [2.0]}),
-            weights=f"{entity}_weight",
-        )
+def _frame(entity: str) -> MicroDataFrame:
+    return MicroDataFrame(
+        pd.DataFrame({f"{entity}_id": [1], f"{entity}_weight": [2.0]}),
+        weights=f"{entity}_weight",
+    )
+
 
+def _uk_year_data() -> UKYearData:
     return UKYearData(
-        person=frame("person"), benunit=frame("benunit"), household=frame("household")
+        person=_frame("person"),
+        benunit=_frame("benunit"),
+        household=_frame("household"),
     )
 
 
-def test_country_year_data_has_entity_frames_and_a_household_frame():
-    year_data = _uk_year_data()
+def _us_year_data() -> USYearData:
+    return USYearData(
+        **{
+            entity: _frame(entity)
+            for entity in (
+                "person",
+                "marital_unit",
+                "family",
+                "spm_unit",
+                "tax_unit",
+                "household",
+            )
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    "year_data_factory", [_uk_year_data, _us_year_data], ids=["UK", "US"]
+)
+def test_country_year_data_has_entity_frames_and_a_household_frame(
+    year_data_factory,
+):
+    year_data = year_data_factory()
 
     assert isinstance(year_data, EntityFrames)
     assert isinstance(year_data, HouseholdFrame)
 
 
+def test_country_model_versions_resolve_entity_variables():
+    # policyengine's US and UK model versions both derive from this class.
+    assert issubclass(MicrosimulationModelVersion, EntityVariableResolver)
+
+
 def test_data_with_frames_by_attribute_matches_the_shapes():
     assert isinstance(SimpleNamespace(entity_data={}), EntityFrames)
     assert isinstance(SimpleNamespace(household=None), HouseholdFrame)

From 55c6abb41ecef42fc243a767b7b64d7ba8541c9b Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 20:27:31 +0400
Subject: [PATCH 13/15] Point the observability and type-checking guides at
 scopes.operation

Both guides prescribed scopes.span only, and the observability guide
still said to pass a received context to runtime.operation, which brings
back the Pyright errors the executor is held to zero on. They now name
scopes.operation for operations, and the type-checking guide says how to
narrow a simulation's output data without hiding the read from the
microdata checks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/engineering/skills/observability.md | 12 +++++++-----
 docs/engineering/skills/type-checking.md | 16 +++++++++++-----
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/docs/engineering/skills/observability.md b/docs/engineering/skills/observability.md
index 2134d896f..d0d08d303 100644
--- a/docs/engineering/skills/observability.md
+++ b/docs/engineering/skills/observability.md
@@ -52,7 +52,7 @@ add `_observability_context` to it.
 
 At the receiving function, validate the separate argument with
 `normalize_observability_context` and pass the result to
-`runtime.operation(..., remote_context=context)`. The operation scope must
+`scopes.operation(runtime, ..., remote_context=context)`. The operation scope must
 cover the worker call. Child dispatches made inside that scope call
 `runtime.capture_context()` again and pass the resulting context through their
 own keyword-only argument.
@@ -79,10 +79,12 @@ All span names for supported run configurations are defined in
 appropriate `StagePlan` and call `plan.name(Stage.VALUE)`. Do not add span name
 string literals in services or workers.
 
-Open a span with `scopes.span(runtime, plan.name(Stage.VALUE))` from
-`policyengine_simulation_observability.scopes` rather than `runtime.span(...)`.
-It returns the same scope, typed so that Pyright knows an exception raised in
-the block propagates; see `docs/engineering/skills/type-checking.md`.
+Open a span with `scopes.span(runtime, plan.name(Stage.VALUE))`, and an
+operation with `scopes.operation(runtime, ...)`, from
+`policyengine_simulation_observability.scopes` rather than `runtime.span(...)`
+or `runtime.operation(...)`. Each returns the runtime's own scope, typed so
+that Pyright knows an exception raised in the block propagates; see
+`docs/engineering/skills/type-checking.md`.
 
 When adding a run configuration or runtime stage:
 
diff --git a/docs/engineering/skills/type-checking.md b/docs/engineering/skills/type-checking.md
index bef001d5b..c4f16d4d8 100644
--- a/docs/engineering/skills/type-checking.md
+++ b/docs/engineering/skills/type-checking.md
@@ -114,8 +114,14 @@ Baselines are under `rules/`, so changing one needs code-owner approval
   same line: `# pyright: ignore[reportAttributeAccessIssue]  # <reason>`.
 - If an error shows a real bug, first write a test that demonstrates it, then
   fix it.
-- Open observability spans with `scopes.span(runtime, ...)` from
-  `policyengine_simulation_observability.scopes`, not `runtime.span(...)`.
-  policyengine-observability types its scope managers as possibly suppressing
-  exceptions, so after `with runtime.span(...)` Pyright reports each variable
-  the block assigned as possibly unbound.
+- Open observability spans and operations with `scopes.span(runtime, ...)` and
+  `scopes.operation(runtime, ...)` from
+  `policyengine_simulation_observability.scopes`, not `runtime.span(...)` or
+  `runtime.operation(...)`. policyengine-observability types its scope
+  managers as possibly suppressing exceptions, so after `with runtime.span(...)`
+  Pyright reports each variable the block assigned as possibly unbound.
+- Before reading a simulation's output frames, narrow
+  `<simulation>.output_dataset.data` with the `EntityFrames` or
+  `HouseholdFrame` protocols in `policyengine_simulation_executor.output_frames`,
+  and keep reading through that expression: the microdata weight checks
+  recognise weighted frames by it.

From 1c382ad34a0e37f301f043e0e0c3c84524df4508 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 20:58:06 +0400
Subject: [PATCH 14/15] Open tracked_stage's scopes through the typed scopes
 helpers

tracked_stage called runtime.operation and runtime.span directly; it now
calls scopes.operation and scopes.span, which call the same runtime
methods with the same name and attributes and no remote trace context.
The test stand-in's operation accepts remote_context, as the real
ObservabilityRuntime.operation does, and asserts it is None.

No behaviour change: test_tracked_stage_scopes_and_records_the_stage_with_its_role
passes before and after.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../stage12_runtime/failure_causes.py                        | 5 +++--
 .../tests/test_stage12_failure_causes.py                     | 5 ++++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py
index 27aac2e99..8a60a4d19 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_causes.py
@@ -35,6 +35,7 @@
 from urllib.parse import urlsplit
 
 from policyengine_simulation_contract.stage12_execution import ErrorSummary
+from policyengine_simulation_observability import scopes
 from pydantic import TypeAdapter, ValidationError
 from sqlalchemy.exc import StatementError
 
@@ -186,9 +187,9 @@ def tracked_stage(
     if runtime is None:
         scope = nullcontext()
     elif operation:
-        scope = runtime.operation(stage_name, attributes=attributes)
+        scope = scopes.operation(runtime, stage_name, attributes=attributes)
     else:
-        scope = runtime.span(stage_name, attributes=attributes)
+        scope = scopes.span(runtime, stage_name, attributes=attributes)
     with (
         tracking.progress.stage(stage_name, role=_simulation_role(attributes)),
         scope,
diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py
index fb651aa53..b908bbc84 100644
--- a/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py
+++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_causes.py
@@ -491,7 +491,10 @@ def span(self, name, *, attributes=None):
         yield
 
     @contextmanager
-    def operation(self, name, *, attributes=None):
+    def operation(self, name, *, attributes=None, remote_context=None):
+        # Mirrors ObservabilityRuntime.operation; a Stage 12 stage continues
+        # no remote trace.
+        assert remote_context is None
         self.scopes.append(("operation", name, attributes))
         yield
 

From 57e3267829c65ba9ab771e56f19f4f15188fb42c Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 21 Oct 2026 20:59:09 +0400
Subject: [PATCH 15/15] Test the UK weight-matrix check's TypeError for a
 dataset with no entity data

_require_certified_uk_weight_matrix raises TypeError("the '<name>' dataset
has no entity data loaded") when dataset.data is None or has no
entity_data mapping. Nothing tested that branch. The new test covers both
inputs; without the isinstance check they fail with AttributeError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../tests/test_uk_weight_matrix_guard.py      | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py
index 7242afaa2..0bc66cca7 100644
--- a/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py
+++ b/projects/policyengine-simulation-executor/tests/test_uk_weight_matrix_guard.py
@@ -172,3 +172,22 @@ def test_guard_rejects_a_year_the_matrix_does_not_cover(monkeypatch, tmp_path) -
 
     with pytest.raises(ValueError, match=r"has no weights for 2026 \(it covers 2025\)"):
         sr._require_certified_uk_weight_matrix(_STRATEGY, _dataset(3, 2026), _SELECTION)
+
+
+@pytest.mark.parametrize(
+    "unloaded_data",
+    [None, SimpleNamespace(tables={})],
+    ids=["no-data", "no-entity-data"],
+)
+def test_guard_rejects_a_dataset_with_no_entity_data(
+    monkeypatch, tmp_path, unloaded_data
+) -> None:
+    _bundle(monkeypatch)
+    _assets(monkeypatch, tmp_path)
+    dataset = SimpleNamespace(year=2025, data=unloaded_data)
+
+    with pytest.raises(
+        TypeError,
+        match=r"the 'enhanced_frs_2024_25' dataset has no entity data loaded",
+    ):
+        sr._require_certified_uk_weight_matrix(_STRATEGY, dataset, _SELECTION)