187 lines
7.2 KiB
Python
187 lines
7.2 KiB
Python
"""A completed/cancelled graph must not leave SSE alive on metadata I/O."""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.runtime import journal as journal_module
|
|
from deerflow.runtime.runs import worker
|
|
from deerflow.runtime.runs.manager import RunManager
|
|
from deerflow.runtime.runs.schemas import RunStatus
|
|
from deerflow.runtime.stream_bridge.memory import MemoryStreamBridge
|
|
|
|
|
|
class FinishedAgent:
|
|
async def astream(self, *args, **kwargs):
|
|
yield {"messages": []}
|
|
|
|
|
|
async def make_run(monkeypatch, phase, *, graph=None, ignore_cancellation=False):
|
|
monkeypatch.setattr(worker, "_RUN_FINALIZATION_TIMEOUT_SECONDS", 0.03)
|
|
entered = asyncio.Event()
|
|
cancelled = asyncio.Event()
|
|
release = asyncio.Event()
|
|
drained = asyncio.Event()
|
|
|
|
async def blocked(*args, **kwargs):
|
|
entered.set()
|
|
try:
|
|
await release.wait()
|
|
except asyncio.CancelledError:
|
|
cancelled.set()
|
|
if ignore_cancellation:
|
|
await release.wait()
|
|
else:
|
|
raise
|
|
finally:
|
|
drained.set()
|
|
|
|
fake_journal = SimpleNamespace(flush=AsyncMock(), get_completion_data=lambda: {})
|
|
monkeypatch.setattr(journal_module, "RunJournal", lambda **kwargs: fake_journal)
|
|
manager = RunManager()
|
|
record = await manager.create("cleanup-test")
|
|
bridge = MemoryStreamBridge()
|
|
# Preserve the stream for assertions and avoid leaving a 60s sleeper.
|
|
monkeypatch.setattr(bridge, "cleanup", AsyncMock())
|
|
thread_store = SimpleNamespace(update_display_name=AsyncMock(), update_status=AsyncMock())
|
|
checkpoint_calls = 0
|
|
|
|
async def checkpoint(*args):
|
|
nonlocal checkpoint_calls
|
|
checkpoint_calls += 1
|
|
# Snapshot and terminal inspection precede final metadata cleanup.
|
|
if (phase == "checkpoint" and checkpoint_calls == 3) or (phase == "terminal_checkpoint" and checkpoint_calls == 2):
|
|
await blocked()
|
|
return SimpleNamespace(checkpoint={"channel_values": {"title": "Test"}}, config={})
|
|
|
|
checkpointer = SimpleNamespace(aget_tuple=checkpoint)
|
|
if phase == "journal":
|
|
fake_journal.flush = blocked
|
|
elif phase == "completion":
|
|
monkeypatch.setattr(manager, "update_run_completion", blocked)
|
|
elif phase == "title":
|
|
thread_store.update_display_name = blocked
|
|
elif phase == "thread_status":
|
|
thread_store.update_status = blocked
|
|
elif phase == "terminal_status":
|
|
async def persist_status(run_id, status, **kwargs):
|
|
if status != "running":
|
|
await blocked()
|
|
manager._store = SimpleNamespace(update_status=persist_status, update_run_completion=AsyncMock())
|
|
task = asyncio.create_task(worker.run_agent(
|
|
bridge, manager, record,
|
|
ctx=worker.RunContext(checkpointer=checkpointer, event_store=object(), thread_store=thread_store),
|
|
agent_factory=lambda **kwargs: graph or FinishedAgent(),
|
|
graph_input={}, config={},
|
|
))
|
|
record.task = task
|
|
return SimpleNamespace(
|
|
task=task, manager=manager, record=record, bridge=bridge,
|
|
entered=entered, cancelled=cancelled, release=release, drained=drained,
|
|
journal=fake_journal, thread_store=thread_store,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("phase", ["journal", "completion", "checkpoint", "title", "thread_status", "terminal_status", "terminal_checkpoint"])
|
|
async def test_metadata_stall_always_ends_sse(monkeypatch, phase, caplog):
|
|
run = await make_run(monkeypatch, phase)
|
|
await asyncio.wait_for(run.task, 1)
|
|
await asyncio.wait_for(run.cancelled.wait(), 1)
|
|
assert run.entered.is_set()
|
|
assert run.record.status == (RunStatus.error if phase == "terminal_checkpoint" else RunStatus.success)
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
assert "RUN_COMPLETION_TIMEOUT" in caplog.text
|
|
if phase == "terminal_checkpoint":
|
|
assert any(event.event == "error" for event in run.bridge._streams[run.record.run_id].events)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delayed_cancellation_acknowledgement_does_not_hold_stream(monkeypatch):
|
|
run = await make_run(monkeypatch, "journal", ignore_cancellation=True)
|
|
try:
|
|
await asyncio.wait_for(run.task, 1)
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
await asyncio.wait_for(run.cancelled.wait(), 1)
|
|
finally:
|
|
run.release.set()
|
|
await asyncio.wait_for(run.drained.wait(), 1)
|
|
await asyncio.sleep(0)
|
|
run.thread_store.update_display_name.assert_not_awaited()
|
|
run.thread_store.update_status.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_graph_ends_sse_even_when_cleanup_stalls(monkeypatch):
|
|
graph_started = asyncio.Event()
|
|
|
|
class WaitingAgent:
|
|
async def astream(self, *args, **kwargs):
|
|
graph_started.set()
|
|
await asyncio.Event().wait()
|
|
yield {}
|
|
|
|
run = await make_run(monkeypatch, "journal", graph=WaitingAgent())
|
|
await asyncio.wait_for(graph_started.wait(), 1)
|
|
await run.manager.cancel(run.record.run_id)
|
|
await asyncio.wait_for(run.task, 1)
|
|
assert run.record.status == RunStatus.interrupted
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancellation_during_cleanup_preserves_terminator(monkeypatch):
|
|
run = await make_run(monkeypatch, "journal")
|
|
await asyncio.wait_for(run.entered.wait(), 1)
|
|
run.task.cancel()
|
|
await asyncio.wait_for(run.task, 1)
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
await asyncio.wait_for(run.cancelled.wait(), 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_healthy_completion_still_persists_metadata(monkeypatch):
|
|
run = await make_run(monkeypatch, "healthy")
|
|
await asyncio.wait_for(run.task, 1)
|
|
run.journal.flush.assert_awaited_once()
|
|
run.thread_store.update_display_name.assert_awaited_once_with("cleanup-test", "Test")
|
|
run.thread_store.update_status.assert_awaited_once_with("cleanup-test", "idle")
|
|
assert run.record.status == RunStatus.success
|
|
assert run.bridge._streams[run.record.run_id].ended
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_older_cleanup_does_not_overwrite_newer_run_metadata(monkeypatch):
|
|
run = await make_run(monkeypatch, "journal")
|
|
await asyncio.wait_for(run.entered.wait(), 1)
|
|
await run.manager.create(run.record.thread_id)
|
|
run.release.set()
|
|
await asyncio.wait_for(run.task, 1)
|
|
run.thread_store.update_display_name.assert_not_awaited()
|
|
run.thread_store.update_status.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completion_budget_does_not_time_out_executing_graph(monkeypatch):
|
|
graph_started = asyncio.Event()
|
|
graph_release = asyncio.Event()
|
|
|
|
class LongAgent:
|
|
async def astream(self, *args, **kwargs):
|
|
graph_started.set()
|
|
await graph_release.wait()
|
|
yield {}
|
|
|
|
run = await make_run(monkeypatch, "healthy", graph=LongAgent())
|
|
try:
|
|
await asyncio.wait_for(graph_started.wait(), 1)
|
|
done, _ = await asyncio.wait({run.task}, timeout=0.06)
|
|
assert not done
|
|
assert run.record.status == RunStatus.running
|
|
assert not run.bridge._streams[run.record.run_id].ended
|
|
finally:
|
|
graph_release.set()
|
|
await asyncio.wait_for(run.task, 1)
|