mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +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
94 lines
4.0 KiB
Python
94 lines
4.0 KiB
Python
"""Third-party completion stores must opt into the strict metadata contract."""
|
|
|
|
import pytest
|
|
|
|
from deerflow.config.run_ownership_config import RunOwnershipConfig
|
|
from deerflow.runtime import RunManager, RunStatus
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
|
|
class AbsorbingStore(MemoryRunStore):
|
|
async def update_run_completion(self, run_id, *, status, **kwargs):
|
|
self.completion_calls = getattr(self, "completion_calls", 0) + 1
|
|
self._runs[run_id].update(status=status, **kwargs)
|
|
return True
|
|
|
|
|
|
class FlagOnlyChild(AbsorbingStore):
|
|
supports_terminal_completion_kwargs = True
|
|
|
|
|
|
class OptedInStore(MemoryRunStore):
|
|
supports_terminal_completion_kwargs = True
|
|
|
|
async def update_run_completion(self, run_id, *, status, **kwargs):
|
|
self.completion_calls = getattr(self, "completion_calls", 0) + 1
|
|
return await super().update_run_completion(run_id, status=status, **kwargs)
|
|
|
|
|
|
class InheritedOptIn(OptedInStore):
|
|
pass
|
|
|
|
|
|
class UnoptedOverride(OptedInStore):
|
|
async def update_run_completion(self, run_id, *, status, **kwargs):
|
|
self.completion_calls = getattr(self, "completion_calls", 0) + 1
|
|
return await MemoryRunStore.update_run_completion(self, run_id, status=status, **kwargs)
|
|
|
|
|
|
class ExplicitStore(MemoryRunStore):
|
|
async def update_run_completion(self, run_id, *, status, require_terminal=False, owner_worker_id=None, **kwargs):
|
|
self.completion_calls = getattr(self, "completion_calls", 0) + 1
|
|
return await super().update_run_completion(run_id, status=status, require_terminal=require_terminal, owner_worker_id=owner_worker_id, **kwargs)
|
|
|
|
|
|
class PartialContract(AbsorbingStore):
|
|
async def update_run_completion(self, run_id, *, status, require_terminal=False, **kwargs):
|
|
return await super().update_run_completion(run_id, status=status, require_terminal=require_terminal, **kwargs)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("store_type", [AbsorbingStore, FlagOnlyChild, UnoptedOverride, PartialContract])
|
|
async def test_unproven_kwargs_completion_is_skipped(store_type, caplog):
|
|
store = store_type()
|
|
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True))
|
|
record = await manager.create_or_reject("thread")
|
|
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
|
|
proof = record.terminal_commit_proof
|
|
assert proof is not None
|
|
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
|
|
assert getattr(store, "completion_calls", 0) == 0
|
|
row = await store.get(record.run_id)
|
|
assert row.get("total_tokens", 0) == 0
|
|
assert "require_terminal" not in row
|
|
assert record.terminal_commit_proof is proof
|
|
assert not record.ownership_lost
|
|
assert "require_terminal" in caplog.text
|
|
assert "owner_worker_id" in caplog.text
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("store_type", [MemoryRunStore, ExplicitStore, OptedInStore, InheritedOptIn])
|
|
async def test_declared_completion_contract_updates_metadata(store_type):
|
|
store = store_type()
|
|
manager = RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True))
|
|
record = await manager.create_or_reject("thread")
|
|
await manager.set_status_if_not_cancelled(record.run_id, RunStatus.success)
|
|
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
|
|
row = await store.get(record.run_id)
|
|
assert row["total_tokens"] == 7
|
|
assert "require_terminal" not in row
|
|
assert row["owner_worker_id"] == "owner"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_single_worker_preserves_legacy_kwargs_completion():
|
|
store = AbsorbingStore()
|
|
manager = RunManager(store=store)
|
|
record = await manager.create_or_reject("thread")
|
|
await manager.update_run_completion(record.run_id, status="success", total_tokens=7)
|
|
assert store.completion_calls == 1
|
|
row = await store.get(record.run_id)
|
|
assert row["total_tokens"] == 7
|
|
assert "require_terminal" not in row
|