314 lines
13 KiB
Python
314 lines
13 KiB
Python
"""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"
|