Bind the metered owner route to worker leases and sandbox lifecycle

Assistant: codex
Assistant-Model: gpt-6-astra
Assistant-Session: 01a07ff8-19d0-7820-b4d0-1353833cb7fc

tests/test_messages_owner.py

@@ -0,0 +1,191 @@
+"""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)

tests/test_repository_artifact_bwrap.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 import json
+from datetime import UTC, datetime, timedelta
 import os
 from pathlib import Path
 from unittest.mock import MagicMock
@@ -31,6 +32,10 @@ from rein_aharness.ops_run_client import (
 from test_repository_artifact import commit, git
 from rein_aharness.repository_grant import RepositoryGrant
 from rein_aharness.spend_admission import SpendLedger, SpendPolicy, digest
+from rein_aharness.messages_owner import MessagesOwner
+from rein_aharness.request_admission import RequestLedger
+from llm_connect.messages_gate import MessagesPolicy
+from test_request_admission import fake_provider as fake_provider
 
 pytestmark = pytest.mark.skipif(
     os.environ.get("REIN_REAL_BWRAP") != "1", reason="opt-in kernel namespace proof"
@@ -42,8 +47,9 @@ class DeterministicRein(Rein):
     workspace = None
     head = None
 
-    def __init__(self, source, baseline):
+    def __init__(self, source, baseline, metered=False):
         self.source, self.baseline = source, baseline
+        self.metered = metered
 
     def start_session(self, profile, inputs, sandbox):
         transport = transport_from_sandbox(sandbox)
@@ -52,6 +58,14 @@ class DeterministicRein(Rein):
 
     def dispatch_tool(self, transport, tool_call):
         self.calls += 1
+        if self.metered:
+            response = transport.run(["python3", "-c", """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':'fixture only'}]}),{'Content-Type':'application/json','x-api-key':os.environ['ANTHROPIC_API_KEY']})
+r=c.getresponse(); assert r.status==200; assert b'message_stop' in r.read(); c.close()
+"""], timeout=10)
+            assert response.returncode == 0, response.stderr
         code = """from pathlib import Path
 import subprocess
 Path('result.txt').write_text('sandbox-result-only\\n')
@@ -79,8 +93,8 @@ print(git('rev-parse','HEAD'))
         assert git(self.source, "rev-parse", "HEAD") == self.baseline
 
 
-@pytest.mark.parametrize("with_spend", [False, True])
-def test_real_bwrap_worker_import_and_close_replay(tmp_path, monkeypatch, with_spend):
+@pytest.mark.parametrize("with_spend", [False, True, "requests"])
+def test_real_bwrap_worker_import_and_close_replay(tmp_path, monkeypatch, with_spend, fake_provider):
     monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "data"))
     monkeypatch.setenv("REIN_AHARNESS_STATE_DIR", str(tmp_path / "state"))
     monkeypatch.setenv("SANDBOXER_NO_STATE_HUB", "1")
@@ -108,6 +122,7 @@ def test_real_bwrap_worker_import_and_close_replay(tmp_path, monkeypatch, with_s
         repository_grant=grant,
         claim_owner="fixture-worker",
         attempt=1,
+        lease_until=(datetime.now(UTC) + timedelta(seconds=90)).isoformat(),
     )
     client = MagicMock(spec=ActivityCoreOpsClient)
     client.config = OpsRunConfig(
@@ -141,6 +156,12 @@ def test_real_bwrap_worker_import_and_close_replay(tmp_path, monkeypatch, with_s
             )
         }
     )
+    if with_spend == "requests":
+        admitted, _ = catalog.resolve(run.harness_profile_ref)
+        catalog.profiles()[(admitted.id, admitted.version)] = admitted.model_copy(update={
+            "model": admitted.model.model_copy(update={"model": "fixture-model"}),
+            "sandbox_profile": "profile.bwrap-local", "credential_route_refs": [],
+        })
     spend = None
     if with_spend:
         private = tmp_path / "spend-state"
@@ -166,12 +187,20 @@ def test_real_bwrap_worker_import_and_close_replay(tmp_path, monkeypatch, with_s
         client.config.spend_policy_path = str(policy_path)
         client.config.spend_ledger_path = str(spend.path)
         monkeypatch.setattr("glas_harness.profiles.ProfileCatalog", lambda: catalog)
+    if with_spend == "requests":
+        RequestLedger(spend).initialize()
+        provider, seen, _ = fake_provider
+        client.config.messages_owner = MessagesOwner(
+            MessagesPolicy("fixture:upper-rate", "fixture-model", 1000, 1000, 1000, 1000),
+            "fixture-owner-key", f"http://127.0.0.1:{provider.server_port}", True,
+        )
+        client.config.require_request_admission = True
     manager = SandboxManager(store=SandboxStore(tmp_path / "sandboxes.json"))
-    rein = DeterministicRein(source, baseline)
+    rein = DeterministicRein(source, baseline, metered=with_spend == "requests")
     monkeypatch.setattr(
         "glas_harness.gateway.run_execution",
         lambda request, **kwargs: run_execution(
-            request, catalog=kwargs.pop("catalog", catalog), rein=rein, manager=manager, **kwargs
+            request, catalog=kwargs.pop("catalog", catalog), rein=rein, manager=kwargs.pop("manager", manager), **kwargs
         ),
     )
     outbox = CloseOutbox(state_dir=tmp_path / "state")
@@ -205,3 +234,10 @@ def test_real_bwrap_worker_import_and_close_replay(tmp_path, monkeypatch, with_s
     assert client.complete.call_count == 2
     assert client.heartbeat.call_count >= 1
     assert git(source, "rev-list", "--count", "HEAD") == "2"
+
+    if with_spend == "requests":
+        assert len(seen) == 1 and seen[0]["key"] == "fixture-owner-key"
+        assert RequestLedger(spend).status()[0]["state"] == "charged"
+        with spend._db() as db:
+            route = db.execute("SELECT * FROM request_routes").fetchone()
+            assert route["revoked"] == 1 and route["expires_at"] == run.lease_until