test/network/test_expectations.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
"""Expectation evaluators, violation handlers, and audit log tests.
Three layers covered:
* **Evaluators** (unit) — ``AcksWithinEvaluator``,
``ReplyWithinEvaluator``, ``MaxSilenceEvaluator`` against synthesised
``ExpectationContext`` inputs.
* **Handlers** (integration) — ``AuditHandler``,
``NotifyChannelHandler``, ``AutoCloseHandler`` driven by the hub's
manual ``_expectation_tick()`` call with a controllable clock.
* **Audit log** — ``register`` / ``unregister`` / ``set_*`` write the
expected records via ``AuditLog.read_all()``.
Time control: tests use a ``_MockClock`` so the threshold logic is
exercised deterministically without sleeping. The sweeper interval is
disabled (``expectation_sweep_interval=0``); tests call
``hub._expectation_tick()`` explicitly.
"""
import asyncio
from datetime import datetime
import pytest
from ag2 import Agent
from ag2.knowledge import MemoryKnowledgeStore
from ag2.network import (
EV_EXPECTATION_VIOLATED,
EV_TEXT,
Envelope,
Hub,
Resume,
Rule,
)
from ag2.network.adapters.consulting import CONSULTING_TYPE
from ag2.network.adapters.conversation import (
CONVERSATION_TYPE,
ConversationAdapter,
)
from ag2.network.channel import (
ChannelManifest,
ChannelMetadata,
ChannelState,
Expectation,
Participant,
ParticipantRole,
ParticipantSchema,
)
from ag2.network.hub import (
AUDIT_KIND_AGENT_REGISTERED,
AUDIT_KIND_AGENT_UNREGISTERED,
AUDIT_KIND_EXPECTATION_VIOLATED,
AUDIT_KIND_RESUME_SET,
AUDIT_KIND_RULE_SET,
AUDIT_KIND_SKILL_SET,
AcksWithinEvaluator,
AuditLog,
ExpectationContext,
MaxSilenceEvaluator,
ReplyWithinEvaluator,
)
from ag2.network.views.builtin import WindowedSummary
from ag2.testing import TestConfig
from ._helpers import _MockClock
def _agent(name: str, *events: object) -> Agent:
return Agent(name=name, config=TestConfig(*events))
async def _silent_handler(_envelope: Envelope) -> None:
"""No-op handler that ignores every envelope.
Used to drive expectation paths that depend on a participant
*not* responding (e.g. ``acks_within`` / ``reply_within`` tests).
"""
def _conv_metadata(
*,
state: ChannelState = ChannelState.ACTIVE,
created_at: str = "2026-01-01T00:00:00+00:00",
pending_acks: list[str] | None = None,
) -> ChannelMetadata:
"""Build a 2-party conversation channel metadata for evaluator tests."""
manifest = ConversationAdapter().manifest
return ChannelMetadata(
channel_id="s1",
manifest=manifest,
creator_id="alice",
participants=[
Participant(agent_id="alice", role=ParticipantRole.INITIATOR, order=0),
Participant(agent_id="bob", role=ParticipantRole.RESPONDENT, order=1),
],
state=state,
created_at=created_at,
pending_acks=list(pending_acks or []),
)
def _envelope(
sender: str,
text: str,
created_at: str,
*,
audience: list[str] | None = None,
) -> Envelope:
env = Envelope(
envelope_id=f"env-{sender}-{text}",
channel_id="s1",
sender_id=sender,
audience=audience,
event_type=EV_TEXT,
event_data={"text": text},
created_at=created_at,
)
return env
def _ctx(
metadata: ChannelMetadata,
*,
wal: list[Envelope] | None = None,
now: str = "2026-01-01T00:01:00+00:00",
) -> ExpectationContext:
now_dt = datetime.fromisoformat(now)
return ExpectationContext(
metadata=metadata,
state=None,
wal=wal or [],
now_iso=now,
now_seconds=now_dt.timestamp(),
)
# ── Evaluator unit tests ────────────────────────────────────────────────────
class TestAcksWithinEvaluator:
def test_fires_after_threshold_with_pending_acks(self) -> None:
metadata = _conv_metadata(
state=ChannelState.PENDING,
created_at="2026-01-01T00:00:00+00:00",
pending_acks=["bob"],
)
evaluator = AcksWithinEvaluator()
expectation = Expectation(
name="acks_within",
on_violation="auto_close",
params={"seconds": 30},
)
violation = evaluator.evaluate(expectation, _ctx(metadata, now="2026-01-01T00:01:00+00:00"))
assert violation is not None
assert violation.violator_ids == ["bob"]
assert violation.detail["threshold_seconds"] == 30
def test_silent_within_threshold(self) -> None:
metadata = _conv_metadata(
state=ChannelState.PENDING,
pending_acks=["bob"],
)
evaluator = AcksWithinEvaluator()
expectation = Expectation(
name="acks_within",
on_violation="auto_close",
params={"seconds": 60},
)
# 30s elapsed, threshold 60s → no violation
violation = evaluator.evaluate(expectation, _ctx(metadata, now="2026-01-01T00:00:30+00:00"))
assert violation is None
def test_silent_when_no_pending_acks(self) -> None:
metadata = _conv_metadata(
state=ChannelState.PENDING,
pending_acks=[],
)
evaluator = AcksWithinEvaluator()
expectation = Expectation(name="acks_within", on_violation="auto_close", params={"seconds": 30})
violation = evaluator.evaluate(expectation, _ctx(metadata, now="2026-01-01T00:10:00+00:00"))
assert violation is None
def test_silent_when_channel_active(self) -> None:
# Active channel (acks already collected) — evaluator skips.
metadata = _conv_metadata(state=ChannelState.ACTIVE, pending_acks=[])
evaluator = AcksWithinEvaluator()
expectation = Expectation(name="acks_within", on_violation="auto_close", params={"seconds": 30})
violation = evaluator.evaluate(expectation, _ctx(metadata, now="2026-01-01T00:10:00+00:00"))
assert violation is None
class TestReplyWithinEvaluator:
def test_fires_when_addressed_participant_silent(self) -> None:
metadata = _conv_metadata()
wal = [_envelope("alice", "hi bob", "2026-01-01T00:00:00+00:00", audience=["bob"])]
evaluator = ReplyWithinEvaluator()
expectation = Expectation(
name="reply_within",
on_violation="audit",
params={"seconds": 60},
)
violation = evaluator.evaluate(expectation, _ctx(metadata, wal=wal, now="2026-01-01T00:02:00+00:00"))
assert violation is not None
assert violation.violator_ids == ["bob"]
def test_silent_when_reply_within_threshold(self) -> None:
# Tight timing: bob replied 15s ago; alice's window to follow up
# hasn't expired yet either.
metadata = _conv_metadata()
wal = [
_envelope("alice", "hi bob", "2026-01-01T00:00:00+00:00", audience=["bob"]),
_envelope("bob", "hi back", "2026-01-01T00:00:30+00:00"),
]
evaluator = ReplyWithinEvaluator()
expectation = Expectation(name="reply_within", on_violation="audit", params={"seconds": 60})
# Now = 00:00:45 → bob's reply is 15s old; alice's outstanding
# incoming (bob's broadcast) is 15s old — both within 60s window.
violation = evaluator.evaluate(expectation, _ctx(metadata, wal=wal, now="2026-01-01T00:00:45+00:00"))
assert violation is None
def test_originally_addressed_party_not_violator_after_reply(self) -> None:
# Bob has answered — even past the threshold he is not flagged
# because his outbound is newer than his most recent inbound.
metadata = _conv_metadata()
wal = [
_envelope("alice", "hi bob", "2026-01-01T00:00:00+00:00", audience=["bob"]),
_envelope("bob", "hi back", "2026-01-01T00:00:30+00:00", audience=["alice"]),
]
evaluator = ReplyWithinEvaluator()
expectation = Expectation(name="reply_within", on_violation="audit", params={"seconds": 60})
# Long after — alice will be flagged but bob will not.
violation = evaluator.evaluate(expectation, _ctx(metadata, wal=wal, now="2026-01-01T00:10:00+00:00"))
assert violation is not None
assert "bob" not in violation.violator_ids
assert "alice" in violation.violator_ids
class TestMaxSilenceEvaluator:
def test_fires_when_channel_silent_past_threshold(self) -> None:
metadata = _conv_metadata(created_at="2026-01-01T00:00:00+00:00")
wal = [_envelope("alice", "hello", "2026-01-01T00:00:30+00:00")]
evaluator = MaxSilenceEvaluator()
expectation = Expectation(
name="max_silence",
on_violation="audit",
params={"seconds": 60},
)
violation = evaluator.evaluate(expectation, _ctx(metadata, wal=wal, now="2026-01-01T00:02:00+00:00"))
assert violation is not None
assert violation.violator_ids == [] # channel-wide
def test_silent_when_recent_activity(self) -> None:
metadata = _conv_metadata(created_at="2026-01-01T00:00:00+00:00")
wal = [_envelope("alice", "hello", "2026-01-01T00:01:30+00:00")]
evaluator = MaxSilenceEvaluator()
expectation = Expectation(name="max_silence", on_violation="audit", params={"seconds": 60})
violation = evaluator.evaluate(expectation, _ctx(metadata, wal=wal, now="2026-01-01T00:02:00+00:00"))
assert violation is None
# ── Handler integration tests ───────────────────────────────────────────────
@pytest.mark.asyncio
async def test_auto_close_handler_terminates_channel_with_audit() -> None:
"""Consulting's ``acks_within(30s, auto_close)`` fires after 30s,
transitions channel to CLOSED, records to audit log."""
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,
invite_ack_timeout=300.0, # long enough that the sweeper closes first
)
alice = await hub.register(_agent("alice"))
# Bob registers without the auto-ack default handler — install a
# silent handler explicitly so invites are never acked.
bob = await hub.register(_agent("bob"), attach_plugin=False)
bob.on_envelope(_silent_handler)
# Open in background; the sweeper auto_closes via ProtocolError on the waiter.
open_task = asyncio.create_task(alice.open(type=CONSULTING_TYPE, target=bob.agent_id))
# Let the invite dispatch.
await asyncio.sleep(0.05)
# Advance past the 30s acks_within threshold and tick.
clock.advance(45)
await hub._expectation_tick()
# Open fails because the channel was auto-closed.
with pytest.raises(Exception):
await open_task
# Find the channel id from cached state.
channels = list(hub._channels.values())
assert len(channels) == 1
channel_id = channels[0].channel_id
final = await hub.get_channel(channel_id)
assert final.state == ChannelState.CLOSED
assert "expectation_violated:acks_within" in final.close_reason
audit = await hub._audit_log.read_all()
kinds = [r["kind"] for r in audit]
assert AUDIT_KIND_EXPECTATION_VIOLATED in kinds
await hub.close()
@pytest.mark.asyncio
async def test_audit_handler_records_without_envelope_or_close() -> None:
"""``audit`` handler logs the violation and nothing else."""
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_TYPE, target=bob.agent_id)
pre_audit = len(await hub._audit_log.read_all())
# Conversation declares max_silence(3600s, audit). Advance 1h+.
clock.advance(3700)
await hub._expectation_tick()
# Channel still ACTIVE; no EV_EXPECTATION_VIOLATED in WAL.
state = await hub.get_channel(channel.channel_id)
assert state.state == ChannelState.ACTIVE
wal = await hub.read_wal(channel.channel_id)
assert not any(e.event_type == EV_EXPECTATION_VIOLATED for e in wal)
# Audit log has one new violation entry.
audit = await hub._audit_log.read_all()
new_records = audit[pre_audit:]
violation_records = [r for r in new_records if r["kind"] == AUDIT_KIND_EXPECTATION_VIOLATED]
assert len(violation_records) == 1
assert violation_records[0]["expectation"] == "max_silence"
await hub.close()
@pytest.mark.asyncio
async def test_notify_channel_handler_broadcasts_envelope() -> None:
"""``notify_channel`` handler audits + posts EV_EXPECTATION_VIOLATED."""
# Use a custom channel with a notify_channel expectation we can drive.
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)
# Register a custom adapter with notify_channel on max_silence so we
# don't have to wait for conversation's 1h default.
class _NotifyAdapter(ConversationAdapter):
def __init__(self) -> None:
super().__init__()
self.manifest = ChannelManifest(
type="conversation_notify",
version=1,
participants=ParticipantSchema(min=2, max=2, roles=["initiator", "respondent"]),
knobs_schema={},
default_view_policy=WindowedSummary.name,
expectations=[
Expectation(
name="max_silence",
on_violation="notify_channel",
params={"seconds": 60},
),
],
)
hub.register_adapter(_NotifyAdapter())
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
channel = await alice.open(type="conversation_notify", target=bob.agent_id)
clock.advance(120)
await hub._expectation_tick()
wal = await hub.read_wal(channel.channel_id)
violation_envelopes = [e for e in wal if e.event_type == EV_EXPECTATION_VIOLATED]
assert len(violation_envelopes) == 1
assert violation_envelopes[0].event_data["expectation"] == "max_silence"
# Channel still ACTIVE — notify_channel does not close.
state = await hub.get_channel(channel.channel_id)
assert state.state == ChannelState.ACTIVE
await hub.close()
@pytest.mark.asyncio
async def test_violation_dedup_within_channel_lifetime() -> None:
"""Two consecutive ticks fire the same violation only once."""
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"))
await alice.open(type=CONVERSATION_TYPE, target=bob.agent_id)
pre_audit = len(await hub._audit_log.read_all())
clock.advance(3700)
await hub._expectation_tick()
await hub._expectation_tick()
await hub._expectation_tick()
audit = await hub._audit_log.read_all()
violations = [r for r in audit[pre_audit:] if r["kind"] == AUDIT_KIND_EXPECTATION_VIOLATED]
assert len(violations) == 1 # deduped across 3 ticks
await hub.close()
# ── Audit log tests for identity changes ────────────────────────────────────
@pytest.mark.asyncio
async def test_audit_log_records_register_and_unregister() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"))
await hub.unregister(alice.agent_id)
audit = await hub._audit_log.read_all()
kinds = [r["kind"] for r in audit]
assert AUDIT_KIND_AGENT_REGISTERED in kinds
assert AUDIT_KIND_AGENT_UNREGISTERED in kinds
await hub.close()
@pytest.mark.asyncio
async def test_audit_log_records_set_resume_skill_rule() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"))
pre_audit = len(await hub._audit_log.read_all())
await hub.set_resume(alice.agent_id, Resume(summary="updated"))
await hub.set_skill(alice.agent_id, "# alice's skill")
await hub.set_rule(alice.agent_id, Rule())
audit = await hub._audit_log.read_all()
new_kinds = [r["kind"] for r in audit[pre_audit:]]
assert AUDIT_KIND_RESUME_SET in new_kinds
assert AUDIT_KIND_SKILL_SET in new_kinds
assert AUDIT_KIND_RULE_SET in new_kinds
await hub.close()
@pytest.mark.asyncio
async def test_audit_log_writer_round_trips() -> None:
"""``AuditLog`` writes JSON lines that ``read_all`` round-trips."""
store = MemoryKnowledgeStore()
log = AuditLog(store)
await log.append({"at": "2026-01-01T00:00:00", "kind": "test", "n": 1})
await log.append({"at": "2026-01-01T00:00:01", "kind": "test", "n": 2})
records = await log.read_all()
assert records == [
{"at": "2026-01-01T00:00:00", "kind": "test", "n": 1},
{"at": "2026-01-01T00:00:01", "kind": "test", "n": 2},
]