atlas-iac/testing/tests/test_hermes_execution_pool_protocol_v2.py

"""Version-2 protocol reliability and exact-attempt fencing contracts."""

from __future__ import annotations

import hashlib
import hmac
import json
import sys
import threading
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

import pytest


SCRIPTS = Path(__file__).parents[2] / "services/hermes/scripts"
sys.path.insert(0, str(SCRIPTS))

import execution_pool_protocol as protocol  # noqa: E402


MASTER = b"m" * 32


def binding(**changes):
    value = {
        "board": "metis",
        "task_id": "t_deadbeef",
        "run_id": "42",
        "worker_ordinal": 0,
        "attempt": 1,
    }
    value.update(changes)
    return value


def resign(envelope, key=MASTER):
    unsigned = dict(envelope)
    unsigned.pop("signature", None)
    envelope["signature"] = hmac.new(
        key, protocol.canonical_json(unsigned), hashlib.sha256
    ).hexdigest()
    return envelope


def test_atomic_json_is_private_durable_and_rejects_symlink_parent(tmp_path):
    target = tmp_path / "state/value.json"
    protocol.atomic_json(target, {"safe": True})
    assert json.loads(target.read_text()) == {"safe": True}
    assert target.stat().st_mode & 0o777 == 0o600

    outside = tmp_path / "outside"
    outside.mkdir()
    (tmp_path / "linked").symlink_to(outside, target_is_directory=True)
    with pytest.raises(protocol.ProtocolError, match="directory"):
        protocol.atomic_json(tmp_path / "linked/value.json", {"safe": False})


def test_ordinal_derivation_and_selector_fail_closed():
    keys = [protocol.derive_ordinal_key(MASTER, ordinal) for ordinal in range(3)]
    assert len(set(keys)) == 3
    for master, ordinal in ((b"short", 0), (MASTER, -1), (MASTER, 3)):
        with pytest.raises(protocol.ProtocolError, match="derivation"):
            protocol.derive_ordinal_key(master, ordinal)
    for value in (None, [], {}, {"worker_ordinal": True}, {"worker_ordinal": 3}):
        with pytest.raises(protocol.ProtocolError):
            protocol.envelope_ordinal(value)


def test_key_reader_bounds_content_and_type(tmp_path):
    for content in (b"x" * 31, b"x" * 4097):
        path = tmp_path / f"key-{len(content)}"
        path.write_bytes(content)
        path.chmod(0o600)
        with pytest.raises(protocol.ProtocolError, match="length"):
            protocol.read_key(path)
    directory = tmp_path / "directory"
    directory.mkdir(mode=0o700)
    with pytest.raises(protocol.ProtocolError, match="regular"):
        protocol.read_key(directory)


def test_envelope_rejects_kind_fields_numeric_lifetime_and_digest():
    with pytest.raises(protocol.ProtocolError, match="unsupported"):
        protocol.sign_envelope(MASTER, "admin", binding(), {})
    valid = protocol.sign_envelope(MASTER, "heartbeat", binding(), {})

    extra = {**valid, "extra": True}
    with pytest.raises(protocol.ProtocolError, match="fields"):
        protocol.verify_envelope(MASTER, extra)
    wrong_version = resign({**valid, "version": 1})
    with pytest.raises(protocol.ProtocolError, match="version"):
        protocol.verify_envelope(MASTER, wrong_version)
    with pytest.raises(protocol.ProtocolError, match="unexpected"):
        protocol.verify_envelope(MASTER, valid, expected_kind="result")

    numeric = resign({**valid, "attempt": "not-a-number"})
    with pytest.raises(protocol.ProtocolError, match="numeric"):
        protocol.verify_envelope(MASTER, numeric)
    numeric_text = resign({**valid, "attempt": "1"})
    with pytest.raises(protocol.ProtocolError, match="numeric"):
        protocol.verify_envelope(MASTER, numeric_text)
    typed_identifier = resign({**valid, "run_id": 42})
    with pytest.raises(protocol.ProtocolError, match="run_id"):
        protocol.verify_envelope(MASTER, typed_identifier)
    lifetime = resign({**valid, "expires_at": valid["issued_at"]})
    with pytest.raises(protocol.ProtocolError, match="lifetime"):
        protocol.verify_envelope(MASTER, lifetime)
    digest = resign({**valid, "payload_digest": "0" * 64})
    with pytest.raises(protocol.ProtocolError, match="digest"):
        protocol.verify_envelope(MASTER, digest)
    with pytest.raises(protocol.ProtocolError, match="oversized"):
        protocol.verify_envelope(MASTER, {"value": "x" * protocol.MAX_WIRE_BYTES})


def test_wire_parser_rejects_non_object_and_empty():
    for body, message in ((b"", "empty"), (b"[]", "object")):
        with pytest.raises(protocol.ProtocolError, match=message):
            protocol.parse_wire(body)


def test_bounded_http_server_sets_timeout_and_releases_slot():
    handled = threading.Event()

    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):  # noqa: N802
            handled.set()
            self.send_response(204)
            self.end_headers()

        def log_message(self, *_args):
            return

    server = protocol.BoundedHTTPServer(("127.0.0.1", 0), Handler, max_workers=1)
    thread = threading.Thread(target=server.handle_request)
    thread.start()
    with urllib.request.urlopen(
        f"http://127.0.0.1:{server.server_port}/", timeout=3
    ) as response:
        assert response.status == 204
    thread.join(timeout=3)
    assert handled.is_set()
    acquired = False
    for _ in range(100):
        acquired = server._slots.acquire(blocking=False)
        if acquired:
            break
        threading.Event().wait(0.01)
    assert acquired
    server._slots.release()
    server.server_close()


def test_bounded_server_releases_slot_when_thread_dispatch_raises(monkeypatch):
    server = protocol.BoundedHTTPServer(
        ("127.0.0.1", 0), BaseHTTPRequestHandler, max_workers=1
    )

    def explode(*_args):
        raise RuntimeError("dispatch failed")

    monkeypatch.setattr(ThreadingHTTPServer, "process_request", explode)
    with pytest.raises(RuntimeError, match="dispatch"):
        server.process_request(object(), ("127.0.0.1", 1))
    assert server._slots.acquire(blocking=False)
    server._slots.release()
    server.server_close()


def test_store_fences_ordinal_attempt_state_and_expired_lease(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db", lease_seconds=60)
    assert protocol.PoolStore._record(None) is None
    store.add(binding(), {"context": "safe"})
    store.offer(0)

    foreign = protocol.sign_envelope(
        MASTER, "heartbeat", binding(worker_ordinal=1), {"note": "foreign"}
    )
    with pytest.raises(protocol.ProtocolError, match="ordinal"):
        store.heartbeat(foreign)

    with store._connect() as connection:
        connection.execute("UPDATE assignments SET lease_until=1")
    expired = protocol.sign_envelope(MASTER, "heartbeat", binding(), {"note": "late"})
    with pytest.raises(protocol.ProtocolError, match="lease expired"):
        store.heartbeat(expired)
    terminal = protocol.sign_envelope(MASTER, "result", binding(), {"structured": {}})
    with pytest.raises(protocol.ProtocolError, match="lease expired"):
        store.accept_result(terminal)


def test_store_pending_result_invalid_state_and_exact_finalize(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    store.add(binding(), {"context": "safe"})
    result = protocol.sign_envelope(MASTER, "result", binding(), {"structured": {}})
    record, duplicate = store.accept_result(result)
    assert duplicate is False and record["state"] == "result"
    assert store.pending_results()[0]["result"] == {"structured": {}}
    with pytest.raises(protocol.ProtocolError, match="terminal"):
        store.finalize(binding(), "invalid")
    store.finalize(binding(attempt=2), "stale")
    assert store.pending_results()[0]["state"] == "result"
    store.finalize(binding(), "finalized")
    heartbeat = protocol.sign_envelope(MASTER, "heartbeat", binding(), {})
    with pytest.raises(protocol.ProtocolError, match="running"):
        store.heartbeat(heartbeat)


def test_expired_attempt_is_reoffered_then_terminally_released(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    store.add(binding(), {"context": "safe"})
    store.offer(0)
    changed = store.expire_leases(now=10_000_000_000, max_attempts=2)
    assert changed[0]["state"] == "assigned" and changed[0]["attempt"] == 2
    stale = protocol.sign_envelope(MASTER, "heartbeat", binding(), {})
    with pytest.raises(protocol.ProtocolError, match="attempt is stale"):
        store.heartbeat(stale)

    offered = store.offer(0)
    assert offered and offered["attempt"] == 2
    changed = store.expire_leases(now=10_000_000_001, max_attempts=2)
    assert changed[0]["state"] == "lease_failed"
    assert store.available_ordinals() == [0, 1, 2]


def test_store_garbage_collection_removes_old_terminal_and_deliveries(tmp_path):
    store = protocol.PoolStore(tmp_path / "pool.db")
    store.add(binding(), {"context": "safe"})
    heartbeat = protocol.sign_envelope(
        MASTER, "heartbeat", binding(), {}, delivery_id="old-delivery"
    )
    store.heartbeat(heartbeat)
    store.finalize(binding(), "stale")
    with store._connect() as connection:
        connection.execute("UPDATE assignments SET updated_at=1")
        connection.execute("UPDATE deliveries SET received_at=1")
    assert store.garbage_collect(3600) == 1
    assert store.active_assignments() == []
    with store._connect() as connection:
        assert connection.execute("SELECT count(*) FROM deliveries").fetchone()[0] == 0