mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
850b3b2547
* fix(runtime): confirm terminal commits and isolate lease renewal * fix(runtime): diagnose unsupported terminal store contracts * fix(runtime): attribute atomic claims and close lease lifecycle gaps * fix(runtime): require explicit terminal completion capability * fix(runtime): retry renewals that race rollback admission
197 lines
9.4 KiB
Python
197 lines
9.4 KiB
Python
"""Execute the terminal/lease contract against memory and real SQL backends."""
|
|
|
|
import os
|
|
import uuid
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
from support.postgres import asyncpg_test_url
|
|
|
|
from deerflow.persistence.run import RunRepository
|
|
from deerflow.persistence.run.model import RunChangeClockRow, RunRow
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
|
|
@pytest.fixture(params=["memory", "sqlite", "postgresql"])
|
|
async def store(request, tmp_path):
|
|
if request.param == "memory":
|
|
yield MemoryRunStore()
|
|
return
|
|
schema = None
|
|
if request.param == "postgresql":
|
|
uri = os.getenv("TEST_POSTGRES_URI")
|
|
if not uri:
|
|
pytest.skip("set TEST_POSTGRES_URI for real PostgreSQL terminal/lease contracts")
|
|
schema = f"r2_{uuid.uuid4().hex}"
|
|
engine = create_async_engine(asyncpg_test_url(uri), connect_args={"server_settings": {"search_path": schema}})
|
|
else:
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'r2.db'}")
|
|
try:
|
|
async with engine.begin() as conn:
|
|
if schema:
|
|
await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
|
|
await conn.run_sync(lambda c: RunRow.metadata.create_all(c, tables=[RunRow.__table__, RunChangeClockRow.__table__]))
|
|
yield RunRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
finally:
|
|
if schema:
|
|
async with engine.begin() as conn:
|
|
await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))
|
|
await engine.dispose()
|
|
|
|
|
|
async def put(store, *, expired=False):
|
|
expiry = (datetime.now(UTC) + timedelta(seconds=-1 if expired else 30)).isoformat()
|
|
await store.put("run", thread_id="thread", status="running", owner_worker_id="owner", lease_expires_at=expiry)
|
|
return expiry
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_expired_lease_cannot_be_resurrected(store):
|
|
expired = await put(store, expired=True)
|
|
renewal = await store.renew_lease("run", owner_worker_id="owner", lease_expires_at=(datetime.now(UTC) + timedelta(seconds=60)).isoformat())
|
|
assert not renewal.renewed
|
|
assert datetime.fromisoformat((await store.get("run"))["lease_expires_at"]) == datetime.fromisoformat(expired)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("expired,owner", [(True, "owner"), (False, "peer")])
|
|
async def test_terminal_cas_requires_live_owner(store, expired, owner):
|
|
await put(store, expired=expired)
|
|
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id=owner)
|
|
assert not result.finalized
|
|
assert (await store.get("run"))["status"] == "running"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancel_wins_and_metadata_cannot_bypass_cas(store):
|
|
await put(store)
|
|
assert await store.request_cancel("run", action="rollback") == "rollback"
|
|
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")
|
|
assert not result.finalized
|
|
assert result.cancel_action == "rollback"
|
|
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", total_tokens=5) is False
|
|
assert (await store.get("run"))["status"] == "running"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_confirmed_terminal_allows_only_matching_metadata(store):
|
|
await put(store)
|
|
assert (await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")).finalized
|
|
assert await store.update_run_completion("run", status="error", require_terminal=True, owner_worker_id="owner") is False
|
|
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="peer") is False
|
|
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", total_tokens=5) is True
|
|
assert (await store.get("run"))["total_tokens"] == 5
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_terminal_metadata_needs_matching_owner_and_outcome_but_not_live_lease(store):
|
|
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
|
|
await store.put("run", thread_id="thread", status="success", owner_worker_id="owner", lease_expires_at=expired)
|
|
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", total_tokens=5)
|
|
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="peer", total_tokens=10) is False
|
|
assert await store.update_run_completion("run", status="error", require_terminal=True, owner_worker_id="owner", total_tokens=10) is False
|
|
row = await store.get("run")
|
|
assert row["status"] == "success"
|
|
assert row["total_tokens"] == 5
|
|
assert datetime.fromisoformat(row["lease_expires_at"]) == datetime.fromisoformat(expired)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("method", ["renew_lease", "update_lease", "finalize_if_not_cancelled", "update_status"])
|
|
async def test_postgres_checks_database_time_after_row_lock(store, monkeypatch, method):
|
|
"""A transaction begun before expiry cannot write after waiting past it."""
|
|
if not isinstance(store, RunRepository) or store._sf.kw["bind"].dialect.name != "postgresql":
|
|
pytest.skip("requires PostgreSQL row locks and clock_timestamp")
|
|
import asyncio
|
|
|
|
from sqlalchemy import func, select
|
|
|
|
expiry = datetime.now(UTC) + timedelta(milliseconds=300)
|
|
await store.put("locked", thread_id="locked-thread", status="running", owner_worker_id="owner", lease_expires_at=expiry.isoformat())
|
|
entered = asyncio.Event()
|
|
lock = store._lock_owned_row
|
|
|
|
async def observe_transaction_start(session, run_id):
|
|
assert await session.scalar(select(func.clock_timestamp())) < expiry
|
|
entered.set()
|
|
await lock(session, run_id)
|
|
|
|
monkeypatch.setattr(store, "_lock_owned_row", observe_transaction_start)
|
|
operation = None
|
|
async with store._sf() as blocker:
|
|
await blocker.execute(select(RunRow).where(RunRow.run_id == "locked").with_for_update())
|
|
try:
|
|
if method in ("renew_lease", "update_lease"):
|
|
operation = asyncio.create_task(getattr(store, method)("locked", owner_worker_id="owner", lease_expires_at=(expiry + timedelta(seconds=30)).isoformat()))
|
|
elif method == "update_status":
|
|
operation = asyncio.create_task(store.update_status("locked", "error", owner_worker_id="owner"))
|
|
else:
|
|
operation = asyncio.create_task(store.finalize_if_not_cancelled("locked", status="success", owner_worker_id="owner"))
|
|
await asyncio.wait_for(entered.wait(), 1)
|
|
async with asyncio.timeout(2):
|
|
while await blocker.scalar(select(func.clock_timestamp())) <= expiry:
|
|
await asyncio.sleep(0.01)
|
|
assert not operation.done()
|
|
finally:
|
|
await blocker.rollback()
|
|
assert operation is not None
|
|
result = await asyncio.wait_for(operation, 1)
|
|
if method == "renew_lease":
|
|
assert not result.renewed
|
|
elif method == "finalize_if_not_cancelled":
|
|
assert not result.finalized
|
|
else:
|
|
assert result is False
|
|
row = await store.get("locked")
|
|
assert row["status"] == "running"
|
|
assert datetime.fromisoformat(row["lease_expires_at"]) == expiry
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_durable_cancel_does_not_authorize_expired_terminal_writer(store):
|
|
await put(store, expired=True)
|
|
assert await store.request_cancel("run", action="rollback") == "rollback"
|
|
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")
|
|
assert not result.finalized
|
|
assert result.cancel_action is None
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_renewal_cannot_write_an_already_expired_target(store):
|
|
original = await put(store)
|
|
renewal = await store.renew_lease("run", owner_worker_id="owner", lease_expires_at=(datetime.now(UTC) - timedelta(seconds=1)).isoformat())
|
|
assert not renewal.renewed
|
|
assert datetime.fromisoformat((await store.get("run"))["lease_expires_at"]) == datetime.fromisoformat(original)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_strict_completion_cannot_rewrite_terminal_outcome_fields(store):
|
|
await put(store)
|
|
verdict = {"satisfied": True}
|
|
assert (await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner", goal_verdict=verdict)).finalized
|
|
assert await store.update_run_completion("run", status="success", require_terminal=True, owner_worker_id="owner", error="changed", goal_verdict={"satisfied": False}, total_tokens=5)
|
|
row = await store.get("run")
|
|
assert row.get("error") is None
|
|
assert row["goal_verdict"] == verdict
|
|
assert row["total_tokens"] == 5
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_terminal_cancel_row_does_not_report_active_cancellation_win(store):
|
|
await put(store)
|
|
assert await store.request_cancel("run", action="interrupt") == "interrupt"
|
|
assert await store.update_status("run", "interrupted", owner_worker_id="owner")
|
|
result = await store.finalize_if_not_cancelled("run", status="success", owner_worker_id="owner")
|
|
assert not result.finalized
|
|
assert result.cancel_action is None
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_owned_status_success_cannot_bypass_durable_cancel(store):
|
|
await put(store)
|
|
assert await store.request_cancel("run", action="interrupt") == "interrupt"
|
|
assert await store.update_status("run", "success", owner_worker_id="owner") is False
|
|
assert (await store.get("run"))["status"] == "running"
|