perf(pre-commit): fetch basedpyright base counts from CI artifacts (#35970)

.github/workflows/publish-basedpyright-base-counts.yml

@@ -0,0 +1,70 @@
+name: Publish basedpyright base counts
+
+# Every commit on litellm_internal_staging is some branch's future merge-base.
+# Publishing its per-rule basedpyright counts as an artifact lets
+# scripts/type_check_gate.py download them in seconds instead of paying a
+# 60-110s second basedpyright pass on every fresh worktree or moved merge-base.
+# No concurrency group on purpose: runs must never cancel each other, because
+# every sha's artifact matters (any of them can become a merge-base).
+
+on:
+  push:
+    branches:
+      - litellm_internal_staging
+  workflow_dispatch:
+    inputs:
+      ref:
+        description: "Ref to compute and publish base counts for"
+        required: false
+        default: litellm_internal_staging
+
+permissions:
+  contents: read
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+
+    steps:
+      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
+        with:
+          ref: ${{ inputs.ref || github.sha }}
+          clean: true
+          persist-credentials: false
+
+      - name: Set up Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+
+      - name: Set up uv
+        uses: ./.github/actions/setup-uv-with-retries
+        with:
+          version: "0.10.9"
+
+      - name: Install dependencies
+        run: |
+          uv sync --frozen --group proxy-dev --group e2e-dev
+
+      # Mirrors test-linting.yml's lint job: basedpyright resolves Prisma's
+      # generated client only after `prisma generate`, and the published counts
+      # must match what that job would measure for the same tree.
+      - name: Generate Prisma client
+        env:
+          PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
+        run: |
+          uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
+
+      - name: Emit basedpyright counts for HEAD
+        run: |
+          uv run --no-sync python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts"
+          counts_file=$(ls "$RUNNER_TEMP"/basedpyright-counts/basedpyright-counts-*.json)
+          echo "COUNTS_ARTIFACT_NAME=$(basename "$counts_file" .json)" >> "$GITHUB_ENV"
+
+      - name: Upload counts artifact
+        uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
+        with:
+          name: ${{ env.COUNTS_ARTIFACT_NAME }}
+          path: ${{ runner.temp }}/basedpyright-counts/
+          if-no-files-found: error

.github/workflows/test-linting.yml

@@ -15,6 +15,12 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     timeout-minutes: 15
+    # actions: read lets scripts/type_check_gate.py download the base-counts
+    # artifact published by publish-basedpyright-base-counts.yml instead of
+    # re-running basedpyright over the merge-base tree.
+    permissions:
+      contents: read
+      actions: read
 
     steps:
       - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
@@ -107,6 +113,8 @@ jobs:
           uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')"
 
       - name: Check basedpyright budget (delta vs base)
+        env:
+          GH_TOKEN: ${{ github.token }}
         run: |
           uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA"
 

scripts/type_check_gate.py

@@ -23,7 +23,11 @@ detached worktree at the merge-base, run under the same environment so import
 resolution matches, and its per-rule counts are cached under the repo's git
 common dir keyed by merge-base commit,
 ``pyrightconfig.json``, and ``uv.lock``, so re-runs against the same branch
-point pay for it once. ``--update`` ratchets each rule's ``limit`` down by the
+point pay for it once. A CI workflow publishes every staging commit's counts as
+an artifact (``--emit-counts-dir`` is its entry point), and on a disk-cache miss
+the gate first tries to download the merge-base's artifact through the ``gh``
+CLI; any fetch failure falls back silently to the local base pass, so the gate
+never gets worse than it was without CI. ``--update`` ratchets each rule's ``limit`` down by the
 number of errors this branch fixed relative to its branch point (the merge-base),
 so the headroom you were granted shrinks by exactly what you cleared and never
 grows.
@@ -37,12 +41,15 @@ carries an unambiguous ``rule`` field.
 import argparse
 import contextlib
 import hashlib
+import io
 import json
 import os
+import re
 import shutil
 import subprocess
 import sys
 import tempfile
+import zipfile
 from collections import Counter
 from collections.abc import Callable, Iterator, Mapping
 from pathlib import Path
@@ -54,6 +61,8 @@ PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json"
 UV_LOCK = REPO_ROOT / "uv.lock"
 DEFAULT_BASE = "origin/litellm_internal_staging"
 CACHE_FILE_PREFIX = "basedpyright-base-"
+ARTIFACT_NAME_PREFIX = "basedpyright-counts-"
+GH_TIMEOUT_SECONDS = 10
 
 # basedpyright's node process needs more than the ~4 GB default heap on this
 # repo; appended last so it wins node's last-flag-wins resolution over any
@@ -225,12 +234,8 @@ def default_cache_dir() -> Path:
     return resolved / "litellm-lint-cache"
 
 
-def load_cached_counts(path: Path) -> dict[str, int] | None:
-    try:
-        data = json.loads(path.read_text())
-    except (OSError, json.JSONDecodeError):
-        return None
-    counts = data.get("counts") if isinstance(data, dict) else None
+def validated_counts(data: object) -> dict[str, int] | None:
+    counts: Final = data.get("counts") if isinstance(data, dict) else None
     if not isinstance(counts, dict):
         return None
     if not all(
@@ -241,6 +246,14 @@ def load_cached_counts(path: Path) -> dict[str, int] | None:
     return counts
 
 
+def load_cached_counts(path: Path) -> dict[str, int] | None:
+    try:
+        data = json.loads(path.read_text())
+    except (OSError, json.JSONDecodeError):
+        return None
+    return validated_counts(data)
+
+
 def scratch_path(path: Path) -> Path:
     """In-flight scratch for the tmp+rename write. Dot-prefixed so the prune
     glob in `store_counts` can never match it (a concurrent run would otherwise
@@ -249,6 +262,16 @@ def scratch_path(path: Path) -> Path:
     return path.with_name(f".{path.name}.{os.getpid()}.tmp")
 
 
+def counts_payload(base_point: str, counts: Mapping[str, int]) -> str:
+    return (
+        json.dumps(
+            {"base_point": base_point, "counts": dict(sorted(counts.items()))},
+            indent=2,
+        )
+        + "\n"
+    )
+
+
 def store_counts(
     directory: Path, path: Path, base_point: str, counts: Mapping[str, int]
 ) -> None:
@@ -257,30 +280,141 @@ def store_counts(
         if stale != path:
             stale.unlink(missing_ok=True)
     scratch = scratch_path(path)
-    scratch.write_text(
-        json.dumps(
-            {"base_point": base_point, "counts": dict(sorted(counts.items()))},
-            indent=2,
-        )
-        + "\n"
-    )
+    scratch.write_text(counts_payload(base_point, counts))
     scratch.replace(path)
 
 
+def parse_origin_slug(url: str) -> str | None:
+    match: Final = re.fullmatch(
+        r"(?:git@github\.com:|https://github\.com/)([^/]+/[^/]+?)(?:\.git)?/?",
+        url.strip(),
+    )
+    return match.group(1) if match else None
+
+
+def origin_slug() -> str | None:
+    proc: Final = subprocess.run(
+        ["git", "remote", "get-url", "origin"],
+        cwd=REPO_ROOT,
+        capture_output=True,
+        text=True,
+    )
+    if proc.returncode != 0:
+        return None
+    return parse_origin_slug(proc.stdout)
+
+
+def artifact_name(base_point: str) -> str:
+    return f"{ARTIFACT_NAME_PREFIX}{cache_key(base_point, environment_fingerprints())}"
+
+
+def _gh_output(args: list[str]) -> bytes | None:
+    try:
+        proc = subprocess.run(
+            ["gh", *args], capture_output=True, timeout=GH_TIMEOUT_SECONDS
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    return proc.stdout if proc.returncode == 0 else None
+
+
+def _parsed_json(raw: bytes) -> object | None:
+    try:
+        return json.loads(raw)
+    except ValueError:
+        return None
+
+
+def _artifact_download_url(listing: object) -> str | None:
+    artifacts: Final = listing.get("artifacts") if isinstance(listing, dict) else None
+    if not isinstance(artifacts, list) or not artifacts:
+        return None
+    newest: Final = artifacts[0]
+    if not isinstance(newest, dict) or newest.get("expired"):
+        return None
+    url: Final = newest.get("archive_download_url")
+    return url if isinstance(url, str) else None
+
+
+def _counts_json_from_zip(zip_bytes: bytes) -> object | None:
+    try:
+        with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive:
+            members: Final = [
+                name for name in archive.namelist() if name.endswith(".json")
+            ]
+            if len(members) != 1:
+                return None
+            return json.loads(archive.read(members[0]))
+    except (zipfile.BadZipFile, ValueError, OSError):
+        return None
+
+
+def counts_for_base(payload: object, base_point: str) -> dict[str, int] | None:
+    if not isinstance(payload, dict) or payload.get("base_point") != base_point:
+        return None
+    counts: Final = validated_counts(payload)
+    return counts if counts else None
+
+
+def _fetch_fallback(reason: str) -> None:
+    sys.stderr.write(f"{reason}; computing base counts locally\n")
+
+
+def fetch_ci_base_counts(
+    base_point: str,
+    gh_output: Callable[[list[str]], bytes | None] = _gh_output,
+) -> dict[str, int] | None:
+    """Base counts from the CI artifact published for `base_point`, or None.
+
+    Every failure mode (no gh, no auth, offline, expired or missing artifact,
+    malformed payload, counts for a different commit) returns None so the
+    caller falls back to the local base pass; the fetch is an optimization and
+    must never make the gate less available than local compute alone."""
+    slug: Final = origin_slug()
+    if slug is None:
+        return _fetch_fallback("origin remote is not a github.com URL")
+    name: Final = artifact_name(base_point)
+    listing: Final = gh_output(
+        ["api", f"repos/{slug}/actions/artifacts?name={name}&per_page=1"]
+    )
+    if listing is None:
+        return _fetch_fallback(f"could not list CI artifacts named {name}")
+    url: Final = _artifact_download_url(_parsed_json(listing))
+    if url is None:
+        return _fetch_fallback(f"no usable CI artifact named {name}")
+    zip_bytes: Final = gh_output(["api", url])
+    if zip_bytes is None:
+        return _fetch_fallback(f"download failed for CI artifact {name}")
+    counts: Final = counts_for_base(_counts_json_from_zip(zip_bytes), base_point)
+    if counts is None:
+        return _fetch_fallback(
+            f"CI artifact {name} is not valid base counts for {base_point[:12]}"
+        )
+    sys.stderr.write(f"base counts fetched from CI artifact {name}\n")
+    return counts
+
+
 def base_counts_cached(
     base_point: str,
     cache_dir: Path | None = None,
     compute: Callable[[str], dict[str, int]] = base_counts,
+    fetch: Callable[[str], dict[str, int] | None] = fetch_ci_base_counts,
 ) -> dict[str, int]:
     """`base_counts` memoized on disk. The base tree at a given commit is
     immutable, so its counts are a pure function of the merge-base plus the
     environment fingerprints in the cache key; an empty result is never stored
-    because it is the signature of a crashed pass, not a clean tree."""
+    because it is the signature of a crashed pass, not a clean tree. On a disk
+    miss the counts CI already published for the merge-base are fetched before
+    the expensive local base pass; a fetch miss of any kind computes locally."""
     directory = default_cache_dir() if cache_dir is None else cache_dir
     path = cache_path(directory, base_point, environment_fingerprints())
     cached = load_cached_counts(path)
     if cached is not None:
         return cached
+    fetched: Final = fetch(base_point)
+    if fetched:
+        store_counts(directory, path, base_point, fetched)
+        return fetched
     counts = compute(base_point)
     if counts:
         store_counts(directory, path, base_point, counts)
@@ -353,6 +487,29 @@ def cmd_update(current: Mapping[str, int], base_ref: str = DEFAULT_BASE) -> None
     )
 
 
+def cmd_emit_counts(head: Mapping[str, int], directory: Path, head_sha: str) -> None:
+    """Write HEAD's per-rule counts as the file the publisher workflow uploads.
+
+    The filename stem is exactly the artifact name `fetch_ci_base_counts` will
+    later look up for this commit, so emit and fetch cannot drift apart. Empty
+    counts are refused for the same reason `is_vacuous_run` exists: a pass that
+    produced nothing almost certainly crashed, and publishing it would poison
+    every branch that fetches it."""
+    if not head:
+        print(
+            "FAIL: basedpyright produced no errors; refusing to publish empty base "
+            "counts because the pass almost certainly crashed or emitted nothing."
+        )
+        raise SystemExit(1)
+    name: Final = artifact_name(head_sha)
+    directory.mkdir(parents=True, exist_ok=True)
+    (directory / f"{name}.json").write_text(counts_payload(head_sha, head))
+    print(
+        f"Emitted base counts for {head_sha} as {name}.json "
+        f"({sum(head.values())} errors total)"
+    )
+
+
 def cmd_check(head: Mapping[str, int], base_ref: str) -> None:
     budget = json.loads(BUDGET_PATH.read_text())
     if is_vacuous_run(head, budget):
@@ -401,9 +558,14 @@ def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--base", default=DEFAULT_BASE)
     parser.add_argument("--update", action="store_true")
+    parser.add_argument("--emit-counts-dir", type=Path)
     args = parser.parse_args()
     head = count_basedpyright(run_basedpyright())
-    if args.update:
+    if args.emit_counts_dir is not None:
+        cmd_emit_counts(
+            head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip()
+        )
+    elif args.update:
         cmd_update(head, args.base)
     else:
         cmd_check(head, args.base)

tests/test_litellm/test_type_check_gate.py

@@ -295,16 +295,27 @@ def test_store_prunes_entries_for_other_branch_points(tmp_path):
     assert gate.load_cached_counts(new) == {"reportAny": 2}
 
 
+def _no_fetch(ref):
+    return None
+
+
+def _never(reason):
+    def callback(ref):
+        raise AssertionError(reason)
+
+    return callback
+
+
 def test_base_counts_cached_returns_the_hit_without_recomputing(tmp_path):
     path = gate.cache_path(tmp_path, "abc123", gate.environment_fingerprints())
     gate.store_counts(tmp_path, path, "abc123", {"reportAny": 7})
 
-    def explode(ref):
-        raise AssertionError("a cache hit must not re-run the base pass")
-
-    assert gate.base_counts_cached("abc123", cache_dir=tmp_path, compute=explode) == {
-        "reportAny": 7
-    }
+    assert gate.base_counts_cached(
+        "abc123",
+        cache_dir=tmp_path,
+        compute=_never("a cache hit must not re-run the base pass"),
+        fetch=_never("a cache hit must not reach for CI"),
+    ) == {"reportAny": 7}
 
 
 def test_base_counts_cached_computes_once_then_hits(tmp_path):
@@ -314,8 +325,12 @@ def test_base_counts_cached_computes_once_then_hits(tmp_path):
         calls.append(ref)
         return {"reportAny": 4}
 
-    first = gate.base_counts_cached("abc123", cache_dir=tmp_path, compute=fake)
-    second = gate.base_counts_cached("abc123", cache_dir=tmp_path, compute=fake)
+    first = gate.base_counts_cached(
+        "abc123", cache_dir=tmp_path, compute=fake, fetch=_no_fetch
+    )
+    second = gate.base_counts_cached(
+        "abc123", cache_dir=tmp_path, compute=fake, fetch=_no_fetch
+    )
     assert first == second == {"reportAny": 4}
     assert calls == ["abc123"]
 
@@ -327,12 +342,204 @@ def test_an_empty_base_pass_is_never_cached(tmp_path):
         calls.append(ref)
         return {}
 
-    assert gate.base_counts_cached("abc123", cache_dir=tmp_path, compute=crashed) == {}
-    assert gate.base_counts_cached("abc123", cache_dir=tmp_path, compute=crashed) == {}
+    assert (
+        gate.base_counts_cached(
+            "abc123", cache_dir=tmp_path, compute=crashed, fetch=_no_fetch
+        )
+        == {}
+    )
+    assert (
+        gate.base_counts_cached(
+            "abc123", cache_dir=tmp_path, compute=crashed, fetch=_no_fetch
+        )
+        == {}
+    )
     assert calls == ["abc123", "abc123"]
     assert list(tmp_path.iterdir()) == []
 
 
+def test_base_counts_cached_uses_fetched_counts_and_persists_them(tmp_path):
+    counts = gate.base_counts_cached(
+        "abc123",
+        cache_dir=tmp_path,
+        compute=_never("fetched counts must skip the local base pass"),
+        fetch=lambda ref: {"reportAny": 9},
+    )
+    assert counts == {"reportAny": 9}
+    path = gate.cache_path(tmp_path, "abc123", gate.environment_fingerprints())
+    assert gate.load_cached_counts(path) == {"reportAny": 9}
+    assert gate.base_counts_cached(
+        "abc123",
+        cache_dir=tmp_path,
+        compute=_never("the persisted fetch must satisfy later runs"),
+        fetch=_never("the persisted fetch must satisfy later runs"),
+    ) == {"reportAny": 9}
+
+
+def test_base_counts_cached_falls_back_to_compute_on_a_fetch_miss(tmp_path):
+    calls = []
+
+    def local(ref):
+        calls.append(ref)
+        return {"reportAny": 4}
+
+    assert gate.base_counts_cached(
+        "abc123", cache_dir=tmp_path, compute=local, fetch=_no_fetch
+    ) == {"reportAny": 4}
+    assert calls == ["abc123"]
+
+
+def test_base_counts_cached_treats_empty_fetched_counts_as_a_miss(tmp_path):
+    assert gate.base_counts_cached(
+        "abc123",
+        cache_dir=tmp_path,
+        compute=lambda ref: {"reportAny": 2},
+        fetch=lambda ref: {},
+    ) == {"reportAny": 2}
+    path = gate.cache_path(tmp_path, "abc123", gate.environment_fingerprints())
+    assert gate.load_cached_counts(path) == {"reportAny": 2}
+
+
+def test_origin_slug_parsing_supports_ssh_and_https_github_forms():
+    assert gate.parse_origin_slug("git@github.com:BerriAI/litellm.git") == "BerriAI/litellm"
+    assert gate.parse_origin_slug("git@github.com:BerriAI/litellm") == "BerriAI/litellm"
+    assert gate.parse_origin_slug("https://github.com/BerriAI/litellm.git") == "BerriAI/litellm"
+    assert gate.parse_origin_slug("https://github.com/BerriAI/litellm") == "BerriAI/litellm"
+    assert gate.parse_origin_slug("https://github.com/BerriAI/litellm/") == "BerriAI/litellm"
+
+
+def test_origin_slug_parsing_rejects_non_github_urls():
+    assert gate.parse_origin_slug("https://gitlab.com/BerriAI/litellm.git") is None
+    assert gate.parse_origin_slug("git@bitbucket.org:BerriAI/litellm.git") is None
+    assert gate.parse_origin_slug("not a url") is None
+    assert gate.parse_origin_slug("") is None
+
+
+def _artifact_zip(payload):
+    import io
+    import zipfile
+
+    buffer = io.BytesIO()
+    with zipfile.ZipFile(buffer, "w") as archive:
+        archive.writestr("basedpyright-counts.json", json.dumps(payload))
+    return buffer.getvalue()
+
+
+def _gh_stub(listing, zip_bytes):
+    def gh_output(args):
+        if args[-1].startswith("repos/"):
+            return json.dumps(listing).encode()
+        return zip_bytes
+
+    return gh_output
+
+
+def _live_listing():
+    return {
+        "artifacts": [
+            {"expired": False, "archive_download_url": "https://api.github.com/x/zip"}
+        ]
+    }
+
+
+def test_fetcher_returns_counts_from_a_matching_artifact(capsys):
+    payload = {"base_point": "abc123", "counts": {"reportAny": 3}}
+    fetched = gate.fetch_ci_base_counts(
+        "abc123", gh_output=_gh_stub(_live_listing(), _artifact_zip(payload))
+    )
+    assert fetched == {"reportAny": 3}
+    assert "fetched from CI artifact" in capsys.readouterr().err
+
+
+def test_fetcher_rejects_an_artifact_for_a_different_base_point():
+    payload = {"base_point": "someothersha", "counts": {"reportAny": 3}}
+    assert (
+        gate.fetch_ci_base_counts(
+            "abc123", gh_output=_gh_stub(_live_listing(), _artifact_zip(payload))
+        )
+        is None
+    )
+
+
+def test_fetcher_rejects_empty_or_misshapen_artifact_counts():
+    for counts in ({}, {"reportAny": "three"}, {"reportAny": True}):
+        payload = {"base_point": "abc123", "counts": counts}
+        assert (
+            gate.fetch_ci_base_counts(
+                "abc123", gh_output=_gh_stub(_live_listing(), _artifact_zip(payload))
+            )
+            is None
+        )
+
+
+def test_fetcher_rejects_an_expired_artifact():
+    listing = {
+        "artifacts": [
+            {"expired": True, "archive_download_url": "https://api.github.com/x/zip"}
+        ]
+    }
+    payload = {"base_point": "abc123", "counts": {"reportAny": 3}}
+    assert (
+        gate.fetch_ci_base_counts(
+            "abc123", gh_output=_gh_stub(listing, _artifact_zip(payload))
+        )
+        is None
+    )
+
+
+def test_fetcher_misses_when_no_artifact_is_published():
+    assert (
+        gate.fetch_ci_base_counts(
+            "abc123", gh_output=_gh_stub({"artifacts": []}, b"")
+        )
+        is None
+    )
+
+
+def test_fetcher_misses_when_gh_is_unusable(capsys):
+    assert gate.fetch_ci_base_counts("abc123", gh_output=lambda args: None) is None
+    assert "computing base counts locally" in capsys.readouterr().err
+
+
+def test_fetcher_misses_on_a_corrupt_artifact_archive():
+    assert (
+        gate.fetch_ci_base_counts(
+            "abc123", gh_output=_gh_stub(_live_listing(), b"not a zip")
+        )
+        is None
+    )
+
+
+def test_emit_writes_the_artifact_json_named_by_the_head_key(tmp_path, capsys):
+    gate.cmd_emit_counts({"reportAny": 3, "aRule": 1}, tmp_path, "deadbeef")
+    key = gate.cache_key("deadbeef", gate.environment_fingerprints())
+    path = tmp_path / f"basedpyright-counts-{key}.json"
+    assert json.loads(path.read_text()) == {
+        "base_point": "deadbeef",
+        "counts": {"aRule": 1, "reportAny": 3},
+    }
+    summary = capsys.readouterr().out
+    assert "deadbeef" in summary
+    assert key in summary
+    assert "4" in summary
+
+
+def test_emit_refuses_to_publish_empty_counts(tmp_path):
+    import pytest
+
+    with pytest.raises(SystemExit):
+        gate.cmd_emit_counts({}, tmp_path, "deadbeef")
+    assert list(tmp_path.iterdir()) == []
+
+
+def test_emitted_file_round_trips_through_the_fetch_validation(tmp_path):
+    gate.cmd_emit_counts({"reportAny": 3}, tmp_path, "deadbeef")
+    key = gate.cache_key("deadbeef", gate.environment_fingerprints())
+    payload = json.loads((tmp_path / f"basedpyright-counts-{key}.json").read_text())
+    assert gate.counts_for_base(payload, "deadbeef") == {"reportAny": 3}
+    assert gate.counts_for_base(payload, "someothersha") is None
+
+
 def _git(cwd, *args):
     proc = subprocess.run(["git", *args], cwd=cwd, capture_output=True, text=True)
     assert proc.returncode == 0, proc.stderr