atlas-iac/testing/tests/test_hermes_execution_pool.py

"""Adversarial contracts for the fenced three-node Hermes execution pool."""

from __future__ import annotations

import json
import subprocess
import sys
import threading
import time
from pathlib import Path

import pytest
import yaml


ROOT = Path(__file__).parents[2]
HERMES = ROOT / "services/hermes"
SCRIPTS = HERMES / "scripts"
sys.path[:0] = [str(SCRIPTS), str(HERMES / "scm-common/scripts")]

import execution_pool_coordinator as coordinator  # noqa: E402
import execution_pool_client as pool_client  # noqa: E402
import execution_pool_protocol as protocol  # noqa: E402
import execution_pool_scm as scm  # noqa: E402
import execution_pool_worker as worker  # noqa: E402


KEY = b"k" * 32


def binding(**values):
    result = {
        "board": "atlas",
        "task_id": "t_deadbeef",
        "run_id": "run-1234",
        "worker_ordinal": 0,
        "attempt": 1,
    }
    result.update(values)
    return result


def assignment_payload(**values):
    result = {
        "context": "Implement the bounded objective.",
        "assignee": "cli-auto",
        "repo_url": "https://scm.bstein.dev/atlas/titan-iac.git",
        "branch": "feature/hermes-safe-pool",
        "base_branch": "main",
        "max_runtime_seconds": 3600,
        "deadline_unix": int(time.time()) + 3600,
    }
    result.update(values)
    return result


def test_envelope_binds_run_ordinal_attempt_and_digest():
    signed = protocol.sign_envelope(KEY, "heartbeat", binding(), {"note": "active"})
    verified = protocol.verify_envelope(KEY, signed, expected_kind="heartbeat")

    assert verified["payload_digest"] == protocol.payload_digest({"note": "active"})
    for name, value in binding().items():
        assert verified[name] == value

    signed["attempt"] = 2
    with pytest.raises(protocol.ProtocolError, match="authentication"):
        protocol.verify_envelope(KEY, signed)


def test_empty_poll_acknowledgement_is_authenticated():
    empty = {
        "board": "", "task_id": "", "run_id": "",
        "worker_ordinal": 2, "attempt": 0,
    }
    signed = protocol.sign_envelope(KEY, "ack", empty, {"assignment": None})
    assert protocol.verify_envelope(KEY, signed, expected_kind="ack")["payload"] == {
        "assignment": None
    }


@pytest.mark.parametrize(
    "mutation,error",
    [
        ({"board": "../atlas"}, "invalid board"),
        ({"worker_ordinal": 3}, "outside the pool"),
        ({"expires_at": 1}, "validity window"),
    ],
)
def test_malformed_or_traversal_bindings_fail_closed(mutation, error):
    signed = protocol.sign_envelope(KEY, "heartbeat", binding(), {"note": "active"})
    signed.update(mutation)
    unsigned = dict(signed)
    unsigned.pop("signature")
    import hashlib
    import hmac

    signed["signature"] = hmac.new(
        KEY, protocol.canonical_json(unsigned), hashlib.sha256
    ).hexdigest()
    with pytest.raises(protocol.ProtocolError, match=error):
        protocol.verify_envelope(KEY, signed)


def test_oversized_and_malformed_payloads_are_rejected():
    with pytest.raises(protocol.ProtocolError, match="wire limit"):
        protocol.sign_envelope(
            KEY, "result", binding(), {"output": "x" * protocol.MAX_WIRE_BYTES}
        )
    with pytest.raises(protocol.ProtocolError, match="malformed"):
        protocol.parse_wire(b"{not-json")
    with pytest.raises(protocol.ProtocolError, match="oversized"):
        protocol.parse_wire(b"x" * (protocol.MAX_WIRE_BYTES + 1))


def test_key_requires_private_regular_file_and_rejects_symlink(tmp_path):
    key = tmp_path / "key"
    key.write_bytes(KEY)
    key.chmod(0o644)
    with pytest.raises(protocol.ProtocolError, match="private"):
        protocol.read_key(key)
    key.chmod(0o600)
    assert protocol.read_key(key) == KEY
    link = tmp_path / "link"
    link.symlink_to(key)
    with pytest.raises(protocol.ProtocolError, match="unavailable"):
        protocol.read_key(link)
    with pytest.raises(protocol.ProtocolError, match="unavailable"):
        protocol.read_key(tmp_path / "missing")


def test_local_signing_boundary_rejects_foreign_or_unassigned_result(monkeypatch):
    monkeypatch.setattr(pool_client, "ORDINAL", 0)
    boundary = pool_client.ClientBoundary(KEY)
    request = {"binding": binding(), "payload": {"note": "active"}}
    with pytest.raises(protocol.ProtocolError, match="does not own"):
        boundary.heartbeat(request)
    boundary.current = binding()
    request["binding"] = binding(worker_ordinal=1)
    with pytest.raises(protocol.ProtocolError, match="does not own"):
        boundary.heartbeat(request)


def test_simultaneous_claim_materialization_has_one_winner(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    barrier = threading.Barrier(8)
    outcomes = []

    def add(index):
        barrier.wait()
        try:
            outcome = store.add(
                binding(task_id=f"t_{index}", run_id=f"run-{index}"),
                assignment_payload(),
            )
        except protocol.ProtocolError:
            outcome = False
        outcomes.append(outcome)

    threads = [threading.Thread(target=add, args=(index,)) for index in range(8)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join()

    assert outcomes.count(True) == 1
    assert store.available_ordinals() == [1, 2]


def test_duplicate_assignment_is_idempotent_but_conflict_is_rejected(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    assert store.add(binding(), assignment_payload()) is True
    assert store.add(binding(), assignment_payload()) is False
    with pytest.raises(protocol.ProtocolError, match="conflicting"):
        store.add(binding(), assignment_payload(context="different"))
    with pytest.raises(protocol.ProtocolError, match="conflicting"):
        store.add(binding(worker_ordinal=1), assignment_payload())


def test_duplicate_heartbeat_and_result_are_idempotent(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    store.add(binding(), assignment_payload())
    store.offer(0)
    heartbeat = protocol.sign_envelope(
        KEY, "heartbeat", binding(), {"note": "active"}, delivery_id="delivery-1"
    )
    assert store.heartbeat(heartbeat) == (True, False)
    assert store.heartbeat(heartbeat) == (True, True)
    result = protocol.sign_envelope(
        KEY, "result", binding(), {"structured": {"status": "completed"}}
    )
    _, duplicate = store.accept_result(result)
    assert duplicate is False
    _, duplicate = store.accept_result(result)
    assert duplicate is True


def test_reused_delivery_and_stale_attempt_cannot_cross_runs(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    store.add(binding(), assignment_payload())
    first = protocol.sign_envelope(
        KEY, "heartbeat", binding(), {"note": "one"}, delivery_id="delivery-1"
    )
    store.heartbeat(first)
    reused = protocol.sign_envelope(
        KEY, "heartbeat", binding(), {"note": "two"}, delivery_id="delivery-1"
    )
    with pytest.raises(protocol.ProtocolError, match="reused"):
        store.heartbeat(reused)
    store.finalize(binding(), "finalized")
    store.add(binding(run_id="other-run"), assignment_payload())
    cross_run = protocol.sign_envelope(
        KEY, "heartbeat", binding(run_id="other-run"), {"note": "one"},
        delivery_id="delivery-1",
    )
    with pytest.raises(protocol.ProtocolError, match="reused"):
        store.heartbeat(cross_run)
    stale = protocol.sign_envelope(
        KEY, "result", binding(attempt=2), {"structured": {"status": "completed"}}
    )
    with pytest.raises(protocol.ProtocolError, match="attempt is stale"):
        store.accept_result(stale)
    replacement = protocol.sign_envelope(
        KEY,
        "result",
        binding(run_id="replacement-run"),
        {"structured": {"status": "completed"}},
    )
    with pytest.raises(protocol.ProtocolError, match="unknown or stale"):
        store.accept_result(replacement)


def test_restart_heartbeat_loss_and_node_replacement_reuse_durable_assignment(tmp_path):
    database = tmp_path / "pool.db"
    first = protocol.PoolStore(database)
    first.add(binding(), assignment_payload())
    offered = first.offer(0)
    assert offered and offered["attempt"] == 1
    # Reopening the SQLite ledger models coordinator restart. An expired lease
    # is deliberately not reassigned across ordinals; the StatefulSet replaces
    # ordinal 0 with the same PVC and provider session state.
    restarted = protocol.PoolStore(database)
    with restarted._connect() as connection:
        connection.execute("UPDATE assignments SET lease_until=0")
    resumed = restarted.offer(0)
    assert resumed and resumed["run_id"] == "run-1234"
    assert resumed["attempt"] == 1
    assert restarted.offer(1) is None
    assert len(restarted.active_assignments()) == 1


def test_result_conflict_never_overwrites_first_result(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    store.add(binding(), assignment_payload())
    first = protocol.sign_envelope(KEY, "result", binding(), {"returncode": 0})
    store.accept_result(first)
    conflict = protocol.sign_envelope(KEY, "result", binding(), {"returncode": 1})
    with pytest.raises(protocol.ProtocolError, match="conflicting result"):
        store.accept_result(conflict)


def test_scm_workspace_is_ordinal_contained_and_rejects_symlink(tmp_path, monkeypatch):
    monkeypatch.setattr(scm, "WORKSPACE_ROOT", tmp_path)
    monkeypatch.setattr(scm, "ORDINAL", 0)
    signed = protocol.sign_envelope(KEY, "assignment", binding(), assignment_payload())
    path = scm.workspace_path(signed)
    assert path == tmp_path / "runs/atlas/t_deadbeef/run-1234"
    path.parent.mkdir(parents=True, exist_ok=True)
    path.symlink_to(tmp_path / "elsewhere", target_is_directory=True)
    with pytest.raises(protocol.ProtocolError, match="symlink"):
        scm.workspace_path(signed)
    root_case = tmp_path / "root-case"
    root_case.mkdir()
    root = tmp_path / "other-root"
    root.mkdir()
    (root_case / "runs").symlink_to(root, target_is_directory=True)
    monkeypatch.setattr(scm, "WORKSPACE_ROOT", root_case)
    with pytest.raises(protocol.ProtocolError, match="run root"):
        scm.workspace_path(signed)


def test_scm_boundary_rejects_other_repo_branch_and_ordinal(monkeypatch):
    monkeypatch.setattr(scm, "ORDINAL", 0)
    for changes, message in (
        ({"repo_url": "https://evil.example/atlas/titan-iac.git"}, "outside Atlas"),
        ({"branch": "main"}, "reviewed namespace"),
    ):
        signed = protocol.sign_envelope(
            KEY, "assignment", binding(), assignment_payload(**changes)
        )
        with pytest.raises(protocol.ProtocolError, match=message):
            scm._binding(signed)
    signed = protocol.sign_envelope(
        KEY, "assignment", binding(worker_ordinal=1), assignment_payload()
    )
    with pytest.raises(protocol.ProtocolError, match="ordinal"):
        scm._binding(signed)


def test_scm_git_process_has_no_credential_or_ambient_configuration():
    environment = scm._git_environment()
    assert environment == {
        "HOME": "/nonexistent",
        "PATH": "/usr/bin:/bin",
        "GIT_CONFIG_NOSYSTEM": "1",
        "GIT_TERMINAL_PROMPT": "0",
    }
    assert scm._broker_repo("titan-iac").startswith("http://hermes-scm-broker.")
    source = (SCRIPTS / "execution_pool_scm.py").read_text()
    assert "GITEA_TOKEN" not in source and "GIT_ASKPASS" not in source


def test_activity_is_bounded_sanitized_and_nofollow(tmp_path):
    text = coordinator.sanitize_activity(
        'route=codex token=super-secret\nAuthorization: Bearer abcdefghijklmnop\n'
        '"refresh_token":"sk-ant-oat01-abcdefghijklmnopqrstuvwxyz"'
    )
    assert "super-secret" not in text
    assert "abcdefghijklmnop" not in text
    assert "abcdefghijklmnopqrstuvwxyz" not in text
    assert len(text.encode()) <= protocol.MAX_ACTIVITY_BYTES

    class FakeKanban:
        @staticmethod
        def worker_log_path(_task, board):
            assert board == "atlas"
            return str(tmp_path / "worker.log")

    envelope = {**binding(), "payload": {"activity": "visible worker activity\n"}}
    coordinator._append_activity(FakeKanban, envelope)
    assert (tmp_path / "worker.log").read_text() == "visible worker activity\n"
    (tmp_path / "worker.log").unlink()
    (tmp_path / "worker.log").symlink_to(tmp_path / "target")
    with pytest.raises(protocol.ProtocolError, match="symlink"):
        coordinator._append_activity(FakeKanban, envelope)


def _documents(path):
    return [item for item in yaml.safe_load_all(path.read_text()) if item]


def test_three_node_statefulset_contract_and_cross_worker_isolation():
    stateful = _documents(HERMES / "execution-worker-statefulset.yaml")[0]
    pod = stateful["spec"]["template"]["spec"]
    assert stateful["spec"]["replicas"] == 3
    assert stateful["spec"]["podManagementPolicy"] == "Parallel"
    assert pod["automountServiceAccountToken"] is False
    claims = {
        item["metadata"]["name"]: item
        for item in stateful["spec"]["volumeClaimTemplates"]
    }
    assert set(claims) == {"workspace", "tools", "provider-access"}
    # The workspace is co-mounted by this ordinal's mediator, so it must be
    # ReadWriteMany: a ReadWriteOnce claim deadlocks on Multi-Attach when a drain
    # or preemption separates the two Pods. Everything only the worker mounts
    # stays single-mounter.
    assert claims["workspace"]["spec"]["accessModes"] == ["ReadWriteMany"]
    assert claims["tools"]["spec"]["accessModes"] == ["ReadWriteOnce"]
    assert claims["provider-access"]["spec"]["accessModes"] == ["ReadWriteOnce"]
    assert all(
        item["spec"]["storageClassName"] == "astreae" for item in claims.values()
    )
    anti = pod["affinity"]["podAntiAffinity"]["requiredDuringSchedulingIgnoredDuringExecution"]
    assert anti[0]["topologyKey"] == "kubernetes.io/hostname"
    spread = pod["topologySpreadConstraints"][0]
    assert spread["whenUnsatisfiable"] == "DoNotSchedule"
    assert spread["maxSkew"] == 1


def test_worker_placement_prefers_accelerators_and_preserves_exclusions():
    stateful = _documents(HERMES / "execution-worker-statefulset.yaml")[0]
    pod = stateful["spec"]["template"]["spec"]
    assert pod["priorityClassName"] == "scavenger"
    affinity = stateful["spec"]["template"]["spec"]["affinity"]["nodeAffinity"]
    terms = affinity["requiredDuringSchedulingIgnoredDuringExecution"]["nodeSelectorTerms"]
    required = {item["key"]: item for item in terms[0]["matchExpressions"]}
    assert required["node-role.kubernetes.io/worker"]["values"] == ["true"]
    excluded = set(required["kubernetes.io/hostname"]["values"])
    assert {"titan-04", "titan-13", "titan-14", "titan-17", "titan-18", "titan-19", "titan-22", "titan-24"} <= excluded
    assert affinity["preferredDuringSchedulingIgnoredDuringExecution"][0]["weight"] == 100
    assert affinity["preferredDuringSchedulingIgnoredDuringExecution"][1]["weight"] == 50
    worker_container = pod["containers"][0]
    assert worker_container["resources"] == {
        "requests": {"cpu": "5m", "memory": "128Mi", "ephemeral-storage": "1Gi"},
        "limits": {"cpu": "2", "memory": "4Gi", "ephemeral-storage": "8Gi"},
    }
    mediators = [
        item
        for item in _documents(HERMES / "execution-mediator.yaml")
        if item["kind"] == "Deployment"
    ]
    assert all(
        item["spec"]["template"]["spec"]["containers"][0]["resources"][
            "requests"
        ]
        == {"cpu": "2m", "memory": "64Mi"}
        for item in mediators
    )


def test_model_worker_has_no_scm_or_cluster_credential_mount():
    stateful = _documents(HERMES / "execution-worker-statefulset.yaml")[0]
    pod = stateful["spec"]["template"]
    containers = {item["name"]: item for item in pod["spec"]["containers"]}
    worker = containers["execution-worker"]
    environment = {item["name"] for item in worker["env"]}
    assert not {"GITEA_TOKEN", "OPENAI_API_KEY", "ANTHROPIC_API_KEY", "KUBECONFIG"} & environment
    mounts = {item["mountPath"] for item in worker["volumeMounts"]}
    assert "/vault/secrets" not in mounts
    assert "/pool-access" not in mounts
    assert "/provider-access" in mounts
    assert not any(item["name"] == "vault-auth-token" for item in worker["volumeMounts"])
    assert pod["metadata"]["annotations"]["vault.hashicorp.com/agent-inject-containers"] == "stage-worker-access"
    assert set(containers) == {"execution-worker"}
    assert "execution-pool-key" not in str(pod["metadata"]["annotations"])
    mediators = [
        item for item in _documents(HERMES / "execution-mediator.yaml")
        if item["kind"] == "Deployment"
    ]
    assert len(mediators) == 3
    for mediator in mediators:
        spec = mediator["spec"]["template"]["spec"]
        assert spec["automountServiceAccountToken"] is False
        privileged = spec["containers"][0]
        assert {mount["mountPath"] for mount in privileged["volumeMounts"]} >= {
            "/pool-access", "/scm-state", "/opt/scm", "/workspace"
        }
        assert "/provider-access" not in {
            mount["mountPath"] for mount in privileged["volumeMounts"]
        }


def test_worker_service_account_has_no_kubernetes_permissions():
    documents = _documents(HERMES / "execution-worker-rbac.yaml")
    assert len(documents) == 1
    account = documents[0]
    assert account["kind"] == "ServiceAccount"
    assert account["automountServiceAccountToken"] is False


def test_coordinator_remains_single_state_owner_and_workers_do_not_mount_home():
    agent = _documents(HERMES / "agent-deployment.yaml")[0]
    assert agent["kind"] == "Deployment"
    assert agent["spec"]["replicas"] == 1
    assert agent["spec"]["strategy"]["type"] == "Recreate"
    assert any(
        item["name"] == "home"
        and item["persistentVolumeClaim"]["claimName"] == "hermes-agent-home"
        for item in agent["spec"]["template"]["spec"]["volumes"]
    )
    worker_text = (HERMES / "execution-worker-statefulset.yaml").read_text()
    assert "hermes-agent-home" not in worker_text
    assert "kanban.db" not in worker_text


def test_worker_protocol_preserves_switchyard_fallback_and_visible_evidence():
    source = (SCRIPTS / "execution_pool_worker.py").read_text()
    coordinator_source = (SCRIPTS / "execution_pool_coordinator.py").read_text()
    server_source = (SCRIPTS / "execution_pool_server.py").read_text()
    assert "cli_lane_runner.select_route" in source
    assert 'alternate = "claude" if route.provider == "codex" else "codex"' in source
    assert "codex_thread_id" in source and "claude_session_id" in source
    assert "final_activity" in source
    assert "worker_ordinal" in coordinator_source and "provider_sessions" in coordinator_source
    assert "coordinator.reconcile" in server_source


def test_retention_gc_removes_only_clean_terminal_workspace(tmp_path, monkeypatch):
    monkeypatch.setattr(worker, "ROOT", tmp_path)
    monkeypatch.setattr(worker, "RETENTION_SECONDS", 3600)
    workspace = tmp_path / "runs/atlas/t_deadbeef/run-1234"
    workspace.mkdir(parents=True)
    subprocess.run(["git", "init", "-q", str(workspace)], check=True)
    subprocess.run(["git", "-C", str(workspace), "config", "user.email", "test@example.com"], check=True)
    subprocess.run(["git", "-C", str(workspace), "config", "user.name", "Test"], check=True)
    (workspace / "tracked").write_text("safe\n")
    subprocess.run(["git", "-C", str(workspace), "add", "tracked"], check=True)
    subprocess.run(["git", "-C", str(workspace), "commit", "-qm", "initial"], check=True)
    state_file = tmp_path / "session-state/atlas/t_deadbeef/run-1234.json"
    state_file.parent.mkdir(parents=True)
    state_file.write_text(
        json.dumps({"terminal_at": time.time() - 7200, "workspace": str(workspace)})
    )
    assert worker.garbage_collect() == 1
    assert not workspace.exists()
    assert not state_file.exists()