test/agent/test_run.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
"""Contract for ``Agent.run`` — the observable counterpart to ``Agent.ask``.
``run`` opens a turn *scope* and yields an ``AgentRun`` handle. The turn does
not advance until the caller drives it; while it drives, every event flows
through ``run.stream``, so observation is push-based (subscribe a callback, or
attach ``observers=``).
* ``await run.result()`` — drive inline and get the authoritative, idempotent
``AgentReply``; cancelling that await (e.g. a timeout) cancels the turn inline.
* ``run.start()`` — drive in a scope-owned background task; ``result()`` then
awaits it, and leaving the block cancels it if still running.
* ``run.stream`` — the underlying stream; subscribe to observe the turn live.
See docs/adr/0008-agent-run-turn-is-scoped-to-its-context-manager.md and its
amendment docs/adr/0009-agent-run-start-drives-in-background.md.
"""
import asyncio
import pytest
from pydantic import BaseModel
from ag2 import Agent, Context, tool
from ag2.events import (
BaseEvent,
MessageEnqueued,
ModelMessage,
ModelRequest,
ModelResponse,
TextInput,
ToolCallEvent,
ToolResultEvent,
)
from ag2.stream import MemoryStream
from ag2.testing import TestConfig, TrackingConfig
def _texts(event: BaseEvent) -> list[str]:
"""Text parts of a ``ModelRequest``-like event (empty for anything else)."""
return [p.content for p in getattr(event, "parts", []) if isinstance(p, TextInput)]
@pytest.mark.asyncio
class TestResult:
async def test_returns_reply_body(self) -> None:
agent = Agent("runner", config=TestConfig("hello"))
async with agent.run("Hi!") as run:
result = await run.result()
assert result.body == "hello"
async def test_is_idempotent(self) -> None:
agent = Agent("runner", config=TestConfig("hello"))
async with agent.run("Hi!") as run:
first = await run.result()
second = await run.result()
assert first is second
async def test_reraises_turn_failure(self) -> None:
@tool
async def explode() -> str:
"""Always raises."""
raise RuntimeError("tool blew up")
agent = Agent(
"runner",
config=TestConfig(
ToolCallEvent(name="explode", arguments="{}"),
"unreachable",
),
tools=[explode],
)
with pytest.raises(RuntimeError, match="tool blew up"):
async with agent.run("Hi!") as run:
await run.result()
async def test_content_validates_against_schema(self) -> None:
class Answer(BaseModel):
value: int
agent = Agent(
"runner",
config=TestConfig('{"value": 42}'),
response_schema=Answer,
)
async with agent.run("Hi!") as run:
content = await (await run.result()).content()
assert content == Answer(value=42)
@pytest.mark.asyncio
async def test_reply_run_continues_conversation() -> None:
agent = Agent(
"runner",
config=TestConfig(
"first",
"second",
),
)
async with agent.run("Hi!") as run:
first = await run.result()
assert first.body == "first"
async with first.run("again") as run2:
second = await run2.result()
assert second.body == "second"
# Continuation: same context/stream, history accumulates across both turns.
assert second.context is first.context
@pytest.mark.asyncio
async def test_events_are_observable_on_stream_while_driving() -> None:
tool_call = ToolCallEvent(name="ping", arguments="{}")
agent = Agent(
"runner",
config=TestConfig(tool_call, "done"),
tools=[tool(lambda: "pong", name="ping")],
)
seen: list = []
async def capture(event) -> None:
seen.append(event)
async with agent.run("Hi!") as run:
run.stream.subscribe(capture) # subscribe before driving
result = await run.result() # callbacks fire inline as the turn drives
assert [e for e in seen if isinstance(e, ToolCallEvent)] == [tool_call]
assert any(isinstance(e, ModelResponse) for e in seen)
assert result.body == "done"
@pytest.mark.asyncio
async def test_stream_is_exposed() -> None:
stream = MemoryStream()
agent = Agent("runner", config=TestConfig(ModelResponse(ModelMessage("hi"))))
async with agent.run("Hi!", stream=stream) as run:
assert run.stream is stream
await run.result()
@pytest.mark.asyncio
async def test_turn_runs_on_result_not_on_enter() -> None:
stream = MemoryStream()
agent = Agent("runner", config=TestConfig(ModelResponse(ModelMessage("hi"))))
async with agent.run("Hi!", stream=stream) as run:
await asyncio.sleep(0.05) # a background task, if any, would drive the turn here
before = await stream.history.get_events()
assert not any(isinstance(e, ModelResponse) for e in before), "no model call before result()"
result = await run.result()
after = await stream.history.get_events()
assert any(isinstance(e, ModelResponse) for e in after)
assert result.body == "hi"
@pytest.mark.asyncio
async def test_turn_never_runs_without_result() -> None:
stream = MemoryStream()
agent = Agent("runner", config=TestConfig(ModelResponse(ModelMessage("hi"))))
async with agent.run("Hi!", stream=stream):
await asyncio.sleep(0.05) # a background task, if any, would drive the turn here
events = await stream.history.get_events()
assert not any(isinstance(e, ModelResponse) for e in events), "an undriven run must not call the model"
@pytest.mark.asyncio
class TestEnqueue:
async def test_forwards_to_stream_inbox(self) -> None:
agent = Agent("runner", config=TestConfig("ok"))
async with agent.run("Hi!") as run:
run.enqueue("queued")
via_handle = list(run.stream.pending_messages)
await run.result()
assert via_handle == [ModelRequest.ensure_request(["queued"])]
async def test_before_result_merges_into_first_model_call(self) -> None:
config = TrackingConfig(TestConfig("answer"))
agent = Agent("runner", config=config)
async with agent.run("Hi!") as run:
run.enqueue("extra context")
result = await run.result()
first_seen = config.mock.call_args_list[0].args[0]
assert _texts(first_seen) == ["extra context", "Hi!"]
assert result.body == "answer"
async def test_during_turn_is_consumed_by_the_same_turn(self) -> None:
config = TrackingConfig(TestConfig(ToolCallEvent(name="ping", arguments="{}"), "r1", "r2", "r3"))
agent = Agent("runner", config=config, tools=[tool(lambda: "pong", name="ping")])
injected = False
async with agent.run("Hi!") as run:
@run.stream.where(ToolResultEvent).subscribe
async def inject(_event: BaseEvent) -> None:
nonlocal injected
if not injected:
injected = True
run.enqueue("injected message")
await run.result()
seen = [text for call in config.mock.call_args_list for text in _texts(call.args[0])]
assert "injected message" in seen, "a message enqueued mid-turn must reach a model call in that turn"
async def test_announces_the_message(self) -> None:
agent = Agent("runner", config=TestConfig("ok"))
announced = asyncio.Event()
async with agent.run("Hi!") as run:
@run.stream.where(MessageEnqueued).subscribe
async def on_enqueued(_event: MessageEnqueued) -> None:
announced.set()
run.enqueue("queued")
await asyncio.wait_for(announced.wait(), timeout=1.0)
await run.result()
async def test_announces_from_a_sync_tool(self) -> None:
stream = MemoryStream()
announced = asyncio.Event()
@stream.where(MessageEnqueued).subscribe
async def on_enqueued(_event: MessageEnqueued) -> None:
announced.set()
@tool
def queue_note(ctx: Context) -> str:
"""Enqueue a note from a worker thread."""
ctx.enqueue("note")
return "queued"
agent = Agent(
"runner",
config=TestConfig(ToolCallEvent(name="queue_note", arguments="{}"), "done"),
tools=[queue_note],
)
reply = await agent.ask("Hi!", stream=stream)
await asyncio.wait_for(announced.wait(), timeout=1.0)
assert reply.body == "done"
async def test_announcement_stays_out_of_history(self) -> None:
agent = Agent("runner", config=TestConfig("ok"))
async with agent.run("Hi!") as run:
run.enqueue("queued")
reply = await run.result()
await asyncio.sleep(0)
events = await reply.context.stream.history.get_events()
assert not any(isinstance(e, MessageEnqueued) for e in events)
@pytest.mark.asyncio
async def test_cancelling_result_cancels_turn_inline() -> None:
started = asyncio.Event()
cancelled = asyncio.Event()
@tool
async def block() -> str:
"""Blocks forever, flagging if it is cancelled."""
started.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
cancelled.set()
raise
return "never"
agent = Agent(
"runner",
config=TestConfig(
ToolCallEvent(name="block", arguments="{}"),
"unreachable",
),
tools=[block],
)
async with agent.run("Hi!") as run:
# result() drives the turn inline; cancel it only once the tool is in
# flight, so the assertion is about propagation, not turn-startup speed.
result_task = asyncio.ensure_future(run.result())
await asyncio.wait_for(started.wait(), timeout=5)
result_task.cancel()
with pytest.raises(asyncio.CancelledError):
await result_task
assert cancelled.is_set(), "cancelling the result() await must cancel the in-flight turn"
@pytest.mark.asyncio
class TestStart:
async def test_drives_in_background_and_result_returns_reply(self) -> None:
tool_call = ToolCallEvent(name="ping", arguments="{}")
agent = Agent(
"runner",
config=TestConfig(tool_call, "done"),
tools=[tool(lambda: "pong", name="ping")],
)
seen: list[BaseEvent] = []
async with agent.run("Hi!") as run:
run.start() # drive in the background; pull events meanwhile
with run.stream.join() as events:
async for event in events:
seen.append(event)
if isinstance(event, ModelResponse) and event.content == "done":
break
reply = await run.result()
assert reply.body == "done"
assert [e for e in seen if isinstance(e, ToolCallEvent)] == [tool_call]
async def test_result_returns_same_reply_after_start(self) -> None:
agent = Agent("runner", config=TestConfig("hello"))
async with agent.run("Hi!") as run:
run.start()
first = await run.result()
second = await run.result()
assert first is second
assert first.body == "hello"
async def test_start_is_idempotent(self) -> None:
agent = Agent("runner", config=TestConfig("hello"))
async with agent.run("Hi!") as run:
run.start()
run.start() # second call is a no-op
reply = await run.result()
assert reply.body == "hello"
async def test_leaving_block_cancels_started_turn(self) -> None:
cancelled = asyncio.Event()
running = asyncio.Event()
@tool
async def block() -> str:
"""Blocks forever, flagging when it runs and when it is cancelled."""
running.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
cancelled.set()
raise
return "never"
agent = Agent(
"runner",
config=TestConfig(ToolCallEvent(name="block", arguments="{}"), "unreachable"),
tools=[block],
)
async with agent.run("Hi!") as run:
run.start()
await asyncio.wait_for(running.wait(), timeout=1.0) # let the turn reach the tool
# leave the block without awaiting result()
assert cancelled.is_set(), "leaving the block must cancel a started, still-running turn"
async def test_start_outside_block_raises(self) -> None:
agent = Agent("runner", config=TestConfig("hello"))
run = agent.run("Hi!")
with pytest.raises(RuntimeError):
run.start()