feat(run-task): add ability to fetch arbitrary extra refs

Bug: 2028208

Author: ahal

src/taskgraph/run-task/run-task

@@ -31,7 +31,7 @@ import time
 import urllib.error
 import urllib.request
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 SECRET_BASEURL_TPL = "{}/secrets/v1/secret/{{}}".format(
     os.environ.get("TASKCLUSTER_PROXY_URL", "http://taskcluster").rstrip("/")
@@ -640,6 +640,7 @@ def git_checkout(
     ssh_key_file: Optional[Path],
     ssh_known_hosts_file: Optional[Path],
     shallow: bool = False,
+    extra_refs: Optional[List[str]] = None,
 ):
     assert head_ref or head_rev
 
@@ -734,6 +735,9 @@ def git_checkout(
         env=env,
     )
 
+    for ref in extra_refs or []:
+        git_fetch(destination_path, f"{ref}:{ref}", remote=head_repo, env=env)
+
     args = [
         "git",
         "checkout",
@@ -932,6 +936,8 @@ def collect_vcs_options(args, project, name):
     head_rev = os.environ.get(f"{env_prefix}_HEAD_REV")
     pip_requirements = os.environ.get(f"{env_prefix}_PIP_REQUIREMENTS")
     private_key_secret = os.environ.get(f"{env_prefix}_SSH_SECRET_NAME")
+    if extra_refs := os.environ.get(f"{env_prefix}_EXTRA_REFS"):
+        extra_refs = json.loads(extra_refs)
 
     store_path = os.environ.get("HG_STORE_PATH")
 
@@ -965,6 +971,7 @@ def collect_vcs_options(args, project, name):
         "ssh-secret-name": private_key_secret,
         "pip-requirements": pip_requirements,
         "shallow-clone": shallow_clone,
+        "extra-refs": extra_refs,
     }
 
 
@@ -1014,6 +1021,7 @@ def vcs_checkout_from_args(options):
                 ssh_key_file,
                 ssh_known_hosts_file,
                 shallow=options.get("shallow-clone", False),
+                extra_refs=options.get("extra-refs"),
             )
         elif options["repo-type"] == "hg":
             revision = hg_checkout(

test/test_scripts_run_task.py

@@ -1,4 +1,5 @@
 import functools
+import json
 import os
 import site
 import stat
@@ -179,6 +180,17 @@ def test_install_pip_requirements_with_uv(
             {"shallow-clone": True},
             id="git_with_shallow_clone",
         ),
+        pytest.param(
+            {},
+            {
+                "REPOSITORY_TYPE": "git",
+                "HEAD_REPOSITORY": "https://github.com/example/repo",
+                "HEAD_REV": "abc123",
+                "EXTRA_REFS": json.dumps(["refs/notes/taskgraph", "refs/notes/other"]),
+            },
+            {"extra-refs": ["refs/notes/taskgraph", "refs/notes/other"]},
+            id="git_with_extra_refs",
+        ),
     ],
 )
 def test_collect_vcs_options(
@@ -216,6 +228,7 @@ def test_collect_vcs_options(
         "shallow-clone": False,
         "ssh-secret-name": env.get("SSH_SECRET_NAME"),
         "store-path": env.get("HG_STORE_PATH"),
+        "extra-refs": None,
     }
     if "PIP_REQUIREMENTS" in env:
         expected["pip-requirements"] = os.path.join(
@@ -638,3 +651,36 @@ def test_main_abspath_environment(mocker, run_main):
     assert env.get("MOZ_UV_HOME") == "/builds/worker/dir/uv"
     for key in envvars:
         assert env[key] == "/builds/worker/file"
+
+
+def test_git_checkout_extra_refs(mock_stdin, run_task_mod, mock_git_repo, tmp_path):
+    """extra_refs are fetched into the local repo during checkout."""
+    # Add a notes ref to the source repo
+    rev = mock_git_repo["main"][-1]
+    subprocess.check_call(
+        ["git", "notes", "--ref=refs/notes/taskgraph", "add", "-m", "test", rev],
+        cwd=mock_git_repo["path"],
+    )
+
+    destination = tmp_path / "destination"
+    run_task_mod.git_checkout(
+        destination_path=str(destination),
+        head_repo=mock_git_repo["path"],
+        base_repo=mock_git_repo["path"],
+        base_rev=None,
+        head_ref="main",
+        head_rev=None,
+        ssh_key_file=None,
+        ssh_known_hosts_file=None,
+        extra_refs=["refs/notes/taskgraph"],
+    )
+
+    # Verify the notes ref is available locally
+    result = subprocess.run(
+        ["git", "notes", "--ref=refs/notes/taskgraph", "show", rev],
+        cwd=str(destination),
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0
+    assert "test" in result.stdout