test/tools/test_delegation_events.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
import asyncio
from dataclasses import dataclass, field
from typing import Any
import pytest
from dirty_equals import IsPartialDict
from ag2 import Agent, Context, MemoryStream, Scope, ScopedValue, observer
from ag2.compact import CompactTrigger, TailWindowCompact
from ag2.events import (
BaseEvent,
ChildEvent,
DeliveryFailed,
DownEvent,
HaltEvent,
HumanInputRequest,
HumanMessage,
ModelMessage,
ModelMessageChunk,
ModelResponse,
TaskCancelled,
TaskStarted,
ToolCallEvent,
Usage,
UsageEvent,
)
from ag2.knowledge import KnowledgePlugin, MemoryKnowledgeStore
from ag2.testing import TestConfig
from ag2.tools.subagents import DeliveryError, background_agent_tool
from ag2.tools.subagents.run_task import run_task
from ag2.usage import UsageReport
@dataclass
class Journal:
events: list[BaseEvent] = field(default_factory=list)
async def record(self, event: BaseEvent) -> None:
self.events.append(event)
@pytest.mark.asyncio
async def test_root_observes_original_leaf_and_counts_each_call_once() -> None:
leaf = Agent(
"leaf",
config=TestConfig(
ModelMessageChunk("working"),
ModelResponse(ModelMessage("leaf done"), usage=Usage(total_tokens=30), model="leaf-model"),
),
)
worker = Agent(
"worker",
config=TestConfig(
ToolCallEvent(name="task_leaf", arguments='{"objective": "dig"}'),
ModelResponse(ModelMessage("worker done"), usage=Usage(total_tokens=20), model="worker-model"),
),
tools=[leaf.as_tool(description="Dig")],
)
root = Agent(
"root",
config=TestConfig(
ToolCallEvent(name="task_worker", arguments='{"objective": "organize"}'),
ModelResponse(ModelMessage("root done"), usage=Usage(total_tokens=10), model="root-model"),
),
tools=[worker.as_tool(description="Organize")],
)
journal = Journal()
reply = await root.ask("start", stream=MemoryStream(), observers=[observer(callback=journal.record)])
chunks = [e for e in journal.events if isinstance(e, ChildEvent) and isinstance(e.inner, ModelMessageChunk)]
assert len(chunks) == 1
assert chunks[0].agent_name == "leaf"
assert len(chunks[0].path) == 2
assert chunks[0].origin.agent_name == "leaf"
history = list(await reply.history.get_events())
assert not any(isinstance(e, ChildEvent) and isinstance(e.inner, ModelMessageChunk) for e in history)
assert not any(isinstance(e, UsageEvent) and e.kind == "subtask" for e in history)
report = await reply.usage()
assert report.total.total_tokens == 60
assert {model: usage.total_tokens for model, usage in report.by_model.items()} == {
"leaf-model": 30,
"worker-model": 20,
"root-model": 10,
}
class WorkSignal(BaseEvent):
__conversational__ = False
class Directive(BaseEvent):
__conversational__ = False
text: str
@dataclass
class Guidance:
origins: list[Any] = field(default_factory=list)
async def monitor(self, event: BaseEvent) -> None:
if isinstance(event, ChildEvent) and isinstance(event.inner, WorkSignal):
self.origins.append(event.origin)
await event.origin.send(Directive(text="keep digging"))
async def receive(self, event: BaseEvent, ctx: Context) -> None:
if isinstance(event, Directive):
ctx.variables["guidance"] = event.text
async def seek_guidance(ctx: Context) -> str:
"""Receive guidance during this tool's execution."""
await ctx.send(WorkSignal())
return ctx.variables["guidance"]
@pytest.mark.asyncio
async def test_origin_addresses_the_leaf_without_echo_and_expires_after_completion() -> None:
guidance, journal = Guidance(), Journal()
leaf = Agent(
"leaf",
config=TestConfig(ToolCallEvent(name="seek_guidance", arguments="{}"), "done"),
tools=[seek_guidance],
observers=[observer(callback=guidance.receive)],
)
worker = Agent(
"worker",
config=TestConfig(ToolCallEvent(name="task_leaf", arguments='{"objective":"dig"}'), "done"),
tools=[leaf.as_tool(description="Dig")],
)
root = Agent(
"root",
config=TestConfig(ToolCallEvent(name="task_worker", arguments='{"objective":"organize"}'), "done"),
tools=[worker.as_tool(description="Organize")],
)
stream = MemoryStream()
stream.subscribe(journal.record)
await root.ask("start", stream=stream, observers=[observer(callback=guidance.monitor)])
assert [origin.agent_name for origin in guidance.origins] == ["leaf"]
assert not any(isinstance(e, ChildEvent) and isinstance(e.inner, Directive) for e in journal.events)
with pytest.raises(DeliveryError):
await guidance.origins[0].send(Directive(text="too late"))
assert any(
isinstance(e, ChildEvent) and isinstance(e.inner, DeliveryFailed) or isinstance(e, DeliveryFailed)
for e in journal.events
)
async def long_work(ctx: Context) -> str:
"""A cancellable operation that must finish cleanup before its terminal."""
try:
await ctx.send(WorkSignal())
await asyncio.Event().wait()
finally:
ctx.variables["cleanup"].append("cleaned")
return "unreachable"
async def stop_on_signal(event: BaseEvent, ctx: Context) -> None:
if isinstance(event, ChildEvent) and isinstance(event.inner, WorkSignal):
sibling = ctx.dependencies["sibling_gate"]
await sibling.entered.wait()
try:
await event.origin.send(HaltEvent(reason="budget", source="test"))
finally:
sibling.release.set()
@pytest.mark.asyncio
@pytest.mark.parametrize("same_name", [False, True])
async def test_targeted_halt_cancels_running_tool_without_aborting_parent_or_sibling(same_name: bool) -> None:
child = Agent(
"worker" if same_name else "child",
config=TestConfig(ToolCallEvent(name="long_work", arguments="{}")),
tools=[long_work],
)
sibling_gate = Gate()
sibling = Agent(
"worker" if same_name else "sibling",
config=TestConfig(ToolCallEvent(name="run", arguments="{}"), "sibling done"),
tools=[sibling_gate.run],
)
root = Agent(
"root",
config=TestConfig(
[
ToolCallEvent(name="cancel_work", arguments='{"objective":"work"}'),
ToolCallEvent(name="keep_work", arguments='{"objective":"work"}'),
],
"root survived",
),
tools=[
child.as_tool(description="Work", name="cancel_work"),
sibling.as_tool(description="Work", name="keep_work"),
],
dependencies={"sibling_gate": sibling_gate},
variables={"cleanup": ScopedValue([], scope=Scope.SHARED)},
)
journal = Journal()
reply = await asyncio.wait_for(
root.ask("start", observers=[observer(callback=stop_on_signal), observer(callback=journal.record)]), 2
)
assert reply.body == "root survived"
assert sibling_gate.received == [False]
assert reply.context.variables["cleanup"] == ["cleaned"]
cancelled = [e.inner if isinstance(e, ChildEvent) else e for e in journal.events]
assert [e.agent_name for e in cancelled if isinstance(e, TaskCancelled)] == [child.name]
async def ask_human(ctx: Context) -> str:
"""Ask the owner of the human-input channel."""
return await ctx.input("May I continue?", timeout=1)
def human_answer(event: HumanInputRequest) -> str:
return "yes"
@pytest.mark.asyncio
@pytest.mark.parametrize("own_hook", [False, True])
async def test_root_audits_one_question_and_answer_from_leaf_with_either_hook_owner(own_hook: bool) -> None:
leaf = Agent(
"leaf",
config=TestConfig(ToolCallEvent(name="ask_human", arguments="{}"), "done"),
tools=[ask_human],
hitl_hook=human_answer if own_hook else None,
)
worker = Agent(
"worker",
config=TestConfig(ToolCallEvent(name="task_leaf", arguments='{"objective":"ask"}'), "done"),
tools=[leaf.as_tool(description="Ask")],
)
root = Agent(
"root",
config=TestConfig(ToolCallEvent(name="task_worker", arguments='{"objective":"ask"}'), "done"),
tools=[worker.as_tool(description="Ask")],
hitl_hook=human_answer,
)
journal = Journal()
await root.ask("start", observers=[observer(callback=journal.record)])
requests = [e.inner for e in journal.events if isinstance(e, ChildEvent) and isinstance(e.inner, HumanInputRequest)]
answers = [e.inner for e in journal.events if isinstance(e, ChildEvent) and isinstance(e.inner, HumanMessage)]
assert len(requests) == len(answers) == 1
assert answers[0].parent_id == requests[0].id
assert any(isinstance(e, TaskStarted) and e.agent_name == "leaf" for e in journal.events)
@dataclass
class Gate:
entered: asyncio.Event = field(default_factory=asyncio.Event)
release: asyncio.Event = field(default_factory=asyncio.Event)
received: list[bool] = field(default_factory=list)
async def run(self, ctx: Context) -> str:
"""Wait until the test releases this public tool."""
self.entered.set()
await self.release.wait()
self.received.append("guidance" in ctx.variables)
return "done"
@dataclass
class HaltCleanup:
entered: asyncio.Event = field(default_factory=asyncio.Event)
cleaned: asyncio.Event = field(default_factory=asyncio.Event)
async def work(self) -> str:
"""Finish cleanup after an externally delivered Halt."""
self.entered.set()
try:
await asyncio.Event().wait()
finally:
self.cleaned.set()
return "unreachable"
async def observe(self, event: BaseEvent) -> None:
if isinstance(event, HaltEvent):
await self.cleaned.wait()
@pytest.mark.asyncio
async def test_successfully_delivered_halt_remains_successful_during_recipient_cleanup() -> None:
cleanup, journal = HaltCleanup(), Journal()
ctx = Context(MemoryStream())
ctx.stream.subscribe(journal.record)
child = Agent(
"child",
config=TestConfig(ToolCallEvent(name="work", arguments="{}")),
tools=[cleanup.work],
observers=[observer(callback=cleanup.observe), observer(callback=journal.record)],
)
task = asyncio.create_task(run_task(child, "work", parent_context=ctx, task_id="child"))
await cleanup.entered.wait()
handle = next(event.origin for event in journal.events if isinstance(event, ChildEvent))
await handle.send(HaltEvent(reason="stop", source="test"))
result = await task
assert result.cancelled
assert cleanup.cleaned.is_set()
assert not any(isinstance(event, DeliveryFailed) for event in await ctx.stream.history.get_events())
@pytest.mark.asyncio
async def test_queued_delegation_has_no_origin_and_does_not_capture_active_turn_delivery() -> None:
gate = Gate()
guidance = Guidance()
worker = Agent(
"worker",
config=TestConfig(ToolCallEvent(name="run", arguments="{}"), "done"),
tools=[gate.run],
observers=[observer(callback=guidance.receive)],
)
ctx = Context(MemoryStream())
journal = Journal()
ctx.stream.subscribe(journal.record)
shared_stream = MemoryStream(storage=ctx.stream.history.storage)
first = asyncio.create_task(run_task(worker, "first", parent_context=ctx, stream=shared_stream, task_id="first"))
await gate.entered.wait()
second = asyncio.create_task(run_task(worker, "second", parent_context=ctx, stream=shared_stream, task_id="second"))
await asyncio.sleep(0)
try:
observations = [event for event in journal.events if isinstance(event, ChildEvent)]
assert {event.task_id for event in observations} == {"first"}
await observations[0].origin.send(Directive(text="first only"))
finally:
gate.release.set()
results = await asyncio.gather(first, second)
assert all(result.completed for result in results)
assert gate.received == [True, False]
@pytest.mark.asyncio
async def test_descendant_accounting_survives_compaction_and_serialized_replay() -> None:
child = Agent(
"child", config=TestConfig(ModelResponse(ModelMessage("done"), usage=Usage(total_tokens=30), model="child"))
)
root = Agent(
"root",
config=TestConfig(ToolCallEvent(name="task_child", arguments='{"objective":"work"}'), "done"),
tools=[child.as_tool(description="Work")],
plugins=[
KnowledgePlugin(
store=MemoryKnowledgeStore(),
compact=TailWindowCompact(target=1),
compact_trigger=CompactTrigger(max_events=1),
expose_tool=False,
write_event_log=False,
)
],
)
reply = await root.ask("start")
history = list(await reply.history.get_events())
replayed = [type(event).from_dict(event.to_dict()) for event in history]
assert UsageReport.from_events(replayed).total.total_tokens == 30
assert UsageReport.from_events([*history, *replayed]).total.total_tokens == 30
@dataclass
class CancellationState:
ready: asyncio.Event = field(default_factory=asyncio.Event)
context: Context | None = None
signals: int = 0
async def record(self, event: BaseEvent, ctx: Context) -> None:
inner = event.inner if isinstance(event, ChildEvent) else event
if isinstance(inner, WorkSignal):
self.context = ctx
self.signals += 1
if self.signals == 2:
self.ready.set()
@pytest.mark.asyncio
async def test_parent_cancellation_settles_background_descendants_before_returning() -> None:
child = Agent("child", config=TestConfig(ToolCallEvent(name="long_work", arguments="{}")), tools=[long_work])
root = Agent(
"root",
config=TestConfig([
ToolCallEvent(name="background_task_child", arguments='{"objective":"work"}'),
ToolCallEvent(name="long_work", arguments="{}"),
]),
tools=[background_agent_tool(child, description="Work"), long_work],
variables={"cleanup": ScopedValue([], scope=Scope.SHARED)},
)
state = CancellationState()
turn = asyncio.create_task(root.ask("start", observers=[observer(callback=state.record)]))
await asyncio.wait_for(state.ready.wait(), 2)
turn.cancel()
with pytest.raises(asyncio.CancelledError):
await turn
assert state.context is not None
assert state.context.variables["cleanup"] == ["cleaned", "cleaned"]
@pytest.mark.asyncio
async def test_republishing_a_usage_template_creates_distinct_source_publications() -> None:
template = UsageEvent(Usage(total_tokens=10))
child = Agent("child", config=TestConfig(template, "done"))
root = Agent(
"root",
config=TestConfig(
[
ToolCallEvent(name="task_child", arguments='{"objective":"one"}'),
ToolCallEvent(name="task_child", arguments='{"objective":"two"}'),
],
"done",
),
tools=[child.as_tool(description="Work")],
)
reply = await root.ask("start")
assert (await reply.usage()).total.total_tokens == 20
@dataclass
class TurnRace:
raw: bool = False
first_entered: asyncio.Event = field(default_factory=asyncio.Event)
second_entered: asyncio.Event = field(default_factory=asyncio.Event)
release_first: asyncio.Event = field(default_factory=asyncio.Event)
release_second: asyncio.Event = field(default_factory=asyncio.Event)
async def hold(self, ctx: Context) -> str:
"""Hold each persistent turn at a public tool boundary."""
if ctx.variables["number"] == 1:
self.first_entered.set()
await self.release_first.wait()
else:
self.second_entered.set()
await self.release_second.wait()
return "done"
async def advance_during_audit(self, event: BaseEvent) -> BaseEvent:
if (self.raw and isinstance(event, HaltEvent)) or (
not self.raw and isinstance(event, DownEvent) and event.task_id == "first"
):
self.release_first.set()
await self.second_entered.wait()
return event
@pytest.mark.asyncio
@pytest.mark.parametrize("stage", ["sender", "recipient_audit", "recipient_event"])
async def test_downlink_revalidates_its_recipient_after_an_awaiting_audit(stage: str) -> None:
race = TurnRace(raw=stage == "recipient_event")
ctx = Context(MemoryStream(), variables={"number": 1})
journal = Journal()
ctx.stream.subscribe(journal.record)
stream = MemoryStream(storage=ctx.stream.history.storage)
if stage == "sender":
ctx.stream.subscribe(race.advance_during_audit)
else:
stream.subscribe(race.advance_during_audit, interrupt=True)
worker = Agent("worker", config=TestConfig(ToolCallEvent(name="hold", arguments="{}"), "done"), tools=[race.hold])
first = asyncio.create_task(run_task(worker, "first", parent_context=ctx, stream=stream, task_id="first"))
await race.first_entered.wait()
handle = next(
event.origin for event in journal.events if isinstance(event, TaskStarted) and event.task_id == "first"
)
assert handle is not None
ctx.variables["number"] = 2
second = asyncio.create_task(run_task(worker, "second", parent_context=ctx, stream=stream, task_id="second"))
try:
with pytest.raises(DeliveryError):
await handle.send(HaltEvent(reason="stop first", source="test"))
finally:
race.release_first.set()
race.release_second.set()
results = await asyncio.gather(first, second)
assert all(result.completed for result in results)
async def halt_on_spend(event: UsageEvent) -> None:
await event.origin.send(HaltEvent(reason="budget spent", source="test"))
@pytest.mark.asyncio
@pytest.mark.parametrize("emit_events", [True, False])
async def test_usage_retained_at_source_when_upward_delivery_cancels_it(emit_events: bool) -> None:
ctx = Context(MemoryStream())
ctx.stream.where(UsageEvent).subscribe(halt_on_spend)
child = Agent(
"child", config=TestConfig(ModelResponse(ModelMessage("done"), usage=Usage(total_tokens=30), model="child"))
)
result = await run_task(child, "work", parent_context=ctx, emit_events=emit_events)
assert result.cancelled
assert result.usage == Usage(total_tokens=30)
assert UsageReport.from_events(await result.stream.history.get_events()).total == Usage(total_tokens=30)
assert UsageReport.from_events(await ctx.stream.history.get_events()).total == Usage(total_tokens=30)
@dataclass
class UsagePause:
entered: asyncio.Event = field(default_factory=asyncio.Event)
release: asyncio.Event = field(default_factory=asyncio.Event)
async def pause(self, event: UsageEvent) -> UsageEvent:
self.entered.set()
await self.release.wait()
return event
@pytest.mark.asyncio
async def test_parent_cancellation_retains_usage_while_its_interrupter_is_awaiting() -> None:
pause = UsagePause()
ctx = Context(MemoryStream())
ctx.stream.where(UsageEvent).subscribe(pause.pause, interrupt=True)
stream = MemoryStream(storage=ctx.stream.history.storage)
child = Agent("child", config=TestConfig(UsageEvent(Usage(total_tokens=30)), "done"))
task = asyncio.create_task(run_task(child, "work", parent_context=ctx, stream=stream))
await pause.entered.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert UsageReport.from_events(await stream.history.get_events()).total == Usage(total_tokens=30)
assert UsageReport.from_events(await ctx.stream.history.get_events()).total == Usage(total_tokens=30)
class MiddleSender:
async def guide(self, event: BaseEvent) -> None:
if isinstance(event, ChildEvent) and isinstance(event.inner, WorkSignal):
await event.origin.send(Directive(text="from middle"))
@pytest.mark.asyncio
async def test_nested_delivery_records_the_calling_ancestor_as_author() -> None:
guidance = Guidance()
leaf = Agent(
"leaf",
config=TestConfig(ToolCallEvent(name="seek_guidance", arguments="{}"), "done"),
tools=[seek_guidance],
observers=[observer(callback=guidance.receive)],
)
sender = MiddleSender()
middle = Agent(
"middle",
config=TestConfig(ToolCallEvent(name="task_leaf", arguments='{"objective":"work"}'), "done"),
tools=[leaf.as_tool(description="Work")],
observers=[observer(callback=sender.guide)],
)
ctx = Context(MemoryStream())
result = await run_task(middle, "organize", parent_context=ctx, task_id="middle")
assert result.completed
audits = [
event.inner
for event in await ctx.stream.history.get_events()
if isinstance(event, ChildEvent) and isinstance(event.inner, DownEvent)
]
assert len(audits) == 1
assert audits[0].to_dict() == IsPartialDict({
"origin_task_id": "middle",
"origin_path": ["middle"],
"origin_stream_id": str(result.stream.id),
})