test/stream/test_transient.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
"""Tests for transient event filtering in MemoryStream persistence.
Verifies that events marked ``__transient__ = True`` are:
- Delivered to subscribers (streaming still works)
- NOT persisted to storage by default
- Persisted when ``persist_all=True``
"""
import pytest
from ag2 import Context, MemoryStream
from ag2.events import (
AggregationCompleted,
BaseEvent,
CompactionCompleted,
ModelMessage,
ModelMessageChunk,
ModelReasoning,
ModelRequest,
ModelResponse,
ObserverCompleted,
ObserverStarted,
TaskCompleted,
TaskProgress,
TaskStarted,
TextInput,
ToolCallEvent,
)
class TestTransientFlag:
"""Verify __transient__ is set correctly on event types."""
def test_base_event_not_transient(self):
assert BaseEvent.__transient__ is False
def test_model_request_not_transient(self):
assert ModelRequest.__transient__ is False
def test_model_response_not_transient(self):
assert ModelResponse.__transient__ is False
def test_tool_call_event_not_transient(self):
assert ToolCallEvent.__transient__ is False
def test_task_started_not_transient(self):
assert TaskStarted.__transient__ is False
def test_task_completed_not_transient(self):
assert TaskCompleted.__transient__ is False
def test_model_message_chunk_transient(self):
assert ModelMessageChunk.__transient__ is True
def test_model_message_transient(self):
assert ModelMessage.__transient__ is True
def test_model_reasoning_transient(self):
assert ModelReasoning.__transient__ is True
def test_task_progress_transient(self):
assert TaskProgress.__transient__ is True
def test_observer_started_transient(self):
assert ObserverStarted.__transient__ is True
def test_observer_completed_transient(self):
assert ObserverCompleted.__transient__ is True
def test_compaction_completed_transient(self):
assert CompactionCompleted.__transient__ is True
def test_aggregation_completed_transient(self):
assert AggregationCompleted.__transient__ is True
class TestTransientDelivery:
"""Transient events must still reach subscribers (for real-time streaming)."""
@pytest.mark.asyncio
async def test_chunk_delivered_to_subscriber(self):
stream = MemoryStream()
received = []
stream.where(ModelMessageChunk).subscribe(lambda ev: received.append(ev))
chunk = ModelMessageChunk(content="hello")
await stream.send(chunk, context=Context(stream))
assert len(received) == 1
assert received[0].content == "hello"
@pytest.mark.asyncio
async def test_task_progress_delivered_to_subscriber(self):
stream = MemoryStream()
received = []
stream.where(TaskProgress).subscribe(lambda ev: received.append(ev))
progress = TaskProgress(
task_id="t1",
agent_name="analyzer",
objective="scan logs",
content="working...",
)
await stream.send(progress, context=Context(stream))
assert len(received) == 1
assert received[0].content == "working..."
class TestTransientNotPersisted:
"""Default MemoryStream should NOT store transient events in history."""
@pytest.mark.asyncio
async def test_chunk_not_in_history(self):
stream = MemoryStream()
ctx = Context(stream)
await stream.send(ModelRequest([TextInput("hi")]), ctx)
await stream.send(ModelMessageChunk(content="hell"), ctx)
await stream.send(ModelMessageChunk(content="o"), ctx)
await stream.send(ModelResponse(message=ModelMessage(content="hello")), ctx)
events = list(await stream.history.get_events())
types = [type(e).__name__ for e in events]
assert "ModelRequest" in types
assert "ModelResponse" in types
assert "ModelMessageChunk" not in types
assert "ModelMessage" not in types
@pytest.mark.asyncio
async def test_lifecycle_events_not_in_history(self):
stream = MemoryStream()
ctx = Context(stream)
await stream.send(ObserverStarted(name="loop_detector"), ctx)
await stream.send(ModelRequest([TextInput("hi")]), ctx)
await stream.send(ModelResponse(message=ModelMessage(content="hello")), ctx)
await stream.send(ObserverCompleted(name="loop_detector"), ctx)
events = list(await stream.history.get_events())
types = [type(e).__name__ for e in events]
assert "ObserverStarted" not in types
assert "ObserverCompleted" not in types
assert "ModelRequest" in types
assert "ModelResponse" in types
@pytest.mark.asyncio
async def test_task_progress_not_in_history(self):
stream = MemoryStream()
ctx = Context(stream)
await stream.send(
TaskStarted(task_id="t1", agent_name="researcher", objective="research"),
ctx,
)
await stream.send(
TaskProgress(task_id="t1", agent_name="researcher", objective="research", content="step 1"),
ctx,
)
await stream.send(
TaskProgress(task_id="t1", agent_name="researcher", objective="research", content="step 2"),
ctx,
)
await stream.send(
TaskCompleted(
task_id="t1",
agent_name="researcher",
objective="research",
result="done",
task_stream=stream.id,
),
ctx,
)
events = list(await stream.history.get_events())
types = [type(e).__name__ for e in events]
assert "TaskStarted" in types
assert "TaskCompleted" in types
assert "TaskProgress" not in types
@pytest.mark.asyncio
async def test_compaction_aggregation_not_in_history(self):
stream = MemoryStream()
ctx = Context(stream)
await stream.send(ModelRequest([TextInput("hi")]), ctx)
await stream.send(CompactionCompleted(agent="pilot", strategy="s", events_before=10, events_after=5), ctx)
await stream.send(AggregationCompleted(agent="pilot", strategy="s", event_count=5), ctx)
events = list(await stream.history.get_events())
types = [type(e).__name__ for e in events]
assert "CompactionCompleted" not in types
assert "AggregationCompleted" not in types
@pytest.mark.asyncio
async def test_non_transient_events_all_persisted(self):
"""Ensure conversation events are all stored."""
stream = MemoryStream()
ctx = Context(stream)
await stream.send(ModelRequest([TextInput("hi")]), ctx)
await stream.send(ToolCallEvent(name="search", arguments="q"), ctx)
await stream.send(ModelResponse(message=ModelMessage(content="hello")), ctx)
events = list(await stream.history.get_events())
assert len(events) == 3
class TestPersistAll:
"""When persist_all=True, ALL events including transient ones are stored."""
@pytest.mark.asyncio
async def test_chunks_persisted_when_persist_all(self):
stream = MemoryStream(persist_all=True)
ctx = Context(stream)
await stream.send(ModelRequest([TextInput("hi")]), ctx)
await stream.send(ModelMessageChunk(content="hell"), ctx)
await stream.send(ModelMessageChunk(content="o"), ctx)
await stream.send(ModelResponse(message=ModelMessage(content="hello")), ctx)
events = list(await stream.history.get_events())
types = [type(e).__name__ for e in events]
assert types.count("ModelMessageChunk") == 2
assert "ModelMessage" not in types # ModelMessage is standalone, not sent here as top-level
@pytest.mark.asyncio
async def test_lifecycle_persisted_when_persist_all(self):
stream = MemoryStream(persist_all=True)
ctx = Context(stream)
await stream.send(ObserverStarted(name="monitor"), ctx)
await stream.send(
TaskProgress(task_id="t1", agent_name="researcher", objective="work", content="working"),
ctx,
)
events = list(await stream.history.get_events())
types = [type(e).__name__ for e in events]
assert "ObserverStarted" in types
assert "TaskProgress" in types