Agentswarm/scripts/test-audit-trace.py

"""Audit / lineage trace tests (issue #17).

Covers the pure assembler (orchestrator/audit.py) and the orchestrator wiring
(build_audit_trace_for_run) that reconstructs a replayable per-step trail —
who / when / which model / which tool / which approval / result — from the
event stream + task lineage + approvals, carrying NO prompt content or secrets.

Run from agent_swarm_v6 (install deps first — needs fakeredis):
    pip install -r orchestrator/requirements.txt
    REDIS_FAKE=1 python scripts/test-audit-trace.py
"""
import asyncio
import json
import os
import sys
from pathlib import Path

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

from orchestrator import audit as audit_mod
from orchestrator.audit import TaskAuditFacts, build_audit_trace, replay, classify_result

failures = []


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


def test_pure():
    events = [
        {"event_type": "task.created", "task_id": "t1", "occurred_at": "T0", "payload": {"title": "x"}},
        {"event_type": "task.claimed", "task_id": "t1", "agent_instance_id": "agent-7", "occurred_at": "T1", "payload": {}},
        {"event_type": "task.completed", "task_id": "t1", "agent_instance_id": "agent-7", "occurred_at": "T2", "payload": {}},
        {"event_type": "approval.requested", "occurred_at": "T3", "payload": {"approval_id": "ap1", "operation": "git.push", "risk_level": "high"}},
        {"event_type": "approval.rejected", "occurred_at": "T4", "payload": {"approval_id": "ap1"}},
        {"event_type": "swarm.failed", "occurred_at": "T5", "payload": {"status": "failed"}},
    ]
    task_facts = {
        "t1": TaskAuditFacts(agent_role="impl", assigned_agent_id="agent-7",
                             model_id="gpt-x", tool_count=3, source="agent_proposed"),
    }
    approvals = {"ap1": {"decision": "rejected", "reason": "blocked by policy"}}
    lineage = {"manager_deployment_id": "m1", "deployment_id": "d1", "swarm_id": "s1", "correlation_id": "c1"}

    recs = build_audit_trace(events, task_facts, approvals, lineage=lineage)
    check("one record per event, in order", len(recs) == len(events))
    check("audit_id is deterministic/sequential", [r["audit_id"] for r in recs][:2] == ["aud_000001", "aud_000002"])
    check("every record redacted + carries full lineage",
          all(r["redacted"] is True and r["lineage"]["swarm_id"] == "s1" for r in recs))

    # who
    claimed = recs[1]
    check("who: actor resolves agent_instance_id", claimed["actor"]["agent_instance_id"] == "agent-7")
    check("who: actor resolves agent_role from task facts", claimed["actor"]["agent_role"] == "impl")
    # which model / tool
    check("which model: surfaced from task usage", claimed["model"] == {"model_id": "gpt-x"})
    check("which tool: count surfaced from task usage", claimed["tool"] == {"count": 3})
    # which approval
    appr_evt = recs[4]
    check("which approval: id + decision surfaced", appr_evt["approval"] == {"approval_id": "ap1", "decision": "rejected"})
    # result classification
    check("result: completed -> success", recs[2]["result"] == "success")
    check("result: approval.rejected -> blocked", appr_evt["result"] == "blocked")
    check("result: swarm.failed -> failed", recs[5]["result"] == "failed")
    check("result: unknown -> info (no fabrication)", classify_result("some.weird.event") == "info")

    # no content / no secrets anywhere in the serialized trail
    blob = json.dumps(recs)
    check("no secret_ref / credential leaked into audit records",
          "secret_ref" not in blob and "azkv://" not in blob and "credential" not in blob)

    # replay is ordered and human-readable
    lines = replay(recs)
    check("replay yields one line per record, ordered", len(lines) == len(recs) and lines[0].startswith("aud_000001"))
    check("replay surfaces who/model/approval", "agent-7" in lines[1] and "model=gpt-x" in lines[1]
          and "approval=ap1:rejected" in lines[4])


async def test_integration():
    from orchestrator.redis_client import redis_client
    from orchestrator.swarm_runtime import swarm_runtime
    from orchestrator.task_queue import task_queue
    from orchestrator import main as orch

    await redis_client.connect()
    body = {"mode": "swarm", "orchestration_plan": {"objective": "audit it"},
            "callback": {"url": "", "subscribed_events": []},
            "metadata": {"manager_deployment_id": "m-audit"}}
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="c-audit")

    task = await task_queue.create_task(
        task_id=f"{run.swarm_id}-t", description="impl", agent_role="impl",
        required_capabilities=["python"], enqueue=True,
    )
    # Agent-reported result carrying model usage + tool calls + a would-be secret (must NOT leak).
    await task_queue.complete_task(task.task_id, result=json.dumps({
        "summary": "done",
        "usage": {"model_id": "gpt-x", "model_tokens": 10},
        "tool_calls": [{"name": "write_file"}, {"name": "run_git"}],
        "secret_ref": "azkv://vault/secrets/x",
    }))
    await swarm_runtime.attach_task(run, task.task_id)
    run.approvals["ap1"] = {"approval_id": "ap1", "decision": "approved"}
    await swarm_runtime.save_run(run)

    await swarm_runtime.emit_event(run, "task.claimed", task_id=task.task_id,
                                   agent_instance_id="agent-1", payload={"agent_role": "impl"})
    await swarm_runtime.emit_event(run, "task.completed", task_id=task.task_id, agent_instance_id="agent-1")

    data = await orch.build_audit_trace_for_run(run)
    recs = data["records"]
    check("integration: trace assembled from real run", data["count"] >= 1 and bool(recs))
    by_type = {r["action"]: r for r in recs}
    check("integration: model_id reconstructed from task usage",
          by_type.get("task.completed", {}).get("model") == {"model_id": "gpt-x"})
    check("integration: tool count reconstructed (2 tool_calls)",
          by_type.get("task.completed", {}).get("tool") == {"count": 2})
    check("integration: no secret leaked into audit trace",
          "secret_ref" not in json.dumps(data) and "azkv://" not in json.dumps(data))


async def test_pagination_no_truncation():
    """A run with >500 events must produce a COMPLETE audit trail (no silent truncation).

    list_events clamps each page to 500; build_audit_trace_for_run must follow next_cursor to
    exhaustion. Regression guard for the #17 'every step replayable' DoD.
    """
    from orchestrator.redis_client import redis_client
    from orchestrator.swarm_runtime import swarm_runtime
    from orchestrator import main as orch

    await redis_client.connect()
    body = {"mode": "swarm", "orchestration_plan": {"objective": "long run"},
            "callback": {"url": "", "subscribed_events": []},
            "metadata": {"manager_deployment_id": "m-long"}}
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="c-long")

    # Count whatever creation already emitted, then add enough to cross two page boundaries.
    baseline = len((await swarm_runtime.list_events(run.swarm_id, limit=500)).get("events", []))
    extra = 612  # > 500 and not a multiple of it, so truncation/off-by-page would show up
    for i in range(extra):
        await swarm_runtime._store_event(run.swarm_id, json.dumps({
            "event_type": "task.heartbeat", "task_id": None, "occurred_at": f"T{i}",
            "swarm_id": run.swarm_id, "payload": {},
        }))
    total = baseline + extra

    data = await orch.build_audit_trace_for_run(run)
    check(f"pagination: audit count == full event count ({total}), not capped at 500",
          data["count"] == total and total > 500)
    check("pagination: replay is not truncated (one line per event)", len(data["replay"]) == total)
    check("pagination: audit_id runs unbroken to the last event",
          data["records"][-1]["audit_id"] == f"aud_{total:06d}")


async def main():
    test_pure()
    await test_integration()
    await test_pagination_no_truncation()
    print()
    if failures:
        print(f"{len(failures)} audit-trace check(s) FAILED: {failures}")
        sys.exit(1)
    print("all audit-trace checks passed")


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