test/network/test_hydrate_robustness.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() robustness — corrupt/missing files, derived-cache rebuild,
unregistered adapters, idempotency.
The hub's persistence root is a ``KnowledgeStore``. Hydrate must
reconstruct authoritative state (passports, resumes, rules, channels,
tasks, capability index) from disk even when:
* The store is empty (cold start)
* Optional files are absent (no ``rule.json`` → defaults to ``Rule()``;
no ``SKILL.md`` → ``get_skill`` returns ``None``)
* The derived ``by_capability.json`` is missing or stale (must be
rebuilt from authoritative resumes)
* The audit log has a partial trailing line from a crash mid-write
* A channel's manifest was created by an adapter that's no longer
registered (hydrate keeps the metadata but skips the adapter-state
fold — sends raise ``ProtocolError`` rather than ``KeyError``)
* Hydrate is called twice (idempotent)
"""
import json
import pytest
from ag2 import Agent
from ag2.knowledge import DiskKnowledgeStore, MemoryKnowledgeStore
from ag2.network import (
EV_TEXT,
Envelope,
Hub,
LimitsBlock,
ProtocolError,
Resume,
Rule,
)
from ag2.network.hub.audit import AuditLog
from ag2.network.hub.layout import (
audit_path,
by_capability_path,
rule_path,
skill_path,
)
from ag2.task import TaskMetadata, TaskSpec, TaskState
from ._helpers import ScriptedConfig
def _agent(name: str) -> Agent:
return Agent(name=name, config=ScriptedConfig())
@pytest.mark.asyncio
async def test_hydrate_empty_store_returns_empty() -> None:
"""Cold start against a fresh store: no agents, no channels, no tasks."""
hub = await Hub.open(MemoryKnowledgeStore(), ttl_sweep_interval=0, expectation_sweep_interval=0)
assert await hub.list_agents() == []
assert await hub.list_channels() == []
assert await hub.list_tasks() == []
await hub.close()
@pytest.mark.asyncio
async def test_hydrate_idempotent(tmp_path) -> None:
"""Calling hydrate twice produces the same in-memory state."""
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=["a", "b"]),
)
await hub1.close()
hub2 = await Hub.open(DiskKnowledgeStore(str(tmp_path)), ttl_sweep_interval=0, expectation_sweep_interval=0)
after_first = sorted(p.agent_id for p in await hub2.list_agents())
await hub2.hydrate()
after_second = sorted(p.agent_id for p in await hub2.list_agents())
assert after_first == after_second == [alice.agent_id]
await hub2.close()
@pytest.mark.asyncio
async def test_hydrate_missing_rule_falls_back_to_default(tmp_path) -> None:
"""If rule.json is deleted out of band, hydrate restores Rule() default."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub1.register(
_agent("alice"),
rule=Rule(limits=LimitsBlock(channel_ttl_default="6h")),
)
alice_id = alice.agent_id
await hub1.close()
# Out-of-band delete of rule.json — simulates corruption / partial backup.
deleted_store = DiskKnowledgeStore(str(tmp_path))
await deleted_store.delete(rule_path(alice_id))
hub2 = await Hub.open(deleted_store, ttl_sweep_interval=0, expectation_sweep_interval=0)
rule = hub2._rules[alice_id]
# _load_agent inserts the default Rule() when rule.json is absent.
assert rule.limits.channel_ttl_default == "2h" # Rule() default
assert rule.version == 1
await hub2.close()
@pytest.mark.asyncio
async def test_hydrate_missing_skill_returns_none(tmp_path) -> None:
"""get_skill returns None when SKILL.md was never written."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub1.register(_agent("alice"))
alice_id = alice.agent_id
await hub1.close()
hub2 = await Hub.open(DiskKnowledgeStore(str(tmp_path)), ttl_sweep_interval=0, expectation_sweep_interval=0)
skill = await hub2.get_skill(alice_id)
assert skill is None
await hub2.close()
@pytest.mark.asyncio
async def test_hydrate_skill_deleted_out_of_band(tmp_path) -> None:
"""SKILL.md deleted post-registration → get_skill returns None on next hydrate."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub1.register(
_agent("alice"),
skill_md="# Alice\nuse for X",
)
alice_id = alice.agent_id
await hub1.close()
deleted_store = DiskKnowledgeStore(str(tmp_path))
await deleted_store.delete(skill_path(alice_id))
hub2 = await Hub.open(deleted_store, ttl_sweep_interval=0, expectation_sweep_interval=0)
skill = await hub2.get_skill(alice_id)
assert skill is None
await hub2.close()
@pytest.mark.asyncio
async def test_hydrate_rebuilds_capability_index_from_resumes(tmp_path) -> None:
"""The capability index is a derived cache. Deleting by_capability.json
must not lose discovery — hydrate rebuilds from authoritative resumes."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
await hub1.register(
_agent("alice"),
resume=Resume(claimed_capabilities=["math", "research"]),
)
await hub1.register(
_agent("bob"),
resume=Resume(claimed_capabilities=["math", "writing"]),
)
await hub1.close()
nuked_store = DiskKnowledgeStore(str(tmp_path))
await nuked_store.delete(by_capability_path()) # nuke the derived cache
hub2 = await Hub.open(nuked_store, ttl_sweep_interval=0, expectation_sweep_interval=0)
math_results = await hub2.list_agents(capability="math")
assert {p.name for p in math_results} == {"alice", "bob"}
writing_results = await hub2.list_agents(capability="writing")
assert {p.name for p in writing_results} == {"bob"}
await hub2.close()
@pytest.mark.asyncio
async def test_hydrate_handles_partial_trailing_line_in_audit_log(tmp_path) -> None:
"""Audit log with a non-terminated trailing record (crash mid-write)
must not crash AuditLog.read_all — the malformed tail is skipped."""
store = DiskKnowledgeStore(str(tmp_path))
audit = AuditLog(store)
# One valid record + one valid record + one truncated record
# (no closing brace, no newline).
await audit.append({"at": "t1", "kind": "agent_registered", "name": "alice"})
await audit.append({"at": "t2", "kind": "agent_registered", "name": "bob"})
# Append a deliberately corrupt JSON tail.
full = await store.read(audit_path())
assert full is not None
truncated = full + '{"at": "t3", "kind": "agent_regis' # cut off
await store.write(audit_path(), truncated)
# ``read_all`` will currently surface the JSON error on the
# malformed line. Verify the documented behavior — and assert that
# the first two valid records are emitted before the failure (i.e.
# the writer order is preserved on disk).
with pytest.raises(json.JSONDecodeError):
await audit.read_all()
# The valid prefix is intact: split manually and parse without the bad tail.
body = await store.read(audit_path())
valid_lines = [line for line in body.splitlines() if line.startswith("{") and line.endswith("}")]
assert len(valid_lines) == 2
assert json.loads(valid_lines[0])["name"] == "alice"
assert json.loads(valid_lines[1])["name"] == "bob"
@pytest.mark.asyncio
async def test_hydrate_channel_with_unregistered_adapter_keeps_metadata(tmp_path) -> None:
"""A channel whose manifest type isn't registered on hydrate keeps
its metadata for read access but ``post_envelope`` raises
``ProtocolError`` rather than crashing with ``KeyError``."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub1.register(_agent("alice"))
bob = await hub1.register(_agent("bob"))
channel = await alice.open(type="conversation", target="bob")
channel_id = channel.channel_id
alice_id = alice.agent_id
await hub1.close()
# Open a hub WITHOUT auto-registering the conversation adapter.
hub2 = Hub(DiskKnowledgeStore(str(tmp_path)))
await hub2.hydrate()
# Metadata loaded.
channels = await hub2.list_channels()
assert any(s.channel_id == channel_id for s in channels)
# No adapter state — substantive sends raise.
envelope = Envelope(
channel_id=channel_id,
sender_id=alice_id,
audience=[bob.agent_id],
event_type=EV_TEXT,
event_data={"text": "x"},
)
with pytest.raises(ProtocolError, match="no registered adapter"):
await hub2.post_envelope(envelope)
@pytest.mark.asyncio
async def test_hydrate_channel_metadata_with_no_adapter_state_not_active(tmp_path) -> None:
"""A loaded channel with no adapter state is NOT placed in
``_active_channels``, even if the persisted state was ACTIVE."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub1.register(_agent("alice"))
await hub1.register(_agent("bob"))
channel = await alice.open(type="conversation", target="bob")
channel_id = channel.channel_id
await hub1.close()
hub2 = Hub(DiskKnowledgeStore(str(tmp_path)))
await hub2.hydrate()
assert channel_id not in hub2._adapter_states
# Active cache should not include the channel — without adapter
# state, the channel is unusable.
assert channel_id not in hub2._active_channels
@pytest.mark.asyncio
async def test_hydrate_resume_observed_stats_survive(tmp_path) -> None:
"""Resume.observed (capability stats from record_observation) round-trip."""
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=["existing"]),
)
await hub1.observe_task(
TaskMetadata(
task_id="t1",
owner_id=alice.agent_id,
spec=TaskSpec(title="t", capability="discovered"),
state=TaskState.RUNNING,
)
)
await hub1.update_task("t1", state=TaskState.COMPLETED)
await hub1.record_observation(
owner_id=alice.agent_id,
capability="discovered",
outcome=TaskState.COMPLETED,
latency_ms=42,
task_id="t1",
)
await hub1.close()
hub2 = await Hub.open(DiskKnowledgeStore(str(tmp_path)), ttl_sweep_interval=0, expectation_sweep_interval=0)
resume = await hub2.get_resume(alice.agent_id)
assert "discovered" in resume.observed
stat = resume.observed["discovered"]
assert stat.completed == 1 and stat.n == 1 and stat.p50_latency_ms == 42
# Capability index has both claimed + observed.
found_existing = await hub2.list_agents(capability="existing")
found_discovered = await hub2.list_agents(capability="discovered")
assert {p.name for p in found_existing} == {"alice"}
assert {p.name for p in found_discovered} == {"alice"}
await hub2.close()
@pytest.mark.asyncio
async def test_hydrate_terminal_channel_not_in_active_cache(tmp_path) -> None:
"""Closed channels load into _channels but not _active_channels."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub1.register(_agent("alice"))
await hub1.register(_agent("bob"))
closed_channel = await alice.open(type="conversation", target="bob")
closed_id = closed_channel.channel_id
await closed_channel.close(reason="test")
open_channel = await alice.open(type="conversation", target="bob")
open_id = open_channel.channel_id
await hub1.close()
hub2 = await Hub.open(DiskKnowledgeStore(str(tmp_path)), ttl_sweep_interval=0, expectation_sweep_interval=0)
assert closed_id in hub2._channels
assert closed_id not in hub2._active_channels
assert open_id in hub2._channels
assert open_id in hub2._active_channels
await hub2.close()