test/test_alert_halt.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
"""Unit tests for AlertPolicy and HaltCheckMiddleware.
AlertPolicy is the assembly-chain replacement for the old Signal delivery
system. HaltCheckMiddleware catches HaltEvent and short-circuits LLM calls.
Most tests wire both through the deprecated ``Agent(assembly=...)``, which
stays covered until its removal in 1.4.0.
These tests isolate both components from the full Agent stack to verify
edge cases not covered by Agent integration tests.
"""
from collections.abc import Sequence
from typing import Any, Self
from unittest.mock import MagicMock
import pytest
from ag2 import Agent, Context
from ag2.assembly import AssemblerMiddleware
from ag2.config import LLMClient, ModelConfig
from ag2.events import (
BaseEvent,
HaltEvent,
ModelMessage,
ModelResponse,
ObserverAlert,
Severity,
ToolCallEvent,
ToolCallsEvent,
)
from ag2.middleware import HaltCheckMiddleware, Middleware
from ag2.observers import BaseObserver
from ag2.policies import AlertPolicy
from ag2.stream import MemoryStream
from ag2.tools.final import tool
from ag2.watch import EventWatch
pytestmark = pytest.mark.filterwarnings("ignore:Agent\\(assembly=:DeprecationWarning")
@tool
def echo_tool(value: str) -> str:
"""Echoes back input."""
return f"echo: {value}"
class _RecordingClient(LLMClient):
"""LLM client that records calls and returns canned responses."""
def __init__(self, *responses: ModelResponse | ToolCallEvent | str) -> None:
self._responses = list(responses)
self._call_count = 0
self.calls: list[tuple[list[BaseEvent], list[str]]] = []
async def __call__(
self,
messages: Sequence[BaseEvent],
context: Context,
**kwargs: Any,
) -> ModelResponse:
self.calls.append((list(messages), list(context.prompt)))
resp = self._responses[self._call_count] if self._call_count < len(self._responses) else "done"
self._call_count += 1
if isinstance(resp, str):
return ModelResponse(message=ModelMessage(content=resp))
if isinstance(resp, ToolCallEvent):
return ModelResponse(tool_calls=ToolCallsEvent(calls=[resp]))
return resp
class _RecordingConfig(ModelConfig):
__test__ = False
def __init__(self, *responses: ModelResponse | ToolCallEvent | str) -> None:
self._responses = responses
self.client: _RecordingClient | None = None
def copy(self) -> Self:
return self
def create(self) -> _RecordingClient:
self.client = _RecordingClient(*self._responses)
return self.client
class TestAlertPolicyUnit:
"""AlertPolicy in isolation (no Agent, no middleware)."""
@pytest.mark.asyncio
async def test_no_alerts_is_noop(self) -> None:
"""When no ObserverAlert events are present, policy is a no-op."""
policy = AlertPolicy()
ctx = MagicMock()
prompts = ["System prompt"]
events: list[BaseEvent] = [ModelMessage(content="hello")]
result_prompts, result_events = await policy.apply(prompts, events, ctx)
assert result_prompts == ["System prompt"]
assert result_events == events
ctx.send.assert_not_called()
@pytest.mark.asyncio
async def test_single_warning_injected(self) -> None:
"""A single non-fatal alert is injected as prompt text."""
policy = AlertPolicy()
ctx = MagicMock()
alert = ObserverAlert(source="monitor", severity=Severity.WARNING, message="Token budget low")
prompts = ["System prompt"]
events: list[BaseEvent] = [ModelMessage(content="hi"), alert]
result_prompts, _ = await policy.apply(prompts, events, ctx)
assert len(result_prompts) == 2
assert "[OBSERVER MONITORING ALERTS]" in result_prompts[1]
assert "Token budget low" in result_prompts[1]
assert "WARNING" in result_prompts[1]
ctx.send.assert_not_called() # No HaltEvent for non-fatal
@pytest.mark.asyncio
async def test_multiple_severities_formatted(self) -> None:
"""Multiple alerts with different severities are all included."""
policy = AlertPolicy()
ctx = MagicMock()
alerts = [
ObserverAlert(source="a", severity=Severity.INFO, message="Info msg"),
ObserverAlert(source="b", severity=Severity.WARNING, message="Warn msg"),
ObserverAlert(source="c", severity=Severity.CRITICAL, message="Critical msg"),
]
prompts = ["System"]
events: list[BaseEvent] = list(alerts)
result_prompts, _ = await policy.apply(prompts, events, ctx)
alert_text = result_prompts[1]
assert "Info msg" in alert_text
assert "Warn msg" in alert_text
assert "Critical msg" in alert_text
@pytest.mark.asyncio
async def test_fatal_emits_halt_event(self) -> None:
"""FATAL alert emits HaltEvent on the stream via context.send()."""
policy = AlertPolicy()
sent: list[BaseEvent] = []
async def mock_send(event: BaseEvent) -> None:
sent.append(event)
ctx = MagicMock()
ctx.send = mock_send
alert = ObserverAlert(source="guard", severity=Severity.FATAL, message="Budget exceeded")
prompts = ["System"]
events: list[BaseEvent] = [alert]
result_prompts, _ = await policy.apply(prompts, events, ctx)
# HaltEvent should have been emitted
assert len(sent) == 1
assert isinstance(sent[0], HaltEvent)
assert "Budget exceeded" in sent[0].reason
assert sent[0].source == "guard"
# Halt note should be in prompts
assert any("[FATAL ALERT]" in p for p in result_prompts)
@pytest.mark.asyncio
async def test_mixed_fatal_and_nonfatal(self) -> None:
"""Both fatal and non-fatal alerts are handled correctly together."""
policy = AlertPolicy()
sent: list[BaseEvent] = []
async def mock_send(event: BaseEvent) -> None:
sent.append(event)
ctx = MagicMock()
ctx.send = mock_send
alerts = [
ObserverAlert(source="a", severity=Severity.WARNING, message="Just a warning"),
ObserverAlert(source="b", severity=Severity.FATAL, message="Total failure"),
]
prompts = ["System"]
events: list[BaseEvent] = list(alerts)
result_prompts, _ = await policy.apply(prompts, events, ctx)
# Non-fatal should be injected
assert any("Just a warning" in p for p in result_prompts)
# Fatal should emit HaltEvent
assert any(isinstance(e, HaltEvent) for e in sent)
# Fatal note should be in prompts
assert any("[FATAL ALERT]" in p for p in result_prompts)
@pytest.mark.asyncio
async def test_dedup_across_calls(self) -> None:
"""Same alerts are not re-injected on subsequent apply() calls."""
policy = AlertPolicy()
ctx = MagicMock()
alert = ObserverAlert(source="mon", severity=Severity.INFO, message="Noted")
prompts = ["System"]
events: list[BaseEvent] = [alert]
# First call — alert delivered
result_prompts_1, _ = await policy.apply(prompts, events, ctx)
assert len(result_prompts_1) == 2
# Second call — same alert should be skipped
result_prompts_2, _ = await policy.apply(prompts, events, ctx)
assert len(result_prompts_2) == 1 # Only original prompt, no alerts
@pytest.mark.asyncio
async def test_new_alerts_delivered_after_dedup(self) -> None:
"""New alerts are delivered even after previous ones were deduped."""
policy = AlertPolicy()
ctx = MagicMock()
alert1 = ObserverAlert(source="a", severity=Severity.INFO, message="First")
alert2 = ObserverAlert(source="b", severity=Severity.WARNING, message="Second")
# First call — deliver alert1
prompts1, _ = await policy.apply(["System"], [alert1], ctx)
assert "First" in prompts1[1]
# Second call — alert1 deduped, alert2 delivered
prompts2, _ = await policy.apply(["System"], [alert1, alert2], ctx)
assert len(prompts2) == 2
assert "Second" in prompts2[1]
assert "First" not in prompts2[1]
@pytest.mark.asyncio
async def test_dedup_survives_history_replacement(self) -> None:
"""A re-constructed alert with identical content is still treated as a duplicate.
Compaction rebuilds events from disk into fresh objects; an
``id(event)``-based dedup would re-deliver the alert. Content
dedup (source, severity, message) must hold.
"""
policy = AlertPolicy()
ctx = MagicMock()
original = ObserverAlert(source="mon", severity=Severity.WARNING, message="watch out")
await policy.apply(["System"], [original], ctx)
replaced = ObserverAlert(source="mon", severity=Severity.WARNING, message="watch out")
assert original is not replaced
prompts, _ = await policy.apply(["System"], [replaced], ctx)
assert all("watch out" not in p for p in prompts)
@pytest.mark.asyncio
async def test_multiple_fatals_uses_first_for_halt(self) -> None:
"""When multiple FATAL alerts arrive, the first one drives the HaltEvent."""
policy = AlertPolicy()
sent: list[BaseEvent] = []
async def mock_send(event: BaseEvent) -> None:
sent.append(event)
ctx = MagicMock()
ctx.send = mock_send
alerts = [
ObserverAlert(source="obs-a", severity=Severity.FATAL, message="Crash A"),
ObserverAlert(source="obs-b", severity=Severity.FATAL, message="Crash B"),
]
await policy.apply(["System"], list(alerts), ctx)
assert len(sent) == 1
halt = sent[0]
assert isinstance(halt, HaltEvent)
assert halt.source == "obs-a"
assert "Crash A" in halt.reason
# Both fatals should be in the alerts list
assert len(halt.alerts) == 2
class TestHaltCheckMiddleware:
"""HaltCheckMiddleware short-circuits LLM on HaltEvent."""
@pytest.mark.asyncio
async def test_no_halt_passes_through(self) -> None:
"""Without any FATAL alert, LLM is called normally."""
config = _RecordingConfig("normal response")
agent = Agent("test", config=config, assembly=[AlertPolicy()])
reply = await agent.ask("Hi")
assert reply.body == "normal response"
assert config.client is not None
assert config.client._call_count == 1
@pytest.mark.asyncio
async def test_halt_short_circuits_second_llm_call(self) -> None:
"""FATAL alert after first LLM call prevents second call.
Flow: LLM call 1 returns tool call -> ModelResponse fires FATAL observer
-> tool executes -> LLM call 2 intercepted by HaltCheckMiddleware.
"""
config = _RecordingConfig(
ToolCallEvent(name="echo_tool", arguments='{"value": "test"}'),
"should-not-reach",
)
class _FatalOnFirst(BaseObserver):
def __init__(self) -> None:
super().__init__("fatal-first", watch=EventWatch(ModelResponse))
self.fired = False
async def process(self, events: list[BaseEvent], ctx: Context) -> ObserverAlert | None:
if not self.fired:
self.fired = True
return ObserverAlert(
source="fatal-first",
severity=Severity.FATAL,
message="Stop now",
)
return None
observer = _FatalOnFirst()
agent = Agent(
"test",
config=config,
observers=[observer],
tools=[echo_tool],
assembly=[AlertPolicy()],
)
reply = await agent.ask("Hi")
assert observer.fired
assert reply.body is not None
assert "HALTED" in reply.body
# Second LLM call was short-circuited
assert config.client is not None
assert config.client._call_count == 1
@pytest.mark.asyncio
async def test_registered_as_middleware_halts_the_same_way(self) -> None:
"""The middleware= replacement for assembly= stops the second LLM call."""
config = _RecordingConfig(
ToolCallEvent(name="echo_tool", arguments='{"value": "test"}'),
"should-not-reach",
)
class _Fatal(BaseObserver):
def __init__(self) -> None:
super().__init__("fatal", watch=EventWatch(ModelResponse))
async def process(self, events: list[BaseEvent], ctx: Context) -> ObserverAlert | None:
return ObserverAlert(source="fatal", severity=Severity.FATAL, message="Stop now")
agent = Agent(
"test",
config=config,
observers=[_Fatal()],
tools=[echo_tool],
middleware=[
Middleware(AssemblerMiddleware, policies=[AlertPolicy()]),
Middleware(HaltCheckMiddleware),
],
)
reply = await agent.ask("Hi")
assert reply.body == "HALTED: FATAL: Stop now"
assert config.client is not None
assert config.client._call_count == 1
@pytest.mark.asyncio
async def test_halt_event_on_stream(self) -> None:
"""HaltEvent is observable on the stream."""
config = _RecordingConfig(
ToolCallEvent(name="echo_tool", arguments='{"value": "x"}'),
"after",
)
class _Fatal(BaseObserver):
def __init__(self) -> None:
super().__init__("fatal", watch=EventWatch(ModelResponse))
async def process(self, events: list[BaseEvent], ctx: Context) -> ObserverAlert | None:
return ObserverAlert(
source="fatal",
severity=Severity.FATAL,
message="halt",
)
stream = MemoryStream()
halts: list[HaltEvent] = []
stream.subscribe(lambda e: halts.append(e))
agent = Agent(
"test",
config=config,
observers=[_Fatal()],
tools=[echo_tool],
assembly=[AlertPolicy()],
)
await agent.ask("Hi", stream=stream)
halt_events = [e for e in halts if isinstance(e, HaltEvent)]
assert len(halt_events) >= 1
assert "halt" in halt_events[0].reason
@pytest.mark.asyncio
async def test_nonfatal_does_not_halt(self) -> None:
"""Non-fatal alerts don't trigger HaltCheckMiddleware."""
config = _RecordingConfig(
ToolCallEvent(name="echo_tool", arguments='{"value": "x"}'),
"final result",
)
class _Warning(BaseObserver):
def __init__(self) -> None:
super().__init__("warner", watch=EventWatch(ModelResponse))
async def process(self, events: list[BaseEvent], ctx: Context) -> ObserverAlert | None:
return ObserverAlert(
source="warner",
severity=Severity.WARNING,
message="be careful",
)
agent = Agent(
"test",
config=config,
observers=[_Warning()],
tools=[echo_tool],
assembly=[AlertPolicy()],
)
reply = await agent.ask("Hi")
# Should complete normally — second LLM call happens
assert reply.body == "final result"
assert config.client is not None
assert config.client._call_count == 2
@pytest.mark.asyncio
async def test_concurrent_fatal_from_two_observers(self) -> None:
"""Two observers both emit FATAL in the same turn — handled gracefully."""
config = _RecordingConfig(
ToolCallEvent(name="echo_tool", arguments='{"value": "x"}'),
"should-not-reach",
)
class _FatalObs(BaseObserver):
def __init__(self, name: str, msg: str) -> None:
super().__init__(name, watch=EventWatch(ModelResponse))
self._msg = msg
async def process(self, events: list[BaseEvent], ctx: Context) -> ObserverAlert | None:
return ObserverAlert(
source=self.name,
severity=Severity.FATAL,
message=self._msg,
)
agent = Agent(
"test",
config=config,
observers=[_FatalObs("obs-a", "crash-a"), _FatalObs("obs-b", "crash-b")],
tools=[echo_tool],
assembly=[AlertPolicy()],
)
reply = await agent.ask("Hi")
assert reply.body is not None
assert "HALTED" in reply.body
class TestAlertPolicyOrdering:
"""Assembly policy ordering validation recognizes AlertPolicy as injection."""
def test_alert_before_reduction_no_warning(self) -> None:
"""AlertPolicy before SlidingWindowPolicy produces no warnings."""
from ag2.assembly import AssemblerMiddleware
from ag2.policies import SlidingWindowPolicy
policies = [AlertPolicy(), SlidingWindowPolicy(50)]
warnings = AssemblerMiddleware.validate_order(policies)
assert warnings == []
def test_reduction_before_alert_warns(self) -> None:
"""SlidingWindowPolicy before AlertPolicy produces a warning."""
from ag2.assembly import AssemblerMiddleware
from ag2.policies import SlidingWindowPolicy
policies = [SlidingWindowPolicy(50), AlertPolicy()]
warnings = AssemblerMiddleware.validate_order(policies)
assert len(warnings) == 1
assert "alert" in warnings[0]