Agentswarm/scripts/test-contract-events.py

"""Manager event-contract test.

Validates that Swarm's emitted callback events conform to the Heicode Manager handler
(`heicode/controller/agent_callback.go`):
- each event_type carries HM's required payload fields,
- the HMAC signing canonical string / headers match HM's verification.

Hermetic: REDIS_FAKE in-memory store; callbacks are captured in-process (not sent).
Run from agent_swarm_v6:  python scripts/test-contract-events.py
"""
import asyncio
import hashlib
import hmac
import json
import os
import sys
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"
os.environ["AGENT_CALLBACK_SIGNING_SECRET"] = "test-signing-secret"
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator.redis_client import redis_client
from orchestrator import swarm_runtime as sr_mod
from orchestrator.swarm_runtime import swarm_runtime

# HM required payload fields per event_type (agent_callback.go registry).
REQUIRED = {
    "deployment.status_changed": ["status"],
    "task.created": ["task_id", "title"],
    "task.claimed": ["task_id", "agent_role"],
    "task.running": ["task_id", "agent_role"],
    "task.heartbeat": ["task_id", "agent_role"],
    "task.blocked": ["task_id", "reason"],
    "task.retried": ["task_id", "attempt"],
    "task.released": ["task_id", "agent_role"],
    "task.failed": ["task_id", "reason"],
    "task.completed": ["task_id"],
    "handoff.requested": ["task_id", "from_role", "to_role"],
    "handoff.completed": ["task_id", "from_role", "to_role"],
    "approval.requested": ["approval_id", "operation", "risk_level"],
    "artifact.created": ["artifact_id"],
    "timeline.updated": ["title"],
    "budget.alert": ["threshold_pct"],
}

SECRET = os.environ["AGENT_CALLBACK_SIGNING_SECRET"]
captured = []  # (raw_body, headers)
failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


def validate_envelope(raw_body: str, headers: dict) -> None:
    env = json.loads(raw_body)
    et = env["event_type"]
    # required fields
    req = REQUIRED.get(et, [])
    payload = env.get("payload") or {}
    missing = [f for f in req if f not in payload or payload.get(f) in (None, "")]
    check(f"{et}: required fields present {req}", not missing)
    # HMAC canonical string == ts + "." + event_id + "." + raw_body
    ts = headers.get("X-Agent-Timestamp")
    sig = headers.get("X-Agent-Signature", "")
    eid = headers.get("X-Agent-Event-Id")
    expect = "sha256=" + hmac.new(SECRET.encode(), f"{ts}.{eid}.{raw_body}".encode(), hashlib.sha256).hexdigest()
    check(f"{et}: HMAC signature matches HM canonical string", sig == expect)
    check(f"{et}: event_id header == body event_id", eid == env.get("event_id"))


async def _capture(self, swarm_id, url, raw_body, headers, event_type, event_id):
    captured.append((raw_body, headers))


async def emit(run, event_type, **payload):
    captured.clear()
    await swarm_runtime.emit_event(run, event_type, payload=payload)
    await asyncio.sleep(0)  # let the create_task'd callback run
    assert captured, f"no callback captured for {event_type}"
    return captured[-1]


async def main():
    await redis_client.connect()
    sr_mod.SwarmRuntime._post_callback = _capture  # capture instead of HTTP POST

    body = {
        "mode": "swarm",
        "requirement": {"objective": "contract test"},
        "callback": {"url": "http://127.0.0.1:9/cb", "subscribed_events": []},
        "metadata": {"manager_deployment_id": "m-contract"},
    }
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="corr-contract")

    # 1) centralized normalization (the 2 emit_event fixes)
    raw, hdr = await emit(run, "timeline.updated", summary="did the thing")
    check("timeline.updated normalized to include title", json.loads(raw)["payload"].get("title") == "did the thing")
    validate_envelope(raw, hdr)

    raw, hdr = await emit(run, "budget.alert", threshold=0.8)
    check("budget.alert normalized to threshold_pct=80", json.loads(raw)["payload"].get("threshold_pct") == 80.0)
    validate_envelope(raw, hdr)

    # 2) the handler-built payloads for the other two fixed types
    raw, hdr = await emit(run, "handoff.completed", task_id="t1", from_role="implementation", to_role="testing")
    validate_envelope(raw, hdr)
    raw, hdr = await emit(run, "task.released", task_id="t1", agent_role="testing", reason="at_capacity")
    validate_envelope(raw, hdr)

    # 3) representative coverage of the common event types
    samples = {
        "deployment.status_changed": {"status": "running"},
        "task.created": {"task_id": "t1", "title": "impl"},
        "task.claimed": {"task_id": "t1", "agent_role": "implementation"},
        "task.completed": {"task_id": "t1"},
        "task.failed": {"task_id": "t1", "reason": "boom"},
        "task.retried": {"task_id": "t1", "attempt": 1},
        "approval.requested": {"approval_id": "ap1", "operation": "git.write", "risk_level": "high"},
        "artifact.created": {"artifact_id": "art1"},
    }
    for et, pl in samples.items():
        raw, hdr = await emit(run, et, **pl)
        validate_envelope(raw, hdr)

    print()
    if failures:
        print(f"{len(failures)} contract check(s) FAILED: {failures}")
        sys.exit(1)
    print("all event-contract checks passed")


if __name__ == "__main__":
    asyncio.run(main())