mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-13 08:18:57 +00:00
1c7531242c
* feat(runtime): record terminal artifact delivery receipts (#4272) * fix(runtime): persist delivery receipts across recovery * test(runtime): cover delivery receipt invariants * fix(runtime): preserve terminal status on receipt outages
363 lines
13 KiB
Python
363 lines
13 KiB
Python
"""Worker-level regression tests for the terminal run.delivery event (#4272 slice 1)."""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, ToolMessage
|
|
from langgraph.types import Command
|
|
|
|
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
|
from deerflow.runtime.runs.manager import RunManager
|
|
from deerflow.runtime.runs.schemas import RunStatus
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
from deerflow.runtime.runs.worker import RunContext, run_agent
|
|
|
|
|
|
def _make_bridge():
|
|
return SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock())
|
|
|
|
|
|
async def _delivery_events(store: MemoryRunEventStore, thread_id: str, run_id: str) -> list[dict]:
|
|
events = await store.list_events(thread_id, run_id)
|
|
return [e for e in events if e["event_type"] == "run.delivery"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_event_records_present_files_paths_on_success():
|
|
run_manager = RunManager()
|
|
record = await run_manager.create("thread-1")
|
|
store = MemoryRunEventStore()
|
|
|
|
class DummyAgent:
|
|
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
|
journal = config["context"]["__run_journal"]
|
|
ai = AIMessage(content="", tool_calls=[{"id": "call_1", "name": "present_files", "args": {}}])
|
|
journal._remember_current_run_tool_calls(ai, caller="lead_agent")
|
|
journal.on_tool_end(
|
|
Command(
|
|
update={
|
|
"artifacts": ["/mnt/user-data/outputs/report.md"],
|
|
"messages": [ToolMessage("Successfully presented files", tool_call_id="call_1")],
|
|
}
|
|
),
|
|
run_id=uuid4(),
|
|
)
|
|
yield {"messages": []}
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=store),
|
|
agent_factory=lambda *, config: DummyAgent(),
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
delivery = await _delivery_events(store, "thread-1", record.run_id)
|
|
assert len(delivery) == 1
|
|
assert delivery[0]["content"]["presented"] == 1
|
|
assert delivery[0]["content"]["paths"] == ["/mnt/user-data/outputs/report.md"]
|
|
assert delivery[0]["content"]["by_tool"] == {"present_files": ["/mnt/user-data/outputs/report.md"]}
|
|
fetched = await run_manager.get(record.run_id)
|
|
assert fetched.status == RunStatus.success
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_event_presented_zero_without_artifact_production():
|
|
run_manager = RunManager()
|
|
record = await run_manager.create("thread-1")
|
|
store = MemoryRunEventStore()
|
|
|
|
class DummyAgent:
|
|
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
|
yield {"messages": []}
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=store),
|
|
agent_factory=lambda *, config: DummyAgent(),
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
delivery = await _delivery_events(store, "thread-1", record.run_id)
|
|
assert len(delivery) == 1
|
|
assert delivery[0]["content"] == {"presented": 0, "paths": [], "by_tool": {}}
|
|
fetched = await run_manager.get(record.run_id)
|
|
assert fetched.status == RunStatus.success
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_event_is_singleton_across_goal_continuations(monkeypatch):
|
|
run_manager = RunManager()
|
|
record = await run_manager.create("thread-1")
|
|
store = MemoryRunEventStore()
|
|
stream_calls = 0
|
|
continuation_calls = 0
|
|
|
|
class ContinuingAgent:
|
|
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
|
nonlocal stream_calls
|
|
stream_calls += 1
|
|
journal = config["context"]["__run_journal"]
|
|
tool_call_id = f"call_{stream_calls}"
|
|
journal._remember_current_run_tool_calls(
|
|
AIMessage(content="", tool_calls=[{"id": tool_call_id, "name": "present_files", "args": {}}]),
|
|
caller="lead_agent",
|
|
)
|
|
artifacts = ["/mnt/user-data/outputs/report.md"]
|
|
if stream_calls == 2:
|
|
artifacts.append("/mnt/user-data/outputs/appendix.md")
|
|
journal.on_tool_end(
|
|
Command(
|
|
update={
|
|
"artifacts": artifacts,
|
|
"messages": [ToolMessage("Successfully presented files", tool_call_id=tool_call_id)],
|
|
}
|
|
),
|
|
run_id=uuid4(),
|
|
)
|
|
yield {"messages": []}
|
|
|
|
async def prepare_continuation(**kwargs):
|
|
nonlocal continuation_calls
|
|
continuation_calls += 1
|
|
if continuation_calls == 1:
|
|
return {"messages": []}
|
|
return None
|
|
|
|
monkeypatch.setattr("deerflow.runtime.runs.worker._prepare_goal_continuation_input", prepare_continuation)
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=store),
|
|
agent_factory=lambda *, config: ContinuingAgent(),
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
|
|
delivery = await _delivery_events(store, "thread-1", record.run_id)
|
|
assert stream_calls == 2
|
|
assert len(delivery) == 1
|
|
assert delivery[0]["content"] == {
|
|
"presented": 2,
|
|
"paths": [
|
|
"/mnt/user-data/outputs/report.md",
|
|
"/mnt/user-data/outputs/appendix.md",
|
|
],
|
|
"by_tool": {
|
|
"present_files": [
|
|
"/mnt/user-data/outputs/report.md",
|
|
"/mnt/user-data/outputs/appendix.md",
|
|
]
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_event_emitted_exactly_once_on_error_path():
|
|
run_manager = RunManager()
|
|
record = await run_manager.create("thread-1")
|
|
store = MemoryRunEventStore()
|
|
|
|
class FailingAgent:
|
|
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
|
raise RuntimeError("boom")
|
|
yield # pragma: no cover - make this an async generator
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=store),
|
|
agent_factory=lambda *, config: FailingAgent(),
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
await asyncio.sleep(0)
|
|
|
|
delivery = await _delivery_events(store, "thread-1", record.run_id)
|
|
assert len(delivery) == 1
|
|
assert delivery[0]["content"]["presented"] == 0
|
|
fetched = await run_manager.get(record.run_id)
|
|
assert fetched.status == RunStatus.error
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_is_durable_before_terminal_run_status():
|
|
events = MemoryRunEventStore()
|
|
|
|
class OrderingRunStore(MemoryRunStore):
|
|
async def update_status(self, run_id, status, *, error=None, stop_reason=None):
|
|
if status not in {"pending", "running"}:
|
|
receipt = await events.list_events("thread-1", run_id, event_types=["run.delivery"])
|
|
assert len(receipt) == 1
|
|
return await super().update_status(run_id, status, error=error, stop_reason=stop_reason)
|
|
|
|
run_store = OrderingRunStore()
|
|
run_manager = RunManager(store=run_store)
|
|
record = await run_manager.create("thread-1")
|
|
|
|
class DummyAgent:
|
|
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
|
yield {"messages": []}
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=events),
|
|
agent_factory=lambda *, config: DummyAgent(),
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
|
|
assert (await run_store.get(record.run_id))["status"] == "success"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_write_retries_before_persisting_success():
|
|
class FlakyReceiptStore(MemoryRunEventStore):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.attempts = 0
|
|
|
|
async def put_if_absent(self, **kwargs):
|
|
self.attempts += 1
|
|
if self.attempts == 1:
|
|
raise RuntimeError("transient event store outage")
|
|
return await super().put_if_absent(**kwargs)
|
|
|
|
event_store = FlakyReceiptStore()
|
|
run_store = MemoryRunStore()
|
|
run_manager = RunManager(store=run_store)
|
|
record = await run_manager.create("thread-1")
|
|
|
|
class DummyAgent:
|
|
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
|
yield {"messages": []}
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=event_store),
|
|
agent_factory=lambda *, config: DummyAgent(),
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
|
|
assert event_store.attempts == 2
|
|
assert len(await _delivery_events(event_store, "thread-1", record.run_id)) == 1
|
|
assert (await run_store.get(record.run_id))["status"] == "success"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_write_failure_preserves_real_durable_terminal_status():
|
|
class FailingReceiptStore(MemoryRunEventStore):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.attempts = 0
|
|
|
|
async def put_if_absent(self, **kwargs):
|
|
self.attempts += 1
|
|
raise RuntimeError("event store unavailable")
|
|
|
|
run_store = MemoryRunStore()
|
|
run_manager = RunManager(store=run_store)
|
|
record = await run_manager.create("thread-1")
|
|
event_store = FailingReceiptStore()
|
|
|
|
class DummyAgent:
|
|
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
|
|
yield {"messages": []}
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=event_store),
|
|
agent_factory=lambda *, config: DummyAgent(),
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
|
|
# A receipt outage must not let lease recovery rewrite a genuine success
|
|
# as an error. After bounded retries, preserve the worker's real outcome.
|
|
assert event_store.attempts > 1
|
|
assert record.status == RunStatus.success
|
|
assert (await run_store.get(record.run_id))["status"] == "success"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_event_emitted_when_checkpoint_preflight_fails(monkeypatch):
|
|
run_manager = RunManager()
|
|
run_manager.update_run_completion = AsyncMock(wraps=run_manager.update_run_completion)
|
|
record = await run_manager.create("thread-1")
|
|
store = MemoryRunEventStore()
|
|
compatibility_check = AsyncMock(side_effect=RuntimeError("incompatible checkpoint"))
|
|
monkeypatch.setattr("deerflow.runtime.runs.worker.aensure_checkpoint_mode_compatible", compatibility_check)
|
|
|
|
def unexpected_agent_factory(**kwargs):
|
|
raise AssertionError("agent must not be built after preflight failure")
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=object(), event_store=store),
|
|
agent_factory=unexpected_agent_factory,
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
|
|
delivery = await _delivery_events(store, "thread-1", record.run_id)
|
|
assert len(delivery) == 1
|
|
assert delivery[0]["content"] == {"presented": 0, "paths": [], "by_tool": {}}
|
|
fetched = await run_manager.get(record.run_id)
|
|
assert fetched.status == RunStatus.error
|
|
run_manager.update_run_completion.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delivery_event_emitted_when_cancelled_waiting_for_prior_finalization(monkeypatch):
|
|
run_manager = RunManager()
|
|
run_manager.update_run_completion = AsyncMock(wraps=run_manager.update_run_completion)
|
|
record = await run_manager.create("thread-1")
|
|
store = MemoryRunEventStore()
|
|
monkeypatch.setattr(
|
|
run_manager,
|
|
"wait_for_prior_finalizing",
|
|
AsyncMock(side_effect=asyncio.CancelledError()),
|
|
)
|
|
|
|
def unexpected_agent_factory(**kwargs):
|
|
raise AssertionError("agent must not be built after preflight cancellation")
|
|
|
|
await run_agent(
|
|
_make_bridge(),
|
|
run_manager,
|
|
record,
|
|
ctx=RunContext(checkpointer=None, event_store=store),
|
|
agent_factory=unexpected_agent_factory,
|
|
graph_input={},
|
|
config={},
|
|
)
|
|
|
|
delivery = await _delivery_events(store, "thread-1", record.run_id)
|
|
assert len(delivery) == 1
|
|
assert delivery[0]["content"] == {"presented": 0, "paths": [], "by_tool": {}}
|
|
fetched = await run_manager.get(record.run_id)
|
|
assert fetched.status == RunStatus.interrupted
|
|
run_manager.update_run_completion.assert_not_awaited()
|