diff --git a/agent_assembly/__init__.py b/agent_assembly/__init__.py index 3a8491e4..7d901a9f 100644 --- a/agent_assembly/__init__.py +++ b/agent_assembly/__init__.py @@ -12,9 +12,10 @@ PolicyError, ToolExecutionBlockedError, ) +from agent_assembly.types import AuditEvent, CallStackNode, CallStackNodeKind try: - from agent_assembly._core import ( # type: ignore[attr-defined] + from agent_assembly._core import ( GovernanceEvent, PolicyResult, PolicyTimeoutError, @@ -39,6 +40,9 @@ "AdapterValidationError", "ToolExecutionBlockedError", "MCPToolBlockedError", + "AuditEvent", + "CallStackNode", + "CallStackNodeKind", ] if "RuntimeClient" in globals(): diff --git a/agent_assembly/types.py b/agent_assembly/types.py index e0a41989..01d60b74 100644 --- a/agent_assembly/types.py +++ b/agent_assembly/types.py @@ -16,9 +16,95 @@ from __future__ import annotations +from dataclasses import dataclass, field +from typing import Literal + +# ── Event types (AAASM-1435) ────────────────────────────────────────────────── + +CallStackNodeKind = Literal["llm", "tool", "result"] +"""Discriminator for a [`CallStackNode`] — open-ended on the wire (proto uses +`string kind`) but typed here as a `Literal` for the three values the +dashboard currently renders.""" + + +@dataclass(frozen=True, slots=True) +class CallStackNode: + """One node in the hierarchical call stack attached to an [`AuditEvent`]. + + Mirrors the proto `assembly.audit.v1.CallStackNode` message added in + AAASM-1419. Field names match the wire-protocol's snake_case so the + dataclass can be constructed directly from a decoded payload without an + intermediate naming transform. + + Attributes: + id: Stable identifier for this node within the call stack. + kind: Node category — one of ``"llm"``, ``"tool"``, or ``"result"``. + label: Human-readable label rendered by downstream UI. + latency_ms: Step-local latency in milliseconds, or `None` when the + producer did not record a duration. + children: Recursive descent — nested calls produced by this step. + Empty list (default) when the node has no children. + """ + + id: str + kind: CallStackNodeKind + label: str + latency_ms: int | None = None + children: list[CallStackNode] = field(default_factory=list) + + +@dataclass(frozen=True, slots=True) +class AuditEvent: + """A single governance-relevant occurrence in the gateway audit trail. + + Python-friendly subset of the proto `assembly.audit.v1.AuditEvent` + message. Field names match the wire-protocol's snake_case so the + dataclass can be constructed directly from a decoded payload. + + Out of scope today (tracked under separate follow-up Tasks): + * The ``detail`` oneof (`LLMCallDetail` / `ToolCallDetail` / + `FileOpDetail` / `NetworkCallDetail` / `ProcessExecDetail` / + `PolicyViolation` / `ApprovalEvent`). + * The lineage block (`root_agent_id`, `parent_agent_id`, `team_id`, + `session_id`, `delegation_reason`, `spawned_by_tool`, `depth`). + + This class exists today to surface the new ``call_stack`` field added + by AAASM-1419 to SDK consumers. Round-tripping into the gateway via the + wire encoder remains the Rust FFI layer's responsibility. + + Attributes: + event_id: Unique identifier for this audit record (UUID v7). + agent_id: Identity string of the agent that produced this event. + action_type: High-level action category (e.g. ``"llm_call"``, + ``"tool_call"``, ``"file_op"``). Open-ended on the wire. + decision: Policy engine verdict (e.g. ``"allow"``, ``"deny"``, + ``"redact"``). + trace_id: Distributed tracing run-level identifier. Empty when + unset. + span_id: Distributed tracing action-level identifier. Empty when + unset. + parent_span_id: Distributed tracing parent span identifier. Empty + when this is a root span. + labels: Arbitrary key/value labels attached at event creation. + call_stack: Hierarchical record of LLM / tool / result steps that + led to this event. Empty list (default) when the producer did + not record a stack. Renders inline beneath an expanded Live + Ops row in the dashboard. + """ + + event_id: str + agent_id: str + action_type: str + decision: str + trace_id: str = "" + span_id: str = "" + parent_span_id: str = "" + labels: dict[str, str] = field(default_factory=dict) + call_stack: list[CallStackNode] = field(default_factory=list) + + __all__ = [ - "AgentId", - "PolicyId", - "AssemblyId", - "ConfigDict", + "AuditEvent", + "CallStackNode", + "CallStackNodeKind", ] diff --git a/rust/aa-ffi-python/Cargo.toml b/rust/aa-ffi-python/Cargo.toml index 1cff3d24..b716d9fc 100644 --- a/rust/aa-ffi-python/Cargo.toml +++ b/rust/aa-ffi-python/Cargo.toml @@ -10,8 +10,8 @@ name = "aa_ffi_python" crate-type = ["cdylib"] [dependencies] -aa-core = { git = "https://github.com/AI-agent-assembly/agent-assembly.git", package = "aa-core", features = ["serde"] } -aa-proto = { git = "https://github.com/AI-agent-assembly/agent-assembly.git", package = "aa-proto" } +aa-core = { git = "https://github.com/AI-agent-assembly/agent-assembly.git", rev = "ed4aa11a8c1d1ce1e6f96b08cf2179fd772099b2", package = "aa-core", features = ["serde"] } +aa-proto = { git = "https://github.com/AI-agent-assembly/agent-assembly.git", rev = "ed4aa11a8c1d1ce1e6f96b08cf2179fd772099b2", package = "aa-proto" } once_cell = "1.20" prost = "0.14" pyo3 = { version = "0.20", features = ["extension-module"] } diff --git a/test/unit/test_types_audit_event.py b/test/unit/test_types_audit_event.py new file mode 100644 index 00000000..23fd7eb9 --- /dev/null +++ b/test/unit/test_types_audit_event.py @@ -0,0 +1,98 @@ +"""Unit tests for the `agent_assembly.types` dataclasses (AAASM-1435). + +Covers `AuditEvent` + `CallStackNode` — the Python-friendly mirrors of +the proto messages added in AAASM-1419. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from agent_assembly import AuditEvent, CallStackNode + + +def test_call_stack_node_required_fields() -> None: + node = CallStackNode(id="n0", kind="llm", label="gpt-4o") + assert node.id == "n0" + assert node.kind == "llm" + assert node.label == "gpt-4o" + assert node.latency_ms is None + assert node.children == [] + + +def test_call_stack_node_with_latency_and_children() -> None: + child = CallStackNode(id="n1", kind="tool", label="gmail.send", latency_ms=120) + parent = CallStackNode(id="n0", kind="llm", label="gpt-4o", latency_ms=300, children=[child]) + assert parent.children[0] is child + assert parent.children[0].latency_ms == 120 + + +def test_call_stack_node_is_frozen() -> None: + node = CallStackNode(id="n0", kind="llm", label="gpt-4o") + with pytest.raises(dataclasses.FrozenInstanceError): + node.id = "n2" # type: ignore[misc] + + +def test_audit_event_minimal_construction_defaults_call_stack_to_empty() -> None: + event = AuditEvent( + event_id="evt-1", + agent_id="support-agent", + action_type="llm_call", + decision="allow", + ) + assert event.call_stack == [] + assert event.labels == {} + assert event.trace_id == "" + assert event.span_id == "" + assert event.parent_span_id == "" + + +def test_audit_event_with_populated_call_stack_tree() -> None: + stack = [ + CallStackNode( + id="n0", + kind="llm", + label="gpt-4o", + latency_ms=300, + children=[ + CallStackNode( + id="n1", + kind="tool", + label="gmail.send", + latency_ms=120, + children=[ + CallStackNode(id="n2", kind="result", label="200 OK"), + ], + ), + ], + ), + ] + event = AuditEvent( + event_id="evt-1", + agent_id="support-agent", + action_type="llm_call", + decision="allow", + trace_id="trace-1", + span_id="span-1", + call_stack=stack, + ) + assert event.call_stack[0].kind == "llm" + assert event.call_stack[0].children[0].label == "gmail.send" + assert event.call_stack[0].children[0].children[0].kind == "result" + assert event.call_stack[0].children[0].children[0].latency_ms is None + + +def test_audit_event_is_frozen() -> None: + event = AuditEvent(event_id="e", agent_id="a", action_type="llm_call", decision="allow") + with pytest.raises(dataclasses.FrozenInstanceError): + event.event_id = "e2" # type: ignore[misc] + + +def test_top_level_imports_resolve_to_types_module() -> None: + from agent_assembly.types import AuditEvent as TypesAuditEvent + from agent_assembly.types import CallStackNode as TypesCallStackNode + + assert AuditEvent is TypesAuditEvent + assert CallStackNode is TypesCallStackNode