From e1015eabf29ce10c438092f5df618d98e045b188 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:51:45 +0400
Subject: [PATCH 01/12] Annotate policyengine-fastapi for ruff ANN

Add the missing parameter and return annotations. AppExit forwards its
callback arguments as object rather than Any, the JWT decoder's private
validator returns the same dict type its caller declares, and the
metrics middleware types call_next as Starlette's RequestResponseEndpoint.
Ping endpoints already pass response_model explicitly, so their new
return annotations do not change the generated responses.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../src/policyengine_fastapi/auth/jwt_decoder.py   |  4 ++--
 .../src/policyengine_fastapi/database.py           |  2 +-
 .../src/policyengine_fastapi/exit.py               | 10 +++++-----
 .../src/policyengine_fastapi/health/__init__.py    |  6 +++---
 .../policyengine_fastapi/observability/config.py   |  3 ++-
 .../observability/correlation.py                   |  3 +--
 .../policyengine_fastapi/observability/emitters.py | 14 ++++++++++----
 .../policyengine_fastapi/opentelemetry/console.py  |  4 ++--
 .../src/policyengine_fastapi/opentelemetry/gcp.py  | 12 ++++++------
 .../opentelemetry/instrumentor.py                  |  2 +-
 .../opentelemetry/middleware.py                    | 10 ++++++----
 .../src/policyengine_fastapi/ping/__init__.py      |  6 +++---
 12 files changed, 42 insertions(+), 34 deletions(-)

diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py
index 07331be90..a8ec9b929 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/auth/jwt_decoder.py
@@ -16,7 +16,7 @@ class JWTDecoder:
     the parsed token with all fields.
     """
 
-    def __init__(self, issuer: str, audience: str, auto_error: bool = True):
+    def __init__(self, issuer: str, audience: str, auto_error: bool = True) -> None:
         self.issuer = issuer
         self.audience = audience
         self.auto_error = auto_error
@@ -38,7 +38,7 @@ def __call__(
             LOG.info("Ignoring invalid auth token")
         return None
 
-    def _validate(self, token: HTTPAuthorizationCredentials | None):
+    def _validate(self, token: HTTPAuthorizationCredentials | None) -> dict[str, str]:
         if token is None:
             LOG.info("No bearer token on request")
             raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py b/libs/policyengine-fastapi/src/policyengine_fastapi/database.py
index e5569f372..d1340a26e 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/database.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/database.py
@@ -30,7 +30,7 @@ def create_session_dep(engine: Engine) -> SessionGeneratorFactory:
     routers.
     """
 
-    def session_dep():
+    def session_dep() -> Generator[Session]:
         with Session(engine) as session:
             yield session
 
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py b/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py
index e2821e764..5f7cde3c6 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/exit.py
@@ -4,7 +4,7 @@
 """
 
 import logging
-from collections.abc import Callable
+from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from typing import Any, TypeVar
 
@@ -14,24 +14,24 @@
 
 
 class AppExit:
-    def __init__(self):
+    def __init__(self) -> None:
         self.callbacks: list[Callable[[], None]] = []
 
-    def __call__(self, *args: Any, **kwargs: Any):
+    def __call__(self, *args: object, **kwargs: object) -> Callable[[T], T]:
         def accept(func: T) -> T:
             self.callbacks.append(lambda: func(*args, **kwargs))
             return func
 
         return accept
 
-    def _exit(self):
+    def _exit(self) -> None:
         log.info("Executing exit callbacks")
         for cb in self.callbacks:
             cb()
         log.info("Exit callbacks executed")
 
     @asynccontextmanager
-    async def lifespan(self):
+    async def lifespan(self) -> AsyncIterator[None]:
         yield
         self._exit()
 
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py
index 464826daf..66c9c64bd 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/health/__init__.py
@@ -24,7 +24,7 @@ class HealthStatus(BaseModel):
 
 
 class HealthSystemReporter:
-    def __init__(self, name: str, probes: dict[str, HealthProbe]):
+    def __init__(self, name: str, probes: dict[str, HealthProbe]) -> None:
         self.name = name
         self.probes = probes
 
@@ -45,10 +45,10 @@ class HealthRegistry:
 
     systems: dict[str, HealthSystemReporter]
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.systems = {}
 
-    def register(self, reporter: HealthSystemReporter):
+    def register(self, reporter: HealthSystemReporter) -> None:
         self.systems[reporter.name] = reporter
 
     def report(self) -> HealthStatus:
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py
index fbe58ebd3..465cf198a 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/config.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from typing import Self
 
 from .stages import TracerCaptureMode
 
@@ -44,5 +45,5 @@ class ObservabilityConfig:
     slow_run_threshold_seconds: float = 30.0
 
     @classmethod
-    def disabled(cls, service_name: str = "policyengine-observability"):
+    def disabled(cls, service_name: str = "policyengine-observability") -> Self:
         return cls(enabled=False, service_name=service_name)
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py
index 1dc46157e..4282d0f3b 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/correlation.py
@@ -2,7 +2,6 @@
 
 import json
 from hashlib import sha256
-from typing import Any
 from uuid import uuid4
 
 
@@ -10,7 +9,7 @@ def generate_observability_id() -> str:
     return str(uuid4())
 
 
-def stable_config_hash(payload: Any) -> str:
+def stable_config_hash(payload: object) -> str:
     serialised = json.dumps(
         payload,
         sort_keys=True,
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py
index 3a06f307d..05d4d7d3d 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/observability/emitters.py
@@ -2,24 +2,30 @@
 
 from contextlib import AbstractContextManager
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, Self
 
 from .config import ObservabilityConfig
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
+    from types import TracebackType
 
     from .contracts import SimulationLifecycleEvent, TracerArtifactManifest
 
 
 class NoOpSpan(AbstractContextManager["NoOpSpan"]):
-    def __enter__(self):
+    def __enter__(self) -> Self:
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
         return None
 
-    def set_attribute(self, key: str, value: Any) -> None:
+    def set_attribute(self, key: str, value: object) -> None:
         return None
 
     def add_event(self, name: str, attributes: Mapping[str, Any] | None = None) -> None:
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py
index dc265d56e..31b21b294 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/console.py
@@ -19,7 +19,7 @@
 # Configure opentelemetry
 # 1. to include python logs and
 # 2. export to console (for demo purposes)
-def export_ot_to_console(resource: Resource):
+def export_ot_to_console(resource: Resource) -> None:
     """
     configure opentelemetry to dump messages to console for debugging on desktop
     """
@@ -33,6 +33,6 @@ def export_ot_to_console(resource: Resource):
     metrics.set_meter_provider(meterProvider)
 
     @exit()
-    def flush_ot():
+    def flush_ot() -> None:
         traceProvider.force_flush()
         meterProvider.force_flush()
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py
index 968fed712..91f40ed8f 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/gcp.py
@@ -23,11 +23,11 @@
 
 
 class AddGcpProjectToTraceFilter(logging.Filter):
-    def __init__(self, project: str):
+    def __init__(self, project: str) -> None:
         super().__init__()
         self.project = project
 
-    def filter(self, record: logging.LogRecord):
+    def filter(self, record: logging.LogRecord) -> bool | logging.LogRecord:
         if not hasattr(record, "otelTraceID"):
             return super().filter(record)
         record.otelTraceID = f"projects/{self.project}/traces/{record.otelTraceID}"  # type: ignore
@@ -58,7 +58,7 @@ def _get_project_id_from_metadata() -> str | None:
 
 
 class GCPLoggingInstrumentor:
-    def __init__(self, project_id: str | None = None):
+    def __init__(self, project_id: str | None = None) -> None:
         self.project_id = project_id or _get_project_id()
 
     """
@@ -66,7 +66,7 @@ def __init__(self, project_id: str | None = None):
     json logs as per gcp expectations.
     """
 
-    def instrument(self):
+    def instrument(self) -> None:
         from opentelemetry.instrumentation.logging import LoggingInstrumentor
         from pythonjsonlogger.json import JsonFormatter
 
@@ -93,7 +93,7 @@ def instrument(self):
         LoggingInstrumentor().instrument()
 
 
-def export_ot_to_gcp(resource: Resource):
+def export_ot_to_gcp(resource: Resource) -> None:
     from opentelemetry.exporter.cloud_monitoring import (
         CloudMonitoringMetricsExporter,
     )
@@ -114,7 +114,7 @@ def export_ot_to_gcp(resource: Resource):
     metrics.set_meter_provider(meterProvider)
 
     @exit()
-    def flush_ot():
+    def flush_ot() -> None:
         log.info("Flushing traces and metrics before shutdown")
         traceProvider.force_flush()
         meterProvider.force_flush()
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py
index c91265317..aab1ce0aa 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/instrumentor.py
@@ -10,7 +10,7 @@ class FastAPIEnhancedInstrumenter:
     metrics instead of global for the whole api.
     """
 
-    def instrument(self, app: FastAPI):
+    def instrument(self, app: FastAPI) -> None:
         FastAPIInstrumentor.instrument_app(app)
         middleware = Middleware(app)
         app.middleware("http")(middleware)
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py
index e39ea23e0..f741b51f9 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/opentelemetry/middleware.py
@@ -1,10 +1,10 @@
 from dataclasses import dataclass
 from time import time
-from typing import Any
 
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Response
 from fastapi.routing import APIRoute
 from opentelemetry import metrics
+from starlette.middleware.base import RequestResponseEndpoint
 from starlette.routing import Match
 
 
@@ -16,7 +16,7 @@ class Route:
 
 
 class Middleware:
-    def __init__(self, app: FastAPI):
+    def __init__(self, app: FastAPI) -> None:
         self.app = app
         self.meter = metrics.get_meter("policyengine")
 
@@ -32,7 +32,9 @@ def _create_route(self, route: APIRoute) -> Route:
             latency=self.meter.create_histogram(f"{name}_latency"),
         )
 
-    async def __call__(self, request: Request, call_next) -> Any:
+    async def __call__(
+        self, request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
         route: Route | None = None
         for r in self.routes:
             if r.route.matches(request.scope)[0] != Match.NONE:
diff --git a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py
index bd53cb5ac..ce927959e 100644
--- a/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py
+++ b/libs/policyengine-fastapi/src/policyengine_fastapi/ping/__init__.py
@@ -24,14 +24,14 @@ async def ping(request: PingRequest) -> PingResponse:
         return PingResponse(incremented=request.value + 1)
 
     @router.get("/ping/started", response_model=str)
-    async def started():
+    async def started() -> str:
         """
         Verify the api is running.
         """
         return "alive"
 
     @router.get("/ping/alive", response_model=HealthStatus)
-    async def alive():
+    async def alive() -> JSONResponse:
         """
         Check if the service is healthy. This will always return
         data, but will return a 503 response code if the service is
@@ -46,5 +46,5 @@ async def alive():
     return router
 
 
-def include_all_routers(api: FastAPI, health_registry: HealthRegistry):
+def include_all_routers(api: FastAPI, health_registry: HealthRegistry) -> None:
     api.include_router(create_router(health_registry))

From f066dd38305f836f96184a7f1fada7743bbfbd3f Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:51:45 +0400
Subject: [PATCH 02/12] Annotate the simulation contract library for ruff ANN

Type the SPM resolution and validation entry points, the gateway request
pre-validators (object in, object out) and the Modal Dict store helpers.
The three wrap model serializers in spm.py stay unannotated with a
reasoned noqa: Pydantic turns a serializer's return annotation into the
model's serialization JSON schema, which would rewrite the OpenAPI spec.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../budget_window_state.py                    |  4 +-
 .../gateway_models.py                         | 12 ++---
 .../policyengine_simulation_contract/spm.py   | 44 ++++++++++---------
 .../stage12_manifest.py                       |  2 +-
 4 files changed, 32 insertions(+), 30 deletions(-)

diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py
index 85d3cddf4..1e396ca0b 100644
--- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py
+++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py
@@ -25,11 +25,11 @@
 BUDGET_WINDOW_JOB_SEED_DICT_NAME = "simulation-api-budget-window-job-seeds"
 
 
-def _budget_window_job_store():
+def _budget_window_job_store() -> modal.Dict:
     return modal.Dict.from_name(BUDGET_WINDOW_JOB_DICT_NAME, create_if_missing=True)
 
 
-def _budget_window_job_seed_store():
+def _budget_window_job_seed_store() -> modal.Dict:
     return modal.Dict.from_name(
         BUDGET_WINDOW_JOB_SEED_DICT_NAME, create_if_missing=True
     )
diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py
index f4580ed20..83e9a6ba8 100644
--- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py
+++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py
@@ -41,7 +41,7 @@ class PolicyParameterChanges(RootModel[JsonObject]):
     """Policy parameter names mapped to their dated values."""
 
 
-def _move_internal_telemetry_alias(value):
+def _move_internal_telemetry_alias(value: object) -> object:
     if not isinstance(value, dict):
         return value
     if "_telemetry" in value and "telemetry" not in value:
@@ -50,7 +50,7 @@ def _move_internal_telemetry_alias(value):
     return value
 
 
-def _strip_internal_passthrough_fields(value):
+def _strip_internal_passthrough_fields(value: object) -> object:
     """Drop internal-only fields the gateway adds to payloads in flight.
 
     Parent batch entrypoints attach ``_metadata`` describing resolved routing,
@@ -71,7 +71,7 @@ def _strip_internal_passthrough_fields(value):
     }
 
 
-def _enforce_max_payload_size(value):
+def _enforce_max_payload_size(value: object) -> object:
     if not isinstance(value, dict):
         return value
     try:
@@ -126,17 +126,17 @@ class GatewayRequestBase(BaseModel):
 
     @model_validator(mode="before")
     @classmethod
-    def move_internal_telemetry_alias(cls, value):
+    def move_internal_telemetry_alias(cls, value: object) -> object:
         return _move_internal_telemetry_alias(value)
 
     @model_validator(mode="before")
     @classmethod
-    def strip_internal_passthrough_fields(cls, value):
+    def strip_internal_passthrough_fields(cls, value: object) -> object:
         return _strip_internal_passthrough_fields(value)
 
     @model_validator(mode="before")
     @classmethod
-    def enforce_max_payload_size(cls, value):
+    def enforce_max_payload_size(cls, value: object) -> object:
         return _enforce_max_payload_size(value)
 
 
diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py
index 135bed7b8..160175fa6 100644
--- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py
+++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py
@@ -1,9 +1,9 @@
 """Dependency-light public models for US SPM selection and calculation receipts."""
 
 import re
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from datetime import date
-from typing import Any, Literal
+from typing import Any, Literal, Self
 
 from pydantic import (
     BaseModel,
@@ -17,8 +17,10 @@
     model_validator,
 )
 
+from policyengine_simulation_contract.json_types import JsonObject
 
-def _selection_schema(schema):
+
+def _selection_schema(schema: dict[str, Any]) -> None:
     # An omitted option inherits the selected bundle's default. Advertising
     # Python attribute defaults here makes generated clients send an explicit
     # county selection even when their caller omitted geography entirely.
@@ -45,7 +47,7 @@ class SPMSelection(BaseModel):
     as_of: str | None = None
 
     @model_serializer(mode="wrap")
-    def serialize_selection(
+    def serialize_selection(  # noqa: ANN201  # a return annotation would become the serialization JSON schema
         self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
     ):
         """Preserve explicitly selected options through ordinary and nested JSON.
@@ -75,14 +77,14 @@ def serialize_selection(
 
     @field_validator("as_of")
     @classmethod
-    def validate_as_of(cls, value):
+    def validate_as_of(cls, value: str | None) -> str | None:
         if value is not None:
             if date.fromisoformat(value).isoformat() != value:
                 raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)")
         return value
 
     @model_validator(mode="after")
-    def validate_location(self):
+    def validate_location(self) -> Self:
         # Request options inherit bundle defaults. Validate coupled options
         # once both are explicit (including after defaults are resolved).
         if not {"geography_kind", "geography_id"} <= self.model_fields_set:
@@ -108,7 +110,7 @@ class SPMResolvedConfiguration(BaseModel):
     as_of: str | None
 
     @model_serializer(mode="wrap")
-    def serialize_resolved_configuration(
+    def serialize_resolved_configuration(  # noqa: ANN201  # a return annotation would become the serialization JSON schema
         self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
     ):
         """Keep required nullable fields in completed response bodies."""
@@ -158,7 +160,7 @@ class SPMProvenance(BaseModel):
     runtime_versions: "SPMRuntimeVersions"
 
     @model_serializer(mode="wrap")
-    def serialize_provenance(
+    def serialize_provenance(  # noqa: ANN201  # a return annotation would become the serialization JSON schema
         self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
     ):
         """Keep required nullable configuration fields in response bodies."""
@@ -325,15 +327,15 @@ class SPMErrorDetail(BaseModel):
 class SPMInputError(ValueError):
     """Transportable error shared by the control plane and worker."""
 
-    def __init__(self, code: str, message: str):
+    def __init__(self, code: str, message: str) -> None:
         self.code = code
         self.message = message
         super().__init__(code, message)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.message
 
-    def to_dict(self):
+    def to_dict(self) -> dict[str, str]:
         return {"code": self.code, "message": self.message}
 
 
@@ -356,7 +358,7 @@ class SPMCapability(BaseModel):
     defaults: SPMSelection
 
     @model_validator(mode="after")
-    def pinned_defaults(self):
+    def pinned_defaults(self) -> Self:
         if not self.defaults.forecast_content_sha256 or not self.defaults.scenario:
             raise ValueError(
                 "Certified SPM defaults must pin artifact hash and scenario"
@@ -423,14 +425,14 @@ def build_spm_comparison_provenance(
 
 
 def resolve_spm_selection(
-    country,
-    selection,
+    country: str,
+    selection: SPMSelection | Mapping[str, object] | None,
     *,
-    capability,
-    policyengine_version,
-    model_version,
-    route_provenance=None,
-):
+    capability: SPMCapability | Mapping[str, object] | None,
+    policyengine_version: str | None,
+    model_version: str | None,
+    route_provenance: str | None = None,
+) -> JsonObject | None:
     """Resolve only certified metadata; unrecognized future bundles fail closed."""
     if country.lower() != "us":
         if selection is not None:
@@ -500,8 +502,8 @@ def resolve_spm_selection(
 
 
 def validate_spm_result(
-    result: dict, selection: Any, *, expected_year: int | str | None = None
-):
+    result: dict, selection: object, *, expected_year: int | str | None = None
+) -> SPMComparisonProvenance | None:
     """Do not accept incomplete or mixed-method child/cached output."""
     if "spm_config" in result:
         raise SPMInputError(
diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py
index a048b8f3a..98178064d 100644
--- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py
+++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py
@@ -186,7 +186,7 @@ def publish_v2_manifest(
 class V2ManifestLoader:
     """Load v2 state while retaining only this loader's last valid v2 copy."""
 
-    def __init__(self, store: V2ManifestStore):
+    def __init__(self, store: V2ManifestStore) -> None:
         self._store = store
         self._last_valid: V2VersionManifest | None = None
 

From 501e968dbfcb165b70538f3fa53dcee25c655d3e Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:51:45 +0400
Subject: [PATCH 03/12] Annotate the observability and Stage 12 persistence
 libraries for ruff ANN

Replace Any with object for values that are only inspected or passed
through, and type the persistence row helpers with SQLAlchemy's
RowMapping, Table and ColumnElement (imported for type checking only).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../identifiers.py                                   |  5 ++---
 .../telemetry.py                                     |  2 +-
 .../src/policyengine_stage12_persistence/store.py    | 12 ++++++++----
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py
index dfe54d716..e60ded480 100644
--- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py
+++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/identifiers.py
@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from typing import Any
 from uuid import UUID, uuid4
 
 OBSERVABILITY_ID_HEADER = "X-PolicyEngine-Observability-Id"
@@ -14,13 +13,13 @@ def generate_observability_id() -> str:
     return str(uuid4())
 
 
-def resolve_observability_id(value: Any) -> str:
+def resolve_observability_id(value: object) -> str:
     """Use a valid caller value or create a new diagnostic identifier."""
 
     return normalize_observability_id(value) or generate_observability_id()
 
 
-def normalize_observability_id(value: Any) -> str | None:
+def normalize_observability_id(value: object) -> str | None:
     """Return a canonical UUID string or ``None`` for malformed input.
 
     Observability metadata is best effort. Callers must never reject or alter a
diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py
index 5e1886752..26015dc0a 100644
--- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py
+++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/telemetry.py
@@ -60,7 +60,7 @@ class TelemetryEnvelope(BaseModel):
 
     @model_validator(mode="before")
     @classmethod
-    def discard_noncanonical_context_fields(cls, value: Any) -> Any:
+    def discard_noncanonical_context_fields(cls, value: object) -> object:
         """Ignore old body fields without using them as correlation inputs.
 
         This keeps the receiver compatible with the deployed API during the
diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py
index fb6ca6e32..031a81330 100644
--- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py
+++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py
@@ -45,6 +45,8 @@
     from datetime import datetime
     from uuid import UUID
 
+    from sqlalchemy import ColumnElement, RowMapping, Table
+
 REPORT_IDENTITY_COLUMNS = (
     "environment",
     "calculation_flow",
@@ -155,7 +157,7 @@ def open_connection_once(
     return connect()
 
 
-def _database_value(value: Any) -> Any:
+def _database_value(value: object) -> object:
     if isinstance(value, StrEnum):
         return value.value
     if isinstance(value, tuple):
@@ -168,15 +170,17 @@ def _record_values(record: object) -> dict[str, Any]:
     return {name: _database_value(value) for name, value in dumped.items()}
 
 
-def _report(row: Any) -> ComparisonReportRecord:
+def _report(row: RowMapping) -> ComparisonReportRecord:
     return ComparisonReportRecord.model_validate(dict(row))
 
 
-def _simulation(row: Any) -> ComparisonSimulationRecord:
+def _simulation(row: RowMapping) -> ComparisonSimulationRecord:
     return ComparisonSimulationRecord.model_validate(dict(row))
 
 
-def _identity_predicates(table: Any, values: dict[str, Any], columns: tuple[str, ...]):
+def _identity_predicates(
+    table: Table, values: dict[str, Any], columns: tuple[str, ...]
+) -> tuple[ColumnElement[bool], ...]:
     return tuple(table.c[name] == values[name] for name in columns)
 
 

From 5ec4eae5b591a81d351498df8cf8f378325041e7 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:51:45 +0400
Subject: [PATCH 04/12] Annotate the simulation gateway for ruff ANN

Every annotated route passes response_model explicitly, so FastAPI keeps
using it and the OpenAPI golden file is unchanged. _bundle_payload keeps
Any for **dump_kwargs, forwarded unchanged to BaseModel.model_dump, with
a reasoned noqa. Test fixtures outside tests/ are annotated too.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../fixtures/gateway_endpoints.py             | 32 ++++++++++---------
 .../fixtures/package_imports.py               |  9 ++++--
 .../policyengine_simulation_gateway/app.py    |  8 ++++-
 .../correlation.py                            |  7 ++--
 .../endpoints.py                              | 32 +++++++++++++------
 .../generate_openapi.py                       |  2 +-
 .../responses.py                              |  4 +--
 .../smoke_app.py                              |  2 +-
 8 files changed, 61 insertions(+), 35 deletions(-)

diff --git a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py
index f5ac439a0..3866aff93 100644
--- a/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py
+++ b/projects/policyengine-simulation-gateway/fixtures/gateway_endpoints.py
@@ -1,6 +1,8 @@
 """Fixtures for gateway endpoint tests."""
 
+from collections.abc import ItemsView, KeysView
 from copy import deepcopy
+from typing import Self
 
 import pytest
 from policyengine_simulation_contract.macro_output import (
@@ -159,28 +161,28 @@ def resolve_test_dataset_uri(
 class MockDict:
     """Mock for Modal.Dict to simulate version registry."""
 
-    def __init__(self, data: dict):
+    def __init__(self, data: dict) -> None:
         self._data = data
 
-    def __getitem__(self, key: str):
+    def __getitem__(self, key: str) -> object:
         if key not in self._data:
             raise KeyError(key)
         return self._data[key]
 
-    def __setitem__(self, key: str, value):
+    def __setitem__(self, key: str, value: object) -> None:
         self._data[key] = value
 
-    def get(self, key: str, default=None):
+    def get(self, key: str, default: object = None) -> object:
         return self._data.get(key, default)
 
-    def keys(self):
+    def keys(self) -> KeysView[str]:
         return self._data.keys()
 
-    def items(self):
+    def items(self) -> ItemsView[str, object]:
         return self._data.items()
 
     @classmethod
-    def from_name(cls, name: str):
+    def from_name(cls, name: str) -> Self:
         """Mock from_name that returns a MockDict based on name."""
         raise NotImplementedError("Mock not configured")
 
@@ -191,7 +193,7 @@ class MockFunctionCall:
     registry = {}
     from_id_errors = {}
 
-    def __init__(self, object_id: str = "mock-job-id-123"):
+    def __init__(self, object_id: str = "mock-job-id-123") -> None:
         self.object_id = object_id
         value = BaselineReformValue(baseline=0.0, reform=0.0)
         age_poverty = AgePovertyOutput(
@@ -241,7 +243,7 @@ def __init__(self, object_id: str = "mock-job-id-123"):
         self.running = False
         self.__class__.registry[object_id] = self
 
-    def get(self, timeout: int = 0):
+    def get(self, timeout: int = 0) -> object:
         if self.running:
             raise TimeoutError
         if self.error is not None:
@@ -249,7 +251,7 @@ def get(self, timeout: int = 0):
         return self.result
 
     @classmethod
-    def from_id(cls, object_id: str):
+    def from_id(cls, object_id: str) -> "MockFunctionCall":
         if object_id in cls.from_id_errors:
             raise cls.from_id_errors[object_id]
         if object_id not in cls.registry:
@@ -260,7 +262,7 @@ def from_id(cls, object_id: str):
 class MockFunction:
     """Mock for Modal Function."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.last_payload = None
         self.last_observability_context = None
         self.last_from_name_call = None
@@ -279,7 +281,7 @@ def call_for(self, object_id: str) -> MockFunctionCall:
 class BoundMockFunction:
     """Function handle returned by Modal.Function.from_name."""
 
-    def __init__(self, recorder: MockFunction, app_name: str, func_name: str):
+    def __init__(self, recorder: MockFunction, app_name: str, func_name: str) -> None:
         self.recorder = recorder
         self.app_name = app_name
         self.func_name = func_name
@@ -318,7 +320,7 @@ class OutputExpiredError(Exception):
 
 
 @pytest.fixture
-def mock_modal(monkeypatch):
+def mock_modal(monkeypatch: pytest.MonkeyPatch) -> dict[str, object]:
     """Patch Modal calls in the gateway endpoints module."""
     from policyengine_simulation_contract import budget_window_state, dataset_uri
     from policyengine_simulation_gateway import endpoints
@@ -333,7 +335,7 @@ def mock_modal(monkeypatch):
 
     class MockModalDict:
         @staticmethod
-        def from_name(name: str, create_if_missing: bool = False):
+        def from_name(name: str, create_if_missing: bool = False) -> MockDict:
             if create_if_missing and name not in mock_dicts:
                 mock_dicts[name] = {}
             if name not in mock_dicts:
@@ -342,7 +344,7 @@ def from_name(name: str, create_if_missing: bool = False):
 
     class MockModalFunction:
         @staticmethod
-        def from_name(app_name: str, func_name: str):
+        def from_name(app_name: str, func_name: str) -> BoundMockFunction:
             mock_func.last_from_name_call = (app_name, func_name)
             return mock_func.bind(app_name, func_name)
 
diff --git a/projects/policyengine-simulation-gateway/fixtures/package_imports.py b/projects/policyengine-simulation-gateway/fixtures/package_imports.py
index b6769a469..d508167d1 100644
--- a/projects/policyengine-simulation-gateway/fixtures/package_imports.py
+++ b/projects/policyengine-simulation-gateway/fixtures/package_imports.py
@@ -10,7 +10,8 @@
 import pytest
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Callable, Iterator
+    from types import ModuleType
 
 GATEWAY_MODEL_MODULE = "policyengine_simulation_contract.gateway_models"
 GATEWAY_ENDPOINTS_MODULE = "policyengine_simulation_gateway.endpoints"
@@ -61,10 +62,12 @@ def isolated_gateway_model_import_modules() -> Iterator[None]:
 
 
 @pytest.fixture()
-def import_gateway_models(isolated_gateway_model_import_modules):
+def import_gateway_models(
+    isolated_gateway_model_import_modules: None,
+) -> Callable[[], ModuleType]:
     """Import gateway models from a clean module state."""
 
-    def import_models():
+    def import_models() -> ModuleType:
         return importlib.import_module(GATEWAY_MODEL_MODULE)
 
     return import_models
diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py
index 2b4535033..ee05f39a1 100644
--- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py
+++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/app.py
@@ -8,14 +8,20 @@
 and spawns jobs on those apps.
 """
 
+from __future__ import annotations
+
 import os
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import modal
 from policyengine_simulation_observability.observability import (
     modal_image_environment,
 )
 
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
 # Stable app name - this should rarely change
 app = modal.App("policyengine-simulation-gateway")
 gateway_auth_secret = modal.Secret.from_name("policyengine-gateway-auth")
@@ -64,7 +70,7 @@ def build_gateway_image() -> modal.Image:
 
 @app.function(image=gateway_image, secrets=[gateway_auth_secret])
 @modal.asgi_app()
-def web_app():
+def web_app() -> FastAPI:
     """
     FastAPI gateway for simulation job submission and polling.
 
diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py
index 620e7f774..c5e30a5c4 100644
--- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py
+++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/correlation.py
@@ -10,14 +10,17 @@
 )
 
 if TYPE_CHECKING:
-    from fastapi import FastAPI, Request
+    from fastapi import FastAPI, Request, Response
+    from starlette.middleware.base import RequestResponseEndpoint
 
 
 def install_observability_id_middleware(app: FastAPI) -> None:
     """Store an incoming identifier candidate without creating a workflow."""
 
     @app.middleware("http")
-    async def observability_id_transport(request: Request, call_next):
+    async def observability_id_transport(
+        request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
         request.state.incoming_observability_id = normalize_observability_id(
             request.headers.get(OBSERVABILITY_ID_HEADER)
         )
diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py
index d23832d98..6fd6796ac 100644
--- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py
+++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py
@@ -5,10 +5,11 @@
 import logging
 import re
 from dataclasses import dataclass
-from typing import TypedDict
+from typing import Any, TypedDict
 
 import modal
 from fastapi import APIRouter, Depends, HTTPException, Request, Response
+from fastapi.responses import JSONResponse
 from policyengine_observability import ObservabilityRuntime
 from policyengine_simulation_contract.budget_window_state import (
     build_batch_status_response,
@@ -23,6 +24,7 @@
     BudgetWindowBatchRequest,
     BudgetWindowBatchStatusResponse,
     BudgetWindowBatchSubmitResponse,
+    GatewayRequestBase,
     HealthResponse,
     JobStatusResponse,
     JobSubmitResponse,
@@ -36,6 +38,7 @@
 from policyengine_simulation_contract.hf_dataset import (
     HuggingFaceDatasetReferenceError,
 )
+from policyengine_simulation_contract.json_types import JsonObject
 from policyengine_simulation_contract.spm import (
     SPMCapability,
     SPMInputError,
@@ -91,7 +94,7 @@ class RouteResolution:
     country_model_version: str | None = None
 
 
-def _job_metadata_store():
+def _job_metadata_store() -> "modal.Dict":
     return modal.Dict.from_name(JOB_METADATA_DICT_NAME, create_if_missing=True)
 
 
@@ -158,7 +161,7 @@ def _resolve_dataset_uri_from_app_bundle(
     )
 
 
-def _modal_exception_class(name: str):
+def _modal_exception_class(name: str) -> type[BaseException] | None:
     exception_module = getattr(modal, "exception", None)
     if exception_module is None:
         return None
@@ -176,7 +179,7 @@ def _is_modal_job_not_found(exc: BaseException) -> bool:
     )
 
 
-def _optional_modal_dict(name: str):
+def _optional_modal_dict(name: str) -> "modal.Dict | None":
     try:
         return modal.Dict.from_name(name)
     except KeyError:
@@ -629,7 +632,11 @@ def _certified_model_version(country: str, route: RouteResolution) -> str | None
     return route.country_model_version
 
 
-def _resolve_request_spm(request, bundle, route):
+def _resolve_request_spm(
+    request: GatewayRequestBase,
+    bundle: PolicyEngineBundle,
+    route: RouteResolution,
+) -> JsonObject | None:
     selection = resolve_spm_selection(
         request.country,
         request.spm,
@@ -645,7 +652,10 @@ def _resolve_request_spm(request, bundle, route):
     return selection
 
 
-def _bundle_payload(bundle: PolicyEngineBundle, **dump_kwargs) -> dict:
+def _bundle_payload(
+    bundle: PolicyEngineBundle,
+    **dump_kwargs: Any,  # noqa: ANN401  # forwarded unchanged to BaseModel.model_dump
+) -> dict:
     """Dump a bundle, omitting fields added after the legacy body when unset.
 
     The 202 and 500 job bodies splat this dict in directly, bypassing the
@@ -778,7 +788,7 @@ async def submit_simulation(
     request: SimulationRequest,
     http_request: Request,
     response: Response,
-):
+) -> JobSubmitResponse | JSONResponse:
     """
     Submit a simulation job.
 
@@ -888,7 +898,7 @@ async def submit_budget_window_batch(
     request: BudgetWindowBatchRequest,
     http_request: Request,
     response: Response,
-):
+) -> BudgetWindowBatchSubmitResponse | JSONResponse:
     """
     Submit a budget-window batch job.
     """
@@ -981,7 +991,9 @@ async def submit_budget_window_batch(
     response_model_exclude_none=True,
     dependencies=[Depends(require_auth)],
 )
-async def get_job_status(job_id: str, request: Request, response: Response):
+async def get_job_status(
+    job_id: str, request: Request, response: Response
+) -> JobStatusResponse | JSONResponse:
     """
     Poll for job status.
 
@@ -1064,7 +1076,7 @@ async def get_budget_window_job_status(
     batch_job_id: str,
     request: Request,
     response: Response,
-):
+) -> JSONResponse:
     """
     Poll for budget-window batch status.
     """
diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py
index 7fa67f982..207e5d76d 100644
--- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py
+++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/generate_openapi.py
@@ -137,7 +137,7 @@ async def ping(request: PingRequest) -> PingResponse:
     return app
 
 
-def main():
+def main() -> None:
     """Generate OpenAPI spec and write to artifacts."""
     app = create_openapi_app()
     openapi_spec = app.openapi()
diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py
index c658c0e32..79428a456 100644
--- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py
+++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py
@@ -15,7 +15,7 @@
 class ServerErrorResponse(JSONResponse):
     """Shared 500 JSON response."""
 
-    def __init__(self, content: dict):
+    def __init__(self, content: dict) -> None:
         super().__init__(status_code=500, content=content)
 
 
@@ -38,7 +38,7 @@ def batch_status_response(
     response: BudgetWindowBatchStatusResponse,
     *,
     headers: dict[str, str] | None = None,
-):
+) -> JSONResponse:
     payload = batch_status_payload(response)
     if response.status in {"submitted", "running"}:
         return JSONResponse(status_code=202, content=payload, headers=headers)
diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py
index aecb829c2..b0651d821 100644
--- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py
+++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/smoke_app.py
@@ -62,7 +62,7 @@ def smoke_import_gateway() -> dict:
 
 
 @app.local_entrypoint()
-def main():
+def main() -> None:
     report = smoke_import_gateway.remote()
     print(report)
     assert report["modules_imported"] > 0

From f6deed3a8d211adb60c8a2380cc348b266fd033c Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:53:09 +0400
Subject: [PATCH 05/12] Annotate the simulation entry service for ruff ANN

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../authenticated_tests/test_deployed_auth.py       |  6 +++---
 .../src/policyengine_simulation_entry/app.py        | 13 ++++++++++---
 .../src/policyengine_simulation_entry/auth.py       |  4 ++--
 .../src/policyengine_simulation_entry/backend.py    |  4 ++--
 .../stage12_adapter.py                              |  7 +++++--
 5 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py
index 183b50d9f..1d9609e0d 100644
--- a/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py
+++ b/projects/policyengine-simulation-entry/authenticated_tests/test_deployed_auth.py
@@ -47,7 +47,7 @@ def access_token() -> str:
     return token
 
 
-def test_protected_route_rejects_missing_token(base_url: str):
+def test_protected_route_rejects_missing_token(base_url: str) -> None:
     response = httpx.post(
         f"{base_url}{PROTECTED_PATH}",
         json=PROBE_PAYLOAD,
@@ -57,7 +57,7 @@ def test_protected_route_rejects_missing_token(base_url: str):
     assert response.status_code == 403
 
 
-def test_protected_route_rejects_invalid_token(base_url: str):
+def test_protected_route_rejects_invalid_token(base_url: str) -> None:
     response = httpx.post(
         f"{base_url}{PROTECTED_PATH}",
         json=PROBE_PAYLOAD,
@@ -71,7 +71,7 @@ def test_protected_route_rejects_invalid_token(base_url: str):
 def test_valid_token_reaches_authenticated_old_gateway(
     base_url: str,
     access_token: str,
-):
+) -> None:
     response = httpx.post(
         f"{base_url}{PROTECTED_PATH}",
         json=PROBE_PAYLOAD,
diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py
index ba77098c9..506905399 100644
--- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py
+++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py
@@ -8,7 +8,7 @@
 import uuid
 from collections.abc import Callable
 from contextlib import asynccontextmanager
-from typing import Any, Literal, Protocol
+from typing import TYPE_CHECKING, Any, Literal, Protocol
 from uuid import UUID
 
 from fastapi import Depends, FastAPI, Request
@@ -68,6 +68,11 @@
     TemporaryStage12UnsupportedRequest,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from starlette.middleware.base import RequestResponseEndpoint
+
 logger = logging.getLogger(__name__)
 BACKEND_RESPONSE_HEADER = {
     "X-PolicyEngine-Simulation-Backend": "old_gateway",
@@ -154,7 +159,7 @@ def create_app(
     runtime_settings = settings or Settings.from_env()
 
     @asynccontextmanager
-    async def lifespan(_: FastAPI):
+    async def lifespan(_: FastAPI) -> AsyncIterator[None]:
         runtime_settings.validate()
         await runtime_backend.start()
         try:
@@ -239,7 +244,9 @@ def adopt_backend_observability_id(
         )
 
     @app.middleware("http")
-    async def request_context(request: Request, call_next):
+    async def request_context(
+        request: Request, call_next: RequestResponseEndpoint
+    ) -> Response:
         headers = MutableHeaders(scope=request.scope)
         request_id = headers.get(REQUEST_ID_HEADER) or str(uuid.uuid4())
         incoming_observability_id = normalize_observability_id(
diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py
index 334084a1d..ca6f143f6 100644
--- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py
+++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/auth.py
@@ -24,7 +24,7 @@
 class JWTDecoder:
     """Validate Auth0 RS256 access tokens without database-layer dependencies."""
 
-    def __init__(self, issuer: str, audience: str):
+    def __init__(self, issuer: str, audience: str) -> None:
         self.issuer = issuer
         self.audience = audience
         self.jwks_client = jwt.PyJWKClient(f"{issuer}.well-known/jwks.json")
@@ -75,7 +75,7 @@ def _decoder(issuer: str, audience: str) -> JWTDecoder:
 class CallerAuthenticator:
     """FastAPI dependency that preserves the gateway's JWT contract."""
 
-    def __init__(self, settings: Settings, runtime: ObservabilityRuntime):
+    def __init__(self, settings: Settings, runtime: ObservabilityRuntime) -> None:
         self.settings = settings
         self.runtime = runtime
 
diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py
index 6cbdcbefd..a74f2ef70 100644
--- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py
+++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/backend.py
@@ -77,7 +77,7 @@ class ClientCredentialsTokenProvider:
 
     REFRESH_MARGIN_SECONDS = 60
 
-    def __init__(self, settings: Settings, client: httpx.AsyncClient):
+    def __init__(self, settings: Settings, client: httpx.AsyncClient) -> None:
         self.settings = settings
         self.client = client
         self._token: str | None = None
@@ -132,7 +132,7 @@ def __init__(
         *,
         transport: httpx.AsyncBaseTransport | None = None,
         runtime: ObservabilityRuntime | None = None,
-    ):
+    ) -> None:
         self.settings = settings
         self.transport = transport
         self.observability = runtime
diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py
index 8095728cc..771a97148 100644
--- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py
+++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py
@@ -22,7 +22,10 @@
 
 if TYPE_CHECKING:
     from policyengine_simulation_contract.stage12_bundle import CountryId
-    from policyengine_simulation_contract.stage12_manifest import V2WorkerVersion
+    from policyengine_simulation_contract.stage12_manifest import (
+        V2CountryWorker,
+        V2WorkerVersion,
+    )
 
 
 class ComparisonSkipReason(StrEnum):
@@ -75,7 +78,7 @@ def _skip(reason: ComparisonSkipReason) -> AdaptedComparisonRun:
     return AdaptedComparisonRun(report=None, skip_reason=reason)
 
 
-def _country_worker(worker: V2WorkerVersion, country: str):
+def _country_worker(worker: V2WorkerVersion, country: str) -> V2CountryWorker | None:
     return next(
         (item for item in worker.countries if item.country == country),
         None,

From 01d873a5014918fe54d83283a06ce1c6136d1a1b Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:57:24 +0400
Subject: [PATCH 06/12] Annotate the simulation executor for ruff ANN

Type the runtime, output builders, Modal entrypoints, scripts and the test
fixtures that live outside tests/. PolicyEngine, Google Cloud Storage and
observability types are imported for type checking only, country modules
are ModuleType, and duck-typed analysis values that tests replace with
SimpleNamespace are object. _try_compute_output uses a ParamSpec so its
arguments are checked against the wrapped function.

The in-repo src/modal package shadows the Modal library here, so Modal
types stay quoted or behind TYPE_CHECKING. _image_setup.py may not import
anything at module level, so it uses a local TYPE_CHECKING flag. The
RemoteFunction protocol keeps Any (with a reasoned noqa) because it
mirrors modal.Function.remote and spawn, whose payloads vary per function.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../fixtures/fake_modal.py                    | 57 +++++++++++------
 .../fixtures/identity_stubs.py                |  8 ++-
 .../fixtures/spm_doubles.py                   | 18 +++++-
 .../fixtures/test_modal_scripts.py            | 12 +++-
 .../test_simulation_output_builder.py         |  8 +--
 .../scripts/dump_package_version.py           |  4 +-
 .../src/modal/_image_setup.py                 | 21 +++++-
 .../src/modal/app.py                          |  8 ++-
 .../src/modal/budget_window_results.py        |  2 +-
 .../src/modal/budget_window_scheduler.py      | 13 ++--
 .../src/modal/logging_redaction.py            |  2 +-
 .../src/modal/precompute_app.py               |  2 +-
 .../src/modal/segmented_national.py           | 16 +++--
 .../src/modal/smoke_app.py                    |  2 +-
 .../modal/utils/update_version_registry.py    |  2 +-
 .../src/modal/v2_app.py                       |  7 +-
 .../__init__.py                               |  2 +-
 .../artifact_store.py                         | 19 ++++--
 .../baseline_artifacts.py                     | 15 +++--
 .../policyengine_simulation_executor/main.py  |  3 +-
 .../precompute.py                             | 17 ++++-
 .../precompute_models.py                      |  4 +-
 .../release_bundle.py                         | 17 +++--
 .../segmented_national_reduce.py              | 19 ++++--
 .../settings.py                               |  4 +-
 .../simulation.py                             |  2 +-
 .../simulation_microdata.py                   | 13 +++-
 .../simulation_output_budget.py               | 11 +++-
 .../simulation_output_builder.py              | 32 +++++-----
 .../simulation_output_cliff.py                |  3 +-
 .../simulation_output_common.py               | 33 ++++++----
 .../simulation_output_distribution.py         | 21 ++++--
 .../simulation_output_geographic.py           | 19 ++++--
 .../simulation_output_inequality.py           |  4 +-
 .../simulation_output_labor.py                |  4 +-
 .../simulation_output_poverty.py              | 22 ++++---
 .../simulation_runtime.py                     | 64 ++++++++++++-------
 .../policyengine_simulation_executor/spm.py   | 26 ++++++--
 .../stage12_artifacts.py                      |  2 +-
 .../stage12_runtime/aggregation.py            |  9 ++-
 .../stage12_runtime/coordination.py           |  3 +-
 .../uk_local_authority_metadata.py            |  6 +-
 42 files changed, 372 insertions(+), 184 deletions(-)

diff --git a/projects/policyengine-simulation-executor/fixtures/fake_modal.py b/projects/policyengine-simulation-executor/fixtures/fake_modal.py
index 926967191..99bcc636b 100644
--- a/projects/policyengine-simulation-executor/fixtures/fake_modal.py
+++ b/projects/policyengine-simulation-executor/fixtures/fake_modal.py
@@ -5,83 +5,100 @@
 a per-file copy. The fakes record calls; they run nothing.
 """
 
+from __future__ import annotations
+
 import sys
 from types import ModuleType
+from typing import TYPE_CHECKING, Self
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Mapping
+    from pathlib import Path
+
+    import pytest
 
 
 class FakeImage:
-    def __init__(self):
-        self.calls = []
+    def __init__(self) -> None:
+        self.calls: list[tuple[object, ...]] = []
 
     @classmethod
-    def debian_slim(cls, python_version):
+    def debian_slim(cls, python_version: str) -> Self:
         image = cls()
         image.calls.append(("debian_slim", python_version))
         return image
 
-    def pip_install(self, *packages):
+    def pip_install(self, *packages: str) -> Self:
         self.calls.append(("pip_install", packages))
         return self
 
-    def pip_install_from_requirements(self, requirements_txt, **kwargs):
+    def pip_install_from_requirements(
+        self, requirements_txt: str, **kwargs: object
+    ) -> Self:
         self.calls.append(("pip_install_from_requirements", requirements_txt, kwargs))
         return self
 
-    def uv_sync(self, uv_project_dir="./", **kwargs):
+    def uv_sync(self, uv_project_dir: str = "./", **kwargs: object) -> Self:
         self.calls.append(("uv_sync", uv_project_dir, kwargs))
         return self
 
-    def run_commands(self, *commands, **kwargs):
+    def run_commands(self, *commands: str, **kwargs: object) -> Self:
         self.calls.append(("run_commands", commands, kwargs))
         return self
 
-    def env(self, env):
+    def env(self, env: Mapping[str, str]) -> Self:
         self.calls.append(("env", env))
         return self
 
-    def add_local_python_source(self, *args, **kwargs):
+    def add_local_python_source(self, *args: str, **kwargs: object) -> Self:
         self.calls.append(("add_local_python_source", args, kwargs))
         return self
 
-    def add_local_file(self, local_path, remote_path, **kwargs):
+    def add_local_file(
+        self, local_path: str | Path, remote_path: str, **kwargs: object
+    ) -> Self:
         self.calls.append(("add_local_file", local_path, remote_path, kwargs))
         return self
 
-    def add_local_dir(self, local_path, remote_path, **kwargs):
+    def add_local_dir(
+        self, local_path: str | Path, remote_path: str, **kwargs: object
+    ) -> Self:
         self.calls.append(("add_local_dir", local_path, remote_path, kwargs))
         return self
 
-    def run_function(self, function, **kwargs):
+    def run_function(self, function: Callable[..., object], **kwargs: object) -> Self:
         self.calls.append(("run_function", function.__name__, kwargs))
         return self
 
 
 class FakeSecret:
     @staticmethod
-    def from_name(*args, **kwargs):
+    def from_name(*args: object, **kwargs: object) -> dict[str, object]:
         return {"args": args, "kwargs": kwargs}
 
 
 class FakeApp:
-    def __init__(self, name):
+    def __init__(self, name: str) -> None:
         self.name = name
-        self.function_calls = []
+        self.function_calls: list[tuple[str, dict[str, object]]] = []
 
-    def function(self, **kwargs):
-        def decorator(function):
+    def function[F: Callable[..., object]](self, **kwargs: object) -> Callable[[F], F]:
+        def decorator(function: F) -> F:
             self.function_calls.append((function.__name__, kwargs))
             return function
 
         return decorator
 
-    def local_entrypoint(self, **kwargs):
-        def decorator(function):
+    def local_entrypoint[F: Callable[..., object]](
+        self, **kwargs: object
+    ) -> Callable[[F], F]:
+        def decorator(function: F) -> F:
             return function
 
         return decorator
 
 
-def install_fake_modal(monkeypatch):
+def install_fake_modal(monkeypatch: pytest.MonkeyPatch) -> None:
     modal = ModuleType("modal")
     modal.Image = FakeImage
     modal.Secret = FakeSecret
diff --git a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py
index 0a6c1ca5a..21204db4c 100644
--- a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py
+++ b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py
@@ -15,10 +15,16 @@
 synthetic inputs instead of whichever PolicyEngine release is installed.
 """
 
+from __future__ import annotations
+
 from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import pytest
 
 
-def install_identity_stubs(monkeypatch):
+def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
     """Stub the identity seams with the canonical synthetic version-set.
 
     Returns mutable bundle, receipt, and SPM-capability state, plus a handle
diff --git a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py
index 790e2ac9b..39ac74819 100644
--- a/projects/policyengine-simulation-executor/fixtures/spm_doubles.py
+++ b/projects/policyengine-simulation-executor/fixtures/spm_doubles.py
@@ -1,9 +1,15 @@
 """Small simulation doubles for executor tests that are not model tests."""
 
+from __future__ import annotations
+
 from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
 
 
-def installed_spm_selection():
+def installed_spm_selection() -> dict[str, object] | None:
     """Return the installed US bundle's default SPM selection, if any."""
     from policyengine_simulation_executor.spm import runtime_spm_capability
 
@@ -11,7 +17,9 @@ def installed_spm_selection():
     return None if capability is None else capability.defaults.model_dump()
 
 
-def spm_receipt(selection, year="2026"):
+def spm_receipt(
+    selection: Mapping[str, object], year: int | str = "2026"
+) -> dict[str, object]:
     """Build the smallest calculation record accepted by ``SPMProvenance``."""
     return {
         "forecast_id": "test-only",
@@ -31,7 +39,11 @@ def spm_receipt(selection, year="2026"):
     }
 
 
-def spm_capable_simulation(selection, year="2026", **attributes):
+def spm_capable_simulation(
+    selection: Mapping[str, object] | None,
+    year: int | str = "2026",
+    **attributes: object,
+) -> SimpleNamespace:
     """Return a stand-in for a simulation that measured ``selection``."""
     if selection is None:
         return SimpleNamespace(**attributes)
diff --git a/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py b/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py
index 9e54267f9..c25ae0fd2 100644
--- a/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py
+++ b/projects/policyengine-simulation-executor/fixtures/test_modal_scripts.py
@@ -1,17 +1,23 @@
 """Fixtures and shared constants for Modal script tests."""
 
+from __future__ import annotations
+
 import os
 import tempfile
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 import pytest
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
 REPO_ROOT = Path(__file__).parent.parent.parent.parent
 SCRIPTS_DIR = REPO_ROOT / ".github" / "scripts"
 
 
 @pytest.fixture
-def temp_github_output():
+def temp_github_output() -> Iterator[str]:
     """Create a temporary file to simulate ``GITHUB_OUTPUT``."""
     with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as file:
         yield file.name
@@ -19,7 +25,7 @@ def temp_github_output():
 
 
 @pytest.fixture
-def temp_github_step_summary():
+def temp_github_step_summary() -> Iterator[str]:
     """Create a temporary file to simulate ``GITHUB_STEP_SUMMARY``."""
     with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".md") as file:
         yield file.name
@@ -27,7 +33,7 @@ def temp_github_step_summary():
 
 
 @pytest.fixture
-def all_modal_scripts():
+def all_modal_scripts() -> list[Path]:
     """Return Modal and full-stack simulation deployment scripts."""
     return [
         *SCRIPTS_DIR.glob("modal-*.sh"),
diff --git a/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py
index b593a0be4..ccc591b04 100644
--- a/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py
+++ b/projects/policyengine-simulation-executor/fixtures/test_simulation_output_builder.py
@@ -8,20 +8,20 @@
 
 
 class FakeCollection:
-    def __init__(self, records):
+    def __init__(self, records: list[dict[str, object]]) -> None:
         self.dataframe = pd.DataFrame(records)
 
 
 class FakeModelOutput:
-    def __init__(self, payload):
+    def __init__(self, payload: dict[str, object]) -> None:
         self.payload = payload
 
-    def model_dump(self, *, mode):
+    def model_dump(self, *, mode: str) -> dict[str, object]:
         assert mode == "json"
         return self.payload
 
 
-def fake_analysis():
+def fake_analysis() -> SimpleNamespace:
     return SimpleNamespace(
         program_statistics=FakeCollection(
             [
diff --git a/projects/policyengine-simulation-executor/scripts/dump_package_version.py b/projects/policyengine-simulation-executor/scripts/dump_package_version.py
index afa2d0403..de7a603de 100644
--- a/projects/policyengine-simulation-executor/scripts/dump_package_version.py
+++ b/projects/policyengine-simulation-executor/scripts/dump_package_version.py
@@ -3,7 +3,9 @@
 import tomllib
 
 
-def get_package_version_from_lock(package_name, lock_file_path="poetry.lock"):
+def get_package_version_from_lock(
+    package_name: str, lock_file_path: str = "poetry.lock"
+) -> str | None:
     try:
         with open(lock_file_path, "rb") as f:
             lock_data = tomllib.load(f)
diff --git a/projects/policyengine-simulation-executor/src/modal/_image_setup.py b/projects/policyengine-simulation-executor/src/modal/_image_setup.py
index f8f7183c3..54788bb6a 100644
--- a/projects/policyengine-simulation-executor/src/modal/_image_setup.py
+++ b/projects/policyengine-simulation-executor/src/modal/_image_setup.py
@@ -7,8 +7,23 @@
 so policyengine_simulation_executor is not importable there at all.
 """
 
-
-def fetch_artifacts(bucket: str, manifest, *, client=None):
+# Annotation-only imports sit behind a local TYPE_CHECKING flag (not
+# typing.TYPE_CHECKING) because this module must have no module-level
+# imports; the annotations that use them are therefore quoted.
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+    from typing import Any
+
+    from google.cloud import storage
+
+
+def fetch_artifacts(
+    bucket: str,
+    manifest: "Mapping[str, Any] | None",
+    *,
+    client: "storage.Client | None" = None,
+) -> None:
     """Download every manifest-listed artifact into the image data folder.
 
     Runs as an image-build layer between ``.env(VERSION_ENV)`` and
@@ -155,7 +170,7 @@ def fetch_artifacts(bucket: str, manifest, *, client=None):
         raise RuntimeError(f"Artifact fetch did not produce expected files: {absent}")
 
 
-def snapshot_models():
+def snapshot_models() -> None:
     """Pre-load models at image build time for fast cold starts."""
     import logging
 
diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py
index 4e5786b01..b30670187 100644
--- a/projects/policyengine-simulation-executor/src/modal/app.py
+++ b/projects/policyengine-simulation-executor/src/modal/app.py
@@ -9,6 +9,7 @@
 
 import os
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from policyengine_simulation_executor.hf_access_validation import (
     HF_MODAL_SECRET_NAME,
@@ -38,6 +39,9 @@
 
 import modal
 
+if TYPE_CHECKING:
+    from policyengine_observability import ObservabilityRuntime
+
 
 def _version_from_env_or_local_bundle_package(env_var: str, package: str) -> str:
     value = os.environ.get(env_var)
@@ -241,7 +245,7 @@ def build_base_simulation_image() -> modal.Image:
 def _configure_modal_observability(
     *,
     service_role: str,
-):
+) -> "ObservabilityRuntime":
     return init_process_observability(
         service_name=APP_NAME,
         platform="modal",
@@ -251,7 +255,7 @@ def _configure_modal_observability(
     )
 
 
-def _set_modal_call_attributes(runtime) -> None:
+def _set_modal_call_attributes(runtime: "ObservabilityRuntime") -> None:
     try:
         runtime.set_context(function_call_id=modal.current_function_call_id())
     except Exception:
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 764e38674..02464c0dc 100644
--- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py
+++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py
@@ -37,7 +37,7 @@ def extract_annual_impact(
     *,
     simulation_year: str,
     child_result: dict[str, Any],
-    spm: Any = None,
+    spm: object = None,
 ) -> BudgetWindowAnnualImpact:
     receipt = validate_spm_result(child_result, spm, expected_year=simulation_year)
     budget = child_result.get("budget", {})
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 2fe139224..eb0596a9f 100644
--- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
+++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py
@@ -34,7 +34,12 @@
 import modal
 
 if TYPE_CHECKING:
+    from types import ModuleType
+
     from policyengine_observability import ObservabilityRuntime
+    from policyengine_simulation_contract.gateway_models import (
+        BudgetWindowBatchState,
+    )
 
 # Polling tuning. The runner busy-loops across child FunctionCall.get(timeout=0)
 # probes; when no child resolved we sleep before the next probe to stop the
@@ -54,7 +59,7 @@
 
 
 def serialize_batch_status(
-    state,
+    state: BudgetWindowBatchState,
     *,
     runtime: ObservabilityRuntime,
 ) -> dict[str, Any]:
@@ -68,7 +73,7 @@ def load_or_create_batch_state(
     context: BudgetWindowBatchContext,
     *,
     runtime: ObservabilityRuntime,
-):
+) -> BudgetWindowBatchState:
     with runtime.span(BUDGET_WINDOW_STAGES.name(Stage.BUDGET_WINDOW_STATE_LOAD)):
         state = get_batch_job_seed(context.batch_job_id)
     if state is None:
@@ -92,11 +97,11 @@ def __init__(
         context: BudgetWindowBatchContext,
         *,
         runtime: ObservabilityRuntime,
-        modal_module=None,
+        modal_module: ModuleType | None = None,
         poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS,
         poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS,
         poll_interval_backoff_factor: float = POLL_INTERVAL_BACKOFF_FACTOR,
-    ):
+    ) -> None:
         self.context = context
         self.runtime = runtime
         self.modal = modal if modal_module is None else modal_module
diff --git a/projects/policyengine-simulation-executor/src/modal/logging_redaction.py b/projects/policyengine-simulation-executor/src/modal/logging_redaction.py
index 0325af918..4f2542b51 100644
--- a/projects/policyengine-simulation-executor/src/modal/logging_redaction.py
+++ b/projects/policyengine-simulation-executor/src/modal/logging_redaction.py
@@ -15,7 +15,7 @@
 SENSITIVE_PARAM_KEYS = ("data", "reform", "baseline")
 
 
-def redact_params_for_logging(params) -> dict:
+def redact_params_for_logging(params: object) -> dict:
     """Return a shallow copy of ``params`` safe for observability attributes.
 
     We preserve routing-relevant fields (country, scope, version, time
diff --git a/projects/policyengine-simulation-executor/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py
index 34d3b0833..1314b0f4c 100644
--- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py
+++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py
@@ -121,7 +121,7 @@ def publish_manifest(bucket: str, plan: dict) -> str:
 
 
 @app.local_entrypoint()
-def main(force: bool = False):
+def main(force: bool = False) -> None:
     from policyengine_simulation_executor.artifact_store import resolve_bucket_name
     from policyengine_simulation_executor.precompute import run_precompute
 
diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py
index 6700dedff..a2e25e823 100644
--- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py
+++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py
@@ -46,6 +46,8 @@
 import modal
 
 if TYPE_CHECKING:
+    from types import ModuleType
+
     from policyengine_observability import ObservabilityRuntime
 
 logger = logging.getLogger(__name__)
@@ -144,10 +146,10 @@ def __init__(
         *,
         app_name: str,
         runtime: ObservabilityRuntime,
-        modal_module=None,
+        modal_module: ModuleType | None = None,
         poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS,
         poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS,
-    ):
+    ) -> None:
         self.params = params
         self.app_name = app_name
         self.runtime = runtime
@@ -177,14 +179,16 @@ def run(self) -> dict[str, Any]:
         child_results = self._collect(handles)
         return self._reduce(child_results, country_module=country_module)
 
-    def _country_module(self):
+    def _country_module(self) -> ModuleType:
         from policyengine_simulation_executor.simulation_runtime import (
             _country_module,
         )
 
         return _country_module(self.country)
 
-    def _extend_last_group_with_uncovered_states(self, country_module) -> None:
+    def _extend_last_group_with_uncovered_states(
+        self, country_module: ModuleType
+    ) -> None:
         """Registry state regions missing from the static partition ride in
         the last group, loudly, so a new state-level region added to the
         model cannot silently drop its households from national results.
@@ -307,7 +311,9 @@ def _cancel_all(self, handles: list[tuple[list[str], Any]]) -> None:
             except Exception:
                 pass
 
-    def _reduce(self, child_results: list[dict], *, country_module) -> dict[str, Any]:
+    def _reduce(
+        self, child_results: list[dict], *, country_module: ModuleType
+    ) -> dict[str, Any]:
         from policyengine_simulation_executor.simulation_runtime import (
             _parse_year,
         )
diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py
index 4318b5210..78ecc52a1 100644
--- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py
+++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py
@@ -161,7 +161,7 @@ def smoke_import_executor() -> dict:
 
 
 @app.local_entrypoint()
-def main():
+def main() -> None:
     report = smoke_import_executor.remote()
     print(report)
     print("executor image smoke OK")
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 4d091f342..71641b504 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
@@ -1028,7 +1028,7 @@ def restore_routing_state(
     return next_state
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(
         formatter_class=argparse.RawDescriptionHelpFormatter,
         description=(
diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py
index f79c96578..dab372b87 100644
--- a/projects/policyengine-simulation-executor/src/modal/v2_app.py
+++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py
@@ -41,7 +41,10 @@
 import modal
 
 if TYPE_CHECKING:
-    from policyengine_simulation_contract.stage12_bundle import CountryId
+    from policyengine_simulation_contract.stage12_bundle import (
+        CountryId,
+        Stage12CountryBundle,
+    )
 
 STAGE12_DATA_DIR = "/opt/policyengine/stage12-data"
 _UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "."
@@ -87,7 +90,7 @@ def _external_assertions(environment: dict[str, str]) -> dict[str, str]:
 ]
 
 
-def _country_bundle(country: CountryId):
+def _country_bundle(country: CountryId) -> Stage12CountryBundle:
     return next(
         item for item in RESOLVED_BUNDLE.bundle.countries if item.country == country
     )
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py
index 98f1a49b6..433343dff 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/__init__.py
@@ -9,7 +9,7 @@
 """
 
 
-def initialize(app: FastAPI, runtime: ObservabilityRuntime):
+def initialize(app: FastAPI, runtime: ObservabilityRuntime) -> None:
     """
     attach all routes to the app and configure them to use the provided SQLModel engine
     and jwt settings.
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py
index b16b155f8..9298fd237 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_store.py
@@ -24,12 +24,15 @@
 import os
 from collections.abc import Mapping
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from google.api_core.exceptions import NotFound, PreconditionFailed
 
 from policyengine_simulation_executor import artifact_keys
 
+if TYPE_CHECKING:
+    from google.cloud import storage
+
 ARTIFACT_BUCKET_ENV = "POLICYENGINE_ARTIFACT_BUCKET"
 
 
@@ -42,7 +45,7 @@ def resolve_bucket_name(explicit: str | None = None) -> str:
     return bucket
 
 
-def _default_client():
+def _default_client() -> storage.Client:
     from google.cloud import storage
 
     from policyengine_simulation_executor.simulation_runtime import (
@@ -56,17 +59,19 @@ def _default_client():
 
 
 class ArtifactStore:
-    def __init__(self, bucket_name: str | None = None, *, client: Any = None):
+    def __init__(
+        self, bucket_name: str | None = None, *, client: storage.Client | None = None
+    ) -> None:
         self.bucket_name = resolve_bucket_name(bucket_name)
         self._client = client
 
     @property
-    def client(self) -> Any:
+    def client(self) -> storage.Client:
         if self._client is None:
             self._client = _default_client()
         return self._client
 
-    def _blob(self, path: str):
+    def _blob(self, path: str) -> storage.Blob:
         return self.client.bucket(self.bucket_name).blob(path)
 
     def exists(self, path: str) -> bool:
@@ -123,7 +128,9 @@ def download_file(self, path: str, local_path: str | Path) -> None:
             raise
         partial.replace(destination)
 
-    def _upload_json(self, path: str, payload: Mapping[str, Any], *, overwrite: bool):
+    def _upload_json(
+        self, path: str, payload: Mapping[str, Any], *, overwrite: bool
+    ) -> bool:
         data = json.dumps(payload, sort_keys=True, indent=2)
         blob = self._blob(path)
         if overwrite:
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 e66a59092..ffc8d7091 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
@@ -21,13 +21,16 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from policyengine.core import Simulation
 from pydantic import PrivateAttr
 
 from policyengine_simulation_executor import artifact_keys
 
+if TYPE_CHECKING:
+    from policyengine.core.scoping_strategy import ScopingStrategy
+
 logger = logging.getLogger(__name__)
 
 # Artifact outcomes reported via the `baseline_artifact` observability
@@ -45,9 +48,9 @@ def qualifying_baseline_identity(
     *,
     country: str,
     dataset_is_default: bool,
-    policy: Any,
+    policy: object,
     region_code: str | None,
-    scoping_strategy: Any,
+    scoping_strategy: ScopingStrategy | None,
     year: int,
 ) -> artifact_keys.BaselineArtifactIdentity | None:
     """Full artifact identity for a qualifying baseline, else None.
@@ -103,9 +106,9 @@ def deterministic_baseline_id(
     *,
     country: str,
     dataset_is_default: bool,
-    policy: Any,
+    policy: object,
     region_code: str | None,
-    scoping_strategy: Any,
+    scoping_strategy: ScopingStrategy | None,
     year: int,
 ) -> str | None:
     """The runtime wrapper: id for a qualifying baseline, else None.
@@ -150,7 +153,7 @@ def artifact_outcome(self) -> str | None:
         """Set by ``ensure()``; None until then."""
         return self._artifact_outcome
 
-    def run(self):
+    def run(self) -> None:
         self._computed_this_process = True
         super().run()
 
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py
index b0b0d7383..9e8391583 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/main.py
@@ -1,4 +1,5 @@
 import logging
+from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
@@ -25,7 +26,7 @@
 
 
 @asynccontextmanager
-async def lifespan(app: FastAPI):
+async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     async with exit.lifespan():
         try:
             yield
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 f8e2caa78..42f3c47be 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py
@@ -34,7 +34,7 @@
 
 import logging
 import time
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from policyengine_simulation_executor.artifact_keys import (
     BaselineArtifactIdentity,
@@ -56,6 +56,13 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable
+    from pathlib import Path
+    from types import ModuleType
+
+    from policyengine.core import Simulation
+
+    from policyengine_simulation_executor.artifact_store import ArtifactStore
+    from policyengine_simulation_executor.simulation_runtime import RegionResolution
 
 logger = logging.getLogger(__name__)
 
@@ -85,7 +92,9 @@ def cohort_params(year: int, group: list[str]) -> dict:
     }
 
 
-def _cohort_resolution(year: int, group: list[str]):
+def _cohort_resolution(
+    year: int, group: list[str]
+) -> tuple[dict[str, Any], ModuleType, RegionResolution]:
     """Cohort request params, country module, and region resolution.
 
     The planner's identity (``cohort_identity``) and the worker's
@@ -277,7 +286,9 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR
     )
 
 
-def _prepare_cohort_baseline(bucket: str, expected: BaselinePlanEntry):
+def _prepare_cohort_baseline(
+    bucket: str, expected: BaselinePlanEntry
+) -> tuple[ArtifactStore, Path, Simulation]:
     """Shared by compute and verify: dataset download + simulation build."""
     from pathlib import Path
 
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py
index a6abc66e2..47dc14793 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py
@@ -162,6 +162,6 @@ class DeterminismVerdict(_StrictModel):
 class RemoteFunction(Protocol):
     """The slice of a Modal function handle the orchestration uses."""
 
-    def remote(self, *args: Any) -> Any: ...
+    def remote(self, *args: Any) -> Any: ...  # noqa: ANN401  # mirrors modal.Function.remote; payloads and results vary per function
 
-    def spawn(self, *args: Any) -> Any: ...
+    def spawn(self, *args: Any) -> Any: ...  # noqa: ANN401  # mirrors modal.Function.spawn; payloads vary and the returned FunctionCall is untyped here
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py
index 0d8ce99b5..102cc3366 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py
@@ -14,6 +14,7 @@
 from dataclasses import dataclass
 from functools import lru_cache
 from pathlib import Path
+from typing import TYPE_CHECKING
 
 from policyengine_simulation_contract.dataset_uri import (
     runtime_dataset_uri,
@@ -21,6 +22,12 @@
     split_dataset_revision,
 )
 
+if TYPE_CHECKING:
+    from policyengine.provenance.manifest import (
+        CountryReleaseManifest,
+        DataPackageVersion,
+    )
+
 os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1")
 
 SUPPORTED_COUNTRIES = frozenset({"us", "uk"})
@@ -99,7 +106,7 @@ def _normalise_country(country: str) -> str:
     return country
 
 
-def _artifact_revision(data_package) -> str:
+def _artifact_revision(data_package: DataPackageVersion) -> str:
     return data_package.release_manifest_revision or data_package.version
 
 
@@ -129,7 +136,7 @@ def _with_hf_revision_unvalidated(dataset_uri: str, revision: str) -> str:
 
 def _artifact_hf_uri(
     *,
-    data_package,
+    data_package: DataPackageVersion,
     reference: object,
 ) -> tuple[str | None, str | None]:
     from policyengine.provenance.manifest import build_hf_uri
@@ -152,7 +159,9 @@ def _artifact_hf_uri(
     )
 
 
-def _dataset_uris_from_manifest(manifest) -> tuple[dict[str, str], dict[str, str]]:
+def _dataset_uris_from_manifest(
+    manifest: CountryReleaseManifest,
+) -> tuple[dict[str, str], dict[str, str]]:
     dataset_uris: dict[str, str] = {}
     dataset_repo_types: dict[str, str] = {}
     for name, reference in manifest.datasets.items():
@@ -171,7 +180,7 @@ def _dataset_uris_from_manifest(manifest) -> tuple[dict[str, str], dict[str, str
 def _dataset_uris_from_release(
     *,
     data_release: Mapping,
-    fallback_data_package,
+    fallback_data_package: DataPackageVersion,
 ) -> tuple[dict[str, str], dict[str, str]]:
     from policyengine.provenance.manifest import build_hf_uri
 
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py
index 71fcfbfdf..3fc6de14f 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/segmented_national_reduce.py
@@ -22,13 +22,14 @@
   aggregates drift ~1e-7 relative.
 """
 
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import pandas as pd
 from microdf import MicroDataFrame
 from policyengine.core import Simulation
-from policyengine_observability import ObservabilityRuntime
 
 from policyengine_simulation_executor.simulation_microdata import (
     rebuild_entity_frame,
@@ -37,6 +38,12 @@
     SimulationOutputBuilder,
 )
 
+if TYPE_CHECKING:
+    from types import ModuleType
+
+    from policyengine.tax_benefit_models.us.datasets import PolicyEngineUSDataset
+    from policyengine_observability import ObservabilityRuntime
+
 logger = logging.getLogger(__name__)
 
 
@@ -82,7 +89,9 @@ def concatenate_microdata(
     return sides
 
 
-def _us_dataset_from_frames(frames: dict[str, pd.DataFrame], *, year: int):
+def _us_dataset_from_frames(
+    frames: dict[str, pd.DataFrame], *, year: int
+) -> PolicyEngineUSDataset:
     """Wrap concatenated entity tables as an in-memory US dataset.
 
     Weights need no special handling: they ride as ``<entity>_weight``
@@ -115,7 +124,7 @@ def build_national_output(
     *,
     country: str,
     simulation_params: dict[str, Any],
-    country_module,
+    country_module: ModuleType,
     year: int,
     resolved_data_version: str | None = None,
     runtime: ObservabilityRuntime | None = None,
@@ -137,7 +146,7 @@ def build_national_output(
         for side, frames in sides.items()
     }
 
-    def stand_in(dataset) -> PrecomputedSimulation:
+    def stand_in(dataset: PolicyEngineUSDataset) -> PrecomputedSimulation:
         simulation = PrecomputedSimulation(
             dataset=dataset,
             tax_benefit_model_version=country_module.model,
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py
index fd76d9751..24bfb5547 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/settings.py
@@ -15,7 +15,7 @@ class AppSettings(BaseSettings):
 
     @field_validator("environment", mode="before")
     @classmethod
-    def strip_environment(cls, v):
+    def strip_environment(cls, v: object) -> object:
         if isinstance(v, str):
             return v.strip()
         return v
@@ -33,5 +33,5 @@ def strip_environment(cls, v):
 
 
 @lru_cache
-def get_settings():
+def get_settings() -> AppSettings:
     return AppSettings()
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py
index 78d8b731d..4016de5ec 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py
@@ -14,7 +14,7 @@
 logger = logging.getLogger(__file__)
 
 
-def create_router(runtime: ObservabilityRuntime):
+def create_router(runtime: ObservabilityRuntime) -> APIRouter:
     router = APIRouter()
 
     @router.post("/simulate/economy/comparison", response_model=EconomyComparison)
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 7f7b8370e..c51a28fc2 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
@@ -13,15 +13,22 @@
 bit-exact against a direct run.
 """
 
+from __future__ import annotations
+
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import pandas as pd
 
+if TYPE_CHECKING:
+    from policyengine.core import Simulation
+
 logger = logging.getLogger(__name__)
 
 
-def extract_output_microdata(baseline, reform) -> dict[str, Any]:
+def extract_output_microdata(
+    baseline: Simulation, reform: Simulation
+) -> dict[str, Any]:
     """Columnar dump of each simulation's computed ``output_dataset.data`` per
     entity, for baseline and reform, plus per-column source dtypes.
 
@@ -34,7 +41,7 @@ def extract_output_microdata(baseline, reform) -> dict[str, Any]:
     """
 
     def dump(
-        simulation,
+        simulation: Simulation,
     ) -> tuple[dict[str, dict[str, list]], dict[str, dict[str, str]]]:
         columns: dict[str, dict[str, list]] = {}
         dtypes: dict[str, dict[str, str]] = {}
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py
index e12028e68..7d9a1e95e 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING
 
 from policyengine_simulation_executor.simulation_macro_output import (
     BudgetaryImpact,
@@ -17,8 +17,11 @@
     _sum_output_variable,
 )
 
+if TYPE_CHECKING:
+    from policyengine.core import Simulation
 
-def build_detailed_budget(analysis: Any) -> DetailedBudgetOutput:
+
+def build_detailed_budget(analysis: object) -> DetailedBudgetOutput:
     collection = getattr(analysis, "program_statistics", None)
     if isinstance(collection, DetailedBudgetOutput):
         return collection
@@ -37,7 +40,9 @@ def build_detailed_budget(analysis: Any) -> DetailedBudgetOutput:
     return DetailedBudgetOutput(detailed_budget)
 
 
-def build_budgetary_impact(country: str, baseline, reform) -> BudgetaryImpact:
+def build_budgetary_impact(
+    country: str, baseline: Simulation, reform: Simulation
+) -> BudgetaryImpact:
     tax_revenue_impact = _change_output_variable(
         baseline, reform, "household_tax", entity="household"
     )
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py
index ee368501b..3207f0fa9 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_builder.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from contextlib import nullcontext
+from contextlib import AbstractContextManager, nullcontext
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -60,11 +60,11 @@ class SimulationOutputBuilder:
     def __post_init__(self) -> None:
         self.country = self.country.lower()
 
-    def _span(self, name: str):
+    def _span(self, name: str) -> AbstractContextManager[object]:
         return self.runtime.span(name) if self.runtime is not None else nullcontext()
 
     @property
-    def analysis(self) -> Any:
+    def analysis(self) -> object:
         if self._analysis is None:
             with self._span(ANNUAL_IMPACT_STAGES.name(Stage.ECONOMIC_IMPACT_ANALYSIS)):
                 self._analysis = self.country_module.economic_impact_analysis(
@@ -148,14 +148,14 @@ def _build_intra_decile_output(self) -> IntraDecileOutput:
                 self.baseline, self.reform
             )
 
-    def _build_wealth_decile(self, wealth_decile: Any) -> DecileOutput | None:
+    def _build_wealth_decile(self, wealth_decile: object) -> DecileOutput | None:
         with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_WEALTH_DECILE)):
             return simulation_output_distribution.build_wealth_decile(
                 self.country, wealth_decile
             )
 
     def _build_intra_wealth_decile(
-        self, intra_wealth_decile: Any
+        self, intra_wealth_decile: object
     ) -> IntraDecileOutput | None:
         with self._span(ANNUAL_IMPACT_STAGES.name(Stage.OUTPUT_INTRA_WEALTH_DECILE)):
             return simulation_output_distribution.build_intra_wealth_decile(
@@ -171,15 +171,15 @@ def _build_cliff_impact(self) -> CliffImpactOutput | None:
             return simulation_output_cliff.build_cliff_impact(self.analysis)
 
     def _build_geographic_impact_output(
-        self, value: Any
+        self, value: object
     ) -> GeographicImpactOutput | None:
         return simulation_output_geographic.build_geographic_impact_output(value)
 
-    def _build_decile_output(self, collection: Any) -> DecileOutput:
+    def _build_decile_output(self, collection: object) -> DecileOutput:
         return simulation_output_distribution.build_decile_output(collection)
 
     def _build_intra_decile_output_from_collection(
-        self, collection: Any
+        self, collection: object
     ) -> IntraDecileOutput:
         return simulation_output_distribution.build_intra_decile_output_from_collection(
             collection
@@ -188,10 +188,10 @@ def _build_intra_decile_output_from_collection(
     def _build_poverty_output(
         self,
         *,
-        baseline: Any,
-        reform: Any,
-        baseline_by_age: Any,
-        reform_by_age: Any,
+        baseline: object,
+        reform: object,
+        baseline_by_age: object,
+        reform_by_age: object,
     ) -> PovertyOutput:
         return simulation_output_poverty.build_poverty_output(
             country=self.country,
@@ -204,8 +204,8 @@ def _build_poverty_output(
     def _build_poverty_by_gender_output(
         self,
         *,
-        baseline_by_gender: Any,
-        reform_by_gender: Any,
+        baseline_by_gender: object,
+        reform_by_gender: object,
     ) -> PovertyByGenderOutput:
         return simulation_output_poverty.build_poverty_by_gender_output(
             country=self.country,
@@ -216,8 +216,8 @@ def _build_poverty_by_gender_output(
     def _build_poverty_by_race_output(
         self,
         *,
-        baseline_by_race: Any,
-        reform_by_race: Any,
+        baseline_by_race: object,
+        reform_by_race: object,
     ) -> PovertyByRaceOutput:
         return simulation_output_poverty.build_poverty_by_race_output(
             baseline_by_race=baseline_by_race,
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py
index 0e2953b87..746d1b8f2 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_cliff.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from typing import Any
 
 from policyengine_simulation_executor.simulation_macro_output import (
     CliffImpactInSimulation,
@@ -12,7 +11,7 @@
 from policyengine_simulation_executor.simulation_output_common import _output_model_dump
 
 
-def build_cliff_impact(analysis: Any) -> CliffImpactOutput | None:
+def build_cliff_impact(analysis: object) -> CliffImpactOutput | None:
     cliff_impact = getattr(analysis, "cliff_impact", None)
     if isinstance(cliff_impact, CliffImpactOutput):
         return cliff_impact
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 941b4ab72..e2fd0d7f2 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
@@ -4,16 +4,21 @@
 
 import logging
 import math
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from importlib import import_module
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import pandas as pd
+    from _typeshed import ConvertibleToFloat
+    from policyengine.core import Simulation
 
 import numpy as np
 
 logger = logging.getLogger(__name__)
 
 
-def _number(value: Any, default: float = 0.0) -> float:
+def _number(value: ConvertibleToFloat | None, default: float = 0.0) -> float:
     if value is None:
         return default
     try:
@@ -25,7 +30,7 @@ def _number(value: Any, default: float = 0.0) -> float:
     return result
 
 
-def _collection_records(collection: Any) -> list[dict[str, Any]]:
+def _collection_records(collection: object) -> list[dict[str, Any]]:
     if collection is None:
         return []
     dataframe = getattr(collection, "dataframe", None)
@@ -36,7 +41,7 @@ def _collection_records(collection: Any) -> list[dict[str, Any]]:
     return []
 
 
-def _output_model_dump(value: Any) -> Any:
+def _output_model_dump(value: object) -> object:
     if value is None:
         return None
     if hasattr(value, "model_dump"):
@@ -46,13 +51,13 @@ def _output_model_dump(value: Any) -> Any:
     return None
 
 
-def _entity_data(simulation, entity: str):
+def _entity_data(simulation: Simulation, entity: str) -> pd.DataFrame:
     if simulation.output_dataset is None or simulation.output_dataset.data is None:
         simulation.ensure()
     return getattr(simulation.output_dataset.data, entity)
 
 
-def _count_output_entities(simulation, entity: str, weight: str) -> float:
+def _count_output_entities(simulation: Simulation, entity: str, weight: str) -> float:
     """The weighted number of entities: the plain sum of their weights.
 
     Output frames are microdf ``MicroDataFrame``s weighted by the entity's
@@ -66,7 +71,7 @@ def _count_output_entities(simulation, entity: str, weight: str) -> float:
     return float(np.asarray(data[weight], dtype=float).sum())
 
 
-def _sum_output_variable(simulation, variable: str, entity: str) -> float:
+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())
@@ -83,7 +88,9 @@ def _sum_output_variable(simulation, variable: str, entity: str) -> float:
     return float(output.result)
 
 
-def _change_output_variable(baseline, reform, variable: str, entity: str) -> float:
+def _change_output_variable(
+    baseline: Simulation, reform: Simulation, variable: str, entity: str
+) -> float:
     baseline_data = _entity_data(baseline, entity)
     reform_data = _entity_data(reform, entity)
     if variable in baseline_data.columns and variable in reform_data.columns:
@@ -102,16 +109,18 @@ def _change_output_variable(baseline, reform, variable: str, entity: str) -> flo
     return float(output.result)
 
 
-def _output_module_function(module_name: str, name: str):
+def _output_module_function(module_name: str, name: str) -> Callable[..., object]:
     module = import_module(f"policyengine.outputs.{module_name}")
     return getattr(module, name)
 
 
-def _poverty_module_function(name: str):
+def _poverty_module_function(name: str) -> Callable[..., object]:
     return _output_module_function("poverty", name)
 
 
-def _try_compute_output(label: str, fn, *args, **kwargs):
+def _try_compute_output[**P, T](
+    label: str, fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs
+) -> T | None:
     try:
         return fn(*args, **kwargs)
     except Exception as exc:
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py
index 59be0d9af..32032e308 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_distribution.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING
 
 from policyengine_simulation_executor.simulation_macro_output import (
     DecileOutput,
@@ -14,6 +14,9 @@
     _try_compute_output,
 )
 
+if TYPE_CHECKING:
+    from policyengine.core import Simulation
+
 INTRA_DECILE_COLUMNS = {
     "Lose more than 5%": "lose_more_than_5pct",
     "Lose less than 5%": "lose_less_than_5pct",
@@ -23,11 +26,11 @@
 }
 
 
-def build_decile(analysis: Any) -> DecileOutput:
+def build_decile(analysis: object) -> DecileOutput:
     return build_decile_output(getattr(analysis, "decile_impacts", None))
 
 
-def build_decile_output(collection: Any) -> DecileOutput:
+def build_decile_output(collection: object) -> DecileOutput:
     if isinstance(collection, DecileOutput):
         return collection
     average: dict[str, float] = {}
@@ -45,7 +48,9 @@ def build_decile_output(collection: Any) -> DecileOutput:
     return DecileOutput(average=average, relative=relative)
 
 
-def build_intra_decile_output(baseline, reform) -> IntraDecileOutput:
+def build_intra_decile_output(
+    baseline: Simulation, reform: Simulation
+) -> IntraDecileOutput:
     from policyengine.outputs.intra_decile_impact import (
         compute_intra_decile_impacts,
     )
@@ -61,21 +66,23 @@ def build_intra_decile_output(baseline, reform) -> IntraDecileOutput:
     return build_intra_decile_output_from_collection(collection)
 
 
-def build_wealth_decile(country: str, wealth_decile: Any) -> DecileOutput | None:
+def build_wealth_decile(country: str, wealth_decile: object) -> DecileOutput | None:
     if country != "uk":
         return None
     return build_decile_output(wealth_decile)
 
 
 def build_intra_wealth_decile(
-    country: str, intra_wealth_decile: Any
+    country: str, intra_wealth_decile: object
 ) -> IntraDecileOutput | None:
     if country != "uk":
         return None
     return build_intra_decile_output_from_collection(intra_wealth_decile)
 
 
-def build_intra_decile_output_from_collection(collection: Any) -> IntraDecileOutput:
+def build_intra_decile_output_from_collection(
+    collection: object,
+) -> IntraDecileOutput:
     if isinstance(collection, IntraDecileOutput):
         return collection
     deciles: dict[str, list[float]] = {label: [] for label in INTRA_DECILE_COLUMNS}
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py
index 4deb42a05..2bf92c1b9 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py
@@ -24,6 +24,7 @@
 )
 
 if TYPE_CHECKING:
+    from policyengine.core import Simulation
     from policyengine.data.uk_geography_assets import UKGeographyAssetSpec
     from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata
 
@@ -89,7 +90,7 @@ def _policyengine_us_district_metadata() -> tuple[dict[int, str], frozenset[str]
     )
 
 
-def build_geographic_impact_output(value: Any) -> GeographicImpactOutput | None:
+def build_geographic_impact_output(value: object) -> GeographicImpactOutput | None:
     if isinstance(value, GeographicImpactOutput):
         return value
     records = _output_model_dump(value)
@@ -105,7 +106,7 @@ def build_geographic_impact_output(value: Any) -> GeographicImpactOutput | None:
 
 
 def build_congressional_district_impact_output(
-    value: Any,
+    value: object,
 ) -> CongressionalDistrictImpactOutput | None:
     if isinstance(value, CongressionalDistrictImpactOutput):
         return value
@@ -187,7 +188,11 @@ def _should_build_us_congressional_district_impact(region_code: str | None) -> b
 
 
 def build_congressional_district_impact(
-    country: str, baseline, reform, *, region_code: str | None = None
+    country: str,
+    baseline: Simulation,
+    reform: Simulation,
+    *,
+    region_code: str | None = None,
 ) -> CongressionalDistrictImpactOutput | None:
     if country != "us":
         return None
@@ -207,7 +212,7 @@ def compute_and_format() -> CongressionalDistrictImpactOutput | None:
     return _try_compute_output("congressional district impacts", compute_and_format)
 
 
-def _output_household(simulation) -> pd.DataFrame:
+def _output_household(simulation: Simulation) -> pd.DataFrame:
     return pd.DataFrame(simulation.output_dataset.data.household)
 
 
@@ -230,7 +235,7 @@ def _require_uk_area_column(
 
 
 def build_uk_constituency_impact(
-    country: str, baseline, reform
+    country: str, baseline: Simulation, reform: Simulation
 ) -> GeographicImpactOutput | None:
     if country != "uk":
         return None
@@ -256,8 +261,8 @@ def build_uk_constituency_impact(
 
 def build_uk_local_authority_impact(
     country: str,
-    baseline,
-    reform,
+    baseline: Simulation,
+    reform: Simulation,
     *,
     uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None,
 ) -> GeographicImpactOutput | None:
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py
index a9d8639ce..7bd1366a5 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_inequality.py
@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 from policyengine_simulation_executor.simulation_macro_output import (
     BaselineReformValue,
     InequalityOutput,
@@ -11,7 +9,7 @@
 from policyengine_simulation_executor.simulation_output_common import _number
 
 
-def build_inequality(analysis: Any) -> InequalityOutput:
+def build_inequality(analysis: object) -> InequalityOutput:
     baseline = getattr(analysis, "baseline_inequality", None)
     reform = getattr(analysis, "reform_inequality", None)
     if isinstance(baseline, InequalityOutput):
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py
index 447ba3e1f..7033fe5ab 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_labor.py
@@ -2,15 +2,13 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 from policyengine_simulation_executor.simulation_macro_output import (
     LaborSupplyResponseOutput,
 )
 from policyengine_simulation_executor.simulation_output_common import _output_model_dump
 
 
-def build_labor_supply_response(analysis: Any) -> LaborSupplyResponseOutput | None:
+def build_labor_supply_response(analysis: object) -> LaborSupplyResponseOutput | None:
     labor_supply_response = getattr(analysis, "labor_supply_response", None)
     if isinstance(labor_supply_response, LaborSupplyResponseOutput):
         return labor_supply_response
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py
index de63db21d..fb260ce52 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_poverty.py
@@ -24,6 +24,8 @@
 if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
 
+    from policyengine.core import Simulation
+
 US_POVERTY_TYPES = {
     "spm": "poverty",
     "spm_deep": "deep_poverty",
@@ -110,7 +112,9 @@ def _race_poverty_output(values: dict[str, dict[str, float]]) -> RacePovertyOutp
     )
 
 
-def build_poverty_outputs(country: str, baseline, reform, analysis: Any):
+def build_poverty_outputs(
+    country: str, baseline: Simulation, reform: Simulation, analysis: object
+) -> PovertyModuleOutputs:
     prefix = "us" if country == "us" else "uk"
     baseline_poverty_by_age = _try_compute_output(
         "baseline poverty by age",
@@ -172,10 +176,10 @@ def build_poverty_outputs(country: str, baseline, reform, analysis: Any):
 def build_poverty_output(
     *,
     country: str,
-    baseline: Any,
-    reform: Any,
-    baseline_by_age: Any,
-    reform_by_age: Any,
+    baseline: object,
+    reform: object,
+    baseline_by_age: object,
+    reform_by_age: object,
 ) -> PovertyOutput:
     if isinstance(baseline, PovertyOutput):
         return baseline
@@ -206,8 +210,8 @@ def build_poverty_output(
 def build_poverty_by_gender_output(
     *,
     country: str,
-    baseline_by_gender: Any,
-    reform_by_gender: Any,
+    baseline_by_gender: object,
+    reform_by_gender: object,
 ) -> PovertyByGenderOutput:
     if isinstance(baseline_by_gender, PovertyByGenderOutput):
         return baseline_by_gender
@@ -230,8 +234,8 @@ def build_poverty_by_gender_output(
 
 def build_poverty_by_race_output(
     *,
-    baseline_by_race: Any,
-    reform_by_race: Any,
+    baseline_by_race: object,
+    reform_by_race: object,
 ) -> PovertyByRaceOutput:
     if isinstance(baseline_by_race, PovertyByRaceOutput):
         return baseline_by_race
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py
index 6004f9283..871656ae5 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py
@@ -6,6 +6,8 @@
 No Modal dependencies here.
 """
 
+from __future__ import annotations
+
 import contextlib
 import hashlib
 import json
@@ -13,13 +15,11 @@
 import os
 import tempfile
 import time
-from collections.abc import Iterator
 from dataclasses import dataclass
 from datetime import datetime
 from importlib import import_module
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from policyengine_observability import ObservabilityRuntime
 from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri
 from policyengine_simulation_observability.stages import (
     ANNUAL_IMPACT_STAGES,
@@ -35,6 +35,16 @@
     SimulationOutputBuilder,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from types import ModuleType
+
+    from policyengine.core import Dataset, Policy, Region, Simulation
+    from policyengine.core.scoping_strategy import ScopingStrategy
+    from policyengine.provenance.manifest import ArtifactPathReference
+    from policyengine_observability import ObservabilityRuntime
+    from policyengine_simulation_observability.telemetry import TelemetryEnvelope
+
 logger = logging.getLogger(__name__)
 
 os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1")
@@ -182,7 +192,7 @@ def _parse_year(params: dict[str, Any]) -> int:
     return int(value)
 
 
-def _parse_policy_period(period_key: Any) -> tuple[datetime, datetime | None]:
+def _parse_policy_period(period_key: object) -> tuple[datetime, datetime | None]:
     """Parse effective dates or inclusive ``start.stop`` policy periods."""
     text = str(period_key)
     parts = text.split(".")
@@ -201,7 +211,7 @@ def _parse_policy_period(period_key: Any) -> tuple[datetime, datetime | None]:
     return start, stop
 
 
-def _normalise_period_key(period_key: Any) -> str:
+def _normalise_period_key(period_key: object) -> str:
     """Validate policy period keys without discarding their stop dates."""
     text = str(period_key)
     _parse_policy_period(text)
@@ -224,7 +234,9 @@ def _normalise_policy(policy: dict[str, Any] | None) -> dict[str, Any] | None:
     return normalised
 
 
-def _compile_interval_policy(policy, *, country_module, year: int):
+def _compile_interval_policy(
+    policy: dict[str, Any] | None, *, country_module: ModuleType, year: int
+) -> dict[str, Any] | Policy | None:
     """Use public wrapper values for intervals its dict compiler cannot parse."""
     if not policy or not any(
         isinstance(value, dict) and any("." in str(key) for key in value)
@@ -260,7 +272,9 @@ def _compile_interval_policy(policy, *, country_module, year: int):
     return Policy(name="Reform", parameter_values=parameter_values)
 
 
-def _runtime_policy(policy, *, country_module, year: int):
+def _runtime_policy(
+    policy: dict[str, Any] | None, *, country_module: ModuleType, year: int
+) -> dict[str, Any] | Policy | None:
     """A request's policy as the wrapper must receive it.
 
     Normalises the period keys, then compiles a policy with interval keys
@@ -276,7 +290,7 @@ def _resolve_dataset_reference(country: str) -> str:
     return resolve_bundle_dataset_name(country, None)
 
 
-def _normalise_region_code(country: str, region: Any) -> str:
+def _normalise_region_code(country: str, region: object) -> str:
     if region is None or str(region).strip() == "":
         return country
 
@@ -305,7 +319,7 @@ def _normalise_region_code(country: str, region: Any) -> str:
     return f"{prefix}/{value}"
 
 
-def _build_uk_weight_replacement_region(region_code: str):
+def _build_uk_weight_replacement_region(region_code: str) -> Region | None:
     if "/" not in region_code:
         return None
 
@@ -345,7 +359,9 @@ def _build_uk_weight_replacement_region(region_code: str):
 
 
 def _require_certified_uk_weight_matrix(
-    scoping_strategy, dataset, dataset_selection
+    scoping_strategy: ScopingStrategy | None,
+    dataset: Dataset,
+    dataset_selection: DatasetSelection,
 ) -> None:
     """Bind a weight-matrix region to the release of the dataset it reweights.
 
@@ -385,7 +401,7 @@ def _require_certified_uk_weight_matrix(
     manifest = get_release_manifest("uk")
     package = manifest.data_package
 
-    def release(reference) -> tuple[str, str]:
+    def release(reference: ArtifactPathReference) -> tuple[str, str]:
         return (
             reference.repo_id or package.repo_id,
             reference.revision or package.release_manifest_revision or package.version,
@@ -454,9 +470,9 @@ def release(reference) -> tuple[str, str]:
 
 
 def _region_parent_dataset_reference(
-    country_module,
+    country_module: ModuleType,
     country: str,
-    region,
+    region: Region,
 ) -> str:
     parent_code = getattr(region, "parent_code", None)
     visited: set[str] = set()
@@ -481,7 +497,7 @@ def _region_parent_dataset_reference(
     return _resolve_dataset_reference(country)
 
 
-def _reject_unscoped_us_place_region(region_code: str, region) -> None:
+def _reject_unscoped_us_place_region(region_code: str, region: Region) -> None:
     if not region_code.startswith("place/"):
         return
     if getattr(region, "dataset_path", None) is not None:
@@ -496,7 +512,7 @@ def _reject_unscoped_us_place_region(region_code: str, region) -> None:
 
 def _resolve_region(
     *,
-    country_module,
+    country_module: ModuleType,
     country: str,
     params: dict[str, Any],
 ) -> RegionResolution:
@@ -544,7 +560,7 @@ def _resolve_region(
 
 def _resolve_region_group(
     *,
-    country_module,
+    country_module: ModuleType,
     country: str,
     params: dict[str, Any],
 ) -> RegionResolution:
@@ -588,7 +604,7 @@ def _resolve_region_group(
     )
 
 
-def _country_module(country: str):
+def _country_module(country: str) -> ModuleType:
     country = country.lower()
     if country not in {"us", "uk"}:
         raise ValueError(f"Unsupported country: {country}")
@@ -649,8 +665,8 @@ def _load_dataset(
     params: dict[str, Any],
     *,
     selection: DatasetSelection,
-    country_module=None,
-):
+    country_module: ModuleType | None = None,
+) -> Dataset:
     country = params.get("country", "us").lower()
     year = _parse_year(params)
     country_module = country_module or _country_module(country)
@@ -678,12 +694,12 @@ def _load_dataset(
 def _build_simulation(
     params: dict[str, Any],
     *,
-    dataset,
+    dataset: Dataset,
     dataset_selection: DatasetSelection,
     policy: dict[str, Any] | None,
-    scoping_strategy=None,
+    scoping_strategy: ScopingStrategy | None = None,
     region_code: str | None = None,
-):
+) -> Simulation:
     from policyengine.core import Simulation
 
     from policyengine_simulation_executor.baseline_artifacts import (
@@ -694,7 +710,7 @@ def _build_simulation(
 
     selection = normalize_runtime_spm(params)
     params = {**params, **({"spm": selection} if selection is not None else {})}
-    spm_kwargs = {"spm": selection} if selection is not None else {}
+    spm_kwargs: dict[str, Any] = {"spm": selection} if selection is not None else {}
     country = params.get("country", "us")
     country_module = _country_module(country)
     policy = _compile_interval_policy(
@@ -862,7 +878,7 @@ def _set_runtime_attributes(
     *,
     runtime: ObservabilityRuntime,
     simulation_params: dict[str, Any],
-    telemetry,
+    telemetry: TelemetryEnvelope | None,
     metadata: dict[str, Any],
     country: str,
 ) -> None:
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 d76ad3fcb..8cb5696f8 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py
@@ -1,8 +1,11 @@
 """Certified SPM runtime selection, independent of caller-supplied metadata."""
 
+from __future__ import annotations
+
 import json
 from collections.abc import Mapping
 from functools import lru_cache
+from typing import TYPE_CHECKING, Any
 
 from policyengine_simulation_contract.spm import (
     SPMCapability,
@@ -17,16 +20,19 @@
     validate_spm_calculation_provenance,
 )
 
+if TYPE_CHECKING:
+    from spm_calculator.rolling_forecast import SPMForecast
+
 
 @lru_cache(maxsize=4)
-def _forecast(expected_sha256):
+def _forecast(expected_sha256: str | None) -> SPMForecast:
     from spm_calculator import load_forecast
 
     return load_forecast(expected_sha256=expected_sha256)
 
 
 @lru_cache(maxsize=1)
-def runtime_spm_capability():
+def runtime_spm_capability() -> SPMCapability | None:
     """The installed image's certified SPM capability, or None.
 
     Memoized for the life of the process because every input is a property
@@ -83,7 +89,9 @@ def runtime_spm_capability():
 
 
 @lru_cache(maxsize=64)
-def _prevalidate_selection(selection_json: str, start_year: int, window_size: int):
+def _prevalidate_selection(
+    selection_json: str, start_year: int, window_size: int
+) -> None:
     """Prove a resolved selection can be measured, before anything expensive.
 
     Keyed on the resolved selection and the year range, which is everything
@@ -139,7 +147,7 @@ def _prevalidate_selection(selection_json: str, start_year: int, window_size: in
 )
 
 
-def reset_spm_runtime_caches():
+def reset_spm_runtime_caches() -> None:
     """Forget what was memoized from the installed environment.
 
     Production never needs this: a container's image is fixed for the life
@@ -154,7 +162,7 @@ def reset_spm_runtime_caches():
         clear()
 
 
-def normalize_runtime_spm(params):
+def normalize_runtime_spm(params: dict[str, Any]) -> dict[str, Any] | None:
     """Resolve before dataset loading, artifact lookup, or child submission."""
     from policyengine_simulation_executor.release_bundle import (
         get_country_release_bundle,
@@ -277,7 +285,13 @@ def simulation_spm_calculation_provenance(
         ) from exc
 
 
-def simulation_spm_result(baseline, reform, selection, *, expected_year=None):
+def simulation_spm_result(
+    baseline: object,
+    reform: object,
+    selection: object,
+    *,
+    expected_year: int | None = None,
+) -> dict[str, object]:
     if selection is None:
         return {}
     baseline_provenance = simulation_spm_calculation_provenance(baseline, selection)
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py
index 798c091af..6d88745ce 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_artifacts.py
@@ -239,7 +239,7 @@ def deserialize_uk_local_authority_metadata(
 
 
 class Stage12ArtifactStore:
-    def __init__(self, bucket_name: str, *, store: ArtifactStore | None = None):
+    def __init__(self, bucket_name: str, *, store: ArtifactStore | None = None) -> None:
         if not bucket_name:
             raise ValueError(
                 "Stage 12 requires an explicitly configured private bucket"
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py
index b61a3c0f1..4c8b05993 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py
@@ -15,6 +15,8 @@
     from collections.abc import Mapping
 
     import pandas as pd
+    from policyengine.tax_benefit_models.uk.datasets import PolicyEngineUKDataset
+    from policyengine.tax_benefit_models.us.datasets import PolicyEngineUSDataset
     from policyengine_simulation_contract.stage12_bundle import CountryId
     from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata
     from pydantic import JsonValue
@@ -86,7 +88,7 @@ def _dataset_from_frames(
     *,
     country: CountryId,
     year: int,
-):
+) -> PolicyEngineUSDataset | PolicyEngineUKDataset:
     from microdf import MicroDataFrame
 
     fields = {}
@@ -159,7 +161,10 @@ def build_aggregate_report(
         ),
     }
 
-    def stand_in(dataset, policy: dict[str, JsonValue]):
+    def stand_in(
+        dataset: PolicyEngineUSDataset | PolicyEngineUKDataset,
+        policy: dict[str, JsonValue],
+    ) -> PrecomputedSimulation:
         simulation = PrecomputedSimulation(
             dataset=dataset,
             tax_benefit_model_version=country_module.model,
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py
index 3063ddbd3..1e6aa584a 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py
@@ -60,6 +60,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
+    from contextlib import AbstractContextManager
 
     from policyengine_observability import ObservabilityRuntime
 
@@ -280,7 +281,7 @@ def coordinate_report(
         def stage_scope(
             stage: Stage,
             attributes: Mapping[str, object] | None = None,
-        ):
+        ) -> AbstractContextManager[object]:
             if runtime is None:
                 return nullcontext()
             return runtime.operation(
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py
index 8f97f2237..e30e3c6f2 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/uk_local_authority_metadata.py
@@ -10,6 +10,7 @@
 from functools import lru_cache
 from importlib.resources import files
 from types import MappingProxyType
+from typing import TYPE_CHECKING
 
 import pandas as pd
 from policyengine_simulation_contract.uk_geography import (
@@ -17,6 +18,9 @@
     UKLocalAuthorityMetadata,
 )
 
+if TYPE_CHECKING:
+    from importlib.resources.abc import Traversable
+
 _CODE_PATTERN = re.compile(r"^[A-Z]\d{8}$")
 _RESOURCE_PACKAGE = "policyengine_simulation_executor"
 _RESOURCE_DIRECTORY = ("static_runtime_files", "uk_local_authorities")
@@ -116,7 +120,7 @@ def metadata_for(
         )
 
 
-def _resource_path(filename: str):
+def _resource_path(filename: str) -> Traversable:
     path = files(_RESOURCE_PACKAGE)
     for part in _RESOURCE_DIRECTORY:
         path = path.joinpath(part)

From bb4b718af92d65ba4e0bd8cea66e8b1b9836802a Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:57:24 +0400
Subject: [PATCH 07/12] Enable ruff group 3 (annotations)

Select ANN. Tests are exempt, alongside the docstring, magic-number and
assert rules a later group will enable. scripts/rules findings are left
for the next sync from policyengine-api.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ruff.toml | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/ruff.toml b/ruff.toml
index a22dc8cc2..2aa803cb6 100644
--- a/ruff.toml
+++ b/ruff.toml
@@ -13,8 +13,15 @@ select = [
     "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "UP", "W",
     # Group 2: typing-only imports.
     "TC",
+    # Group 3: annotations.
+    "ANN",
 ]
 
+[lint.per-file-ignores]
+# Tests may use pytest's bare assert, and need no docstrings, annotations or
+# named constants for magic numbers. Every other rule applies to tests.
+"**/tests/**" = ["ANN", "D1", "PLR2004", "S101"]
+
 [lint.flake8-type-checking]
 # Classes whose annotations Pydantic evaluates at runtime, so their imports must
 # stay at runtime. Ruff checks direct bases only, so every project model base

From b4eb91c50325109a6afd909de7c3cfa5b8b58600 Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:58:18 +0400
Subject: [PATCH 08/12] Annotate the package-update script and the Stage 12
 proof fixtures

_json_from_stdin returns object (every caller narrows it), and the proof
fixtures annotate their SQLAlchemy connection and table parameters. The
schema-authority proof reads the fixtures as text, so it is unaffected.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .github/scripts/update_policyengine_package.py           | 2 +-
 .../sim.stage12.no-schema-ddl/bad_stage12_runtime.py     | 9 ++++++++-
 .../sim.stage12.no-schema-ddl/good_stage12_runtime.py    | 9 ++++++++-
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/.github/scripts/update_policyengine_package.py b/.github/scripts/update_policyengine_package.py
index 1d3ab3ce6..56197deb3 100644
--- a/.github/scripts/update_policyengine_package.py
+++ b/.github/scripts/update_policyengine_package.py
@@ -28,7 +28,7 @@ def _mapping(value: object) -> Mapping[str, Any]:
     return value if isinstance(value, Mapping) else {}
 
 
-def _json_from_stdin() -> Any:
+def _json_from_stdin() -> object:
     return json.load(sys.stdin)
 
 
diff --git a/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py b/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py
index f6b4b7262..8e6ea09d6 100644
--- a/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py
+++ b/rules/fixtures/sim.stage12.no-schema-ddl/bad_stage12_runtime.py
@@ -1,7 +1,14 @@
 """Proof fixture: a Stage 12 runtime that creates its own table. The rule must reject it."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from sqlalchemy import text
 
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
+
 
-def ensure_comparison_table(connection) -> None:
+def ensure_comparison_table(connection: Connection) -> None:
     connection.execute(text("CREATE TABLE stage12_evaluation_reports (id uuid)"))
diff --git a/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py b/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py
index 6dddf8588..b8cf2de0c 100644
--- a/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py
+++ b/rules/fixtures/sim.stage12.no-schema-ddl/good_stage12_runtime.py
@@ -1,7 +1,14 @@
 """Proof fixture: a Stage 12 runtime that writes rows with typed DML. The rule must accept it."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from sqlalchemy import insert
 
+if TYPE_CHECKING:
+    from sqlalchemy import Connection, Table
+
 
-def record_report(connection, reports_table, report_id: str) -> None:
+def record_report(connection: Connection, reports_table: Table, report_id: str) -> None:
     connection.execute(insert(reports_table).values(id=report_id))

From 28a9c83e44443423c61e0a3dcdc891880549d56b Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:34:38 +0400
Subject: [PATCH 09/12] Sync the annotated rules tool from policyengine-api

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 scripts/rules/guard.py |  4 ++--
 scripts/rules/load.py  | 21 +++++++++++++--------
 scripts/rules/model.py |  6 +++++-
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/scripts/rules/guard.py b/scripts/rules/guard.py
index 42d4b6a37..c68fce02b 100644
--- a/scripts/rules/guard.py
+++ b/scripts/rules/guard.py
@@ -16,7 +16,7 @@
 )
 from scripts.rules.discover import discover_hidden
 from scripts.rules.load import load_registry
-from scripts.rules.model import Violation
+from scripts.rules.model import Registry, Violation
 from scripts.rules.render import (
     AGENTS_MD,
     RULES_MD,
@@ -42,7 +42,7 @@ class GuardResult:
     notes: list[str] = field(default_factory=list)
 
 
-def _stale_outputs(registry, root: Path) -> list[Violation]:
+def _stale_outputs(registry: Registry, root: Path) -> list[Violation]:
     violations = []
     rules_md = root / RULES_MD
     if not rules_md.exists():
diff --git a/scripts/rules/load.py b/scripts/rules/load.py
index 267818647..83b455e2e 100644
--- a/scripts/rules/load.py
+++ b/scripts/rules/load.py
@@ -8,8 +8,9 @@
 
 import tomllib
 from datetime import date
+from enum import Enum
 from pathlib import Path
-from typing import Any
+from typing import Any, TypeVar
 
 from scripts.rules.model import (
     Agents,
@@ -30,6 +31,8 @@
     Violation,
 )
 
+_EnumT = TypeVar("_EnumT", bound=Enum)
+
 REGISTRY_PATH = Path("rules/registry.toml")
 
 _TOP_KEYS = {
@@ -127,7 +130,9 @@ def string_list(
         )
         return ()
 
-    def enum(self, rule_id: str | None, enum_type: Any, value: Any, where: str) -> Any:
+    def enum(
+        self, rule_id: str | None, enum_type: type[_EnumT], value: object, where: str
+    ) -> _EnumT | None:
         try:
             return enum_type(value)
         except ValueError:
@@ -135,7 +140,7 @@ def enum(self, rule_id: str | None, enum_type: Any, value: Any, where: str) -> A
             self.add(rule_id, f"{where} is {value!r}", f"use one of {allowed}")
             return None
 
-    def day(self, rule_id: str | None, value: Any, where: str) -> date | None:
+    def day(self, rule_id: str | None, value: object, where: str) -> date | None:
         if isinstance(value, date):
             return value
         if isinstance(value, str):
@@ -147,7 +152,7 @@ def day(self, rule_id: str | None, value: Any, where: str) -> date | None:
         return None
 
 
-def _ignored(c: _Collector, raw: Any, key: str, field: str) -> tuple[Ignored, ...]:
+def _ignored(c: _Collector, raw: object, key: str, field: str) -> tuple[Ignored, ...]:
     if raw is None:
         return ()
     items = []
@@ -168,7 +173,7 @@ def _ignored(c: _Collector, raw: Any, key: str, field: str) -> tuple[Ignored, ..
     return tuple(items)
 
 
-def _discovery(c: _Collector, raw: Any) -> Discovery:
+def _discovery(c: _Collector, raw: object) -> Discovery:
     if raw is None:
         return Discovery()
     if not isinstance(raw, dict):
@@ -189,7 +194,7 @@ def _discovery(c: _Collector, raw: Any) -> Discovery:
     )
 
 
-def _agents(c: _Collector, raw: Any) -> Agents | None:
+def _agents(c: _Collector, raw: object) -> Agents | None:
     if raw is None:
         return None
     if not isinstance(raw, dict):
@@ -207,7 +212,7 @@ def _agents(c: _Collector, raw: Any) -> Agents | None:
     )
 
 
-def _ci_scripts(c: _Collector, raw: Any) -> CiScripts | None:
+def _ci_scripts(c: _Collector, raw: object) -> CiScripts | None:
     if raw is None:
         return None
     if not isinstance(raw, dict):
@@ -225,7 +230,7 @@ def _ci_scripts(c: _Collector, raw: Any) -> CiScripts | None:
     )
 
 
-def _rule(c: _Collector, raw: Any, index: int) -> Rule | None:
+def _rule(c: _Collector, raw: object, index: int) -> Rule | None:
     if not isinstance(raw, dict):
         c.add(None, f"rule #{index + 1} must be a table", "write [[rule]]")
         return None
diff --git a/scripts/rules/model.py b/scripts/rules/model.py
index 619138b01..6d8f9c365 100644
--- a/scripts/rules/model.py
+++ b/scripts/rules/model.py
@@ -100,7 +100,11 @@ class Enforcer:
     kind: Kind
     fields: Mapping[str, Any] = field(default_factory=dict)
 
-    def get(self, key: str, default: Any = None) -> Any:
+    def get(
+        self,
+        key: str,
+        default: object = None,
+    ) -> Any:  # noqa: ANN401  # a raw TOML value that callers read by its shape
         return self.fields.get(key, default)
 
     @property

From d9de4e5dc14d7e5c5f6a64e172bd48a5aaa491fb Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Tue, 13 Oct 2026 16:34:38 +0400
Subject: [PATCH 10/12] Register ANN and the test-file ruff exemptions

sim.quality.ruff-lint now includes ANN, and the **/tests/** per-file
ignore list (S101, D1, ANN, PLR2004) is registered as
sim.quality.tests-ruff-exemptions.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 AGENTS.md               |  3 ++-
 changelog_entry.yaml    |  1 +
 docs/generated/rules.md | 17 +++++++++++------
 rules/registry.toml     | 21 ++++++++++++++++++++-
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/AGENTS.md b/AGENTS.md
index d5cb893ea..0bae03c1c 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -138,8 +138,9 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`.
 - `sim.quality.every-project-tested`: Every project and library with unit tests runs them in the PR test matrix. (doc-only, not checked: follow it yourself)
 - `sim.quality.ruff-format-src`: Source code under projects/*/src and libs/*/src is formatted with ruff format. (blocking: ruff format)
 - `sim.quality.ruff-format-tests`: Test code is formatted with ruff format too; CI currently checks only src directories. (doc-only, not checked: follow it yourself)
-- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
+- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W))
 - `sim.quality.tests-pass`: Each project's unit tests pass in its own locked environment. (blocking: CI job `test`)
+- `sim.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**`)
 
 ### CI workflows and scripts
 
diff --git a/changelog_entry.yaml b/changelog_entry.yaml
index 78db3d851..7ed733326 100644
--- a/changelog_entry.yaml
+++ b/changelog_entry.yaml
@@ -4,6 +4,7 @@
       - Made CLAUDE.md and the Copilot instructions thin pointers to AGENTS.md, required workflow logic longer than one line to live in tested and shellchecked scripts, and registered the tracking-issue rule
       - Added a rules registry (rules/registry.toml) with a standard-library tool that validates it, checks every rule is enforced where it claims, and renders docs/generated/rules.md and the Rules block in AGENTS.md
     changed:
+      - Enabled ruff's annotation rules (ANN) and annotated every function outside tests; test files are exempt from ANN, missing docstrings, magic numbers and bare-assert rules
       - Enabled ruff's typing-only import rules (TC) and moved annotation-only imports under TYPE_CHECKING, keeping runtime imports for Pydantic models and FastAPI endpoints
       - Ran ruff check in CI with the first group of rule families (imports, pyupgrade and other safe auto-fixes) and fixed every finding
       - Moved the multi-line `run` blocks of the simulation deployment workflows into single-line steps and tested scripts without changing their behaviour or outputs, leaving only the release-sentinel and routing-rollback blocks in the one-line `run` rule's baseline
diff --git a/docs/generated/rules.md b/docs/generated/rules.md
index 6dac8864c..1d979899d 100644
--- a/docs/generated/rules.md
+++ b/docs/generated/rules.md
@@ -8,11 +8,11 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 
 | Tier | enforced | ratcheted | advisory | doc-only | Total |
 | --- | ---: | ---: | ---: | ---: | ---: |
-| mechanical | 32 | 1 | 0 | 14 | 47 |
+| mechanical | 33 | 1 | 0 | 14 | 48 |
 | metric | 0 | 0 | 0 | 0 | 0 |
 | judgment | 0 | 0 | 0 | 0 | 0 |
 
-**Blocking:** 33 of 47 rules.
+**Blocking:** 34 of 48 rules.
 
 ## Gaps
 
@@ -245,8 +245,9 @@ flowchart LR
 | `sim.quality.every-project-tested` | Every project and library with unit tests runs them in the PR test matrix. | mechanical | doc-only | no | nothing | — | — |
 | `sim.quality.ruff-format-src` | Source code under projects/*/src and libs/*/src is formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — |
 | `sim.quality.ruff-format-tests` | Test code is formatted with ruff format too; CI currently checks only src directories. | mechanical | doc-only | no | nothing | — | — |
-| `sim.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
+| `sim.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — |
 | `sim.quality.tests-pass` | Each project's unit tests pass in its own locked environment. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — |
+| `sim.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**` | `pr.yml:lint` | — |
 
 ```mermaid
 flowchart LR
@@ -257,18 +258,22 @@ flowchart LR
   j_CI_pr_yml_lint["CI pr.yml:lint"]
   r_sim_quality_ruff_format_tests("sim.quality.ruff-format-tests (doc-only)")
   r_sim_quality_ruff_lint("sim.quality.ruff-lint (enforced)")
-  e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"]
+  e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"]
   r_sim_quality_tests_pass("sim.quality.tests-pass (enforced)")
   e_CI_job_test["CI job test"]
   j_CI_pr_yml_test["CI pr.yml:test"]
+  r_sim_quality_tests_ruff_exemptions("sim.quality.tests-ruff-exemptions (enforced)")
+  e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___["ruff per-file ignore of ANN, D1, PLR2004, S101 for **/tests/**"]
   r_sim_quality_every_project_tested -.-> gap_unenforced
   r_sim_quality_ruff_format_src --> e_ruff_format
   e_ruff_format --> j_CI_pr_yml_lint
   r_sim_quality_ruff_format_tests -.-> gap_unenforced
-  r_sim_quality_ruff_lint --> e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_
-  e_ruff_check__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint
+  r_sim_quality_ruff_lint --> e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_
+  e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint
   r_sim_quality_tests_pass --> e_CI_job_test
   e_CI_job_test --> j_CI_pr_yml_test
+  r_sim_quality_tests_ruff_exemptions --> e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___
+  e_ruff_per_file_ignore_of_ANN__D1__PLR2004__S101_for____tests___ --> j_CI_pr_yml_lint
 ```
 
 ## CI workflows and scripts
diff --git a/rules/registry.toml b/rules/registry.toml
index 10879a735..ae72bb9f1 100644
--- a/rules/registry.toml
+++ b/rules/registry.toml
@@ -710,7 +710,26 @@ autofix = true
 [rule.enforcer]
 kind = "ruff"
 mode = "lint"
-codes = ["E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"]
+codes = ["ANN", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"]
+[[rule.ci]]
+workflow = ".github/workflows/pr.yml"
+job = "lint"
+invocation = "ruff check ."
+
+[[rule]]
+id = "sim.quality.tests-ruff-exemptions"
+category = "quality"
+statement = "Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests."
+rationale = "docs/engineering/skills/rules-registry.md#baselines"
+scope = ["projects/*/tests/**/*.py", "libs/*/tests/**/*.py", "scripts/rules/tests/**/*.py"]
+tier = "mechanical"
+status = "enforced"
+autofix = false
+[rule.enforcer]
+kind = "ruff"
+mode = "per-file-ignore"
+pattern = "**/tests/**"
+codes = ["ANN", "D1", "PLR2004", "S101"]
 [[rule.ci]]
 workflow = ".github/workflows/pr.yml"
 job = "lint"

From 9cad8d334172b912ecbe698215f66a5bdf192d0f Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Wed, 14 Oct 2026 20:32:51 +0400
Subject: [PATCH 11/12] Annotate the routing and bundle helpers #743 and #744
 added

#743 and #744 landed on main after this stack branched. ANN, which this
PR enables, wants annotations on #744's routing-store, publication-journal
and active-state helpers and on #743's bundle dataset-digest helper. The
Modal store annotations are quoted, as in the gateway, because the
executor tests replace the modal module with a stub.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../src/modal/utils/update_version_registry.py       | 12 +++++++-----
 .../release_bundle.py                                |  2 +-
 2 files changed, 8 insertions(+), 6 deletions(-)

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 71641b504..b6cf9069b 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
@@ -439,7 +439,7 @@ def validate_routing_plan(plan: object) -> dict:
     return plan
 
 
-def _routing_store(environment: str, *, create_if_missing: bool):
+def _routing_store(environment: str, *, create_if_missing: bool) -> "modal.Dict | None":
     try:
         return modal.Dict.from_name(
             ROUTING_STATE_DICT_NAME,
@@ -501,7 +501,7 @@ def _require_serialized_routing_workflow() -> None:
     )
 
 
-def _unfinished_publication(store) -> dict | None:
+def _unfinished_publication(store: "modal.Dict | None") -> dict | None:
     record = store.get(ROUTING_PUBLICATION_KEY) if store is not None else None
     if record is None:
         return None
@@ -517,7 +517,9 @@ def _unfinished_publication(store) -> dict | None:
     return record["plan"] if record["status"] == "pending" else None
 
 
-def _require_finished_publication(store, *, allowed_plans: Iterable[dict] = ()) -> None:
+def _require_finished_publication(
+    store: "modal.Dict | None", *, allowed_plans: Iterable[dict] = ()
+) -> None:
     pending = _unfinished_publication(store)
     if pending is not None and not any(pending == plan for plan in allowed_plans):
         raise RoutingConflict(
@@ -528,7 +530,7 @@ def _require_finished_publication(store, *, allowed_plans: Iterable[dict] = ())
         )
 
 
-def _write_publication_record(store, *, plan: dict, status: str) -> None:
+def _write_publication_record(store: "modal.Dict", *, plan: dict, status: str) -> None:
     record = {"status": status, "plan": _json_copy(plan)}
     store[ROUTING_PUBLICATION_KEY] = record
     if store.get(ROUTING_PUBLICATION_KEY) != record:
@@ -1175,7 +1177,7 @@ def main() -> None:
     )
     args = parser.parse_args()
 
-    def active_state():
+    def active_state() -> object:
         snapshot = _modal_dict_snapshot(
             name=ROUTING_STATE_DICT_NAME, environment=args.environment
         )
diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py
index 102cc3366..522b6df3e 100644
--- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py
+++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/release_bundle.py
@@ -70,7 +70,7 @@ def _default_dataset_sha256(
     *,
     default_dataset: str,
     data_release: Mapping,
-    manifest,
+    manifest: CountryReleaseManifest,
 ) -> str | None:
     """The certified SHA-256 of the default dataset, from the bundle.
 

From 80d7ab171d6f4bdd8c32cd9ce5c67960cd22f10c Mon Sep 17 00:00:00 2001
From: Anthony Volk <14987227+anth-volk@users.noreply.github.com>
Date: Thu, 15 Oct 2026 00:24:46 +0400
Subject: [PATCH 12/12] Annotate the shared object store fake #787 added

#787 landed on main after this stack branched. ANN, which this PR enables,
wants annotations on the in-memory object store it moved into the executor's
fixtures. They match the byte interface of the real ArtifactStore it stands
in for.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../fixtures/fake_object_store.py                             | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/projects/policyengine-simulation-executor/fixtures/fake_object_store.py b/projects/policyengine-simulation-executor/fixtures/fake_object_store.py
index eb29be88b..23cbca3e9 100644
--- a/projects/policyengine-simulation-executor/fixtures/fake_object_store.py
+++ b/projects/policyengine-simulation-executor/fixtures/fake_object_store.py
@@ -12,7 +12,7 @@ class InMemoryObjectStore:
     def __init__(self) -> None:
         self.objects: dict[str, bytes] = {}
 
-    def upload_bytes(self, path, payload, *, content_type):
+    def upload_bytes(self, path: str, payload: bytes, *, content_type: str) -> bool:
         # The real store uploads with ``if_generation_match=0``: an existing
         # object is left in place and the call reports that it created nothing.
         if path in self.objects:
@@ -20,5 +20,5 @@ def upload_bytes(self, path, payload, *, content_type):
         self.objects[path] = payload
         return True
 
-    def read_bytes(self, path):
+    def read_bytes(self, path: str) -> bytes | None:
         return self.objects.get(path)