From dc95f0afffbeffa03855eca7c19e5b6a766acf0e Mon Sep 17 00:00:00 2001
From: Owen Price Skelly <21372141+OwenPriceSkelly@users.noreply.github.com>
Date: Tue, 1 Sep 2026 10:31:47 -0500
Subject: [PATCH] fetch_checkpoint: fall back to VERIFY_KWARGS like the verify
 path
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The download phase runs setup(checkpoint, "cpu") via _run_download, but
unlike verify_checkpoint it never fell back to the env's declared
VERIFY_KWARGS when no explicit setup_kwargs were given. Since #219 made
multi-head checkpoints raise without a head selection, every such
checkpoint (uma task=..., mace-mh-1 head=...) failed the sync's download
phase with e.g. "download: ValueError: uma-s-1p1 is multi-task and has
no default head" (first seen on the 2026-08-28 Aurora sync, 1.6.3).

Resolve the hosting env first, then default setup_kwargs to
verify_kwargs_for(root, env_name, checkpoint) — explicit kwargs still
win, and envs without VERIFY_KWARGS still get {}.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 rootstock/operations.py          |  8 +++-
 tests/commands/test_add_split.py | 65 ++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/rootstock/operations.py b/rootstock/operations.py
index 73def46..188723f 100644
--- a/rootstock/operations.py
+++ b/rootstock/operations.py
@@ -39,6 +39,7 @@
     parse_checkpoints_dict,
     parse_clusters_list,
     parse_custom_checkpoint_ids,
+    verify_kwargs_for,
 )
 from .exceptions import RootstockError
 from .layout import resolve_cache_root
@@ -1211,12 +1212,17 @@ def fetch_checkpoint(
     manifest's ``last_error``).
     """
     root = Path(root)
-    setup_kwargs = setup_kwargs or {}
     if cache_root is None:
         cache_root = resolve_cache_root(root)
 
     env_name = _resolve_built_env(root, checkpoint, cluster)
 
+    # No explicit kwargs → fall back to the env's declared VERIFY_KWARGS,
+    # exactly like the verify path: setup() runs the download here too, and a
+    # multi-head env (UMA's task, MACE-MH-1's head) raises without a head
+    # selection. Explicit kwargs always win.
+    setup_kwargs = setup_kwargs or verify_kwargs_for(root, env_name, checkpoint)
+
     # Unlocked peek for idempotence; every write below loads fresh inside
     # its own transaction, so a racing writer costs at most a redundant
     # (idempotent) download, never a lost update.
diff --git a/tests/commands/test_add_split.py b/tests/commands/test_add_split.py
index ccd3c48..96ea03f 100644
--- a/tests/commands/test_add_split.py
+++ b/tests/commands/test_add_split.py
@@ -140,6 +140,71 @@ def test_fetch_refresh_knob(fake_root, refresh_calls, monkeypatch):
     assert len(refresh_calls) == 1
 
 
+_MULTIHEAD_ENV_SOURCE = '''\
+"""Multi-head env: setup() requires a head selection, no default."""
+
+CHECKPOINTS = {
+    "uma-s-1p1": "uma-s-1p1",
+}
+
+VERIFY_KWARGS = {
+    "uma-s-1p1": {"task": "omat"},
+}
+
+
+def setup(checkpoint, device="cuda", task=None):
+    return None
+'''
+
+
+@pytest.fixture
+def multihead_root(tmp_path: Path) -> Path:
+    """A root whose env declares VERIFY_KWARGS for its checkpoint."""
+    root = tmp_path
+    env_dir = root / "envs" / "uma"
+    (env_dir / "bin").mkdir(parents=True)
+    (env_dir / "bin" / "python").touch()
+    (env_dir / "env_source.py").write_text(_MULTIHEAD_ENV_SOURCE)
+
+    from rootstock.config import UserConfig
+    from rootstock.manifest import create_manifest, save_manifest
+
+    cfg = UserConfig(name="t", email="t@t.t")
+    save_manifest(create_manifest(root, ["test"], cfg), root)
+    return root
+
+
+def test_fetch_falls_back_to_verify_kwargs(multihead_root, refresh_calls, monkeypatch):
+    """No explicit kwargs → the env's VERIFY_KWARGS reach the download's
+    setup() call, same as the verify path (a multi-head setup() raises
+    without a head selection)."""
+    captured = {}
+
+    def fake_download(root, env_name, checkpoint, setup_kwargs, **kw):
+        captured["setup_kwargs"] = setup_kwargs
+        return True, None
+
+    monkeypatch.setattr(operations, "_run_download", fake_download)
+
+    fetch_checkpoint(multihead_root, "uma-s-1p1")
+
+    assert captured["setup_kwargs"] == {"task": "omat"}
+
+
+def test_fetch_explicit_kwargs_beat_verify_kwargs(multihead_root, refresh_calls, monkeypatch):
+    captured = {}
+
+    def fake_download(root, env_name, checkpoint, setup_kwargs, **kw):
+        captured["setup_kwargs"] = setup_kwargs
+        return True, None
+
+    monkeypatch.setattr(operations, "_run_download", fake_download)
+
+    fetch_checkpoint(multihead_root, "uma-s-1p1", setup_kwargs={"task": "omol"})
+
+    assert captured["setup_kwargs"] == {"task": "omol"}
+
+
 def test_fetch_fails_fast_when_env_not_built(tmp_path, refresh_calls):
     from rootstock.config import UserConfig
     from rootstock.manifest import create_manifest, save_manifest