"""RC-BE-012: claim, lease renew, recovery, two-phase cancel, dispatcher.""" from __future__ import annotations import asyncio from datetime import UTC, datetime, timedelta import pytest from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from app.report_collaboration.execution.dispatcher import ReportCollaborationDispatcher from app.report_collaboration.execution.executor import LeaseParams, ReportCollaborationExecutor from app.report_collaboration.execution.ledger import TaskLedger, TaskResult from app.report_collaboration.execution.member_runner import ScriptedMemberKernel from app.report_collaboration.execution.recovery import recover_orphaned_attempts from deerflow.config.report_collaboration_config import ReportCollaborationConfig from deerflow.persistence.base import Base from deerflow.persistence.report_collaboration import MemoryReportCollaborationStore, make_report_collaboration_store from deerflow.persistence.report_collaboration.codec import iso def _evidence() -> dict: return { "schema_version": 1, "claims": [ { "claim_id": "c1", "text": "2024 年新能源乘用车零售份额上升", "kind": "fact", "source_ids": ["src_1"], "excerpt": "份额上升", } ], "sources": [{"source_id": "src_1", "title": "乘联会月报", "excerpt": "份额上升"}], } def _draft() -> dict: return { "schema_version": 1, "section_id": "overview", "title": "市场趋势", "markdown": "份额上升，来源 c1。", "evidence_ids": ["c1"], "claim_citation_map": {"份额上升": ["c1"]}, } def _plan() -> dict: return { "id": "plan_lease", "proposal_group_id": "grp_lease", "title": "租约", "strategy": "focused_fast", "summary": "研究后写作", "rationale": "012", "recommended": True, "estimated_duration_seconds": 60, "estimated_cost_level": "low", "requirement_revision": 0, "roles": [ {"key": "researcher", "display_name": "研究员", "responsibility": "检索"}, {"key": "writer", "display_name": "撰写员", "responsibility": "成稿"}, ], "nodes": [ { "id": "research", "label": "检索", "role_key": "researcher", "role_display_name": "研究员", "mission": "收集证据", "output_artifact_type": "EvidenceBundle", "allowed_tools": ["web_search"], "acceptance_criteria": ["每条主张有来源"], "depends_on": [], "max_attempts": 3, }, { "id": "write", "label": "写作", "role_key": "writer", "role_display_name": "撰写员", "mission": "成稿", "output_artifact_type": "ReportSectionDraft", "allowed_tools": ["submit_artifact"], "depends_on": ["research"], "max_attempts": 3, }, ], "edges": [{"source": "research", "target": "write"}], "quality_gates": [], "validation": {"ok": True, "errors": []}, "revision": 1, "status": "proposed", } async def _boot(store) -> str: session = await store.create_session(owner_id="u1", title="lease", idempotency_key="ik-s") inserted = await store.insert_plan(session["id"], _plan()) await store.select_plan(session["id"], inserted["id"], idempotency_key="ik-sel", expected_revision=None) run = await store.create_run(session["id"], plan_id=inserted["id"], idempotency_key="ik-run", expected_revision=None) return run["id"] async def _sqlite_store(): engine = create_async_engine("sqlite+aiosqlite:///:memory:") async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) sf = async_sessionmaker(engine, expire_on_commit=False) return make_report_collaboration_store(sf), engine def _kernel() -> ScriptedMemberKernel: kernel = ScriptedMemberKernel() kernel.enqueue("research", TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence(), reply_text="检索完成")) kernel.enqueue("write", TaskResult(kind="artifact_draft", artifact_type="ReportSectionDraft", artifact_draft=_draft(), reply_text="成稿")) return kernel @pytest.mark.asyncio @pytest.mark.parametrize("factory", ["memory", "sqlite"]) async def test_claim_renew_and_expired_reclaim(factory: str) -> None: if factory == "memory": store = MemoryReportCollaborationStore() engine = None else: store, engine = await _sqlite_store() try: run_id = await _boot(store) now = datetime.now(UTC) assert await store.list_claimable(now=now, limit=5) == [run_id] first = await store.claim_run(run_id, lease_owner="w1", lease_until=now + timedelta(seconds=30)) assert first is not None and first["lease_owner"] == "w1" assert await store.claim_run(run_id, lease_owner="w2", lease_until=now + timedelta(seconds=30)) is None assert await store.list_claimable(now=now, limit=5) == [] assert await store.renew_lease(run_id, lease_owner="w1", lease_until=now + timedelta(seconds=60)) is True assert await store.renew_lease(run_id, lease_owner="w2", lease_until=now + timedelta(seconds=60)) is False expired = await store.claim_run(run_id, lease_owner="w2", lease_until=now + timedelta(seconds=30)) assert expired is None past = now - timedelta(seconds=5) record = await store.get_run_record(run_id) assert record is not None if factory == "memory": store._runs[run_id]["lease_until"] = iso(past) else: await store.patch_run(run_id, status="running") async with store._sf() as db: from deerflow.persistence.report_collaboration.model import ReportCollaborationRunRow row = await db.get(ReportCollaborationRunRow, run_id) assert row is not None row.lease_until = past await db.commit() stolen = await store.claim_run(run_id, lease_owner="w2", lease_until=now + timedelta(seconds=30)) assert stolen is not None assert stolen["lease_owner"] == "w2" finally: if engine is not None: await engine.dispose() @pytest.mark.asyncio async def test_recovery_preserves_completed_nodes() -> None: store = MemoryReportCollaborationStore() run_id = await _boot(store) ledger = TaskLedger(store) research = (await ledger.ready_assignments(run_id))[0] await ledger.apply_result( research, TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence()), ) nodes = {item["node_run_id"]: item for item in await store.list_node_runs(run_id)} write = next(item for item in nodes.values() if item["node_id"] == "write") write["status"] = "running" await store.save_node_run(write) report = await recover_orphaned_attempts(store, run_id) assert f"{run_id}-research-attempt1" in report.preserved_completed assert write["node_run_id"] in report.interrupted_node_run_ids latest = {item["node_id"]: item for item in await store.list_node_runs(run_id)} assert latest["research"]["status"] == "completed" assert latest["write"]["attempt"] == 2 assert latest["write"]["status"] == "planned" @pytest.mark.asyncio async def test_executor_completes_claimed_run() -> None: store = MemoryReportCollaborationStore() run_id = await _boot(store) now = datetime.now(UTC) claimed = await store.claim_run(run_id, lease_owner="w1", lease_until=now + timedelta(seconds=30)) assert claimed is not None status = await ReportCollaborationExecutor(store, _kernel()).execute(LeaseParams(run_id=run_id, lease_owner="w1")) assert status == "completed" latest = {item["node_id"]: item for item in await store.list_node_runs(run_id)} assert latest["research"]["status"] == "completed" assert latest["write"]["status"] == "completed" @pytest.mark.asyncio async def test_reclaim_does_not_rerun_completed_research() -> None: store = MemoryReportCollaborationStore() run_id = await _boot(store) ledger = TaskLedger(store) research = (await ledger.ready_assignments(run_id))[0] await ledger.apply_result( research, TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence()), ) write = next(item for item in await store.list_node_runs(run_id) if item["node_id"] == "write") write["status"] = "running" await store.save_node_run(write) now = datetime.now(UTC) await store.claim_run(run_id, lease_owner="w2", lease_until=now + timedelta(seconds=30)) kernel = ScriptedMemberKernel() kernel.enqueue("write", TaskResult(kind="artifact_draft", artifact_type="ReportSectionDraft", artifact_draft=_draft())) status = await ReportCollaborationExecutor(store, kernel).execute(LeaseParams(run_id=run_id, lease_owner="w2")) assert status == "completed" assert kernel.calls == ["write"] artifacts = await store.list_artifacts_full(run_id) assert {item["artifact_type"] for item in artifacts} == {"EvidenceBundle", "ReportSectionDraft"} @pytest.mark.asyncio async def test_lost_lease_stops_local_kernel() -> None: store = MemoryReportCollaborationStore() run_id = await _boot(store) now = datetime.now(UTC) await store.claim_run(run_id, lease_owner="w1", lease_until=now + timedelta(seconds=30)) class _StealAfterResearch(ScriptedMemberKernel): async def run(self, assignment, *, context): result = await super().run(assignment, context=context) store._runs[run_id]["lease_until"] = iso(datetime.now(UTC) - timedelta(seconds=2)) stolen = await store.claim_run(run_id, lease_owner="w2", lease_until=datetime.now(UTC) + timedelta(seconds=30)) assert stolen is not None return result kernel = _StealAfterResearch() kernel.enqueue("research", TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence())) kernel.enqueue("write", TaskResult(kind="artifact_draft", artifact_type="ReportSectionDraft", artifact_draft=_draft())) status = await ReportCollaborationExecutor(store, kernel).execute(LeaseParams(run_id=run_id, lease_owner="w1")) assert status == "running" assert kernel.calls == ["research"] latest = {item["node_id"]: item for item in await store.list_node_runs(run_id)} assert latest["research"]["status"] == "completed" assert latest["write"]["status"] != "completed" @pytest.mark.asyncio async def test_cancel_late_result_is_superseded() -> None: store = MemoryReportCollaborationStore() run_id = await _boot(store) now = datetime.now(UTC) await store.claim_run(run_id, lease_owner="w1", lease_until=now + timedelta(seconds=30)) ledger = TaskLedger(store) research = (await ledger.ready_assignments(run_id))[0] await store.request_cancel_run(run_id, idempotency_key="ik-c") await store.finalize_cancel(run_id, lease_owner="w1") outcome = await ledger.apply_result( research, TaskResult(kind="artifact_draft", artifact_type="EvidenceBundle", artifact_draft=_evidence()), ) assert outcome.run_status == "cancelled" assert outcome.node_status != "completed" artifacts = await store.list_artifacts_full(run_id) assert artifacts assert all(item["validation_status"] == "superseded" for item in artifacts) run = await store.get_run(run_id) assert run["status"] == "cancelled" @pytest.mark.asyncio async def test_dispatch_once_claims_and_start_is_gated() -> None: store = MemoryReportCollaborationStore() run_id = await _boot(store) kernel = _kernel() cfg = ReportCollaborationConfig(enabled=True, worker_enabled=False) executor = ReportCollaborationExecutor(store, kernel, config=cfg) dispatcher = ReportCollaborationDispatcher(store, executor, config=cfg, worker_id="w-test") dispatcher.start() assert dispatcher._task is None started = await dispatcher.dispatch_once() assert started == 1 task = executor._tasks.get(run_id) if task is not None: assert await asyncio.wait_for(task, timeout=5) == "completed" assert (await store.get_run(run_id))["status"] == "completed" cfg_on = ReportCollaborationConfig(enabled=True, worker_enabled=True) live = ReportCollaborationDispatcher(store, executor, config=cfg_on, worker_id="w-live") live.start() assert live._task is not None await live.stop() @pytest.mark.asyncio async def test_dispatch_reaps_cancel_without_running_kernel() -> None: store = MemoryReportCollaborationStore() run_id = await _boot(store) now = datetime.now(UTC) await store.claim_run(run_id, lease_owner="dead", lease_until=now - timedelta(seconds=1)) await store.request_cancel_run(run_id, idempotency_key="ik-dead") kernel = _kernel() dispatcher = ReportCollaborationDispatcher( store, ReportCollaborationExecutor(store, kernel), config=ReportCollaborationConfig(enabled=True, worker_enabled=False), worker_id="reaper", ) started = await dispatcher.dispatch_once() assert started == 0 assert kernel.calls == [] assert (await store.get_run(run_id))["status"] == "cancelled"