rein-aharness/tests/test_messages_owner.py

"""Owner/worker lifecycle; real namespace proof remains explicitly opt-in."""

from dataclasses import replace
from datetime import UTC, datetime, timedelta
import json
import os
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from rein_aharness.execution_cancel import ExecutionCancel
from rein_aharness.messages_owner import MessagesOwner, accepted_lease
from rein_aharness.ops_run_client import ActivityCoreOpsClient, OpsRunConfig
from rein_aharness.request_admission import RequestLedger
from rein_aharness.spend_admission import SpendAdmissionError
from test_spend_admission import ledger as ledger, run, configured
from test_request_admission import policy as policy, fake_provider as fake_provider


def claimed():
    return replace(run(), state="claimed", lease_until=(datetime.now(UTC) + timedelta(seconds=60)).isoformat())


@pytest.mark.parametrize("fields", [
    {"id": "other"}, {"claim_owner": "other"}, {"attempt": 2},
    {"state": "open"}, {"lease_until": None}, {"lease_until": "invalid"},
    {"lease_until": "2020-01-01T00:00:00Z"},
])
def test_refuse_unaccepted_lease(fields):
    item = claimed()
    with pytest.raises(SpendAdmissionError):
        accepted_lease(item, replace(item, **fields), "worker-1")


def test_worker_requires_owner_before_claim(tmp_path, monkeypatch):
    from rein_aharness.claim_loop import process_one
    from rein_aharness.close_outbox import CloseOutbox
    client = MagicMock(spec=ActivityCoreOpsClient)
    client.config = OpsRunConfig(require_request_admission=True)
    result = process_one(client, outbox=CloseOutbox(state_dir=tmp_path / "close"))
    assert not result.claimed and "not configured" in result.reason
    client.claim.assert_not_called()


@pytest.mark.parametrize("cancel_reason", ["lease-loss", "signal", "timeout", None])
def test_owner_closes_and_never_rebinds(ledger, policy, fake_provider, cancel_reason, monkeypatch, tmp_path):
    monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "data"))
    meter = RequestLedger(ledger)
    meter.initialize()
    item = claimed()
    ledger.reserve(item)
    cfg = configured(ledger)
    provider, seen, _ = fake_provider
    owner = MessagesOwner(policy, "dummy-provider-key", f"http://127.0.0.1:{provider.server_port}", True)
    profile = SimpleNamespace(credential_route_refs=[], model=SimpleNamespace(model="fixture-model"),
                              sandbox_profile="profile.bwrap-local")
    cancel = ExecutionCancel()
    with owner.activate(item, cfg, ledger, profile, cancel) as manager:
        binding = manager._messages_route
        assert binding.socket_path.is_socket()
        if cancel_reason:
            cancel.cancel(cancel_reason)
            with pytest.raises(SpendAdmissionError, match="revoked"):
                meter.reserve_request(binding.token, policy.sha256, 1)
    assert not binding.socket_path.exists() and not seen
    with pytest.raises(SpendAdmissionError, match="revoked"):
        meter.reserve_request(binding.token, policy.sha256, 1)
    assert ledger.status()["reservations"][0]["state"] == "held"


@pytest.mark.skipif(os.environ.get("REIN_REAL_BWRAP") != "1", reason="opt-in kernel namespace proof")
def test_real_owner_route_confinement_and_revocation(ledger, policy, fake_provider, tmp_path, monkeypatch):
    from sandboxer.models import SandboxCreateRequest, SandboxExecRequest, Consumer

    monkeypatch.setenv("SANDBOXER_NO_STATE_HUB", "1")
    monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "data"))
    meter = RequestLedger(ledger)
    meter.initialize()
    item = claimed()
    ledger.reserve(item)
    cfg = configured(ledger)
    provider, seen, _ = fake_provider
    owner = MessagesOwner(policy, "synthetic-provider-key-owner-only", f"http://127.0.0.1:{provider.server_port}", True)
    profile = SimpleNamespace(credential_route_refs=[], model=SimpleNamespace(model="fixture-model"),
                              sandbox_profile="profile.bwrap-local")
    consumer = Consumer(actor="agt", project=cfg.execution_project, run_id=item.id)
    cancel = ExecutionCancel()
    source = tmp_path / "source"
    source.mkdir()
    (source / "README.md").write_text("No owner state belongs in this tree.\n")
    with owner.activate(item, cfg, ledger, profile, cancel) as manager:
        binding = manager._messages_route
        status = manager.create(SandboxCreateRequest(profile=profile.sandbox_profile,
                                                     inputs={"repo": str(source)}, consumer=consumer))
        try:
            assert status.state.value == "ready", status.error
            def execute(code):
                return manager.execute(status.sandbox_id, SandboxExecRequest(
                    command=["python3", "-c", code], consumer=consumer, timeout_seconds=10))
            probe = execute("""import os, json, socket
from pathlib import Path
blocked = []
for address in [("127.0.0.1", PROVIDER_PORT), ("1.1.1.1", 443)]:
 try:
  s = socket.create_connection(address, timeout=0.2); s.close(); blocked.append(False)
 except OSError: blocked.append(True)
print(json.dumps({"private_absent": not Path(PRIVATE).exists(),
 "owner_state_through_proc_absent": not Path(OWNER_PROCESS).exists(),
 "pid_namespace": os.readlink("/proc/self/ns/pid"),
 "network_namespace": os.readlink("/proc/self/ns/net"),
 "interfaces": sorted(p.name for p in Path("/sys/class/net").glob("*")) if Path("/sys/class/net").exists() else [line.split(':')[0].strip() for line in Path('/proc/net/dev').read_text().splitlines()[2:]],
 "direct_blocked": blocked, "proxy_absent": not any("proxy" in k.lower() for k in os.environ),
 "token_is_opaque": len(os.environ["ANTHROPIC_API_KEY"]) == 43}))
""".replace("PROVIDER_PORT", str(provider.server_port)).replace("PRIVATE", repr(str(ledger.path.parent))).replace("OWNER_PROCESS", repr(f"/proc/{os.getpid()}/root/{str(ledger.path.parent).lstrip(chr(47))}")))
            assert probe.exit_code == 0, probe.stderr
            evidence = json.loads(probe.stdout)
            assert evidence["private_absent"] and evidence["owner_state_through_proc_absent"], evidence
            assert evidence["direct_blocked"] == [True, True] and evidence["proxy_absent"]
            assert evidence["token_is_opaque"] and evidence["interfaces"] == ["lo"]
            assert evidence["network_namespace"] != os.readlink("/proc/self/ns/net")
            assert evidence["pid_namespace"] != os.readlink("/proc/self/ns/pid")
            request_code = """import os,json,http.client,urllib.parse
u=urllib.parse.urlsplit(os.environ['ANTHROPIC_BASE_URL'])
c=http.client.HTTPConnection(u.hostname,u.port,timeout=5)
c.request('POST','/v1/messages',json.dumps({'model':'fixture-model','max_tokens':1000,'stream':True,'messages':[{'role':'user','content':'synthetic request'}]}),{'Content-Type':'application/json','x-api-key':os.environ['ANTHROPIC_API_KEY']})
r=c.getresponse(); print(json.dumps({'status':r.status,'stopped':b'message_stop' in r.read()})); c.close()
"""
            response = execute(request_code)
            assert response.exit_code == 0, response.stderr
            assert json.loads(response.stdout) == {"status": 200, "stopped": True}
            assert len(seen) == 1 and seen[0]["key"] == "synthetic-provider-key-owner-only"
            assert meter.status()[0]["state"] == "charged"
            # This hook is the same callback used by the worker heartbeat monitor.
            cancel.cancel("lease-loss")
            response = execute(request_code)
            assert json.loads(response.stdout)["status"] == 400
            assert len(seen) == 1
            assert binding.token not in json.dumps(status.model_dump(mode="json"))
            assert "synthetic-provider-key-owner-only" not in probe.stdout + response.stdout
            if directory := os.environ.get("REIN_OWNER_PROOF_OUTPUT"):
                out = Path(directory); out.mkdir(parents=True, exist_ok=True)
                (out / "owner-route.json").write_text(json.dumps({
                    "scope": "real local bwrap owner transport; deterministic provider, synthetic key and lease",
                    **evidence, "upstream_requests": len(seen), "revoked_request_status": 400,
                    "request_reservations": meter.status(), "factory_attempts": 0,
                }, indent=2) + "\n")
        finally:
            manager.destroy(status.sandbox_id)
        assert not Path(status.inputs["workspace_dir"]).exists()
    assert not binding.socket_path.exists()


@pytest.mark.parametrize("valid", [False, True])
def test_claim_loop_uses_accepted_heartbeat_lease(ledger, policy, fake_provider, tmp_path, monkeypatch, valid):
    from rein_aharness.claim_loop import process_one
    from rein_aharness.close_outbox import CloseOutbox
    client = MagicMock(spec=ActivityCoreOpsClient)
    client.config = configured(ledger)
    client.config.messages_owner = MessagesOwner(policy, "dummy")
    item = claimed()
    client.claim.return_value = [item]
    expiry = (datetime.now(UTC) + timedelta(seconds=30)).isoformat() if valid else None
    client.heartbeat.return_value = replace(item, lease_until=expiry)
    process = MagicMock(return_value=SimpleNamespace(claimed=True))
    monkeypatch.setattr("rein_aharness.claim_loop._process_profiled_run", process)
    result = process_one(client, outbox=CloseOutbox(state_dir=tmp_path / "close"))
    if valid:
        assert item.lease_until == expiry
        process.assert_called_once()
    else:
        assert "accepted queue lease" in result.reason
        process.assert_not_called()


def test_lease_expiry_revokes_without_a_heartbeat(ledger, policy, tmp_path, monkeypatch):
    monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "data"))
    meter = RequestLedger(ledger)
    meter.initialize()
    item = claimed()
    item.lease_until = (datetime.now(UTC) + timedelta(seconds=0.5)).isoformat()
    ledger.reserve(item)
    profile = SimpleNamespace(credential_route_refs=[], model=SimpleNamespace(model="fixture-model"),
                              sandbox_profile="profile.bwrap-local")
    cancel = ExecutionCancel()
    with MessagesOwner(policy, "dummy").activate(item, configured(ledger), ledger, profile, cancel) as manager:
        binding = manager._messages_route
        assert cancel.wait(2), "accepted lease deadline was not enforced"
        with pytest.raises(SpendAdmissionError):
            meter.reserve_request(binding.token, policy.sha256, 1)