test/network/test_audit_and_lifecycle.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
"""Audit log + lifecycle invariants.
Covers:
* ``delegation_depth`` enforcement at hub.
* Hydrate when an adapter is unregistered no longer raises ``KeyError``
on the next ``post_envelope`` (clear ``ProtocolError`` instead).
* Sweeper ``auto_close`` does not bleed across channels in one tick.
* Channel lifecycle audit records (``channel_created`` / ``_closed``).
* Task lifecycle audit record (``task_terminated``).
"""
import pytest
from ag2 import Agent
from ag2.knowledge import MemoryKnowledgeStore
from ag2.network import (
EV_TEXT,
AccessDeniedError,
Envelope,
Hub,
ProtocolError,
Rule,
)
from ag2.network.adapters.conversation import ConversationAdapter
from ag2.network.channel import (
ChannelManifest,
Expectation,
ParticipantSchema,
)
from ag2.network.hub.audit import (
AUDIT_KIND_CHANNEL_CLOSED,
AUDIT_KIND_CHANNEL_CREATED,
AUDIT_KIND_CHANNEL_EXPIRED,
AUDIT_KIND_EXPECTATION_VIOLATED,
AUDIT_KIND_TASK_TERMINATED,
)
from ag2.network.rule import LimitsBlock
from ag2.task import TaskMetadata, TaskSpec, TaskState
from ag2.testing import TestConfig
from ._helpers import _MockClock
def _agent(name: str) -> Agent:
return Agent(name=name, config=TestConfig())
# ── delegation_depth ────────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_post_envelope_rejects_envelope_above_delegation_depth() -> None:
"""Hub rejects an envelope whose ``depth`` exceeds the sender's
``Rule.limits.delegation_depth`` cap. ``0`` disables the check."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
capped_rule = Rule(limits=LimitsBlock(delegation_depth=2))
alice = await hub.register(_agent("alice"), rule=capped_rule)
bob = await hub.register(_agent("bob"))
channel = await alice.open(type="conversation", target=bob.agent_id)
# depth=3 exceeds cap of 2.
too_deep = Envelope(
channel_id=channel.channel_id,
sender_id=alice.agent_id,
audience=[bob.agent_id],
event_type=EV_TEXT,
event_data={"text": "hi"},
depth=3,
)
with pytest.raises(AccessDeniedError):
await hub.post_envelope(too_deep)
# depth=2 is at the cap → accepted.
at_cap = Envelope(
channel_id=channel.channel_id,
sender_id=alice.agent_id,
audience=[bob.agent_id],
event_type=EV_TEXT,
event_data={"text": "hi"},
depth=2,
)
await hub.post_envelope(at_cap)
await hub.close()
@pytest.mark.asyncio
async def test_delegation_depth_zero_disables_cap() -> None:
"""``delegation_depth=0`` means no cap; arbitrary depth is allowed."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
no_cap = Rule(limits=LimitsBlock(delegation_depth=0))
alice = await hub.register(_agent("alice"), rule=no_cap)
bob = await hub.register(_agent("bob"))
channel = await alice.open(type="conversation", target=bob.agent_id)
deep = Envelope(
channel_id=channel.channel_id,
sender_id=alice.agent_id,
audience=[bob.agent_id],
event_type=EV_TEXT,
event_data={"text": "hi"},
depth=999,
)
await hub.post_envelope(deep) # accepted
await hub.close()
# ── Hydrate with missing adapter ────────────────────────────────────────────
@pytest.mark.asyncio
async def test_post_envelope_after_hydrate_without_adapter_state_raises_protocol_error() -> None:
"""If a channel is loaded by ``hydrate()`` while its adapter is
missing, the channel is kept dormant — its metadata is read-only,
not added to ``_active_channels``, and ``_adapter_states`` has no
entry. If an adapter is later registered (so ``_adapter_for``
succeeds) and a stale envelope arrives, the hub raises a clear
``ProtocolError`` instead of bare ``KeyError``."""
store = MemoryKnowledgeStore()
custom_manifest = ChannelManifest(
type="custom_recovery",
version=1,
participants=ParticipantSchema(min=2),
expectations=[],
)
class _CustomAdapter(ConversationAdapter):
def __init__(self) -> None:
super().__init__()
self.manifest = custom_manifest
# First boot: register custom adapter, open channel, persist.
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
hub1.register_adapter(_CustomAdapter())
alice = await hub1.register(_agent("alice"))
bob = await hub1.register(_agent("bob"))
channel = await alice.open(type="custom_recovery", target=bob.agent_id)
channel_id = channel.channel_id
alice_id = alice.agent_id
await hub1.close()
# Second boot: hydrate WITHOUT the custom adapter. Channel is
# dormant: metadata loaded, but no adapter state, not active.
hub2 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
assert channel_id in hub2._channels
assert channel_id not in hub2._active_channels
assert channel_id not in hub2._adapter_states
# Now register the adapter post-hydrate (e.g. a recovery script).
# ``_adapter_for`` will find it; ``_adapter_states`` is still empty.
hub2.register_adapter(_CustomAdapter())
# The channel metadata says it's ACTIVE (carried over from boot 1)
# but no fold ran. ``post_envelope`` must surface a clear error.
envelope = Envelope(
channel_id=channel_id,
sender_id=alice_id,
audience=None,
event_type=EV_TEXT,
event_data={"text": "hi"},
)
with pytest.raises(ProtocolError, match="no adapter state"):
await hub2.post_envelope(envelope)
await hub2.close()
# ── Sweeper auto_close cross-channel bleed ──────────────────────────────────
@pytest.mark.asyncio
async def test_expectation_tick_processes_all_channels_when_one_auto_closes() -> None:
"""When channel A's expectation fires ``auto_close``, channel B's
expectations on the same tick are still evaluated.
Pre-fix: the sweeper ``return``ed out of the whole tick after the
first ``auto_close``, so channels later in the iteration order
waited for the next 10s tick.
"""
clock = _MockClock("2026-01-01T00:00:00+00:00")
store = MemoryKnowledgeStore()
hub = await Hub.open(store, clock=clock, ttl_sweep_interval=0, expectation_sweep_interval=0)
# Custom conversation manifest with an aggressive max_silence that
# auto_closes — so two channels both violate at the same tick.
aggressive_manifest = ChannelManifest(
type="conversation_aggressive",
version=1,
participants=ParticipantSchema(min=2),
expectations=[
Expectation(
name="max_silence",
on_violation="auto_close",
params={"seconds": 60},
),
],
)
class _AggressiveAdapter(ConversationAdapter):
def __init__(self) -> None:
super().__init__()
self.manifest = aggressive_manifest
hub.register_adapter(_AggressiveAdapter())
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
carol = await hub.register(_agent("carol"))
dave = await hub.register(_agent("dave"))
sess_ab = await alice.open(type="conversation_aggressive", target=bob.agent_id)
sess_cd = await carol.open(type="conversation_aggressive", target=dave.agent_id)
# Advance past max_silence — both should violate.
clock.advance(120)
await hub._expectation_tick()
audit = await hub._audit_log.read_all()
closed_channel_ids = {r["channel_id"] for r in audit if r["kind"] == AUDIT_KIND_CHANNEL_CLOSED}
# Both channels auto-closed in a single tick.
assert sess_ab.channel_id in closed_channel_ids
assert sess_cd.channel_id in closed_channel_ids
# Both expectation violations were logged.
violations = [r for r in audit if r["kind"] == AUDIT_KIND_EXPECTATION_VIOLATED]
violation_channel_ids = {r["channel_id"] for r in violations}
assert sess_ab.channel_id in violation_channel_ids
assert sess_cd.channel_id in violation_channel_ids
await hub.close()
# ── Channel + task lifecycle audit ──────────────────────────────────────────
@pytest.mark.asyncio
async def test_audit_log_records_channel_created_and_closed() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
pre_audit = len(await hub._audit_log.read_all())
channel = await alice.open(type="conversation", target=bob.agent_id)
await hub.close_channel(channel.channel_id, reason="explicit")
audit = await hub._audit_log.read_all()
new_records = audit[pre_audit:]
kinds = [r["kind"] for r in new_records]
assert AUDIT_KIND_CHANNEL_CREATED in kinds
assert AUDIT_KIND_CHANNEL_CLOSED in kinds
closed_record = next(r for r in new_records if r["kind"] == AUDIT_KIND_CHANNEL_CLOSED)
assert closed_record["channel_id"] == channel.channel_id
assert closed_record["reason"] == "explicit"
await hub.close()
@pytest.mark.asyncio
async def test_audit_log_records_channel_expired_on_ttl_sweep() -> None:
"""``EXPIRED`` transitions emit ``channel_expired`` (not ``channel_closed``)."""
clock = _MockClock("2026-01-01T00:00:00+00:00")
store = MemoryKnowledgeStore()
hub = await Hub.open(store, clock=clock, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
channel = await alice.open(type="conversation", target=bob.agent_id, ttl="60s")
pre_audit = len(await hub._audit_log.read_all())
clock.advance(120)
await hub.expire_due()
audit = await hub._audit_log.read_all()
new_records = audit[pre_audit:]
expired = [
r for r in new_records if r["kind"] == AUDIT_KIND_CHANNEL_EXPIRED and r["channel_id"] == channel.channel_id
]
assert len(expired) == 1
assert expired[0]["reason"] == "ttl_expired"
await hub.close()
@pytest.mark.asyncio
async def test_audit_log_records_task_terminated_on_channel_cascade() -> None:
"""Tasks under a closing channel cascade to ``EXPIRED`` and emit
``task_terminated`` audit records carrying ``capability``."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
channel = await alice.open(type="conversation", target=bob.agent_id)
# Plant a non-terminal capability-tagged task under the channel.
task_meta = TaskMetadata(
task_id="task-cap-1",
owner_id=bob.agent_id,
spec=TaskSpec(title="analysing", capability="analysis"),
state=TaskState.RUNNING,
created_at="2026-01-01T00:00:00+00:00",
channel_id=channel.channel_id,
)
hub._tasks[task_meta.task_id] = task_meta
hub._channel_tasks.setdefault(channel.channel_id, set()).add(task_meta.task_id)
pre_audit = len(await hub._audit_log.read_all())
await hub.close_channel(channel.channel_id, reason="explicit")
audit = await hub._audit_log.read_all()
new_records = audit[pre_audit:]
task_records = [r for r in new_records if r["kind"] == AUDIT_KIND_TASK_TERMINATED]
assert len(task_records) == 1
rec = task_records[0]
assert rec["task_id"] == task_meta.task_id
assert rec["capability"] == "analysis"
assert rec["outcome"] == TaskState.EXPIRED.value
await hub.close()