open-swe/tests/test_repo_binding_isolation.py

"""Repo-binding isolation for the sandbox and GitHub-token caches.

Covers TID-COLLIDE-01: a second repo presenting a colliding thread_id must not
reuse the first repo's in-process sandbox or cached GitHub token.
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from typing import Any

import pytest

from agent import server
from agent.utils import github_token


@pytest.fixture(autouse=True)
def _clear_caches() -> None:
    github_token._GITHUB_TOKEN_CACHE.clear()
    server.SANDBOX_BACKENDS.clear()


# --- token cache repo binding ------------------------------------------------


def test_cached_token_not_reused_across_repos() -> None:
    """A token bound to repo A is refused for a colliding thread_id from repo B."""
    future = (datetime.now(UTC) + timedelta(hours=1)).isoformat()
    github_token.cache_github_token_for_thread(
        "tid", "ghp_repoA", expires_at=future, repo={"owner": "acme", "name": "alpha"}
    )

    # Same thread_id, different repo → must NOT serve repo A's token.
    cfg_b = {"configurable": {"thread_id": "tid", "repo": {"owner": "evil", "name": "beta"}}}
    assert github_token.get_github_token(cfg_b) is None
    # And the poisoned entry is evicted.
    assert "tid" not in github_token._GITHUB_TOKEN_CACHE


def test_cached_token_reused_for_same_repo() -> None:
    future = (datetime.now(UTC) + timedelta(hours=1)).isoformat()
    github_token.cache_github_token_for_thread(
        "tid", "ghp_repoA", expires_at=future, repo={"owner": "acme", "name": "alpha"}
    )
    cfg_a = {"configurable": {"thread_id": "tid", "repo": {"owner": "acme", "name": "alpha"}}}
    assert github_token.get_github_token(cfg_a) == "ghp_repoA"


def test_unbound_token_served_when_repo_unknown() -> None:
    """Legacy entries with no bound repo still resolve when no repo is supplied."""
    github_token.cache_github_token_for_thread("tid", "ghp_legacy")
    assert github_token.get_github_token({"configurable": {"thread_id": "tid"}}) == "ghp_legacy"


def test_cached_token_reused_for_same_repo_case_insensitive() -> None:
    """``Org/Repo`` and ``org/repo`` are the same repo: no spurious refusal."""
    future = (datetime.now(UTC) + timedelta(hours=1)).isoformat()
    github_token.cache_github_token_for_thread(
        "tid", "ghp_repoA", expires_at=future, repo={"owner": "Acme", "name": "Alpha"}
    )
    cfg = {"configurable": {"thread_id": "tid", "repo": {"owner": "acme", "name": "alpha"}}}
    assert github_token.get_github_token(cfg) == "ghp_repoA"


def test_repo_cache_key_normalizes() -> None:
    assert github_token.repo_cache_key({"owner": "o", "name": "r"}) == "o/r"
    assert github_token.repo_cache_key("o/r") == "o/r"
    assert github_token.repo_cache_key({"owner": "o"}) is None
    assert github_token.repo_cache_key(None) is None
    # Casefolded so different casing of the same repo collapses to one key.
    assert github_token.repo_cache_key({"owner": "Org", "name": "Repo"}) == "org/repo"
    assert github_token.repo_cache_key("Org/Repo") == "org/repo"


# --- sandbox repo binding ----------------------------------------------------


class _FakeBackend:
    def __init__(self, sandbox_id: str) -> None:
        self.id = sandbox_id
        self.bound_repo: str | None = None

    def execute(self, *_a: Any, **_k: Any) -> Any:  # pragma: no cover - must not run
        raise AssertionError("colliding-repo sandbox must not be pinged/reused")


@pytest.mark.asyncio
async def test_ensure_sandbox_refuses_colliding_repo(monkeypatch: pytest.MonkeyPatch) -> None:
    async def fake_sandbox_id(_tid: str) -> str:
        return "sb-A"

    async def fake_bound_repo(_tid: str) -> str:
        return "acme/alpha"

    monkeypatch.setattr(server, "get_sandbox_id_from_metadata", fake_sandbox_id)
    monkeypatch.setattr(server, "get_bound_repo_from_metadata", fake_bound_repo)
    backend = _FakeBackend("sb-A")
    backend.bound_repo = "acme/alpha"
    server.SANDBOX_BACKENDS["tid"] = backend  # type: ignore[assignment]

    with pytest.raises(server.SandboxRepoMismatchError):
        await server.ensure_sandbox_for_thread("tid", repo={"owner": "evil", "name": "beta"})


@pytest.mark.asyncio
async def test_ensure_sandbox_allows_matching_repo(monkeypatch: pytest.MonkeyPatch) -> None:
    async def fake_sandbox_id(_tid: str) -> str:
        return "sb-A"

    async def fake_bound_repo(_tid: str) -> str:
        return "acme/alpha"

    calls: dict[str, int] = {"git": 0}
    backend = _FakeBackend("sb-A")
    backend.bound_repo = "acme/alpha"

    async def fake_check(b: Any, *_a: Any, **_k: Any) -> Any:
        return b

    async def fake_refresh(b: Any, *_a: Any, **_k: Any) -> Any:
        return b

    async def fake_git(_b: Any) -> None:
        calls["git"] += 1

    monkeypatch.setattr(server, "get_sandbox_id_from_metadata", fake_sandbox_id)
    monkeypatch.setattr(server, "get_bound_repo_from_metadata", fake_bound_repo)
    monkeypatch.setattr(server, "check_or_recreate_sandbox", fake_check)
    monkeypatch.setattr(server, "_refresh_github_proxy_or_recreate", fake_refresh)
    monkeypatch.setattr(server, "set_sandbox_backend", lambda _tid, b, **_k: b)
    monkeypatch.setattr(server, "_configure_git_identity", fake_git)
    server.SANDBOX_BACKENDS["tid"] = backend  # type: ignore[assignment]

    result = await server.ensure_sandbox_for_thread("tid", repo={"owner": "acme", "name": "alpha"})
    assert result is backend
    assert calls["git"] == 1


@pytest.mark.asyncio
async def test_ensure_sandbox_reuses_same_repo_case_insensitive(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A bound sandbox is reused when the current repo differs only by casing."""

    async def fake_sandbox_id(_tid: str) -> str:
        return "sb-A"

    async def fake_bound_repo(_tid: str) -> str:
        return "acme/alpha"

    backend = _FakeBackend("sb-A")
    backend.bound_repo = "acme/alpha"

    async def fake_check(b: Any, *_a: Any, **_k: Any) -> Any:
        return b

    async def fake_refresh(b: Any, *_a: Any, **_k: Any) -> Any:
        return b

    async def fake_git(_b: Any) -> None:
        return None

    monkeypatch.setattr(server, "get_sandbox_id_from_metadata", fake_sandbox_id)
    monkeypatch.setattr(server, "get_bound_repo_from_metadata", fake_bound_repo)
    monkeypatch.setattr(server, "check_or_recreate_sandbox", fake_check)
    monkeypatch.setattr(server, "_refresh_github_proxy_or_recreate", fake_refresh)
    monkeypatch.setattr(server, "set_sandbox_backend", lambda _tid, b, **_k: b)
    monkeypatch.setattr(server, "_configure_git_identity", fake_git)
    server.SANDBOX_BACKENDS["tid"] = backend  # type: ignore[assignment]

    # Different casing of the same repo must not raise and must reuse the sandbox.
    result = await server.ensure_sandbox_for_thread("tid", repo={"owner": "ACME", "name": "Alpha"})
    assert result is backend


@pytest.mark.asyncio
async def test_legacy_unbound_sandbox_recreated_not_reused(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A legacy thread (sandbox_id present, bound_repo absent) must fail closed.

    With no recorded binding, the existing sandbox cannot be confirmed to belong
    to the current repo, so it is never reconnected-and-served: a fresh sandbox is
    created and bound to the requesting repo instead.
    """

    async def fake_sandbox_id(_tid: str) -> str:
        return "sb-legacy"

    async def fake_bound_repo(_tid: str) -> None:
        return None

    def fake_reconnect(*_a: Any, **_k: Any) -> Any:  # pragma: no cover - must not run
        raise AssertionError("must not reconnect to a legacy unbound sandbox")

    fresh = _FakeBackend("sb-fresh")

    async def fake_create_with_proxy(*_a: Any, **_k: Any) -> Any:
        return fresh

    async def fake_git(_b: Any) -> None:
        return None

    class _Threads:
        async def update(self, **_k: Any) -> None:
            return None

        async def get(self, *_a: Any, **_k: Any) -> dict[str, Any]:
            return {}

    class _Client:
        threads = _Threads()

    monkeypatch.setattr(server, "get_sandbox_id_from_metadata", fake_sandbox_id)
    monkeypatch.setattr(server, "get_bound_repo_from_metadata", fake_bound_repo)
    monkeypatch.setattr(server, "create_sandbox", fake_reconnect)
    monkeypatch.setattr(server, "_create_sandbox_with_proxy", fake_create_with_proxy)
    monkeypatch.setattr(server, "_configure_git_identity", fake_git)
    monkeypatch.setattr(server, "client", _Client())

    result = await server.ensure_sandbox_for_thread("tid", repo={"owner": "evil", "name": "beta"})

    assert result.id == "sb-fresh"
    assert server.SANDBOX_BACKENDS["tid"].bound_repo == "evil/beta"