test/network/test_observation.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
"""Resume observation, capability index, skill render, mutation tests.
Three layers covered:
* **Skill render** (unit) — ``parse_skill_frontmatter`` + ``render_fallback_skill``.
* **Capability index + record_observation** (integration) — register
populates the index; ``record_observation`` updates ``Resume.observed``
and adds the agent to the capability bucket; ``unregister`` removes
the agent and prunes empty buckets; the index round-trips through
``Hub.hydrate()``; the on-disk JSON cache reflects the in-memory state.
* **AgentClient mutation** — ``set_resume`` / ``add_example`` go through
the hub and refresh the local cache.
* **TaskMirror end-to-end** — ``Task(spec=TaskSpec(title="...", capability="X"), ...)`` inside
a notify-handler turn drives ``Hub.record_observation`` automatically.
"""
import json
import tempfile
from pathlib import Path
import pytest
from ag2 import Agent, Context
from ag2.knowledge import DiskKnowledgeStore, MemoryKnowledgeStore
from ag2.network import (
Hub,
HubClient,
LocalLink,
Passport,
ResumeExample,
)
from ag2.network.client.skill_render import (
parse_skill_frontmatter,
render_fallback_skill,
)
from ag2.network.hub.audit import (
AUDIT_KIND_RESUME_SET,
RESUME_SOURCE_OBSERVED,
RESUME_SOURCE_TENANT,
)
from ag2.network.hub.layout import by_capability_path
from ag2.network.identity import (
ObservedStat,
Resume,
)
from ag2.network.task_mirror import TaskMirror
from ag2.stream import MemoryStream
from ag2.task import Task, TaskSpec, TaskState
from ag2.testing import TestConfig
from ._helpers import ScriptedConfig
def _agent(name: str, *events: object) -> Agent:
return Agent(name=name, config=TestConfig(*events))
# ── Skill render unit tests ─────────────────────────────────────────────────
class TestParseSkillFrontmatter:
def test_parses_basic_frontmatter_and_body(self) -> None:
md = "---\nname: alice\ndescription: Senior policy analyst.\n---\n\n## What I do\n\nCost-benefit framing.\n"
parsed = parse_skill_frontmatter(md)
assert parsed.frontmatter == {
"name": "alice",
"description": "Senior policy analyst.",
}
assert "## What I do" in parsed.body
assert parsed.body.startswith("\n## What I do") or parsed.body.startswith("## What I do")
def test_no_frontmatter_returns_full_body(self) -> None:
md = "# alice\n\njust a body\n"
parsed = parse_skill_frontmatter(md)
assert parsed.frontmatter == {}
assert parsed.body == md
def test_unterminated_frontmatter_returns_full_body(self) -> None:
md = "---\nname: alice\nbut no closing fence"
parsed = parse_skill_frontmatter(md)
assert parsed.frontmatter == {}
assert parsed.body == md
def test_skips_empty_and_comment_lines(self) -> None:
md = "---\nname: alice\n\n# a comment\nrole: analyst\n---\nbody"
parsed = parse_skill_frontmatter(md)
assert parsed.frontmatter == {"name": "alice", "role": "analyst"}
class TestRenderFallbackSkill:
def test_includes_capabilities_domains_and_summary(self) -> None:
passport = Passport(name="alice")
resume = Resume(
claimed_capabilities=["debate", "analysis"],
domains=["policy", "economics"],
summary="Senior policy analyst.",
)
rendered = render_fallback_skill(passport, resume)
assert rendered.startswith("---\n")
assert "name: alice" in rendered
assert "description: Senior policy analyst." in rendered
assert "## Capabilities" in rendered
assert "- debate" in rendered
assert "## Domains" in rendered
def test_includes_observed_track_record(self) -> None:
passport = Passport(name="bob")
resume = Resume(
claimed_capabilities=["analysis"],
observed={
"analysis": ObservedStat(n=3, completed=2, failed=1),
},
)
rendered = render_fallback_skill(passport, resume)
assert "## Track record" in rendered
assert "analysis" in rendered
assert "n=3" in rendered
def test_minimal_resume_renders_default_description(self) -> None:
passport = Passport(name="solo")
resume = Resume()
rendered = render_fallback_skill(passport, resume)
assert "name: solo" in rendered
assert "Network-registered agent." in rendered
# ── Capability index integration tests ──────────────────────────────────────
@pytest.mark.asyncio
async def test_register_populates_capability_index() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(
_agent("alice"),
resume=Resume(claimed_capabilities=["debate", "analysis"]),
)
assert hub.agents_with_capability("debate") == [alice.agent_id]
assert hub.agents_with_capability("analysis") == [alice.agent_id]
assert hub.agents_with_capability("missing") == []
await hub.close()
@pytest.mark.asyncio
async def test_unregister_removes_from_capability_index() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"), resume=Resume(claimed_capabilities=["debate"]))
bob = await hub.register(_agent("bob"), resume=Resume(claimed_capabilities=["debate"]))
assert set(hub.agents_with_capability("debate")) == {alice.agent_id, bob.agent_id}
await hub.unregister(alice.agent_id)
assert hub.agents_with_capability("debate") == [bob.agent_id]
await hub.unregister(bob.agent_id)
assert hub.agents_with_capability("debate") == []
# Empty bucket pruned from the index entirely.
assert "debate" not in hub._capability_index
await hub.close()
@pytest.mark.asyncio
async def test_capability_index_persisted_to_disk(tmp_path) -> None:
store = DiskKnowledgeStore(str(tmp_path))
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(
_agent("alice"),
resume=Resume(claimed_capabilities=["debate"]),
)
raw = await store.read(by_capability_path())
assert raw is not None
snapshot = json.loads(raw)
assert snapshot == {"debate": [alice.agent_id]}
await hub.close()
@pytest.mark.asyncio
async def test_capability_index_rebuilt_on_hydrate(tmp_path) -> None:
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=["debate"],
observed={"reviews": ObservedStat(n=5, completed=4, failed=1)},
),
)
await hub1.close()
store2 = DiskKnowledgeStore(str(tmp_path))
hub2 = await Hub.open(store2, ttl_sweep_interval=0, expectation_sweep_interval=0)
# Both claimed and observed capabilities show up in the rebuilt index.
assert hub2.agents_with_capability("debate") == [alice.agent_id]
assert hub2.agents_with_capability("reviews") == [alice.agent_id]
await hub2.close()
# ── record_observation tests ────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_record_observation_updates_resume_observed_counters() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(
_agent("alice"),
resume=Resume(claimed_capabilities=["analysis"]),
)
await hub.record_observation(
owner_id=alice.agent_id,
capability="analysis",
outcome=TaskState.COMPLETED,
latency_ms=420,
)
await hub.record_observation(
owner_id=alice.agent_id,
capability="analysis",
outcome=TaskState.FAILED,
)
await hub.record_observation(
owner_id=alice.agent_id,
capability="analysis",
outcome=TaskState.EXPIRED,
)
resume = await hub.get_resume(alice.agent_id)
stat = resume.observed["analysis"]
assert stat.n == 3
assert stat.completed == 1
assert stat.failed == 1
assert stat.expired == 1
assert stat.p50_latency_ms == 420 # last observed value (V1 placeholder)
await hub.close()
@pytest.mark.asyncio
async def test_record_observation_adds_unclaimed_capability_to_index() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice")) # no claims
assert hub.agents_with_capability("emergent") == []
await hub.record_observation(
owner_id=alice.agent_id,
capability="emergent",
outcome=TaskState.COMPLETED,
)
assert hub.agents_with_capability("emergent") == [alice.agent_id]
await hub.close()
@pytest.mark.asyncio
async def test_record_observation_ignores_non_terminal_state() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"))
await hub.record_observation(
owner_id=alice.agent_id,
capability="x",
outcome=TaskState.RUNNING,
)
resume = await hub.get_resume(alice.agent_id)
assert "x" not in resume.observed
await hub.close()
# ── AgentClient mutation tests ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_agent_client_set_resume_refreshes_local_cache() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"), resume=Resume(summary="initial"))
await alice.set_resume(Resume(summary="updated", claimed_capabilities=["x"]))
assert alice.resume.summary == "updated"
assert "x" in alice.resume.claimed_capabilities
# ``set_resume`` re-indexes claimed_capabilities so newly-claimed
# caps surface under ``peers(action="find", capability=...)``.
assert hub.agents_with_capability("x") == [alice.agent_id]
# Removing a claim drops the agent from that bucket so the index
# stays consistent with the current resume.
await alice.set_resume(Resume(summary="updated2", claimed_capabilities=[]))
assert hub.agents_with_capability("x") == []
await hub.close()
@pytest.mark.asyncio
async def test_agent_client_add_example_appends() -> None:
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
alice = await hub.register(_agent("alice"))
assert alice.resume.examples == []
await alice.add_example(ResumeExample(title="reviewed PR #42", outcome="completed"))
await alice.add_example(ResumeExample(title="triaged incident #7", outcome="completed"))
fresh = await hub.get_resume(alice.agent_id)
titles = [e.title for e in fresh.examples]
assert titles == ["reviewed PR #42", "triaged incident #7"]
await hub.close()
# ── TaskMirror end-to-end via Task lifecycle ────────────────────────────────
@pytest.mark.asyncio
async def test_task_mirror_records_observation_on_capability_tagged_task() -> None:
"""Running ``Task(spec=TaskSpec(title="...", capability=X), ...)`` to completion through a mirror
auto-calls ``Hub.record_observation`` and updates ``Resume.observed[X]``.
Exercises the mirror plumbing directly (without the LLM tool
surface) so this contract is verified independently of the
end-to-end notify-handler integration covered in test_tools.py.
"""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
bob_agent = Agent(name="bob", config=ScriptedConfig("ack"))
bob = await hub.register(
bob_agent,
resume=Resume(claimed_capabilities=["analysis"]),
)
stream = MemoryStream()
mirror = TaskMirror(hub=hub, owner_id=bob.agent_id)
sub_ids = mirror.attach(stream)
try:
async with Task(
owner_id=bob_agent.name,
spec=TaskSpec(title="analysing alice's question", capability="analysis"),
context=Context(stream=stream),
) as task:
await task.complete(result="done")
finally:
mirror.detach(stream, sub_ids)
fresh = await hub.get_resume(bob.agent_id)
assert "analysis" in fresh.observed
stat = fresh.observed["analysis"]
assert stat.n == 1
assert stat.completed == 1
assert bob.agent_id in hub.agents_with_capability("analysis")
await hub.close()
@pytest.mark.asyncio
async def test_record_observation_writes_audit_with_observed_source() -> None:
"""Hub-side observation mutations are auditable as ``resume_set``
records with ``source="observed"``, distinct from tenant-driven
``set_resume`` calls (``source="tenant"``)."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
link = LocalLink(hub)
bob_hc = HubClient(link, hub=hub)
bob = await bob_hc.register(
_agent("bob"),
Passport(name="bob"),
Resume(claimed_capabilities=["analysis"]),
)
pre_audit = len(await hub._audit_log.read_all())
# Tenant-driven update.
await bob_hc._hub.set_resume(bob.agent_id, Resume(summary="updated by tenant"))
# Hub-driven observation.
await hub.record_observation(
owner_id=bob.agent_id,
capability="analysis",
outcome=TaskState.COMPLETED,
latency_ms=42,
)
audit = await hub._audit_log.read_all()
new_records = audit[pre_audit:]
resume_records = [r for r in new_records if r["kind"] == AUDIT_KIND_RESUME_SET]
sources = [r.get("source") for r in resume_records]
assert RESUME_SOURCE_TENANT in sources
assert RESUME_SOURCE_OBSERVED in sources
observed = next(r for r in resume_records if r.get("source") == RESUME_SOURCE_OBSERVED)
assert observed["agent_id"] == bob.agent_id
assert observed["capability"] == "analysis"
assert observed["outcome"] == TaskState.COMPLETED.value
await bob_hc.close()
await hub.close()
@pytest.mark.asyncio
async def test_task_capability_survives_hub_hydrate() -> None:
"""``TaskSpec.capability`` round-trips through hub persistence so an
observation can fire on a terminal event after a hub restart.
Pre-fix: ``_task_metadata_to_dict`` dropped ``capability``; on
hydrate the spec came back with ``capability=None`` and
``record_observation`` wouldn't fire even after the task terminated.
"""
with tempfile.TemporaryDirectory() as tmpdir:
store = DiskKnowledgeStore(Path(tmpdir))
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
bob_agent = Agent(name="bob", config=ScriptedConfig("ack"))
bob = await hub.register(
bob_agent,
resume=Resume(claimed_capabilities=["analysis"]),
)
# Start a capability-tagged task; mirror persists metadata.
stream = MemoryStream()
mirror = TaskMirror(hub=hub, owner_id=bob.agent_id)
sub_ids = mirror.attach(stream)
async with Task(
owner_id=bob_agent.name,
spec=TaskSpec(title="analysing", capability="analysis"),
context=Context(stream=stream),
) as task:
task_id = task.task_id
await task.complete(result="done")
mirror.detach(stream, sub_ids)
await hub.close()
# Restart: new Hub, same store. ``_load_task`` rehydrates
# TaskMetadata; the spec must preserve ``capability``.
hub2 = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
rehydrated = hub2._tasks[task_id]
assert rehydrated.spec.capability == "analysis"
await hub2.close()
@pytest.mark.asyncio
async def test_task_mirror_no_observation_when_capability_absent() -> None:
"""Untagged tasks emit lifecycle events but don't touch ``observed``."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, expectation_sweep_interval=0)
bob_agent = Agent(name="bob", config=ScriptedConfig("ack"))
bob = await hub.register(bob_agent)
stream = MemoryStream()
mirror = TaskMirror(hub=hub, owner_id=bob.agent_id)
sub_ids = mirror.attach(stream)
try:
async with Task(
owner_id=bob_agent.name, spec=TaskSpec(title="untagged work"), context=Context(stream=stream)
) as task:
await task.complete(result="ok")
finally:
mirror.detach(stream, sub_ids)
fresh = await hub.get_resume(bob.agent_id)
assert fresh.observed == {}
await hub.close()