ci: resolve fork PRs missing Castiron run associations

.github/workflows/castiron-custom-code-comment.yml

@@ -145,8 +145,11 @@ jobs:
             const {data: main} = await github.rest.git.getRef({...context.repo, ref: 'heads/main'});
             const base = main.object.sha;
             if (run.event === 'pull_request') {
-              const pulls = run.pull_requests.length ? run.pull_requests : await github.paginate(
+              let pulls = run.pull_requests.length ? run.pull_requests : await github.paginate(
                 github.rest.repos.listPullRequestsAssociatedWithCommit, {...context.repo, commit_sha: head});
+              if (!pulls.length) pulls = await github.paginate(github.rest.pulls.list, {
+                ...context.repo, state: 'open', head: `${run.head_repository.owner.login}:${run.head_branch}`,
+              });
               const current = [];
               for (const pull of pulls) {
                 const {data: pr} = await github.rest.pulls.get({...context.repo, pull_number: pull.number});
@@ -222,7 +225,10 @@ jobs:
             const marker = '<!-- castiron:custom-code-report:v1 -->';
             const run = context.payload.workflow_run;
             if (run.event !== 'pull_request' || run.path !== '.github/workflows/castiron-custom-code.yml') return;
-            const pulls = run.pull_requests?.length ? run.pull_requests : await github.paginate(github.rest.repos.listPullRequestsAssociatedWithCommit, {...context.repo, commit_sha: run.head_sha});
+            let pulls = run.pull_requests?.length ? run.pull_requests : await github.paginate(github.rest.repos.listPullRequestsAssociatedWithCommit, {...context.repo, commit_sha: run.head_sha});
+            if (!pulls.length) pulls = await github.paginate(github.rest.pulls.list, {
+              ...context.repo, state: 'open', head: `${run.head_repository.owner.login}:${run.head_branch}`,
+            });
             for (const pull of pulls) {
               const {data: current} = await github.rest.pulls.get({...context.repo, pull_number: pull.number});
               if (current.state !== 'open' || current.head.sha !== run.head_sha) continue;

.github/workflows/castiron-custom-code.yml

@@ -16,7 +16,7 @@ concurrency:
   cancel-in-progress: false
 
 env:
-  REPORTER_SHA256: ac48ca88e9f7ad57195038157e99f055c0cd3dac8de856e4d102dca807766d4a
+  REPORTER_SHA256: be56ac91a1fc757d5d842586ae364354f11e8f0ec4cc458c06af1d7bd54a12cf
 
 jobs:
   queue-signal:

scripts/castiron/custom_code_budget.py

@@ -300,9 +300,7 @@ def github_evaluate(
         raise ValueError("unexpected or superseded source workflow run")
     head = report.require_sha(run["head_sha"])
     if run["event"] == "pull_request":
-        associated = run["pull_requests"] or report.api(
-            "GET", f"{root}/commits/{head}/pulls?per_page=100"
-        )
+        associated = report.associated_pulls(repository, run)
         current: list[int] = []
         for number in sorted({int(pr["number"]) for pr in associated}):
             if number <= 0:

scripts/castiron/custom_code_report.py

@@ -25,6 +25,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, cast
+from urllib.parse import urlencode
 
 DOMAIN = b"castiron-codegen-v1\0"
 MARKER = "<!-- castiron:custom-code-report:v1 -->"
@@ -766,6 +767,25 @@ def api(method: str, path: str, payload: dict[str, Any] | None = None) -> Any:
     return json.loads(result.stdout)
 
 
+def associated_pulls(repository: str, run: dict[str, Any]) -> list[dict[str, Any]]:
+    root = f"repos/{repository}"
+    associated = run["pull_requests"] or api(
+        "GET", f"{root}/commits/{require_sha(run['head_sha'])}/pulls?per_page=100"
+    )
+    if associated:
+        return cast(list[dict[str, Any]], associated)
+    # Fork runs can be absent from both association endpoints. Discover candidates
+    # by branch; callers still verify the current PR head and target through GitHub.
+    head = f"{run['head_repository']['owner']['login']}:{run['head_branch']}"
+    for page in range(1, 101):
+        query = urlencode({"state": "open", "head": head, "per_page": 100, "page": page})
+        pulls = api("GET", f"{root}/pulls?{query}")
+        associated.extend(pulls)
+        if len(pulls) < 100:
+            return cast(list[dict[str, Any]], associated)
+    raise ReportError("too many pull requests for source branch")
+
+
 def publish_comment(
     report: dict[str, Any],
     repository: str,
@@ -793,10 +813,7 @@ def publish_comment(
         or run["head_sha"] != report["head_sha"]
     ):
         raise ReportError("workflow run does not match report PR/head")
-    associated = run["pull_requests"]
-    if not associated:
-        # GitHub can omit the PR association on workflow runs from forks.
-        associated = api("GET", f"{root}/commits/{require_sha(run['head_sha'])}/pulls?per_page=100")
+    associated = associated_pulls(repository, run)
     if not any(pr["number"] == number for pr in associated):
         raise ReportError("workflow run does not match report PR/head")
     if run["run_attempt"] != run_attempt:
@@ -892,7 +909,7 @@ def trusted_report(
     if run["run_attempt"] != run_attempt:
         return
     head = require_sha(run["head_sha"])
-    associated = run["pull_requests"] or api("GET", f"{root}/commits/{head}/pulls?per_page=100")
+    associated = associated_pulls(repository, run)
     current: list[tuple[int, str]] = []
     for number in sorted({int(pr["number"]) for pr in associated}):
         if number <= 0:

scripts/castiron/test_custom_code_budget.py

@@ -22,6 +22,7 @@ def source_run(head: str, event: str = "pull_request") -> dict[str, Any]:
         "head_sha": head,
         "head_branch": "gh-readonly-queue/main/pr-7-example" if event == "merge_group" else "sdk",
         "repository": {"full_name": "openai/example"},
+        "head_repository": {"owner": {"login": "contributor"}},
         "path": ".github/workflows/castiron-custom-code.yml",
         "status": "completed",
         "run_attempt": 1,
@@ -309,6 +310,7 @@ def publish(
         base_changed: bool = False,
         no_result: bool = False,
         failed_budget: bool = False,
+        fallback_pulls: list[dict[str, int]] | None = None,
         run_overrides: dict[str, Any] | None = None,
     ) -> list[dict[str, Any]]:
         path = (
@@ -321,6 +323,7 @@ def publish(
         base, head = "a" * 40, "b" * 40
         payload = {
             "script": script,
+            "fallback_pulls": fallback_pulls,
             "context": {
                 "eventName": "workflow_run",
                 "repo": {"owner": "openai", "repo": "example"},
@@ -352,10 +355,16 @@ def publish(
           const data = JSON.parse(fs.readFileSync(0, 'utf8'));
           const published = [];
           const github = {rest: {
-            pulls: {get: async () => ({data: data.current})},
+            pulls: {get: async () => ({data: data.current}), list: 'pulls.list'},
             actions: {getWorkflowRun: async () => ({data: data.run})},
             git: {getRef: async () => ({data: {object: {sha: data.current.base.sha}}})},
-            repos: {createCommitStatus: async value => published.push(value)},
+            repos: {createCommitStatus: async value => published.push(value),
+              listPullRequestsAssociatedWithCommit: 'commits.pulls'},
+          }, paginate: async (method, params) => {
+            if (method === 'commits.pulls') return [];
+            if (method !== 'pulls.list' || params.head !== 'contributor:sdk' || params.state !== 'open')
+              throw new Error('Unexpected fallback lookup');
+            return data.fallback_pulls;
           }};
           const AsyncFunction = Object.getPrototypeOf(async function(){}).constructor;
           new AsyncFunction('github','context','process', data.script)(github, data.context, {env:data.env})
@@ -380,6 +389,23 @@ def test_statuses_attach_to_candidate_not_main(self) -> None:
                     all(r["sha"] == "b" * 40 and r["state"] == "success" for r in results)
                 )
 
+    def test_fork_statuses_with_no_commit_association(self) -> None:
+        options: dict[str, Any] = {"run_overrides": {"pull_requests": []}, "fallback_pulls": [{"number": 3}]}
+        results = self.publish(**options)
+        self.assertEqual(len(results), 2)
+        self.assertTrue(all(r["sha"] == "b" * 40 and r["state"] == "success" for r in results))
+        self.assertEqual(self.publish(**options, head_changed=True), [])
+        self.assertTrue(
+            all(r["state"] == "failure" for r in self.publish(**options, no_result=True))
+        )
+        self.assertEqual(self.publish(run_overrides={"pull_requests": []}, fallback_pulls=[]), [])
+        self.assertEqual(
+            self.publish(
+                run_overrides={"pull_requests": []}, fallback_pulls=[{"number": 3}, {"number": 4}]
+            ),
+            [],
+        )
+
     def test_stale_pr_head_is_not_published(self) -> None:
         self.assertEqual(self.publish(head_changed=True), [])
 
@@ -568,10 +594,12 @@ def test_github_uses_fresh_bare_data_repo_and_checks_current_head(self) -> None:
             "head": {"sha": head},
             "base": {"sha": base, "ref": "main", "repo": {"full_name": "openai/example"}},
         }
-        responses = [
+        responses: list[Any] = [
             {"default_branch": "main", "private": False},
             {"object": {"sha": base}},
-            source_run(head),
+            {**source_run(head), "pull_requests": []},
+            [],
+            [{"number": 3}],
             pull,
             [{"type": "merge_queue"}],
         ]
@@ -622,7 +650,9 @@ def test_stale_pull_and_wrong_target_fail_before_objects_created(self) -> None:
                     side_effect=[
                         {"default_branch": "main"},
                         {"object": {"sha": base}},
-                        source_run(head),
+                        {**source_run(head), "pull_requests": []},
+                        [],
+                        [{"number": 3}],
                         pull,
                     ],
                 ):

scripts/castiron/test_custom_code_report.py

@@ -283,14 +283,17 @@ def test_trusted_failure_publisher_updates_one_current_comment(self) -> None:
         harness = r"""
 const assert = require('node:assert/strict');
 const AsyncFunction = Object.getPrototypeOf(async function(){}).constructor;
-async function check(stale, exists, priorRun, expected) {
+async function check(stale, exists, priorRun, expected, missing = false) {
   const writes = [];
   const event = {id: 20, run_attempt: 1, event: 'pull_request', path: '.github/workflows/castiron-custom-code.yml', head_sha: 'a'.repeat(40), pull_requests: [{number: 1}]};
+  if (missing) Object.assign(event, {pull_requests: [], head_repository: {owner: {login: 'contributor'}}, head_branch: 'sdk'});
   const current = {state: 'open', head: {sha: (stale ? 'c' : 'a').repeat(40)}};
   const previous = {id: 42, user: {type: 'Bot', login: 'github-actions[bot]'},
     body: `<!-- castiron:custom-code-report:v1 -->\n<!-- castiron:run:v1:${priorRun}:1 -->`};
-  const github = {paginate: async () => exists ? [previous] : [], rest: {
-    pulls: {get: async () => ({data: current})},
+  const github = {paginate: async method => method === 'associations' ? []
+    : method === 'pulls' ? [{number: 1}] : exists ? [previous] : [], rest: {
+    repos: {listPullRequestsAssociatedWithCommit: 'associations'},
+    pulls: {get: async () => ({data: current}), list: 'pulls'},
     issues: {listComments() {}, updateComment: async x => writes.push(['update', x]),
       createComment: async x => writes.push(['create', x])}}};
   const context = {payload: {workflow_run: event}, repo: {owner: 'openai', repo: 'example'},
@@ -308,6 +311,8 @@ def test_trusted_failure_publisher_updates_one_current_comment(self) -> None:
   await check(false, false, 10, 'create');
   await check(true, true, 10, null);
   await check(false, true, 21, null);
+  await check(false, true, 10, 'update', true);
+  await check(true, true, 10, null, true);
 })().catch(error => { console.error(error); process.exitCode = 1; });
 """
         subprocess.run(
@@ -450,7 +455,12 @@ def local_git(repo: Path, *args: str, input_bytes: bytes | None = None) -> bytes
             self.assertNotIn("checkout", args)
             return real_git(repo, *args, input_bytes=input_bytes)
 
-        for label, revision in (("genuine", base), ("custom", head), ("broken", broken)):
+        for label, revision in (
+            ("genuine", base),
+            ("custom", head),
+            ("fork", head),
+            ("broken", broken),
+        ):
             with self.subTest(label=label):
                 calls: list[tuple[str, str]] = []
                 bodies: list[str] = []
@@ -466,6 +476,8 @@ def local_git(repo: Path, *args: str, input_bytes: bytes | None = None) -> bytes
                     "head_sha": revision,
                     "run_attempt": 1,
                     "pull_requests": [],
+                    "head_repository": {"owner": {"login": "contributor"}},
+                    "head_branch": "fix/branch",
                 }
                 forged: dict[str, Any] = {**legitimate, "head_sha": revision, "files": []}
                 self.assertIn("Generated baselines verified", report.render_report(forged))
@@ -481,7 +493,10 @@ def fake_api(
                         responses: dict[str, Any] = {
                             "repos/openai/example": {"private": False},
                             "repos/openai/example/actions/runs/2": run,
-                            f"repos/openai/example/commits/{revision}/pulls?per_page=100": [
+                            f"repos/openai/example/commits/{revision}/pulls?per_page=100": (
+                                [] if label == "fork" else [{"number": 1}]
+                            ),
+                            "repos/openai/example/pulls?state=open&head=contributor%3Afix%2Fbranch&per_page=100&page=1": [
                                 {"number": 1}
                             ],
                             "repos/openai/example/pulls/1": pull,
@@ -529,7 +544,7 @@ def fake_api(
                 if label == "broken":
                     self.assertIn("Report unavailable", body)
                     self.assertNotIn("Generated baselines verified", body)
-                elif label == "custom":
+                elif label in ("custom", "fork"):
                     self.assertIn("1 newly customized", body)
                     self.assertIn("generated.py", body)
                     self.assertIn(b"+# custom", (out / "custom-code.patch").read_bytes())
@@ -539,6 +554,24 @@ def fake_api(
                     self.assertIn("Generated baselines verified", body)
                     self.assertIn("--name castiron-custom-code-9-3", body)
 
+    def test_fork_association_fallback_paginates_candidates(self) -> None:
+        run: dict[str, Any] = {
+            "head_sha": "a" * 40,
+            "pull_requests": [],
+            "head_repository": {"owner": {"login": "contributor"}},
+            "head_branch": "fix/branch&other=value",
+        }
+        first = [{"number": number} for number in range(1, 101)]
+        with mock.patch.object(report, "api", side_effect=[[], first, [{"number": 101}]]) as api:
+            self.assertEqual(len(report.associated_pulls("openai/example", run)), 101)
+        self.assertEqual(
+            api.call_args.args,
+            (
+                "GET",
+                "repos/openai/example/pulls?state=open&head=contributor%3Afix%2Fbranch%26other%3Dvalue&per_page=100&page=2",
+            ),
+        )
+
     def test_trusted_report_rejects_invalid_or_stale_association_before_fetch(self) -> None:
         run = {
             "event": "pull_request",
@@ -547,6 +580,8 @@ def test_trusted_report_rejects_invalid_or_stale_association_before_fetch(self)
             "head_sha": "a" * 40,
             "run_attempt": 1,
             "pull_requests": [{"number": 1}],
+            "head_repository": {"owner": {"login": "contributor"}},
+            "head_branch": "fix/branch",
         }
         pull = {
             "state": "open",
@@ -563,7 +598,21 @@ def test_trusted_report_rejects_invalid_or_stale_association_before_fetch(self)
                 [run, {**pull, "base": {"sha": "b" * 40, "repo": {"full_name": "other/repo"}}}],
                 False,
             ),
-            ([{**run, "pull_requests": []}, []], False),
+            ([{**run, "pull_requests": []}, [], []], False),
+            (
+                [{**run, "pull_requests": []}, [], [{"number": 1}], {**pull, "state": "closed"}],
+                False,
+            ),
+            (
+                [
+                    {**run, "pull_requests": []},
+                    [],
+                    [{"number": 1}],
+                    {**pull, "head": {"sha": "c" * 40}},
+                ],
+                False,
+            ),
+            ([{**run, "pull_requests": []}, [], [{"number": 1}, {"number": 2}], pull, pull], True),
             ([{**run, "pull_requests": [{"number": 1}, {"number": 2}]}, pull, pull], True),
         ]
         for responses, raises in cases:
@@ -861,6 +910,8 @@ def test_comment_rejects_older_runs_attempts_and_wrong_pr(self) -> None:
             "head_sha": base,
             "run_attempt": 2,
             "pull_requests": [{"number": 1}],
+            "head_repository": {"owner": {"login": "contributor"}},
+            "head_branch": "fix/branch",
         }
         comment = {
             "id": 8,
@@ -879,7 +930,7 @@ def test_comment_rejects_older_runs_attempts_and_wrong_pr(self) -> None:
             )
             self.assertEqual(api.call_count, 3)
         with (
-            mock.patch.object(report, "api", side_effect=[pull, {**run, "pull_requests": []}, []]),
+            mock.patch.object(report, "api", side_effect=[pull, {**run, "pull_requests": []}, [], []]),
             self.assertRaisesRegex(report.ReportError, "does not match report PR"),
         ):
             report.publish_comment(result, "openai/example", 1, 2, 2)