test/network/test_hydrate_scale.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
"""``Hub.hydrate()`` correctness at scale.
Verifies:
* All channel metadata round-trips through ``DiskKnowledgeStore``.
* Every active channel's WAL is re-folded through its adapter so the
in-memory ``_adapter_states`` cache matches the on-disk truth.
* The capability index rebuilds from loaded resumes.
* Round-trip is deterministic: hydrating twice yields the same state.
Correctness benchmark, not a perf benchmark. Default: 100 channels ×
100 envelopes (~10k total) so the test runs in <2s. Bump
``ENVELOPES_PER_CHANNEL`` locally to exercise larger sweeps.
"""
import pytest
from ag2 import Agent
from ag2.knowledge import DiskKnowledgeStore
from ag2.network import (
Envelope,
Hub,
Resume,
)
from ag2.network.adapters.conversation import (
CONVERSATION_TYPE,
ConversationState,
)
from ag2.network.adapters.discussion import (
DISCUSSION_TYPE,
ORDERING_ROUND_ROBIN,
DiscussionState,
)
from ag2.network.channel import ChannelState
from ag2.network.envelope import EV_TEXT
from ag2.testing import TestConfig
# Scale chosen for unit-run speed: 100 channels × 100 envelopes ≈
# 10k envelopes, ~1s populate, instant hydrate. Bump
# ``ENVELOPES_PER_CHANNEL`` locally for larger sweeps. Hydrate cost
# scales linearly with envelope count and is dwarfed by populate
# (write throughput) in any realistic setup.
CHANNELS = 100
ENVELOPES_PER_CHANNEL = 100
def _agent(name: str) -> Agent:
return Agent(name=name, config=TestConfig())
@pytest.mark.asyncio
async def test_hydrate_round_trips_many_channels(tmp_path) -> None:
"""Populate disk via a live hub, tear down, re-open, verify state matches."""
n_channels = CHANNELS
n_envelopes = ENVELOPES_PER_CHANNEL
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub1.register(
_agent("alice"),
resume=Resume(claimed_capabilities=["analysis"]),
)
bob = await hub1.register(
_agent("bob"),
resume=Resume(claimed_capabilities=["coding"]),
)
# Build conversation channels in parallel batches; each channel
# fills its WAL with N alternating EV_TEXT envelopes.
expected_states: dict[str, dict] = {}
channels = []
for _ in range(n_channels):
channel = await alice.open(type=CONVERSATION_TYPE, target=bob.agent_id)
channels.append(channel)
# Fill each channel's WAL.
for channel in channels:
for i in range(n_envelopes):
sender = alice if (i % 2 == 0) else bob
envelope = Envelope(
channel_id=channel.channel_id,
sender_id=sender.agent_id,
audience=None,
event_type=EV_TEXT,
event_data={"text": f"msg-{i}"},
)
await hub1.post_envelope(envelope)
# Snapshot expected state.
cached = hub1._adapter_states[channel.channel_id]
expected_states[channel.channel_id] = {
"turn_count": cached.turn_count,
"last_speaker_id": cached.last_speaker_id,
}
await hub1.close()
# Round 1: re-open and verify everything matches.
store2 = DiskKnowledgeStore(str(tmp_path))
hub2 = await Hub.open(store2, ttl_sweep_interval=0, expectation_sweep_interval=0)
# Identities preserved.
assert alice.agent_id in hub2._passports
assert bob.agent_id in hub2._passports
# Capability index rebuilt from resumes.
assert hub2.agents_with_capability("analysis") == [alice.agent_id]
assert hub2.agents_with_capability("coding") == [bob.agent_id]
# Every channel round-tripped.
for channel in channels:
meta = await hub2.get_channel(channel.channel_id)
assert meta.manifest.type == CONVERSATION_TYPE
assert meta.state == ChannelState.ACTIVE
cached = hub2._adapter_states[channel.channel_id]
assert isinstance(cached, ConversationState)
assert cached.turn_count == expected_states[channel.channel_id]["turn_count"]
assert cached.last_speaker_id == expected_states[channel.channel_id]["last_speaker_id"]
# Round 2: hydrating twice from the same store is idempotent.
await hub2.hydrate()
for channel in channels:
cached = hub2._adapter_states[channel.channel_id]
assert cached.turn_count == expected_states[channel.channel_id]["turn_count"]
await hub2.close()
@pytest.mark.asyncio
async def test_hydrate_refolds_discussion_round_robin_state(tmp_path) -> None:
"""Multi-party discussion rotation survives hub restart deterministically."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
names = ["alice", "bob", "carol", "dave", "erin"]
clients = []
for name in names:
client = await hub1.register(_agent(name))
clients.append(client)
alice = clients[0]
channel = await alice.open(
type=DISCUSSION_TYPE,
target=[c.agent_id for c in clients[1:]],
knobs={"ordering": ORDERING_ROUND_ROBIN},
)
# 50 turns through the rotation (10 full cycles of 5 speakers).
for i in range(50):
sender = clients[i % 5]
envelope = Envelope(
channel_id=channel.channel_id,
sender_id=sender.agent_id,
audience=None,
event_type=EV_TEXT,
event_data={"text": f"turn-{i}"},
)
await hub1.post_envelope(envelope)
expected = hub1._adapter_states[channel.channel_id]
expected_turn_count = expected.turn_count
expected_next = expected.expected_next_speaker
expected_last = expected.last_speaker_id
await hub1.close()
store2 = DiskKnowledgeStore(str(tmp_path))
hub2 = await Hub.open(store2, ttl_sweep_interval=0, expectation_sweep_interval=0)
rebuilt = hub2._adapter_states[channel.channel_id]
assert isinstance(rebuilt, DiscussionState)
assert rebuilt.turn_count == expected_turn_count
assert rebuilt.expected_next_speaker == expected_next
assert rebuilt.last_speaker_id == expected_last
assert rebuilt.participant_order == [c.agent_id for c in clients]
await hub2.close()