test/test_task.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
"""Tests for the framework-core ``Task`` primitive.
Covers standalone usage with no hub or network — pure
``ag2.task`` lifecycle, event emission, dependency stamping
for ``TaskInject``, and TTL field population.
"""
import asyncio
from typing import Any
import pytest
from dirty_equals import IsList
from ag2 import Agent, Context, TaskInject
from ag2.config import AnthropicConfig
from ag2.context import ConversationContext
from ag2.events import (
BaseEvent,
ModelResponse,
TaskCompleted,
TaskExpired,
TaskFailed,
TaskProgress,
TaskStarted,
ToolCallEvent,
ToolCallsEvent,
ToolResultEvent,
ToolResultsEvent,
)
from ag2.middleware import ToolExecution, ToolResultType
from ag2.stream import MemoryStream
from ag2.task import TERMINAL_TASK_STATES, Task, TaskSpec, TaskState
from ag2.testing import TestConfig, TrackingConfig
from ag2.tools import tool
def _agent() -> Agent:
"""A bare Agent — no model calls, just a name."""
return Agent(name="researcher", config=AnthropicConfig(model="claude-sonnet-5"))
async def _persisted(stream: MemoryStream) -> list[BaseEvent]:
"""Read durably-persisted events from a MemoryStream's storage.
Note: transient events (``TaskProgress``, ``ModelMessageChunk``) are
not persisted — use ``_subscribe`` to capture them live.
"""
return list(await stream.history.storage.get_history(stream.id))
def _subscribe(stream: MemoryStream) -> list[BaseEvent]:
"""Subscribe to ``stream`` and return a list mutated as events flow.
Captures all events (including transient). Subscription must happen
before any event of interest fires.
"""
events: list[BaseEvent] = []
stream.subscribe(lambda ev: events.append(ev), sync_to_thread=False)
return events
class TestTaskState:
def test_terminal_set(self) -> None:
assert TaskState.COMPLETED in TERMINAL_TASK_STATES
assert TaskState.FAILED in TERMINAL_TASK_STATES
assert TaskState.EXPIRED in TERMINAL_TASK_STATES
assert TaskState.CREATED not in TERMINAL_TASK_STATES
assert TaskState.RUNNING not in TERMINAL_TASK_STATES
class TestTaskSpec:
def test_defaults(self) -> None:
spec = TaskSpec(title="hello")
assert spec.title == "hello"
assert spec.description == ""
assert spec.payload == {}
def test_payload_isolated(self) -> None:
# Default factory produces a fresh dict each construction.
a = TaskSpec(title="a")
b = TaskSpec(title="b")
a.payload["k"] = 1
assert "k" not in b.payload
class TestStandaloneLifecycle:
"""Task with no bound context — creates its own MemoryStream."""
@pytest.mark.asyncio
async def test_clean_exit_auto_completes(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="research X")) as task:
assert task.state == TaskState.RUNNING
assert task.metadata.spec.title == "research X"
stream = task.context.stream
assert task.state == TaskState.COMPLETED
events = await _persisted(stream)
types = [type(e) for e in events]
assert TaskStarted in types
assert TaskCompleted in types
@pytest.mark.asyncio
async def test_explicit_complete_records_result(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="compute")) as task:
await task.complete({"answer": 42})
assert task.state == TaskState.COMPLETED
assert task.metadata.result == {"answer": 42}
stream = task.context.stream
events = await _persisted(stream)
completed = [e for e in events if isinstance(e, TaskCompleted)]
assert len(completed) == 1
assert completed[0].result == {"answer": 42}
assert completed[0].agent_name == "researcher"
assert completed[0].objective == "compute"
@pytest.mark.asyncio
async def test_complete_is_idempotent(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="once")) as task:
await task.complete("first")
await task.complete("second") # no-op
assert task.metadata.result == "first"
stream = task.context.stream
events = await _persisted(stream)
assert sum(1 for e in events if isinstance(e, TaskCompleted)) == 1
@pytest.mark.asyncio
async def test_exception_emits_task_failed_and_propagates(self) -> None:
agent = _agent()
boom = ValueError("kaboom")
captured_stream = None
captured_task = None
with pytest.raises(ValueError) as exc_info:
async with Task(owner_id=agent.name, spec=TaskSpec(title="flaky")) as task:
captured_task = task
captured_stream = task.context.stream
raise boom
assert exc_info.value is boom
assert captured_task is not None
assert captured_task.state == TaskState.FAILED
assert captured_task.metadata.error == "kaboom"
assert captured_stream is not None
events = await _persisted(captured_stream)
failed = [e for e in events if isinstance(e, TaskFailed)]
assert len(failed) == 1
assert failed[0].error is boom
@pytest.mark.asyncio
async def test_explicit_fail_with_string(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="rejected")) as task:
await task.fail("not authorized")
assert task.state == TaskState.FAILED
assert task.metadata.error == "not authorized"
stream = task.context.stream
events = await _persisted(stream)
failed = [e for e in events if isinstance(e, TaskFailed)]
assert len(failed) == 1
assert isinstance(failed[0].error, RuntimeError)
assert str(failed[0].error) == "not authorized"
class TestProgress:
@pytest.mark.asyncio
async def test_progress_accumulates_into_metadata(self) -> None:
agent = _agent()
stream = MemoryStream()
ctx = ConversationContext(stream=stream)
events = _subscribe(stream)
async with Task(owner_id=agent.name, spec=TaskSpec(title="pipeline"), context=ctx) as task:
await task.progress({"step": "search"})
await task.progress({"hits": 12})
assert task.metadata.progress == {"step": "search", "hits": 12}
assert task.metadata.last_progress_at is not None
progress_events = [e for e in events if isinstance(e, TaskProgress)]
assert len(progress_events) == 2
assert progress_events[0].payload == {"step": "search"}
assert progress_events[1].payload == {"hits": 12}
@pytest.mark.asyncio
async def test_progress_after_terminal_is_noop(self) -> None:
agent = _agent()
stream = MemoryStream()
ctx = ConversationContext(stream=stream)
events = _subscribe(stream)
async with Task(owner_id=agent.name, spec=TaskSpec(title="done-fast"), context=ctx) as task:
await task.complete("ok")
await task.progress({"too": "late"})
assert "too" not in task.metadata.progress
assert not [e for e in events if isinstance(e, TaskProgress)]
class TestExpire:
@pytest.mark.asyncio
async def test_expire_emits_task_expired(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="slow")) as task:
await task.expire()
assert task.state == TaskState.EXPIRED
stream = task.context.stream
events = await _persisted(stream)
expired = [e for e in events if isinstance(e, TaskExpired)]
assert len(expired) == 1
assert expired[0].objective == "slow"
class TestBoundContext:
"""Task fed an explicit ConversationContext — events flow on caller's stream."""
@pytest.mark.asyncio
async def test_events_fire_on_supplied_stream(self) -> None:
agent = _agent()
stream = MemoryStream()
ctx = ConversationContext(stream=stream)
events = _subscribe(stream)
async with Task(owner_id=agent.name, spec=TaskSpec(title="with-ctx"), context=ctx) as task:
await task.progress({"k": "v"})
await task.complete("done")
types = [type(e) for e in events]
assert types.count(TaskStarted) == 1
assert types.count(TaskProgress) == 1
assert types.count(TaskCompleted) == 1
@pytest.mark.asyncio
async def test_task_inject_stamped_during_block(self) -> None:
agent = _agent()
stream = MemoryStream()
ctx = ConversationContext(stream=stream)
assert "ag2.task" not in ctx.dependencies
async with Task(owner_id=agent.name, spec=TaskSpec(title="inject-test"), context=ctx) as task:
assert ctx.dependencies["ag2.task"] is task
assert "ag2.task" not in ctx.dependencies
@pytest.mark.asyncio
async def test_nested_tasks_restore_previous(self) -> None:
agent = _agent()
stream = MemoryStream()
ctx = ConversationContext(stream=stream)
async with Task(owner_id=agent.name, spec=TaskSpec(title="outer"), context=ctx) as outer:
assert ctx.dependencies["ag2.task"] is outer
async with Task(owner_id=agent.name, spec=TaskSpec(title="inner"), context=ctx) as inner:
assert ctx.dependencies["ag2.task"] is inner
assert ctx.dependencies["ag2.task"] is outer
assert "ag2.task" not in ctx.dependencies
class TestTtl:
@pytest.mark.asyncio
async def test_ttl_seconds_populates_expires_at(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="with-ttl"), ttl_seconds=60) as task:
assert task.metadata.expires_at is not None
assert task.metadata.started_at is not None
# expires_at should be later than started_at
assert task.metadata.expires_at > task.metadata.started_at
@pytest.mark.asyncio
async def test_no_ttl_leaves_expires_at_none(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="no-ttl")) as task:
assert task.metadata.expires_at is None
class TestPropertiesBeforeEntry:
def test_state_returns_created(self) -> None:
agent = _agent()
task = Task(owner_id=agent.name, spec=TaskSpec(title="not-yet-entered"))
assert task.state == TaskState.CREATED
def test_task_id_raises_before_entry(self) -> None:
agent = _agent()
task = Task(owner_id=agent.name, spec=TaskSpec(title="not-yet-entered"))
with pytest.raises(RuntimeError, match="before __aenter__"):
_ = task.task_id
def test_metadata_raises_before_entry(self) -> None:
agent = _agent()
task = Task(owner_id=agent.name, spec=TaskSpec(title="not-yet-entered"))
with pytest.raises(RuntimeError, match="before __aenter__"):
_ = task.metadata
@pytest.mark.asyncio
async def test_double_enter_raises(self) -> None:
agent = _agent()
task = Task(owner_id=agent.name, spec=TaskSpec(title="once-only"))
async with task:
pass
with pytest.raises(RuntimeError, match="already entered"):
async with task:
pass
class TestMetadata:
@pytest.mark.asyncio
async def test_owner_id_is_agent_name(self) -> None:
agent = _agent()
async with Task(owner_id=agent.name, spec=TaskSpec(title="ownership")) as task:
assert task.metadata.owner_id == "researcher"
assert task.metadata.spec.title == "ownership"
@pytest.mark.asyncio
async def test_payload_passed_through(self) -> None:
agent = _agent()
async with Task(
owner_id=agent.name, spec=TaskSpec(title="payload-test", payload={"capability": "search"})
) as task:
assert task.metadata.spec.payload == {"capability": "search"}
async def _report(note: str, task: TaskInject) -> str:
"""Report progress on the active task."""
if task is None:
return "no active task"
await task.progress({"note": note})
return "reported"
def _report_call(note: str) -> ToolCallEvent:
return ToolCallEvent(id=note, name="report", arguments=f'{{"note": "{note}"}}')
def _reporter(note: str) -> Agent:
"""An agent whose one turn calls ``report(note)``."""
return Agent(
"reporter",
config=TestConfig(ModelResponse(tool_calls=ToolCallsEvent([_report_call(note)])), "Done"),
tools=[tool(_report, name="report")],
)
def _progress(events: list[BaseEvent]) -> list[tuple[str, dict[str, Any]]]:
"""``(task title, payload)`` of every ``TaskProgress`` seen, in order."""
return [(event.objective, event.payload) for event in events if isinstance(event, TaskProgress)]
@pytest.mark.asyncio
class TestTaskInject:
"""A ``@tool`` reports progress through ``TaskInject`` to the task whose
``Task(...)`` block encloses the call."""
async def test_parallel_tool_calls_report_to_their_own_task(self) -> None:
"""Tool middleware opens a task per call; with every block open at once,
each call's progress lands on its own task, not on the last one opened."""
notes = ["A", "B", "C"]
opened: set[str] = set()
all_open = asyncio.Event()
async def task_per_call(call_next: ToolExecution, event: ToolCallEvent, ctx: Context) -> ToolResultType:
async with Task(owner_id=agent.name, spec=TaskSpec(title=event.id), context=ctx):
opened.add(event.id)
if opened == set(notes):
all_open.set()
await all_open.wait()
return await call_next(event, ctx)
agent = Agent(
"planner",
config=TestConfig(ModelResponse(tool_calls=ToolCallsEvent([_report_call(n) for n in notes])), "Done"),
tools=[tool(_report, name="report", middleware=[task_per_call])],
)
stream = MemoryStream()
events = _subscribe(stream)
await agent.ask("Run the steps.", stream=stream)
assert _progress(events) == IsList(*((n, {"note": n}) for n in notes), check_order=False)
@pytest.mark.parametrize("capability", [None, "reporting"])
async def test_turn_inside_a_block_reports_to_it(self, capability: str | None) -> None:
stream = MemoryStream()
events = _subscribe(stream)
async with Task(owner_id="researcher", spec=TaskSpec(title="job"), context=ConversationContext(stream=stream)):
await _reporter("progress").ask("Report.", stream=stream, capability=capability)
assert _progress(events) == [("job", {"note": "progress"})]
async def test_nested_blocks_report_to_the_innermost(self) -> None:
owner = _agent()
stream = MemoryStream()
ctx = ConversationContext(stream=stream)
events = _subscribe(stream)
async with Task(owner_id=owner.name, spec=TaskSpec(title="outer"), context=ctx):
async with Task(owner_id=owner.name, spec=TaskSpec(title="inner"), context=ctx):
await _reporter("inside inner").ask("Report.", stream=stream)
await _reporter("after inner").ask("Report.", stream=stream)
await _reporter("after outer").ask("Report.", stream=stream)
assert _progress(events) == [
("inner", {"note": "inside inner"}),
("outer", {"note": "after inner"}),
]
async def test_without_a_block_there_is_no_task(self) -> None:
call = _report_call("progress")
tracking = TrackingConfig(TestConfig(ModelResponse(tool_calls=ToolCallsEvent([call])), "Done"))
agent = Agent("reporter", config=tracking, tools=[tool(_report, name="report")])
stream = MemoryStream()
events = _subscribe(stream)
await agent.ask("Report.", stream=stream)
tracking.mock.assert_called_with(ToolResultsEvent([ToolResultEvent.from_call(call, "no active task")]))
assert _progress(events) == []
async def test_a_turn_on_another_stream_does_not_see_the_block(self) -> None:
stream = MemoryStream()
events = _subscribe(stream)
async with Task(owner_id="researcher", spec=TaskSpec(title="job"), context=ConversationContext(stream=stream)):
await _reporter("elsewhere").ask("Report.", stream=MemoryStream())
assert _progress(events) == []
async def test_a_subagent_does_not_see_the_parent_task(self) -> None:
"""The parent opens a task around the delegation; the worker runs its own
conversation and does not report to it."""
async def task_per_call(call_next: ToolExecution, event: ToolCallEvent, ctx: Context) -> ToolResultType:
async with Task(owner_id=parent.name, spec=TaskSpec(title="delegation"), context=ctx):
return await call_next(event, ctx)
worker = _reporter("from worker")
parent = Agent(
"parent",
config=TestConfig(
ModelResponse(
tool_calls=ToolCallsEvent([
ToolCallEvent(name="task_reporter", arguments='{"objective": "Report."}')
])
),
"Done",
),
tools=[worker.as_tool(description="Reports on the active task.", middleware=[task_per_call])],
)
stream = MemoryStream()
events = _subscribe(stream)
await parent.ask("Delegate.", stream=stream)
assert _progress(events) == []
async def test_a_task_passed_as_a_dependency_is_injected(self) -> None:
task_stream = MemoryStream()
events = _subscribe(task_stream)
async with Task(
owner_id="researcher", spec=TaskSpec(title="given"), context=ConversationContext(stream=task_stream)
) as given:
await _reporter("progress").ask("Report.", dependencies={"ag2.task": given})
assert _progress(events) == [("given", {"note": "progress"})]