Upgrade the simulation runtime to PolicyEngine 6.1.2

.github/scripts/modal-extract-versions.sh

@@ -1,9 +1,10 @@
 #!/bin/bash
 # Extract policyengine.py, policyengine-core, country package versions, and
-# country data release versions.
+# SPM calculator, and country data release versions.
 # Usage: ./modal-extract-versions.sh <project-dir>
 # Outputs: Sets policyengine_version, policyengine_core_version, us_version,
-# us_data_version, uk_version, and uk_data_version in GITHUB_OUTPUT
+# spm_calculator_version, us_data_version, uk_version, and uk_data_version in
+# GITHUB_OUTPUT
 
 set -euo pipefail
 

.github/scripts/update-policyengine-package.sh

@@ -32,12 +32,109 @@ PROJECT_DIR="${PROJECT_DIR:-projects/policyengine-simulation-executor}"
 PROJECT_PATH="${ROOT_DIR}/${PROJECT_DIR}"
 PYPROJECT="${PROJECT_PATH}/pyproject.toml"
 LOCKFILE="${PROJECT_PATH}/uv.lock"
+REPOSITORY="${GITHUB_REPOSITORY:-PolicyEngine/policyengine-sim-api}"
+ISSUE_NUMBER=""
+
+ensure_update_issue() {
+  local issue_details
+  local issue_title
+  local issue_url
+
+  issue_title="Update policyengine to ${LATEST}"
+  ISSUE_NUMBER=$(
+    gh api --paginate --slurp \
+      "repos/${REPOSITORY}/issues?state=open&per_page=100" \
+      | python3 -c '
+import json
+import sys
+
+title = sys.argv[1]
+pages = json.load(sys.stdin)
+matches = sorted(
+    item["number"]
+    for page in pages
+    for item in page
+    if "pull_request" not in item and item.get("title") == title
+)
+if matches:
+    print(matches[0])
+' "$issue_title"
+  )
+
+  if [[ -z "$ISSUE_NUMBER" ]]; then
+    issue_url=$(gh issue create \
+      --repo "$REPOSITORY" \
+      --title "$issue_title" \
+      --body "Track the automated simulation runtime update to policyengine ${LATEST}.")
+    ISSUE_NUMBER="${issue_url##*/}"
+  fi
+
+  if [[ ! "$ISSUE_NUMBER" =~ ^[0-9]+$ ]]; then
+    echo "ERROR: Could not resolve an issue for policyengine ${LATEST}." >&2
+    exit 1
+  fi
+
+  issue_details=$(gh issue view "$ISSUE_NUMBER" \
+    --repo "$REPOSITORY" \
+    --json number,state,title)
+  printf '%s' "$issue_details" | python3 -c '
+import json
+import sys
+
+expected_number = int(sys.argv[1])
+expected_title = sys.argv[2]
+issue = json.load(sys.stdin)
+if issue.get("number") != expected_number:
+    raise SystemExit("Resolved update issue has an unexpected number")
+if issue.get("state") != "OPEN":
+    raise SystemExit("Resolved update issue is not open")
+if issue.get("title") != expected_title:
+    raise SystemExit("Resolved update issue has an unexpected title")
+' "$ISSUE_NUMBER" "$issue_title"
+}
+
+verify_update_pr() {
+  local pr_details
+
+  pr_details=$(gh pr view "$BRANCH" \
+    --repo "$REPOSITORY" \
+    --json isDraft,headRepositoryOwner,headRepository)
+  printf '%s' "$pr_details" | python3 -c '
+import json
+import sys
+
+expected_repository = sys.argv[1]
+pr = json.load(sys.stdin)
+if pr.get("isDraft") is not True:
+    raise SystemExit("Automated policyengine update PR is not a draft")
+head_repository = pr.get("headRepository") or {}
+if head_repository.get("nameWithOwner") != expected_repository:
+    raise SystemExit("Automated policyengine update PR is not from the canonical repository")
+' "$REPOSITORY"
+}
+
+create_update_pr() {
+  local pr_body_file
+
+  ensure_update_issue
+  pr_body_file="$(create_pr_body_file)"
+  gh pr create \
+    --draft \
+    --repo "$REPOSITORY" \
+    --base main \
+    --head "$BRANCH" \
+    --title "chore(deps): update policyengine to ${LATEST}" \
+    --body-file "$pr_body_file"
+  verify_update_pr
+}
 
 create_pr_body_file() {
   local pr_body_file
 
   pr_body_file="$(mktemp)"
   {
+    echo "Fixes #${ISSUE_NUMBER}"
+    echo
     echo "## Summary"
     echo
     echo "Update policyengine.py from ${CURRENT} to ${LATEST} in the simulation API runtime."
@@ -46,8 +143,11 @@ create_pr_body_file() {
     echo "- policyengine-core: ${BUNDLED_CORE_VERSION:-resolved from bundle during update}"
     echo "- policyengine-us: ${BUNDLED_US_VERSION:-resolved from bundle during update}"
     echo "- policyengine-uk: ${BUNDLED_UK_VERSION:-resolved from bundle during update}"
+    echo "- spm-calculator: ${BUNDLED_SPM_VERSION:-resolved from bundle during update}"
     echo
-    echo "Country data package versions remain manifest-derived at runtime/deploy time rather than independently pinned here."
+    echo "The bundle also selects these certified data releases:"
+    echo "- US: ${BUNDLED_US_DATA_VERSION:-resolved from bundle during update}"
+    echo "- UK: ${BUNDLED_UK_DATA_VERSION:-resolved from bundle during update}"
     echo
     echo "---"
     echo "Generated automatically by GitHub Actions."
@@ -117,24 +217,18 @@ if [[ "$DRY_RUN" == "1" ]]; then
   exit 0
 fi
 
-EXISTING_PR=$(gh pr list \
-  --head "$BRANCH" \
-  --state open \
-  --json number \
-  --jq '.[0].number' 2>/dev/null || true)
+EXISTING_PR=$(gh pr view "$BRANCH" \
+  --repo "$REPOSITORY" \
+  --json number,state \
+  --jq 'select(.state == "OPEN") | .number' 2>/dev/null || true)
 if [[ -n "$EXISTING_PR" ]]; then
   echo "PR #${EXISTING_PR} already exists for ${BRANCH}. Skipping."
   exit 0
 fi
 
 if git ls-remote --exit-code --heads origin "$BRANCH" >/dev/null 2>&1; then
   echo "Remote branch '${BRANCH}' already exists without an open PR. Creating PR."
-  PR_BODY_FILE="$(create_pr_body_file)"
-  gh pr create \
-    --base main \
-    --head "$BRANCH" \
-    --title "chore(deps): update policyengine to ${LATEST}" \
-    --body-file "$PR_BODY_FILE"
+  create_update_pr
   echo "PR created for existing branch ${BRANCH}"
   exit 0
 fi
@@ -144,18 +238,51 @@ git config user.email "github-actions[bot]@users.noreply.github.com"
 git checkout -b "$BRANCH"
 
 python3 - "$PYPROJECT" "$PACKAGE" "$CURRENT" "$LATEST" <<'PY'
+import os
 import sys
+import tempfile
+import tomllib
 from pathlib import Path
 
 pyproject_path, package, current, latest = sys.argv[1:]
 
 pyproject = Path(pyproject_path)
 pyproject_text = pyproject.read_text(encoding="utf-8")
+parsed = tomllib.loads(pyproject_text)
+dependency_lists = {
+    "project.dependencies": parsed.get("project", {}).get("dependencies", []),
+    "dependency-groups.policyengine-models": parsed.get("dependency-groups", {}).get(
+        "policyengine-models", []
+    ),
+}
 old_pin = f'"{package}=={current}"'
 new_pin = f'"{package}=={latest}"'
-if old_pin not in pyproject_text:
-    raise SystemExit(f"Could not find {old_pin} in {pyproject}")
-pyproject.write_text(pyproject_text.replace(old_pin, new_pin), encoding="utf-8")
+for location, dependencies in dependency_lists.items():
+    matches = [
+        dependency
+        for dependency in dependencies
+        if isinstance(dependency, str) and dependency.startswith(f"{package}==")
+    ]
+    if matches != [f"{package}=={current}"]:
+        raise SystemExit(
+            f"Expected {package}=={current} in {location}; found {matches!r}"
+        )
+updated_text = pyproject_text.replace(old_pin, new_pin)
+if pyproject_text.count(old_pin) != len(dependency_lists):
+    raise SystemExit(
+        f"Expected {old_pin} {len(dependency_lists)} times in {pyproject}; "
+        f"found {pyproject_text.count(old_pin)}"
+    )
+with tempfile.NamedTemporaryFile(
+    mode="w",
+    encoding="utf-8",
+    dir=pyproject.parent,
+    prefix=f".{pyproject.name}.",
+    delete=False,
+) as temporary:
+    temporary.write(updated_text)
+    temporary_path = temporary.name
+os.replace(temporary_path, pyproject)
 PY
 
 # The PyPI Simple index (which uv resolves from) can lag the JSON API right
@@ -182,59 +309,133 @@ BUNDLE_OUTPUT=$(
 BUNDLED_US_VERSION=$(printf '%s\n' "$BUNDLE_OUTPUT" | awk -F= '$1 == "us_version" {print $2}')
 BUNDLED_UK_VERSION=$(printf '%s\n' "$BUNDLE_OUTPUT" | awk -F= '$1 == "uk_version" {print $2}')
 BUNDLED_CORE_VERSION=$(printf '%s\n' "$BUNDLE_OUTPUT" | awk -F= '$1 == "policyengine_core_version" {print $2}')
+BUNDLED_POLICYENGINE_VERSION=$(printf '%s\n' "$BUNDLE_OUTPUT" | awk -F= '$1 == "policyengine_version" {print $2}')
+BUNDLED_SPM_VERSION=$(printf '%s\n' "$BUNDLE_OUTPUT" | awk -F= '$1 == "spm_calculator_version" {print $2}')
+BUNDLED_US_DATA_VERSION=$(printf '%s\n' "$BUNDLE_OUTPUT" | awk -F= '$1 == "us_data_version" {print $2}')
+BUNDLED_UK_DATA_VERSION=$(printf '%s\n' "$BUNDLE_OUTPUT" | awk -F= '$1 == "uk_data_version" {print $2}')
 
-if [[ -z "$BUNDLED_CORE_VERSION" || -z "$BUNDLED_US_VERSION" || -z "$BUNDLED_UK_VERSION" ]]; then
+if [[ -z "$BUNDLED_POLICYENGINE_VERSION" || -z "$BUNDLED_CORE_VERSION" || -z "$BUNDLED_US_VERSION" || -z "$BUNDLED_UK_VERSION" || -z "$BUNDLED_SPM_VERSION" || -z "$BUNDLED_US_DATA_VERSION" || -z "$BUNDLED_UK_DATA_VERSION" ]]; then
   echo "ERROR: Could not resolve bundled runtime package versions." >&2
   echo "$BUNDLE_OUTPUT" >&2
   exit 1
 fi
+if [[ "$BUNDLED_POLICYENGINE_VERSION" != "$LATEST" ]]; then
+  echo "ERROR: Installed policyengine.py reports bundle ${BUNDLED_POLICYENGINE_VERSION}, expected ${LATEST}." >&2
+  exit 1
+fi
 
 echo "Bundled runtime pins:"
+echo "  policyengine==${BUNDLED_POLICYENGINE_VERSION}"
 echo "  policyengine-core==${BUNDLED_CORE_VERSION}"
 echo "  policyengine-us==${BUNDLED_US_VERSION}"
 echo "  policyengine-uk==${BUNDLED_UK_VERSION}"
+echo "  spm-calculator==${BUNDLED_SPM_VERSION}"
+echo "Certified data releases:"
+echo "  us=${BUNDLED_US_DATA_VERSION}"
+echo "  uk=${BUNDLED_UK_DATA_VERSION}"
 
-python3 - "$PYPROJECT" "$BUNDLED_CORE_VERSION" "$BUNDLED_US_VERSION" "$BUNDLED_UK_VERSION" <<'PY'
+python3 - "$PYPROJECT" "$BUNDLED_POLICYENGINE_VERSION" "$BUNDLED_CORE_VERSION" "$BUNDLED_US_VERSION" "$BUNDLED_UK_VERSION" "$BUNDLED_SPM_VERSION" <<'PY'
+import os
 import re
 import sys
+import tempfile
+import tomllib
 from pathlib import Path
 
-pyproject_path, core_version, us_version, uk_version = sys.argv[1:]
+(
+    pyproject_path,
+    policyengine_version,
+    core_version,
+    us_version,
+    uk_version,
+    spm_version,
+) = sys.argv[1:]
 pyproject = Path(pyproject_path)
 text = pyproject.read_text(encoding="utf-8")
 pins = {
+    "policyengine": policyengine_version,
     "policyengine-core": core_version,
     "policyengine-us": us_version,
     "policyengine-uk": uk_version,
+    "spm-calculator": spm_version,
 }
+
+# These two dependency lists are installed in different places: project
+# dependencies supply local development/tests, while policyengine-models is
+# the only source for the Modal image. Require one exact pin in each list so a
+# release cannot update one runtime while leaving the other on an older bundle.
+parsed = tomllib.loads(text)
+dependency_lists = {
+    "project.dependencies": parsed.get("project", {}).get("dependencies", []),
+    "dependency-groups.policyengine-models": parsed.get("dependency-groups", {}).get(
+        "policyengine-models", []
+    ),
+}
+for location, dependencies in dependency_lists.items():
+    for package in pins:
+        matches = [
+            dependency
+            for dependency in dependencies
+            if isinstance(dependency, str) and dependency.startswith(f"{package}==")
+        ]
+        if len(matches) != 1:
+            raise SystemExit(
+                f"Expected one exact {package} pin in {location}; found {len(matches)}"
+            )
+
 for package, version in pins.items():
     pattern = rf'"{re.escape(package)}==[^"]+"'
     replacement = f'"{package}=={version}"'
-    text, count = re.subn(pattern, replacement, text, count=1)
-    if count != 1:
-        raise SystemExit(f"Could not update {package} in {pyproject}")
-pyproject.write_text(text, encoding="utf-8")
+    text, count = re.subn(pattern, replacement, text)
+    if count != len(dependency_lists):
+        raise SystemExit(
+            f"Expected to update {package} {len(dependency_lists)} times in "
+            f"{pyproject}; updated {count}"
+        )
+
+updated = tomllib.loads(text)
+for location, dependencies in {
+    "project.dependencies": updated["project"]["dependencies"],
+    "dependency-groups.policyengine-models": updated["dependency-groups"][
+        "policyengine-models"
+    ],
+}.items():
+    for package, version in pins.items():
+        if f"{package}=={version}" not in dependencies:
+            raise SystemExit(f"Updated {package} pin is missing from {location}")
+
+with tempfile.NamedTemporaryFile(
+    mode="w",
+    encoding="utf-8",
+    dir=pyproject.parent,
+    prefix=f".{pyproject.name}.",
+    delete=False,
+) as temporary:
+    temporary.write(text)
+    temporary_path = temporary.name
+os.replace(temporary_path, pyproject)
 PY
 
 (
   cd "$PROJECT_PATH"
   uv lock
+  uv lock --check
+  uv run --extra test pytest \
+    tests/test_bundle_version_export.py \
+    tests/test_policyengine_dependency_source.py \
+    tests/test_modal_bundle_image.py \
+    -q
 )
 
 if git diff --quiet -- "$PYPROJECT" "$LOCKFILE"; then
   echo "No changes after update. Nothing to do."
   exit 0
 fi
 
-PR_BODY_FILE="$(create_pr_body_file)"
-
 git add "$PYPROJECT" "$LOCKFILE"
 git commit -m "chore(deps): update policyengine to ${LATEST}"
 git push -u origin "$BRANCH"
 
-gh pr create \
-  --base main \
-  --title "chore(deps): update policyengine to ${LATEST}" \
-  --body-file "$PR_BODY_FILE"
+create_update_pr
 
 echo "PR created for policyengine ${CURRENT} -> ${LATEST}"

.github/workflows/check-policyengine-updates.yml

@@ -16,6 +16,7 @@ concurrency:
 
 permissions:
   contents: write
+  issues: write
   pull-requests: write
 
 jobs:
@@ -31,6 +32,9 @@ jobs:
         with:
           app-id: ${{ secrets.APP_ID }}
           private-key: ${{ secrets.APP_PRIVATE_KEY }}
+          permission-contents: write
+          permission-issues: write
+          permission-pull-requests: write
 
       - name: Checkout code
         uses: actions/checkout@v6