tests/test_llama_runtime.py

"""The runtime fetcher must be as careful as the weights fetcher.

It downloads an executable from another project and puts it on a volunteer's
machine, so every one of these is a property somebody would be harmed by losing:
the digest is checked before anything is installed, an archive cannot write
outside the runtime directory, and an archive cannot fill the disk.
"""

from __future__ import annotations

import hashlib
import io
import json
import zipfile
from pathlib import Path

import pytest

from distinct_agent import runtime


def _archive(names=("llama-server.exe", "ggml.dll"), size=32, deflate=False) -> bytes:
    buf = io.BytesIO()
    mode = zipfile.ZIP_DEFLATED if deflate else zipfile.ZIP_STORED
    with zipfile.ZipFile(buf, "w", mode) as handle:
        for name in names:
            handle.writestr(f"build/bin/{name}", b"x" * size)
    return buf.getvalue()


def _opener(payload: bytes):
    class _Response:
        headers = {"Content-Length": str(len(payload))}

        def __init__(self) -> None:
            self._buffer = io.BytesIO(payload)

        def read(self, size=-1):
            return self._buffer.read(size)

        def __enter__(self):
            return self

        def __exit__(self, *exc):
            return False

    return lambda request, timeout=None: _Response()


def _pins(tmp_path: Path, sha: str, key: str = "test-x64") -> Path:
    path = tmp_path / "pins.json"
    path.write_text(
        json.dumps(
            {
                "builds": {
                    key: {
                        "tag": "bTEST",
                        "asset": "a.zip",
                        "url": "https://github.com/ggml-org/llama.cpp/releases/download/bTEST/a.zip",
                        "sha256": sha,
                        "bytes": 1,
                    }
                }
            }
        )
    )
    return path


def test_a_verified_archive_installs_and_is_not_fetched_twice(tmp_path: Path) -> None:
    payload = _archive()
    root = tmp_path / "proj"
    root.mkdir()
    pins = _pins(tmp_path, hashlib.sha256(payload).hexdigest())

    installed = runtime.ensure(root=root, key="test-x64", opener=_opener(payload), pins_path=pins)
    assert Path(installed).is_file()
    assert Path(installed).name.startswith("llama-server")

    def refuse(*args, **kwargs):
        raise AssertionError("an installed runtime must not be downloaded again")

    assert runtime.ensure(root=root, key="test-x64", opener=refuse, pins_path=pins) == installed


def test_a_wrong_digest_installs_nothing(tmp_path: Path) -> None:
    payload = _archive()
    root = tmp_path / "proj"
    root.mkdir()
    with pytest.raises(runtime.RuntimeVerificationError):
        runtime.ensure(
            root=root, key="test-x64", opener=_opener(payload), pins_path=_pins(tmp_path, "0" * 64)
        )
    assert not list((root / "runtime").glob("*")) if (root / "runtime").exists() else True


def test_an_unrecorded_platform_is_never_downloaded(tmp_path: Path) -> None:
    """The bar the catalogue sets for a model applies to the thing that runs it."""

    payload = _archive()
    root = tmp_path / "proj"
    root.mkdir()
    with pytest.raises(runtime.RuntimeUnavailable) as caught:
        runtime.ensure(
            root=root,
            key="mystery-arch",
            opener=_opener(payload),
            pins_path=_pins(tmp_path, hashlib.sha256(payload).hexdigest()),
        )
    assert "will not download one" in str(caught.value)


def test_path_traversal_is_refused(tmp_path: Path) -> None:
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as handle:
        handle.writestr("../../escaped.txt", b"x")
    payload = buf.getvalue()
    root = tmp_path / "proj"
    root.mkdir()
    with pytest.raises(runtime.RuntimeUnavailable) as caught:
        runtime.ensure(
            root=root,
            key="test-x64",
            opener=_opener(payload),
            pins_path=_pins(tmp_path, hashlib.sha256(payload).hexdigest()),
        )
    assert "unsafe path" in str(caught.value)
    assert not (tmp_path / "escaped.txt").exists()


def test_a_decompression_bomb_is_refused(tmp_path: Path) -> None:
    """Small on the wire, enormous on disk. The digest would have matched."""

    payload = _archive(names=("llama-server",), size=runtime.MAX_UNPACKED_BYTES + 1, deflate=True)
    assert len(payload) < 5 * 1024 * 1024
    root = tmp_path / "proj"
    root.mkdir()
    with pytest.raises(runtime.RuntimeUnavailable) as caught:
        runtime.ensure(
            root=root,
            key="test-x64",
            opener=_opener(payload),
            pins_path=_pins(tmp_path, hashlib.sha256(payload).hexdigest()),
        )
    assert "unpacks to more than" in str(caught.value)


def test_an_archive_without_a_server_is_refused(tmp_path: Path) -> None:
    payload = _archive(names=("readme.txt",))
    root = tmp_path / "proj"
    root.mkdir()
    with pytest.raises(runtime.RuntimeUnavailable):
        runtime.ensure(
            root=root,
            key="test-x64",
            opener=_opener(payload),
            pins_path=_pins(tmp_path, hashlib.sha256(payload).hexdigest()),
        )


@pytest.mark.parametrize(
    "url,allowed",
    [
        ("https://github.com/ggml-org/llama.cpp/releases/download/b1/a.zip", True),
        ("https://objects.githubusercontent.com/whatever", True),
        ("http://github.com/ggml-org/llama.cpp/a.zip", False),
        ("https://github.com.evil.example/a.zip", False),
        ("https://example.com/a.zip", False),
    ],
)
def test_redirects_leave_github_nowhere(url: str, allowed: bool) -> None:
    assert runtime._RuntimeRedirects._permitted(url) is allowed


def test_the_shipped_pin_file_is_readable_and_pinned() -> None:
    """Whatever is recorded must be a real 64-character digest over https."""

    for key, pin in runtime.load_pins().items():
        assert len(pin.sha256) == 64, key
        assert int(pin.sha256, 16) >= 0, key
        assert pin.url.startswith("https://"), key
        assert pin.tag and pin.asset, key


def test_platform_key_is_stable() -> None:
    assert runtime.platform_key("Windows", "AMD64") == "windows-x64"
    assert runtime.platform_key("Linux", "x86_64") == "linux-x64"
    assert runtime.platform_key("Darwin", "arm64") == "macos-arm64"