test/network/test_task_durability.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
"""Task durability primitives — cancellation + checkpoint/resume.
The framework-core ``Task`` now supports owner-driven cancellation
and opt-in restart recovery through the :class:`CheckpointStore`
Protocol. The hub ships :class:`HubBackedCheckpointStore` as the
canonical default; tenants may plug in any compatible store.
"""
import asyncio
import pytest
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult, SpanExporter
from ag2 import Agent, Context
from ag2.context import ConversationContext
from ag2.events import TaskCancelled
from ag2.knowledge import MemoryKnowledgeStore
from ag2.network import (
Hub,
HubBackedCheckpointStore,
NotFoundError,
Passport,
Resume,
Rule,
)
from ag2.network.task_mirror import TaskMirror
from ag2.stream import MemoryStream
from ag2.task import (
TERMINAL_TASK_STATES,
CheckpointStore,
Task,
TaskSpec,
TaskState,
)
from ag2.testing import TestConfig
class _InMemoryCheckpointStore:
"""Test double satisfying :class:`CheckpointStore` without a hub."""
def __init__(self) -> None:
self.data: dict[str, dict] = {}
async def write(self, task_id: str, state: dict) -> None:
self.data[task_id] = dict(state)
async def read(self, task_id: str) -> dict | None:
snap = self.data.get(task_id)
return dict(snap) if snap is not None else None
class TestCancellationState:
def test_cancelled_is_a_terminal_state(self) -> None:
assert TaskState.CANCELLED in TERMINAL_TASK_STATES
def test_task_state_values_include_cancelled(self) -> None:
assert TaskState.CANCELLED.value == "cancelled"
class TestTaskCancel:
@pytest.mark.asyncio
async def test_cancel_transitions_to_cancelled_and_emits_event(self) -> None:
stream = MemoryStream()
events: list = []
stream.subscribe(lambda ev: events.append(ev))
task = Task(
owner_id="alice",
spec=TaskSpec(title="thing"),
context=ConversationContext(stream=stream),
)
async with task:
await task.cancel("ran too long")
assert task.state == TaskState.CANCELLED
assert task.metadata.error == "ran too long"
# Captured TaskCancelled with the reason.
cancels = [e for e in events if isinstance(e, TaskCancelled)]
assert len(cancels) == 1
assert cancels[0].reason == "ran too long"
@pytest.mark.asyncio
async def test_cancel_is_idempotent_on_terminal_task(self) -> None:
stream = MemoryStream()
events: list = []
stream.subscribe(lambda ev: events.append(ev))
task = Task(
owner_id="alice",
spec=TaskSpec(title="thing"),
context=ConversationContext(stream=stream),
)
async with task:
await task.complete("done")
assert task.state == TaskState.COMPLETED
# Second terminal call must not flip the state or emit another event.
await task.cancel("late")
assert task.state == TaskState.COMPLETED
assert not any(isinstance(e, TaskCancelled) for e in events)
@pytest.mark.asyncio
async def test_cancel_owned_task(self) -> None:
agent = Agent(name="alice", config=TestConfig())
async with Task(owner_id=agent.name, spec=TaskSpec(title="work")) as task:
await task.cancel()
assert task.state == TaskState.CANCELLED
class TestCheckpointStandalone:
@pytest.mark.asyncio
async def test_checkpoint_writes_via_store_and_resume_reads_it_back(self) -> None:
store = _InMemoryCheckpointStore()
first = Task(
owner_id="alice",
spec=TaskSpec(title="work"),
context=ConversationContext(stream=MemoryStream()),
checkpoint_store=store,
)
prior_task_id: str
async with first:
await first.checkpoint({"step": 3, "scratch": [1, 2, 3]})
prior_task_id = first.task_id
second = Task(
owner_id="alice",
spec=TaskSpec(title="work"),
context=ConversationContext(stream=MemoryStream()),
checkpoint_store=store,
resume_from=prior_task_id,
)
async with second:
assert second.resumed_state == {"step": 3, "scratch": [1, 2, 3]}
@pytest.mark.asyncio
async def test_checkpoint_without_store_is_a_silent_noop(self) -> None:
"""Standalone agents that never wire a store can still call
``Task.checkpoint`` — the call is just dropped."""
task = Task(
owner_id="alice",
spec=TaskSpec(title="work"),
context=ConversationContext(stream=MemoryStream()),
)
async with task:
await task.checkpoint({"step": 1}) # must not raise
@pytest.mark.asyncio
async def test_resume_from_unknown_task_yields_none(self) -> None:
store = _InMemoryCheckpointStore()
task = Task(
owner_id="alice",
spec=TaskSpec(title="work"),
context=ConversationContext(stream=MemoryStream()),
checkpoint_store=store,
resume_from="never-checkpointed",
)
async with task:
assert task.resumed_state is None
@pytest.mark.asyncio
async def test_checkpoint_after_terminal_is_a_noop(self) -> None:
store = _InMemoryCheckpointStore()
task = Task(
owner_id="alice",
spec=TaskSpec(title="work"),
context=ConversationContext(stream=MemoryStream()),
checkpoint_store=store,
)
async with task:
tid = task.task_id
await task.complete("done")
await task.checkpoint({"step": "should-not-persist"})
assert tid not in store.data
class TestHubBackedCheckpointStore:
@pytest.mark.asyncio
async def test_round_trip_through_hub(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
store = HubBackedCheckpointStore(hub)
await store.write("task-42", {"step": 7, "buffer": ["a", "b"]})
assert await store.read("task-42") == {"step": 7, "buffer": ["a", "b"]}
finally:
await hub.close()
@pytest.mark.asyncio
async def test_read_unknown_task_returns_none(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
store = HubBackedCheckpointStore(hub)
assert await store.read("never-written") is None
finally:
await hub.close()
@pytest.mark.asyncio
async def test_last_write_wins(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
store = HubBackedCheckpointStore(hub)
await store.write("task-1", {"v": 1})
await store.write("task-1", {"v": 2})
assert await store.read("task-1") == {"v": 2}
finally:
await hub.close()
@pytest.mark.asyncio
async def test_full_resume_cycle_through_hub(self) -> None:
"""End-to-end: first run checkpoints via hub-backed store; a
fresh task that resumes from the prior id sees the snapshot."""
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
store = HubBackedCheckpointStore(hub)
agent = Agent(name="alice", config=TestConfig())
async with Task(owner_id=agent.name, spec=TaskSpec(title="work"), checkpoint_store=store) as first:
await first.checkpoint({"phase": "research", "found": 42})
prior_id = first.task_id
async with Task(
owner_id=agent.name, spec=TaskSpec(title="work"), checkpoint_store=store, resume_from=prior_id
) as second:
assert second.resumed_state == {"phase": "research", "found": 42}
await second.complete("resumed and finished")
assert second.state == TaskState.COMPLETED
finally:
await hub.close()
def test_protocol_conformance_runtime_checkable(self) -> None:
# CheckpointStore is runtime_checkable; HubBackedCheckpointStore
# should satisfy isinstance even without an actual hub instance.
# Use a sentinel object for the hub param since the constructor
# only stores it without using it during the isinstance check.
store = HubBackedCheckpointStore.__new__(HubBackedCheckpointStore)
assert isinstance(store, CheckpointStore)
class _CapturingExporter(SpanExporter):
"""Minimal in-memory span sink for asserting on emitted spans."""
def __init__(self) -> None:
self.spans: list[ReadableSpan] = []
def export(self, spans) -> SpanExportResult:
self.spans.extend(spans)
return SpanExportResult.SUCCESS
def shutdown(self) -> None:
self.spans.clear()
def _events_named(span: ReadableSpan, name: str) -> list:
return [e for e in span.events if e.name == name]
class TestCheckpointTracing:
"""``HubBackedCheckpointStore`` pins span-events on the active span so
checkpoint save/restore — which bypass the envelope path — still surface
in a trace."""
@pytest.mark.asyncio
async def test_write_and_read_emit_span_events(self) -> None:
exporter = _CapturingExporter()
provider = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
tracer = provider.get_tracer("test")
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
store = HubBackedCheckpointStore(hub)
with tracer.start_as_current_span("run"):
await store.write("task-42", {"step": 7, "buffer": ["a", "b"]})
# Hit: a checkpoint exists for this id.
assert await store.read("task-42") == {"step": 7, "buffer": ["a", "b"]}
# Miss: nothing written for this id.
assert await store.read("never-written") is None
finally:
await hub.close()
[span] = exporter.spans
[write_ev] = _events_named(span, "checkpoint.write")
assert write_ev.attributes["ag2.checkpoint.task_id"] == "task-42"
assert write_ev.attributes["ag2.checkpoint.bytes"] > 0
read_evs = _events_named(span, "checkpoint.read")
assert len(read_evs) == 2
hit, miss = read_evs
assert hit.attributes == {"ag2.checkpoint.task_id": "task-42", "ag2.checkpoint.hit": True}
assert miss.attributes == {"ag2.checkpoint.task_id": "never-written", "ag2.checkpoint.hit": False}
@pytest.mark.asyncio
async def test_no_provider_is_safe_noop(self) -> None:
"""With no tracer provider configured, write/read still work — the
span-event call lands on a non-recording span and is a no-op."""
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
store = HubBackedCheckpointStore(hub)
await store.write("task-1", {"v": 1})
assert await store.read("task-1") == {"v": 1}
finally:
await hub.close()
def test_in_memory_double_also_satisfies_protocol(self) -> None:
assert isinstance(_InMemoryCheckpointStore(), CheckpointStore)
@pytest.mark.asyncio
class TestMirrorCancellation:
"""``TaskMirror`` forwards owner-driven cancellation to the hub.
Without this bridge, an agent calling ``Task.cancel`` would emit a
``TaskCancelled`` event on its stream but the hub's task cache
would stay at ``RUNNING`` and capability-tagged stats would never
record the cancellation.
"""
async def test_mirror_forwards_cancellation_to_hub_metadata(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
bob_passport = await hub.register_identity(Passport(name="bob"), Resume())
agent = Agent(name="bob", config=TestConfig())
stream = MemoryStream()
mirror = TaskMirror(hub=hub, owner_id=bob_passport.agent_id)
sub_ids = mirror.attach(stream)
try:
async with Task(
owner_id=agent.name, spec=TaskSpec(title="indexing"), context=Context(stream=stream)
) as task:
task_id = task.task_id
await task.cancel("wrap-up")
# Stream callbacks run via subscribe(sync_to_thread=False) — yield
# the loop so the mirror's _on_cancelled completes before assert.
await asyncio.sleep(0)
finally:
mirror.detach(stream, sub_ids)
meta = await hub.get_task(task_id)
assert meta.state == TaskState.CANCELLED
assert meta.error == "wrap-up"
assert meta.completed_at # terminal stamp
finally:
await hub.close()
async def test_mirror_records_observation_on_capability_tagged_cancellation(
self,
) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
bob_passport = await hub.register_identity(
Passport(name="bob"),
Resume(claimed_capabilities=["indexing"]),
)
agent = Agent(name="bob", config=TestConfig())
stream = MemoryStream()
mirror = TaskMirror(hub=hub, owner_id=bob_passport.agent_id)
sub_ids = mirror.attach(stream)
try:
async with Task(
owner_id=agent.name,
spec=TaskSpec(title="indexing", description="search & summarise", capability="indexing"),
context=Context(stream=stream),
) as task:
await task.cancel("wrap-up")
await asyncio.sleep(0)
finally:
mirror.detach(stream, sub_ids)
resume = await hub.get_resume(bob_passport.agent_id)
stat = resume.observed.get("indexing")
assert stat is not None
assert stat.n == 1
finally:
await hub.close()
@pytest.mark.asyncio
class TestHubAccessors:
"""``Hub.find_agent_id`` + ``Hub.get_rule`` are the public-surface
primitives ``HubClient.attach`` uses to look up persisted identity
without reaching into hub internals.
"""
async def test_find_agent_id_returns_id_for_registered_name(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
passport = await hub.register_identity(Passport(name="alice"), Resume())
assert hub.find_agent_id("alice") == passport.agent_id
finally:
await hub.close()
async def test_find_agent_id_returns_none_for_unknown_name(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
assert hub.find_agent_id("ghost") is None
finally:
await hub.close()
async def test_get_rule_returns_default_rule_for_registered_agent(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
passport = await hub.register_identity(Passport(name="alice"), Resume())
rule = await hub.get_rule(passport.agent_id)
assert isinstance(rule, Rule)
finally:
await hub.close()
async def test_get_rule_raises_not_found_for_unknown_agent(self) -> None:
hub = await Hub.open(
MemoryKnowledgeStore(),
ttl_sweep_interval=0,
expectation_sweep_interval=0,
)
try:
with pytest.raises(NotFoundError):
await hub.get_rule("never-registered")
finally:
await hub.close()