test/network/test_context_vars.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
"""Channel-scoped context variables on ``WorkflowState``.
Pins the contract that ``EV_CONTEXT_SET`` envelopes:
* mutate ``WorkflowState.context_vars`` via ``fold`` (set + delete),
* survive ``Hub.hydrate()`` because they live on the WAL,
* don't advance ``turn_count`` (non-substantive),
* may be sent by any participant regardless of turn order (loose semantics),
* drive transitions via the new ``ContextEquals`` condition.
"""
import pytest
from ag2 import Agent
from ag2.knowledge import MemoryKnowledgeStore
from ag2.network import (
EV_CHANNEL_CLOSED,
EV_CONTEXT_SET,
WORKFLOW_TYPE,
AgentTarget,
ContextEquals,
Envelope,
FromSpeaker,
Hub,
TerminateTarget,
Transition,
TransitionGraph,
)
from ag2.testing import TestConfig
def _agent(name: str, *replies: str) -> Agent:
return Agent(name=name, config=TestConfig(*replies))
async def _post_context_set(hub: Hub, *, sender: str, channel_id: str, **event_data) -> None:
"""Helper: post an EV_CONTEXT_SET envelope through the hub."""
envelope = Envelope(
channel_id=channel_id,
sender_id=sender,
audience=None,
event_type=EV_CONTEXT_SET,
event_data=event_data,
)
await hub.post_envelope(envelope)
@pytest.mark.asyncio
class TestContextVars:
"""``ag2.context.set`` envelopes mutate ``WorkflowState.context_vars``."""
async def test_set_persists_across_fold(self) -> None:
"""One EV_CONTEXT_SET envelope merges into context_vars."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
graph = TransitionGraph(
initial_speaker=alice.agent_id,
transitions=[
Transition(when=FromSpeaker(alice.agent_id), then=AgentTarget(bob.agent_id)),
],
default_target=TerminateTarget(reason="done"),
max_turns=10,
)
channel = await alice.open(
type=WORKFLOW_TYPE,
target=[bob.agent_id],
knobs={"graph": graph.to_dict()},
)
await _post_context_set(
hub,
sender=alice.agent_id,
channel_id=channel.channel_id,
set={"priority": "high", "ticket_id": "T-481"},
)
state = hub._adapter_states[channel.channel_id]
assert state.context_vars == {"priority": "high", "ticket_id": "T-481"}
await hub.close()
async def test_set_then_delete(self) -> None:
"""Subsequent envelopes can both set and remove keys."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
graph = TransitionGraph(
initial_speaker=alice.agent_id,
transitions=[
Transition(when=FromSpeaker(alice.agent_id), then=AgentTarget(bob.agent_id)),
],
default_target=TerminateTarget(reason="done"),
max_turns=10,
)
channel = await alice.open(
type=WORKFLOW_TYPE,
target=[bob.agent_id],
knobs={"graph": graph.to_dict()},
)
await _post_context_set(
hub,
sender=alice.agent_id,
channel_id=channel.channel_id,
set={"priority": "high", "temporary": "1"},
)
await _post_context_set(
hub,
sender=alice.agent_id,
channel_id=channel.channel_id,
delete=["temporary"],
set={"resolved": True},
)
state = hub._adapter_states[channel.channel_id]
assert state.context_vars == {"priority": "high", "resolved": True}
await hub.close()
async def test_loose_semantics_any_participant(self) -> None:
"""Non-current speaker can still post EV_CONTEXT_SET."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
# alice is initial speaker — bob is NOT expected_next_speaker.
graph = TransitionGraph(
initial_speaker=alice.agent_id,
transitions=[
Transition(when=FromSpeaker(alice.agent_id), then=AgentTarget(bob.agent_id)),
],
default_target=TerminateTarget(reason="done"),
max_turns=10,
)
channel = await alice.open(
type=WORKFLOW_TYPE,
target=[bob.agent_id],
knobs={"graph": graph.to_dict()},
)
# bob writes context out-of-turn — should succeed.
await _post_context_set(
hub,
sender=bob.agent_id,
channel_id=channel.channel_id,
set={"observer_flag": True},
)
state = hub._adapter_states[channel.channel_id]
assert state.context_vars == {"observer_flag": True}
await hub.close()
async def test_context_equals_drives_transition(self) -> None:
"""ContextEquals condition routes to a different agent based on
a value set via EV_CONTEXT_SET."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
triage = await hub.register(_agent("triage", "triaged"))
security = await hub.register(_agent("security", "handled"))
legal = await hub.register(_agent("legal", "legal-reviewed"))
# Terminate transitions are listed first so the post-handoff
# speaker hits them before the ContextEquals rule (which would
# otherwise re-fire and loop). Then the ContextEquals override
# for triage's turn, then the default triage→legal route.
graph = TransitionGraph(
initial_speaker=triage.agent_id,
transitions=[
Transition(
when=FromSpeaker(security.agent_id),
then=TerminateTarget(reason="security_done"),
),
Transition(
when=FromSpeaker(legal.agent_id),
then=TerminateTarget(reason="legal_done"),
),
Transition(
when=ContextEquals(key="route", value="security"),
then=AgentTarget(security.agent_id),
),
Transition(
when=FromSpeaker(triage.agent_id),
then=AgentTarget(legal.agent_id),
),
],
default_target=TerminateTarget(reason="fall_through"),
max_turns=10,
)
channel = await triage.open(
type=WORKFLOW_TYPE,
target=[security.agent_id, legal.agent_id],
knobs={"graph": graph.to_dict()},
)
# Triage writes route=security BEFORE its own substantive turn
# so the condition fires when triage's text is folded.
await _post_context_set(
hub,
sender=triage.agent_id,
channel_id=channel.channel_id,
set={"route": "security"},
)
await channel.send("kickoff")
close_env = await triage.wait_for_channel_event(
channel_id=channel.channel_id,
predicate=lambda e: e.event_type == EV_CHANNEL_CLOSED,
timeout=5.0,
)
# Routed to security, not legal.
assert close_env.event_data.get("reason") == "security_done"
await hub.close()
async def test_hydrate_replays_context_vars(self) -> None:
"""Closing and re-opening a hub against the same store
reconstructs ``context_vars`` via WAL replay."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
graph = TransitionGraph(
initial_speaker=alice.agent_id,
transitions=[
Transition(when=FromSpeaker(alice.agent_id), then=AgentTarget(bob.agent_id)),
],
default_target=TerminateTarget(reason="done"),
max_turns=10,
)
channel = await alice.open(
type=WORKFLOW_TYPE,
target=[bob.agent_id],
knobs={"graph": graph.to_dict()},
)
channel_id = channel.channel_id
await _post_context_set(
hub,
sender=alice.agent_id,
channel_id=channel_id,
set={"k1": "v1", "k2": "v2"},
)
await _post_context_set(
hub,
sender=alice.agent_id,
channel_id=channel_id,
delete=["k1"],
set={"k3": "v3"},
)
# Snapshot before close.
before = dict(hub._adapter_states[channel_id].context_vars)
await hub.close()
# Re-open against the same store and re-fold the WAL.
hub2 = await Hub.open(store, ttl_sweep_interval=0)
try:
after = dict(hub2._adapter_states[channel_id].context_vars)
assert after == before
assert after == {"k2": "v2", "k3": "v3"}
finally:
await hub2.close()
async def test_context_set_does_not_advance_turn(self) -> None:
"""EV_CONTEXT_SET must not increment turn_count or rotate the speaker."""
store = MemoryKnowledgeStore()
hub = await Hub.open(store, ttl_sweep_interval=0)
alice = await hub.register(_agent("alice"))
bob = await hub.register(_agent("bob"))
graph = TransitionGraph(
initial_speaker=alice.agent_id,
transitions=[
Transition(when=FromSpeaker(alice.agent_id), then=AgentTarget(bob.agent_id)),
],
default_target=TerminateTarget(reason="done"),
max_turns=10,
)
channel = await alice.open(
type=WORKFLOW_TYPE,
target=[bob.agent_id],
knobs={"graph": graph.to_dict()},
)
before_turn = hub._adapter_states[channel.channel_id].turn_count
before_speaker = hub._adapter_states[channel.channel_id].expected_next_speaker
await _post_context_set(
hub,
sender=alice.agent_id,
channel_id=channel.channel_id,
set={"k": "v"},
)
after = hub._adapter_states[channel.channel_id]
assert after.turn_count == before_turn
assert after.expected_next_speaker == before_speaker
await hub.close()