Agentswarm/orchestrator/audit.py

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

Reconstructs a **replayable, ordered** audit trail for one swarm run from data the runtime
already persists — the event stream, task lineage, per-task model usage, and approvals — and
normalizes it into a single frozen record shape that answers, per step:
**who / when / which model / which tool / which approval / result / lineage**.

Design (mirrors convergence.py / cross_review.py):
  * **Pure** — no Redis / WebSocket / FastAPI / model calls. The caller fetches events/tasks/
    approvals and passes them in; this module only transforms. That makes it hermetically testable
    and keeps the audit logic free of I/O side effects.
  * **No content, no secrets** (rule: audit records carry only non-content metadata). Prompt text,
    code, and conversation are deliberately NOT recorded — only `model_id` + token counts + tool
    counts. Every record is marked `redacted: true`. `secret_ref` is never copied in.
  * **Deterministic** — `audit_id` is derived from the ordered position, not a random id, so the
    same run replays byte-identically.

Frozen record schema: see docs/integration/audit-trace-schema.md §5 (FROZEN v1).
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


# event_type -> coarse result classification (the "结果如何" column).
_RESULT_BY_EVENT = {
    "task.completed": "success",
    "swarm.completed": "success",
    "approval.approved": "success",
    "artifact.created": "success",
    "task.failed": "failed",
    "swarm.failed": "failed",
    "task.blocked": "blocked",
    "approval.rejected": "blocked",
    "swarm.stopped": "stopped",
}


def classify_result(event_type: str) -> str:
    """Map an event_type to a coarse result. Unknown -> 'info' (no fabrication)."""
    return _RESULT_BY_EVENT.get(event_type, "info")


@dataclass
class TaskAuditFacts:
    """Per-task facts the caller extracts from task state + parsed result.usage.

    model_id / tool_count come from the agent-reported `usage` block; left None/0 when the
    agent did not report them (never fabricated).
    """
    agent_role: Optional[str] = None
    assigned_agent_id: Optional[str] = None
    model_id: Optional[str] = None
    tool_count: int = 0
    source: Optional[str] = None
    parent_task_id: Optional[str] = None
    root_task_id: Optional[str] = None


@dataclass
class AuditRecord:
    audit_id: str
    occurred_at: Optional[str]
    action: str                                  # = event_type
    actor: Dict[str, Optional[str]]              # who: agent_role / agent_instance_id
    task_id: Optional[str]
    model: Optional[Dict[str, Any]]              # which model: {model_id} (+counts), or None
    tool: Optional[Dict[str, Any]]               # which tool: {count} (no SK/MCP layer yet), or None
    approval: Optional[Dict[str, Any]]           # which approval: {approval_id, decision}, or None
    result: str
    lineage: Dict[str, Optional[str]]
    redacted: bool = True

    def to_dict(self) -> Dict[str, Any]:
        return {
            "audit_id": self.audit_id,
            "occurred_at": self.occurred_at,
            "action": self.action,
            "actor": self.actor,
            "task_id": self.task_id,
            "model": self.model,
            "tool": self.tool,
            "approval": self.approval,
            "result": self.result,
            "lineage": self.lineage,
            "redacted": self.redacted,
        }


def build_audit_trace(
    events: List[Dict[str, Any]],
    task_facts: Dict[str, TaskAuditFacts],
    approvals: Dict[str, Dict[str, Any]],
    *,
    lineage: Dict[str, Optional[str]],
) -> List[Dict[str, Any]]:
    """Assemble an ordered, replayable audit trail.

    `events`     : the persisted event envelopes, in stored (chronological) order.
    `task_facts` : task_id -> TaskAuditFacts (who/model/tool/source), extracted by the caller.
    `approvals`  : run.approvals (approval_id -> {decision, ...}); only id+decision are surfaced.
    `lineage`    : the 4-id lineage (manager_deployment_id/deployment_id/swarm_id/correlation_id).

    Returns a list of audit-record dicts (frozen schema), one per event, in order.
    """
    records: List[Dict[str, Any]] = []
    for idx, ev in enumerate(events):
        event_type = ev.get("event_type") or "unknown"
        task_id = ev.get("task_id")
        facts = task_facts.get(task_id) if task_id else None
        payload = ev.get("payload") or {}

        # who: prefer the event's instance id, fall back to the task's assigned agent / role.
        actor = {
            "agent_role": (facts.agent_role if facts else None) or payload.get("agent_role"),
            "agent_instance_id": ev.get("agent_instance_id")
            or (facts.assigned_agent_id if facts else None),
        }

        # which model / tool — only when the task reported them (no fabrication).
        model = None
        tool = None
        if facts:
            if facts.model_id:
                model = {"model_id": facts.model_id}
            if facts.tool_count:
                tool = {"count": facts.tool_count}

        # which approval — for approval.* events, surface id + decision (never the secret_ref).
        approval = None
        approval_id = payload.get("approval_id")
        if approval_id:
            decision = (approvals.get(approval_id) or {}).get("decision") or payload.get("decision")
            approval = {"approval_id": approval_id, "decision": decision}

        records.append(
            AuditRecord(
                audit_id=f"aud_{idx + 1:06d}",
                occurred_at=ev.get("occurred_at"),
                action=event_type,
                actor=actor,
                task_id=task_id,
                model=model,
                tool=tool,
                approval=approval,
                result=classify_result(event_type),
                lineage=dict(lineage),
            ).to_dict()
        )
    return records


def replay(records: List[Dict[str, Any]]) -> List[str]:
    """Render the ordered trail as human-readable step lines (for replay / diagnostics).

    Pure formatting over the assembled records — no re-derivation, so the replay cannot diverge
    from the audited records.
    """
    lines: List[str] = []
    for r in records:
        who = r["actor"].get("agent_instance_id") or r["actor"].get("agent_role") or "system"
        model = (r.get("model") or {}).get("model_id")
        appr = r.get("approval") or {}
        bits = [r["audit_id"], r.get("occurred_at") or "-", who, r["action"], f"->{r['result']}"]
        if r.get("task_id"):
            bits.append(f"task={r['task_id']}")
        if model:
            bits.append(f"model={model}")
        if appr.get("approval_id"):
            bits.append(f"approval={appr['approval_id']}:{appr.get('decision')}")
        lines.append(" ".join(str(b) for b in bits))
    return lines