From 75388a210ac605b4c5342b95c07d589f87f4c262 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 b7c728e873f57a18146b39c3cfb9112c274c3f6a Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Mon, 26 Oct 2026 15:47:18 +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, and the Stage 12 fan-out and
comparison 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 output builders already open their spans through main's shared
output_span helper, which uses scopes.span, so they are unchanged. The
Stage 12 failure-cause scope already declared its type
(AbstractContextManager[object]) and is unchanged too. The executor's
baseline drops from 66 entries to 49, and docs/generated/rules.md is
rendered again.

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 ++++--
 .../stage12_runtime/comparison.py             |  3 +-
 .../stage12_runtime/fanout.py                 | 12 +++--
 rules/baselines/pyright-executor.txt          | 17 -------
 rules/registry.toml                           |  2 +-
 10 files changed, 59 insertions(+), 54 deletions(-)

diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 4268bb317..fc8464433 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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`: 66 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 49 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/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 91f8199eb..f8f57d72e 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 19720c3be..d5142f52d 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 = 66
+count = 49
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From f9a4212ae96316b68d905891bb14ee083f8e3c4f Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Mon, 26 Oct 2026 15:47:55 +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 49 entries to 41.

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 fc8464433..825930135 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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`: 49 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 41 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 1499ba3c3..ed71dec46 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
@@ -70,8 +70,11 @@ class BundleManifestMetadata(TypedDict):
     local_area_breakdown: 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
@@ -93,7 +96,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:
@@ -468,7 +475,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))
 
 
@@ -508,7 +516,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,
     }
 
 
@@ -822,9 +830,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
@@ -1060,8 +1071,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)
@@ -1384,7 +1396,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:")
@@ -1449,9 +1461,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:
@@ -1530,9 +1542,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 f8f57d72e..57d83d84c 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 d5142f52d..c09ef4743 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 = 49
+count = 41
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From c123ba1e8bfb5c397d9fd02a8989fc10e0a6df95 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Mon, 26 Oct 2026 15:49:32 +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, and the output
dataset after ensure() in _entity_data (10 errors).

output_frames adds a runtime-checkable protocol, EntityFrames. The
microdata export and the determinism check now check
`isinstance(<simulation>.output_dataset.data, EntityFrames)` first and
raise 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.
The check needs only the attribute, so the tests' stand-in datasets
still pass. The protocol types frames as pandas DataFrames, which a
MicroDataFrame is, since microdf types a column read as
MicroSeries | MicroDataFrame. _entity_data raises RuntimeError when
ensure() leaves no output dataset.

The executor's baseline drops from 41 entries to 31, and
docs/generated/rules.md is rendered again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/generated/rules.md                       |  2 +-
 .../output_frames.py                          | 36 +++++++++++++++++++
 .../precompute.py                             | 10 ++++++
 .../simulation_microdata.py                   |  7 ++++
 .../simulation_output_common.py               |  3 ++
 .../tests/test_output_frames.py               | 35 ++++++++++++++++++
 rules/baselines/pyright-executor.txt          | 10 ------
 rules/registry.toml                           |  2 +-
 8 files changed, 93 insertions(+), 12 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 825930135..8bd07f306 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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`: 41 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..963520e95
--- /dev/null
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/output_frames.py
@@ -0,0 +1,36 @@
+"""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 protocol is 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)."""
+        ...
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 ae5c73308..6b8ddb9e4 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
@@ -60,6 +60,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)
 
 
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..744e4fff5
--- /dev/null
+++ b/projects/policyengine-simulation-executor/tests/test_output_frames.py
@@ -0,0 +1,35 @@
+"""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
+
+
+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():
+    year_data = _uk_year_data()
+
+    assert isinstance(year_data, EntityFrames)
+
+
+def test_data_with_frames_by_attribute_matches_the_shape():
+    assert isinstance(SimpleNamespace(entity_data={}), EntityFrames)
+
+
+def test_missing_output_data_does_not_match_the_shape():
+    assert not isinstance(None, EntityFrames)
+    assert not isinstance(SimpleNamespace(household=None), EntityFrames)
diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt
index 57d83d84c..03ffc022c 100644
--- a/rules/baselines/pyright-executor.txt
+++ b/rules/baselines/pyright-executor.txt
@@ -14,17 +14,8 @@ 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__"
@@ -32,7 +23,6 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/s
 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/simulation_output_common.py :: reportOptionalMemberAccess :: "data" 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 c09ef4743..2e8ddc381 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 = 41
+count = 31
 [rule.proof]
 bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad"
 good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"

From 7d9d66867dd51ac8f9da449b688c40f850d84897 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 8bd07f306..260049334 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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 6b8ddb9e4..c4b579465 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
@@ -50,8 +50,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
@@ -118,7 +119,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,
@@ -132,6 +137,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)
 
 
@@ -142,7 +150,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,
@@ -157,6 +172,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 2e8ddc381..b9d43ca59 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 9b848c6628f4663c2b113fd42fe2cecc677b100a 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 260049334..945f527b7 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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 b9d43ca59..7de7d86f1 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 3b6830b5fd952fd4bbcc04ad725230fb000e16a9 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 945f527b7..c098af5ed 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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 728331273..2409ed33c 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 7de7d86f1..c7cc7fcd7 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 9cc433c825bad31a428c97761de910486dac7410 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 c098af5ed..a377b2caa 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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 c7cc7fcd7..5fecaf020 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 abe88d777dcd34d619faa69387c7facd016c6903 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 a377b2caa..701231032 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -67,7 +67,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`: 0 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 5fecaf020..dd19e0d58 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 272713d5d034fe91ff96c7f0af4ea868cd71309c 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 da3550a264bd8fe9d220bd175c8c9476c8466f7d Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Mon, 26 Oct 2026 15:50:46 +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.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../simulation_output_common.py               |  6 ++---
 .../tests/test_output_frames.py               | 24 +++++++++++++++++++
 2 files changed, 27 insertions(+), 3 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 c4b579465..eeb419290 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
@@ -61,9 +61,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/tests/test_output_frames.py b/projects/policyengine-simulation-executor/tests/test_output_frames.py
index 744e4fff5..d5e240427 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
+from policyengine_simulation_executor.simulation_output_common import _entity_data
 
 
 def _uk_year_data() -> UKYearData:
@@ -33,3 +35,25 @@ def test_data_with_frames_by_attribute_matches_the_shape():
 def test_missing_output_data_does_not_match_the_shape():
     assert not isinstance(None, EntityFrames)
     assert not isinstance(SimpleNamespace(household=None), EntityFrames)
+
+
+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 b237fe3bf816d53b2ffdc106aa135b91913db7e7 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Mon, 26 Oct 2026 15:51:23 +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 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               | 50 +++++++++++++++----
 1 file changed, 41 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 d5e240427..d229ab863 100644
--- a/projects/policyengine-simulation-executor/tests/test_output_frames.py
+++ b/projects/policyengine-simulation-executor/tests/test_output_frames.py
@@ -5,29 +5,61 @@
 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
 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 _us_year_data() -> USYearData:
+    return USYearData(
+        **{
+            entity: _frame(entity)
+            for entity in (
+                "person",
+                "marital_unit",
+                "family",
+                "spm_unit",
+                "tax_unit",
+                "household",
+            )
+        }
     )
 
 
-def test_country_year_data_has_entity_frames():
-    year_data = _uk_year_data()
+@pytest.mark.parametrize(
+    "year_data_factory", [_uk_year_data, _us_year_data], ids=["UK", "US"]
+)
+def test_country_year_data_has_entity_frames(year_data_factory):
+    year_data = year_data_factory()
 
     assert isinstance(year_data, EntityFrames)
 
 
+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_shape():
     assert isinstance(SimpleNamespace(entity_data={}), EntityFrames)
 

From 3334f2135bb25a62c819995a80529b7915417bd0 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 1ce8d6153f94121fa721ce095a27df46030646c8 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 83ad566c409e114cfce0398beda215e573bb7b04 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Mon, 26 Oct 2026 16:01:23 +0400
Subject: [PATCH 15/15] Import pandas at runtime in the output helpers again

_sum_output_variable and _change_output_variable check a column read
with isinstance(column, pd.Series), so the module needs pandas at
runtime. Main had moved `import pandas as pd` under TYPE_CHECKING, since
nothing there used it at runtime any more, and after this branch was
rebased onto main both helpers raised NameError: name 'pd' is not
defined whenever the variable was already a column of the output frame.
Pyright reads TYPE_CHECKING as true, so it did not report it; ruff
reported it as TC004.

The existing tests caught it: before this change 11 executor tests
failed with that NameError, in test_simulation_output_budget.py,
test_simulation_output_builder.py and
test_simulation_output_household_bounds.py. They pass now.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../simulation_output_common.py                                 | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 eeb419290..ab6ae10b1 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
@@ -8,10 +8,10 @@
 from importlib import import_module
 from typing import TYPE_CHECKING, Any
 
+import pandas as pd
 from policyengine_simulation_contract.spm import spm_error_detail
 
 if TYPE_CHECKING:
-    import pandas as pd
     from _typeshed import ConvertibleToFloat
     from policyengine.core import Simulation