test/network/test_consulting.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
"""Consulting adapter integration tests.
Covers:
* Single-recipient consulting handshake (invite → auto-ack → ACTIVE).
* Full LLM-driven turn: initiator sends prompt → respondent's notify
handler runs ``Agent.ask`` → respondent replies → adapter
``on_accepted`` returns CLOSING → channel auto-closes with
``close_reason="consulting_complete"``.
* Adapter rejects out-of-order sends (``ProtocolError``).
* ``Hub.hydrate()`` re-folds an active channel's WAL deterministically
through ``adapter.fold`` so the in-memory ``AdapterState`` matches
what's on disk.
This suite uses ``TestConfig`` so it runs offline and fast.
"""
import pytest
from ag2 import Agent
from ag2.events import ToolCallEvent
from ag2.knowledge import DiskKnowledgeStore, MemoryKnowledgeStore
from ag2.network import (
EV_CHANNEL_CLOSED,
EV_CHANNEL_INVITE,
EV_CHANNEL_INVITE_ACK,
EV_CHANNEL_OPENED,
EV_TEXT,
Envelope,
Hub,
)
from ag2.network.adapters.consulting import ConsultingAdapter, ConsultingState
from ag2.network.channel import ChannelState
from ag2.network.errors import ProtocolError
from ag2.testing import TestConfig
def _agent(name: str, *events: object) -> Agent:
return Agent(name=name, config=TestConfig(*events))
@pytest.mark.asyncio
async def test_consulting_handshake_transitions_to_active() -> None:
"""alice.open → hub posts invite → bob auto-acks → channel ACTIVE."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
await hub.register(_agent("bob"))
channel = await alice.open(type="consulting", target="bob")
assert channel.state == ChannelState.ACTIVE
assert len(channel.metadata.participants) == 2
assert channel.metadata.creator_id == alice.agent_id
assert channel.metadata.pending_acks == []
wal = await hub.read_wal(channel.channel_id)
event_types = [e.event_type for e in wal]
assert EV_CHANNEL_INVITE in event_types
assert EV_CHANNEL_INVITE_ACK in event_types
assert EV_CHANNEL_OPENED in event_types
await hub.close()
@pytest.mark.asyncio
async def test_consulting_full_flow_auto_closes() -> None:
"""Initiator sends prompt → respondent replies via default handler → CLOSED."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(
_agent("bob", "Hi, this is bob's reply."),
)
channel = await alice.open(type="consulting", target="bob")
await channel.send("Hello bob, can you help?", audience=[bob.agent_id])
# Wait for channel close to propagate to alice's inbox.
close_envelope = await alice.wait_for_channel_event(
channel_id=channel.channel_id,
predicate=lambda e: e.event_type == EV_CHANNEL_CLOSED,
timeout=5.0,
)
assert close_envelope.event_data.get("reason") == "consulting_complete"
final = await hub.get_channel(channel.channel_id)
assert final.state == ChannelState.CLOSED
assert final.close_reason == "consulting_complete"
wal = await hub.read_wal(channel.channel_id)
text_events = [e for e in wal if e.event_type == EV_TEXT]
assert len(text_events) == 2
assert text_events[0].sender_id == alice.agent_id
assert text_events[0].event_data["text"] == "Hello bob, can you help?"
assert text_events[1].sender_id == bob.agent_id
assert text_events[1].event_data["text"] == "Hi, this is bob's reply."
await hub.close()
@pytest.mark.asyncio
async def test_consulting_rejects_out_of_order_send() -> None:
"""Respondent cannot send before initiator's first envelope."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
channel = await alice.open(type="consulting", target="bob")
# Bob tries to send first — adapter should reject.
bad = Envelope(
channel_id=channel.channel_id,
sender_id=bob.agent_id,
audience=[alice.agent_id],
event_type=EV_TEXT,
event_data={"text": "I'm jumping the queue"},
)
with pytest.raises(ProtocolError, match="initiator"):
await hub.post_envelope(bad)
await hub.close()
@pytest.mark.asyncio
async def test_consulting_rejects_send_after_complete() -> None:
"""Adapter rejects any send after both turns have happened."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(
_agent("bob", "ok"),
)
channel = await alice.open(type="consulting", target="bob")
await channel.send("the question", audience=[bob.agent_id])
await alice.wait_for_channel_event(
channel_id=channel.channel_id,
predicate=lambda e: e.event_type == EV_CHANNEL_CLOSED,
timeout=5.0,
)
# Channel is now closed — any further post should fail.
extra = Envelope(
channel_id=channel.channel_id,
sender_id=alice.agent_id,
audience=[bob.agent_id],
event_type=EV_TEXT,
event_data={"text": "follow-up"},
)
with pytest.raises(ProtocolError):
await hub.post_envelope(extra)
await hub.close()
@pytest.mark.asyncio
async def test_hub_hydrate_refolds_active_channel(tmp_path) -> None:
"""Close hub mid-flight, reopen, verify adapter state survives."""
store = DiskKnowledgeStore(str(tmp_path))
hub1 = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub1.register(_agent("alice"))
bob = await hub1.register(_agent("bob"))
channel = await alice.open(type="consulting", target="bob")
await channel.send("the question", audience=[bob.agent_id])
# Don't wait for bob's reply — close hub mid-flight.
await hub1.close()
# Reopen with a fresh hub against the same store.
store2 = DiskKnowledgeStore(str(tmp_path))
hub2 = await Hub.open(store2, ttl_sweep_interval=0)
# Channel metadata reloaded.
refreshed = await hub2.get_channel(channel.channel_id)
assert refreshed.channel_id == channel.channel_id
assert refreshed.manifest.type == "consulting"
# Adapter state cache rebuilt by re-folding the WAL.
state = hub2._adapter_states[channel.channel_id]
assert isinstance(state, ConsultingState)
# Alice has sent her prompt; bob has not replied yet.
assert state.initiator_sent is True
assert state.respondent_replied is False
await hub2.close()
@pytest.mark.asyncio
async def test_consulting_invite_timeout_when_no_handler() -> None:
"""If bob has no auto-ack handler, hub times out and fails the open."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0, invite_ack_timeout=0.2)
alice = await hub.register(_agent("alice"))
# Bob registers with attach_plugin=False to skip the default handler.
bob = await hub.register(
_agent("bob"),
attach_plugin=False,
)
# And clear the auto-installed default handler — bob ignores invites.
bob.on_envelope(_silent_handler)
with pytest.raises(ProtocolError, match="ack timeout"):
await alice.open(type="consulting", target="bob")
await hub.close()
async def _silent_handler(_envelope: Envelope) -> None:
"""No-op handler — used to test invite timeout behaviour."""
@pytest.mark.asyncio
async def test_default_consulting_adapter_registered_on_open() -> None:
"""``Hub.open`` auto-registers ConsultingAdapter for ``consulting@v1``."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
adapter = hub._adapters.get(("consulting", 1))
assert isinstance(adapter, ConsultingAdapter)
await hub.close()
@pytest.mark.asyncio
async def test_delegate_tool_end_to_end() -> None:
"""End-to-end: Alice's LLM uses ``delegate`` to consult Bob.
Alice's TestConfig delivers a ``delegate`` tool call followed by a
final user-facing response. Bob's TestConfig delivers a single
string reply. The full chain — open consulting, invite/ack, send
prompt, run Bob's LLM via the default handler, return Bob's reply
to Alice, Alice's second LLM call to incorporate the reply — runs
without any real LLM calls.
"""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice_agent = Agent(
name="alice",
config=TestConfig(
[
ToolCallEvent(
name="delegate",
arguments='{"target": "bob", "prompt": "what is 2+2?"}',
),
],
"The answer is 4.",
),
)
bob_agent = Agent(name="bob", config=TestConfig("4"))
await hub.register(alice_agent)
await hub.register(bob_agent)
reply = await alice_agent.ask("ask bob to do math for me")
assert reply.body == "The answer is 4."
await hub.close()