fix(colab): verify model access before loading the pipeline

tests/test_notebook_contract.py

@@ -37,3 +37,189 @@ def test_abliterate_notebook_stage_callback_uses_stage_result_contract(capsys):
 
     assert "STAGE: PROBE" in output
     assert "loading model" in output
+
+
+def _notebook_code(marker):
+    notebook = json.loads(Path("notebooks/abliterate.ipynb").read_text())
+    return next(
+        "".join(cell["source"])
+        for cell in notebook["cells"]
+        if cell["cell_type"] == "code" and marker in "".join(cell["source"])
+    )
+
+
+def _run_harness(monkeypatch, *, token=None, failure=None):
+    """Execute real notebook cells with an offline Hub and observable pipeline."""
+    import sys
+    from unittest.mock import Mock
+
+    monkeypatch.delenv("HF_TOKEN", raising=False)
+    events = []
+
+    def download(**kwargs):
+        events.append(("access", kwargs))
+        if failure:
+            raise failure
+        return "offline-config.json"
+
+    def pipeline(**kwargs):
+        import os
+        events.append(("construct", kwargs, os.environ.get("HF_TOKEN")))
+        return SimpleNamespace(run=lambda: events.append(("run",)) or "output")
+
+    hub = SimpleNamespace(
+        get_token=Mock(return_value=token), hf_hub_download=download,
+        HfApi=Mock(side_effect=AssertionError("Unexpected upload")),
+        notebook_login=Mock(side_effect=AssertionError("Unexpected login widget")),
+    )
+    secret = Mock(return_value=None)
+    monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
+    monkeypatch.setitem(sys.modules, "obliteratus.credential_sources", SimpleNamespace(resolve_secret=secret))
+    monkeypatch.setitem(sys.modules, "obliteratus.abliterate", SimpleNamespace(AbliterationPipeline=pipeline))
+    namespace = {}
+    exec(_notebook_code("#@title Abliteration Config"), namespace)
+    return namespace, events, hub, secret
+
+
+def test_default_gated_model_stops_before_pipeline_without_access(monkeypatch, capsys):
+    import traceback
+    import pytest
+
+    namespace, events, hub, _ = _run_harness(
+        monkeypatch, failure=PermissionError("private request details fake-secret-value"),
+    )
+    with pytest.raises(RuntimeError, match="HF_TOKEN") as caught:
+        exec(_notebook_code("def check_model_access"), namespace)
+    assert [event[0] for event in events] == ["access"]
+    assert events[0][1] == {
+        "repo_id": "meta-llama/Llama-3.1-8B-Instruct", "filename": "config.json",
+        "token": False, "force_download": True,
+    }
+    rendered = "".join(traceback.format_exception(caught.type, caught.value, caught.tb))
+    assert "fake-secret-value" not in rendered + capsys.readouterr().out
+    assert "wait for approval" in str(caught.value)
+    hub.notebook_login.assert_not_called()
+
+
+def test_authorized_colab_token_reaches_loader_without_output(monkeypatch, capsys):
+    namespace, events, _, _ = _run_harness(monkeypatch, token="fake-private-read-token")
+    exec(_notebook_code("def check_model_access"), namespace)
+    assert [event[0] for event in events] == ["access", "construct", "run"]
+    assert events[0][1]["token"] == "fake-private-read-token"
+    assert events[1][2] == "fake-private-read-token"
+    assert "fake-private-read-token" not in capsys.readouterr().out
+
+
+def test_ungated_model_runs_anonymously_and_upload_is_opt_in(monkeypatch):
+    namespace, events, hub, _ = _run_harness(monkeypatch)
+    namespace["MODEL"] = "Qwen/Qwen2.5-7B-Instruct"
+    exec(_notebook_code("def check_model_access"), namespace)
+    exec(_notebook_code("UPLOAD_TO_HUB ="), namespace)
+    assert [event[0] for event in events] == ["access", "construct", "run"]
+    assert events[0][1]["token"] is False
+    assert events[1][2] is None
+    hub.HfApi.assert_not_called()
+    hub.notebook_login.assert_not_called()
+
+
+def test_direct_rerun_rechecks_changed_model_and_blocks_expired_access(monkeypatch):
+    import pytest
+
+    namespace, events, hub, _ = _run_harness(monkeypatch, token="fake-token")
+    source = _notebook_code("def check_model_access")
+    exec(source, namespace)
+    namespace["MODEL"] = "other/private-model"
+
+    def denied(**kwargs):
+        events.append(("access", kwargs))
+        raise PermissionError("expired")
+
+    hub.hf_hub_download = denied
+    with pytest.raises(RuntimeError, match="Model access"):
+        exec(source, namespace)
+    assert [event[0] for event in events] == ["access", "construct", "run", "access"]
+    assert events[-1][1]["repo_id"] == "other/private-model"
+
+
+def test_configured_secret_precedes_hub_cached_token(monkeypatch):
+    namespace, events, hub, secret = _run_harness(monkeypatch, token="fake-other-account")
+    secret.return_value = "fake-configured-token"
+    exec(_notebook_code("def check_model_access"), namespace)
+    assert events[0][1]["token"] == "fake-configured-token"
+    hub.get_token.assert_not_called()
+
+
+def test_secret_resolution_failure_cannot_fall_back_or_run(monkeypatch):
+    import pytest
+
+    namespace, events, hub, secret = _run_harness(monkeypatch, token="fake-cached-token")
+    secret.side_effect = RuntimeError("sensitive credential source")
+    with pytest.raises(RuntimeError, match="Model access"):
+        exec(_notebook_code("def check_model_access"), namespace)
+    assert events == []
+    hub.get_token.assert_not_called()
+
+
+def test_notebook_preflight_network_failure_stops_before_loading(monkeypatch):
+    import pytest
+
+    namespace, events, _, _ = _run_harness(monkeypatch, failure=ConnectionError("offline"))
+    with pytest.raises(RuntimeError, match="connection"):
+        exec(_notebook_code("def check_model_access"), namespace)
+    assert [event[0] for event in events] == ["access"]
+
+
+def test_notebook_has_no_persisted_execution_outputs():
+    notebook = json.loads(Path("notebooks/abliterate.ipynb").read_text())
+    for cell in notebook["cells"]:
+        if cell["cell_type"] == "code":
+            assert cell.get("outputs", []) == []
+            assert cell.get("execution_count") is None
+
+
+def test_token_without_gated_approval_cannot_run(monkeypatch):
+    import pytest
+
+    namespace, events, _, _ = _run_harness(
+        monkeypatch, token="fake-unapproved-token", failure=PermissionError("403"),
+    )
+    with pytest.raises(RuntimeError, match="wait for approval"):
+        exec(_notebook_code("def check_model_access"), namespace)
+    assert [event[0] for event in events] == ["access"]
+
+
+def test_explicit_upload_uses_resolved_token_and_destination(monkeypatch):
+    from unittest.mock import Mock
+
+    namespace, _, hub, secret = _run_harness(monkeypatch)
+    secret.return_value = "fake-write-token"
+    api = Mock()
+    hub.HfApi = Mock(return_value=api)
+    namespace["model_dir"] = Path("saved-model")
+    source = _notebook_code("UPLOAD_TO_HUB =").replace(
+        "UPLOAD_TO_HUB = False", "UPLOAD_TO_HUB = True",
+    ).replace(
+        'HF_REPO = "your-username/model-name-abliterated"',
+        'HF_REPO = "researcher/output"',
+    )
+    exec(source, namespace)
+    hub.HfApi.assert_called_once_with(token="fake-write-token")
+    api.create_repo.assert_called_once_with("researcher/output", exist_ok=True)
+    api.upload_folder.assert_called_once_with(
+        folder_path="saved-model", repo_id="researcher/output", repo_type="model",
+    )
+
+
+def test_explicit_upload_stops_without_credentials(monkeypatch):
+    import pytest
+
+    namespace, _, hub, _ = _run_harness(monkeypatch)
+    source = _notebook_code("UPLOAD_TO_HUB =").replace(
+        "UPLOAD_TO_HUB = False", "UPLOAD_TO_HUB = True",
+    ).replace(
+        'HF_REPO = "your-username/model-name-abliterated"',
+        'HF_REPO = "researcher/output"',
+    )
+    with pytest.raises(RuntimeError, match="write token"):
+        exec(source, namespace)
+    hub.HfApi.assert_not_called()