mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-12 14:00:35 +00:00
fix(runtime): keep assembly evidence from reviving a run a peer reclaimed (#8046)
* fix(runtime): keep assembly evidence from reviving a run a peer reclaimed In heartbeat mode every run-row write is conditional on this worker's live lease except assembly evidence, which went through RunStore.put(). That upsert rewrote status, owner, lease, error and stop_reason from the local record. When a peer had already reclaimed the run, it put the row back to running and erased the takeover error. If the owner's local lease was still live, its renewal then succeeded, the peer never reclaimed it again, and the run finished as success after the peer had published its terminal recovery event. If the write landed late, the row came back with an expired lease and was reclaimed a second time. Add RunStore.merge_owned_metadata(), which merges metadata keys only for the owner holding a live lease on a pending/running row and writes no other field. SQL follows the existing clock -> row lock -> conditional UPDATE order; MemoryRunStore mirrors it. With heartbeat enabled, record_assembly_evidence() uses it, refuses a locally fenced run before any write, fences the run when the store rejects the merge, and fails closed on stores without the method instead of falling back to put(). Single-worker mode keeps the full snapshot. * docs(changelog): reference #8046 for the assembly evidence takeover fix * test(runtime): scope the simulated lease expiry to the run's thread --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
cf8c29c6b6
commit
8850f8eb7e
12
CHANGELOG.md
12
CHANGELOG.md
@ -47,6 +47,17 @@ This section accumulates work toward the **2.3.0** milestone
|
||||
|
||||
### Fixed
|
||||
|
||||
- **runtime:** In multi-worker mode (`run_ownership.heartbeat_enabled`), a
|
||||
worker that had lost its run to a peer could put the row back. Its assembly
|
||||
evidence write was an unconditional full-row upsert, so after a peer
|
||||
reclaimed the run it reset the row to `running` and erased the takeover
|
||||
error. The run could then finish as `success` after the peer had already
|
||||
reported it failed, or be reclaimed a second time. The evidence is now merged
|
||||
into the row's metadata only while the worker still holds a live lease on an
|
||||
active run (`RunStore.merge_owned_metadata()`); otherwise the worker stops
|
||||
the run before executing it. Custom `RunStore` implementations used with
|
||||
heartbeats must implement that method. ([#8046])
|
||||
|
||||
- **agents:** A high-risk tool call without a usable id no longer wedges the
|
||||
thread when `hitl_approval` is enabled. A provider that omits a tool-call id
|
||||
(`None`) or repeats one within a response made the middleware pause with an
|
||||
@ -16270,6 +16281,7 @@ with **180 merged pull requests** since the first 2.0 milestone tag.
|
||||
[#7998]: https://github.com/bytedance/deer-flow/issues/7998
|
||||
[#8006]: https://github.com/bytedance/deer-flow/pull/8006
|
||||
[#8043]: https://github.com/bytedance/deer-flow/pull/8043
|
||||
[#8046]: https://github.com/bytedance/deer-flow/pull/8046
|
||||
[#8057]: https://github.com/bytedance/deer-flow/pull/8057
|
||||
[#8064]: https://github.com/bytedance/deer-flow/pull/8064
|
||||
[#8069]: https://github.com/bytedance/deer-flow/issues/8069
|
||||
|
||||
@ -26,6 +26,13 @@
|
||||
|
||||
### 修复
|
||||
|
||||
- **运行时:** 多 Worker 模式(`run_ownership.heartbeat_enabled`)下,已被其他 Worker
|
||||
接管的运行不再被原 Worker 恢复。此前装配证据通过无条件的整行 upsert 写入,在对端接管后会把
|
||||
行改回 `running` 并清除接管错误;运行可能在对端已报告失败后以 `success` 结束,或被再次接管。
|
||||
现在仅当该 Worker 仍持有活跃运行的有效租约时,才将证据合并到行元数据
|
||||
(`RunStore.merge_owned_metadata()`),否则在执行前停止该运行。启用心跳的自定义
|
||||
`RunStore` 实现必须实现该方法。([#8046])
|
||||
|
||||
- **网关:** 同一用户对同一次运行重复提交反馈时,现在返回 409 而不是 500。
|
||||
`POST /api/threads/{id}/runs/{run_id}/feedback` 每个用户每次运行只允许一条反馈;此前重试,
|
||||
或对同一次运行的另一条消息提交反馈,会触发该唯一约束,Gateway 因未处理的数据库错误而失败。
|
||||
@ -12853,4 +12860,5 @@ DeerFlow 2.0 是围绕"超级智能体"框架的彻底重写,核心包含子
|
||||
[#7993]: https://github.com/bytedance/deer-flow/pull/7993
|
||||
[#8006]: https://github.com/bytedance/deer-flow/pull/8006
|
||||
[#8043]: https://github.com/bytedance/deer-flow/pull/8043
|
||||
[#8046]: https://github.com/bytedance/deer-flow/pull/8046
|
||||
[#8064]: https://github.com/bytedance/deer-flow/pull/8064
|
||||
|
||||
@ -365,7 +365,14 @@ while its replacement awaits a Store result. Preserve terminal-attribution and
|
||||
cancellation compensation below; changing lock granularity does not grant a new
|
||||
worker authority over a durable row.
|
||||
|
||||
Assembly evidence persists a full row snapshot. Serialize that write with the
|
||||
With heartbeat enabled, assembly evidence is merged into the row's metadata
|
||||
through `RunStore.merge_owned_metadata()`, which requires this Worker's live
|
||||
lease on a `pending`/`running` row and writes no other field. A full snapshot
|
||||
would also rewrite status, owner, lease and error, reviving a row a peer already
|
||||
reclaimed. A rejection fences the run and fails assembly before execution; a
|
||||
locally fenced run is refused before any write, and a Store without the method
|
||||
fails closed instead of falling back to `put()`. Without heartbeat, evidence
|
||||
still persists a full row snapshot. Serialize either write with the
|
||||
same run's lease renewal and local acknowledgment, so it cannot restore an old
|
||||
lease after renewal. Lock order is thread gate, then snapshot gate, then the
|
||||
short registry lock. Renewal uses only the snapshot gate and must release it
|
||||
|
||||
@ -78,10 +78,10 @@ recheck policy, aborting on activation; retained-snapshot isolation fails closed
|
||||
Strict startup cannot downgrade durable policy.
|
||||
[Offline cutover](../../../../../README.md#configuration).
|
||||
|
||||
`RunRepository.update_status(..., owner_worker_id=...)` and
|
||||
`finalize_if_not_cancelled(..., owner_worker_id=...)` require that owner and a
|
||||
live lease. Increment the change clock before taking the run row lock to retain
|
||||
the existing clock/row lock order. PostgreSQL checks `clock_timestamp()` after
|
||||
`RunRepository.update_status`/`finalize_if_not_cancelled` with an owner, and
|
||||
`merge_owned_metadata`, require that owner and a live lease. Increment the
|
||||
change clock before taking the run row lock to retain the existing clock/row
|
||||
lock order. PostgreSQL checks `clock_timestamp()` after
|
||||
`SELECT FOR UPDATE`, rather than transaction-start `now()`; SQLite evaluates its
|
||||
database clock in the conditional UPDATE after acquiring its writer lock.
|
||||
`renew_lease` and legacy `update_lease` also reject expired existing leases and
|
||||
|
||||
@ -992,6 +992,29 @@ class RunRepository(RunStore):
|
||||
cancel_action=cancel_action,
|
||||
)
|
||||
|
||||
async def merge_owned_metadata(self, run_id: str, metadata: dict[str, Any], *, owner_worker_id: str) -> bool:
|
||||
"""Merge metadata keys only while this owner holds a live lease on an active row."""
|
||||
async with self._sf() as session:
|
||||
change_seq = await self._next_change_seq(session)
|
||||
await self._lock_owned_row(session, run_id)
|
||||
current = await session.scalar(select(RunRow.metadata_json).where(RunRow.run_id == run_id))
|
||||
result = await session.execute(
|
||||
update(RunRow)
|
||||
.where(
|
||||
RunRow.run_id == run_id,
|
||||
RunRow.owner_worker_id == owner_worker_id,
|
||||
RunRow.status.in_(("pending", "running")),
|
||||
self._live_lease_predicate(session),
|
||||
)
|
||||
.values(
|
||||
metadata_json={**(current or {}), **(self._safe_json(metadata) or {})},
|
||||
updated_at=datetime.now(UTC),
|
||||
change_seq=change_seq,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
return result.rowcount != 0
|
||||
|
||||
async def claim_for_takeover(
|
||||
self,
|
||||
run_id: str,
|
||||
|
||||
@ -1942,15 +1942,43 @@ class RunManager:
|
||||
record = self._runs.get(run_id)
|
||||
if record is None or evidence.get("run_id") != run_id or evidence.get("thread_id") != record.thread_id:
|
||||
raise ValueError("Assembly/run binding mismatch")
|
||||
if record.ownership_lost:
|
||||
raise RuntimeError("Unable to persist assembly evidence: this worker lost lease ownership of the run")
|
||||
if EVIDENCE_KEY in record.metadata:
|
||||
if record.metadata[EVIDENCE_KEY] != evidence:
|
||||
raise ValueError("Run assembly evidence is immutable")
|
||||
return
|
||||
record.metadata = {**record.metadata, EVIDENCE_KEY: copy.deepcopy(evidence)}
|
||||
if not await self._persist_to_store(record):
|
||||
owned = self.heartbeat_enabled and self._store is not None and not record.store_only
|
||||
# A full snapshot would also rewrite status, owner, lease and error,
|
||||
# undoing a peer's takeover after this worker's lease expired.
|
||||
persisted = await self._merge_owned_metadata(record, {EVIDENCE_KEY: evidence}) if owned else await self._persist_to_store(record)
|
||||
if not persisted:
|
||||
record.metadata.pop(EVIDENCE_KEY, None)
|
||||
raise RuntimeError("Unable to persist assembly evidence")
|
||||
|
||||
async def _merge_owned_metadata(self, record: RunRecord, metadata: dict[str, Any]) -> bool:
|
||||
"""Write metadata keys only under this worker's live lease; fence on rejection."""
|
||||
try:
|
||||
merged = await self._call_store_with_retry(
|
||||
"merge_owned_metadata",
|
||||
record.run_id,
|
||||
lambda: self._store.merge_owned_metadata(record.run_id, copy.deepcopy(metadata), owner_worker_id=self._worker_id),
|
||||
)
|
||||
except NotImplementedError:
|
||||
logger.warning(
|
||||
"Unable to persist metadata for run %s: heartbeat RunStore.merge_owned_metadata must atomically enforce the owner and a live lease",
|
||||
record.run_id,
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
logger.warning("Failed to persist owned metadata for run %s", record.run_id, exc_info=True)
|
||||
return False
|
||||
if merged is not True:
|
||||
await self._mark_ownership_lost(record, reason="The durable store rejected this worker's metadata write.")
|
||||
return False
|
||||
return not record.ownership_lost
|
||||
|
||||
async def update_model_name(self, run_id: str, model_name: str | None) -> None:
|
||||
"""Update the model name for a run."""
|
||||
async with self._run_registry_lock(run_id):
|
||||
|
||||
@ -439,6 +439,23 @@ class RunStore(abc.ABC):
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def merge_owned_metadata(
|
||||
self,
|
||||
run_id: str,
|
||||
metadata: dict[str, Any],
|
||||
*,
|
||||
owner_worker_id: str,
|
||||
) -> bool:
|
||||
"""Merge top-level keys into an active run's metadata under a live owner lease.
|
||||
|
||||
Atomically require ``owner_worker_id``, an unexpired lease and a
|
||||
``pending``/``running`` row; change no other field. Return True only
|
||||
after the write commits. False means a peer took over, the lease
|
||||
expired or the run is no longer active. Legacy stores cannot enforce
|
||||
these conditions and fail closed instead of writing a full snapshot.
|
||||
"""
|
||||
raise NotImplementedError("Store must implement owner/lease-fenced metadata merge")
|
||||
|
||||
async def finalize_if_not_cancelled(
|
||||
self,
|
||||
run_id: str,
|
||||
|
||||
@ -414,6 +414,15 @@ class MemoryRunStore(RunStore):
|
||||
|
||||
return run.get("owner_worker_id") == owner_worker_id and not is_lease_expired(run.get("lease_expires_at"), grace_seconds=0)
|
||||
|
||||
async def merge_owned_metadata(self, run_id, metadata, *, owner_worker_id):
|
||||
run = self._runs.get(run_id)
|
||||
if run is None or run["status"] not in ("pending", "running") or not self._live_owner(run, owner_worker_id):
|
||||
return False
|
||||
run["metadata"] = {**(run.get("metadata") or {}), **copy.deepcopy(metadata)}
|
||||
run["updated_at"] = datetime.now(UTC).isoformat()
|
||||
self._mark_changed(run)
|
||||
return True
|
||||
|
||||
async def update_lease(
|
||||
self,
|
||||
run_id: str,
|
||||
|
||||
@ -512,6 +512,68 @@ async def test_evidence_persistence_failure_prevents_execution(runtime, monkeypa
|
||||
await runs.configuration_evidence(thread_id=thread, run_id=run.run_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_peer_takeover_during_assembly_is_not_undone_by_evidence(runtime, monkeypatch, tmp_path):
|
||||
"""A peer reclaims the run while this worker assembles; evidence must not revive the row."""
|
||||
import asyncio
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from sqlalchemy import update
|
||||
from test_run_repository import _cleanup, _make_repo
|
||||
|
||||
from deerflow.config.run_ownership_config import RunOwnershipConfig
|
||||
from deerflow.persistence.run.model import RunRow
|
||||
from deerflow.runtime import RunManager
|
||||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||||
|
||||
runs, state, _, _, handles = runtime
|
||||
repository = await _make_repo(tmp_path)
|
||||
state.run_store = repository
|
||||
await state.run_manager.shutdown()
|
||||
ownership = RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=30, grace_seconds=1)
|
||||
state.run_manager = RunManager(store=repository, run_ownership_config=ownership, worker_id="worker-a")
|
||||
peer = RunManager(store=repository, run_ownership_config=ownership, worker_id="worker-b")
|
||||
loop = asyncio.get_running_loop()
|
||||
taken_over = []
|
||||
|
||||
async def peer_takeover():
|
||||
# The peer sees an expired lease (clock skew or a stalled owner) that
|
||||
# the owning worker still believes is live.
|
||||
async with repository._sf() as session:
|
||||
await session.execute(update(RunRow).where(RunRow.thread_id == thread).values(lease_expires_at=datetime.now(UTC) - timedelta(minutes=1)))
|
||||
await session.commit()
|
||||
taken_over.extend(await peer.reconcile_orphaned_inflight_runs(error="reclaimed by worker-b", stop_reason="orphan_recovered"))
|
||||
|
||||
import app.gateway.services as services
|
||||
|
||||
assemble = services.resolve_agent_factory("lead_agent")
|
||||
|
||||
def factory_with_takeover(*, config):
|
||||
asyncio.run_coroutine_threadsafe(peer_takeover(), loop).result(timeout=5)
|
||||
return assemble(config=config)
|
||||
|
||||
monkeypatch.setattr("app.gateway.services.resolve_agent_factory", lambda _: factory_with_takeover)
|
||||
try:
|
||||
thread = await runs.create_thread(thread_id="takeover-during-assembly")
|
||||
run = await runs.start(thread_id=thread, input={"messages": [{"role": "user", "content": "must not run"}]})
|
||||
record = await state.run_manager.get(run.run_id)
|
||||
await asyncio.wait({record.task}, timeout=5)
|
||||
assert [record.run_id for record in taken_over] == [run.run_id]
|
||||
row = await repository.get(run.run_id, user_id=None)
|
||||
assert row["status"] == "error"
|
||||
assert row["error"] == "reclaimed by worker-b"
|
||||
assert row["stop_reason"] == "orphan_recovered"
|
||||
assert EVIDENCE_KEY not in row["metadata"]
|
||||
assert record.ownership_lost
|
||||
assert not handles
|
||||
# The owner's renewal stays rejected, so the peer's outcome is final.
|
||||
await state.run_manager._renew_leases()
|
||||
assert (await repository.get(run.run_id, user_id=None))["status"] == "error"
|
||||
finally:
|
||||
await state.run_manager.shutdown()
|
||||
await _cleanup()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_durable_configuration_revision_survives_host_replacement_and_rotation(host, monkeypatch):
|
||||
control, request, load_user = host
|
||||
|
||||
@ -2892,3 +2892,94 @@ async def test_cancellation_during_terminal_receipt_cannot_commit_staged_success
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(asyncio.gather(record.task, return_exceptions=True), 1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Assembly evidence publication under lease ownership
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _evidence(record) -> dict:
|
||||
return {"run_id": record.run_id, "thread_id": record.thread_id, "version": "v1"}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_assembly_evidence_after_peer_takeover_fences_instead_of_reviving_row():
|
||||
"""A heartbeat owner must not rewrite a row a peer already reclaimed."""
|
||||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||||
|
||||
store = MemoryRunStore()
|
||||
config = _lease_config(heartbeat_enabled=True, grace_seconds=1)
|
||||
owner = _make_manager(store=store, run_ownership_config=config, worker_id="worker-a")
|
||||
peer = _make_manager(store=store, run_ownership_config=config, worker_id="worker-b")
|
||||
record = await owner.create_or_reject("thread-1")
|
||||
await owner.set_status(record.run_id, RunStatus.running)
|
||||
# The peer sees an expired lease that the owner still believes is live.
|
||||
store._runs[record.run_id]["lease_expires_at"] = (datetime.now(UTC) - timedelta(minutes=1)).isoformat()
|
||||
assert [r.run_id for r in await peer.reconcile_orphaned_inflight_runs(error="reclaimed", stop_reason=ORPHAN_RECOVERY_STOP_REASON)] == [record.run_id]
|
||||
reclaimed = dict(store._runs[record.run_id])
|
||||
|
||||
with pytest.raises(RuntimeError, match="assembly evidence"):
|
||||
await owner.record_assembly_evidence(record.run_id, _evidence(record))
|
||||
|
||||
assert store._runs[record.run_id] == reclaimed
|
||||
assert record.ownership_lost
|
||||
assert EVIDENCE_KEY not in record.metadata
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_assembly_evidence_is_merged_without_rewriting_the_owned_row():
|
||||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||||
|
||||
store = MemoryRunStore()
|
||||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True), worker_id="worker-a")
|
||||
record = await manager.create_or_reject("thread-1", metadata={"kept": True})
|
||||
await manager.set_status(record.run_id, RunStatus.running)
|
||||
original_put = store.put
|
||||
store.put = AsyncMock(wraps=original_put)
|
||||
|
||||
await manager.record_assembly_evidence(record.run_id, _evidence(record))
|
||||
|
||||
store.put.assert_not_awaited()
|
||||
|
||||
row = await store.get(record.run_id)
|
||||
assert row["metadata"] == {"kept": True, EVIDENCE_KEY: _evidence(record)}
|
||||
assert row["status"] == "running"
|
||||
assert not record.ownership_lost
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_assembly_evidence_fails_closed_on_store_without_owned_metadata_merge():
|
||||
from deerflow.runtime.assembly_evidence import EVIDENCE_KEY
|
||||
from deerflow.runtime.runs.store.base import RunStore
|
||||
|
||||
class LegacyHeartbeatStore(MemoryRunStore):
|
||||
merge_owned_metadata = RunStore.merge_owned_metadata
|
||||
|
||||
store = LegacyHeartbeatStore()
|
||||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True), worker_id="worker-a")
|
||||
record = await manager.create_or_reject("thread-1")
|
||||
# A full snapshot cannot carry the owner/lease conditions.
|
||||
store.put = AsyncMock()
|
||||
|
||||
with pytest.raises(RuntimeError, match="assembly evidence"):
|
||||
await manager.record_assembly_evidence(record.run_id, _evidence(record))
|
||||
|
||||
store.put.assert_not_awaited()
|
||||
|
||||
assert EVIDENCE_KEY not in (await store.get(record.run_id))["metadata"]
|
||||
assert EVIDENCE_KEY not in record.metadata
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_assembly_evidence_is_refused_after_local_fence():
|
||||
store = MemoryRunStore()
|
||||
manager = _make_manager(store=store, run_ownership_config=_lease_config(heartbeat_enabled=True), worker_id="worker-a")
|
||||
record = await manager.create_or_reject("thread-1")
|
||||
await manager.set_status(record.run_id, RunStatus.running)
|
||||
assert await manager._mark_ownership_lost(record, reason="lease expired")
|
||||
store.merge_owned_metadata = AsyncMock(side_effect=AssertionError("a fenced worker must not write"))
|
||||
store.put = AsyncMock(side_effect=AssertionError("a fenced worker must not write"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="lost lease ownership"):
|
||||
await manager.record_assembly_evidence(record.run_id, _evidence(record))
|
||||
|
||||
@ -3,6 +3,8 @@
|
||||
Uses a temp SQLite DB to test ORM-backed CRUD operations.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import copy
|
||||
import os
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from unittest.mock import AsyncMock
|
||||
@ -1604,3 +1606,101 @@ class TestMemoryRunStoreDeleteByThread:
|
||||
# list_by_thread is run-scoped, so the surviving reservation is not listed
|
||||
# even though the row is still tracked.
|
||||
assert await store.list_by_thread("t1", user_id="alice") == []
|
||||
|
||||
|
||||
class TestMergeOwnedMetadata:
|
||||
"""Both stores merge metadata only for the live owner of an active row."""
|
||||
|
||||
@staticmethod
|
||||
@contextlib.asynccontextmanager
|
||||
async def _store(kind, tmp_path):
|
||||
if kind == "memory":
|
||||
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
||||
|
||||
yield MemoryRunStore()
|
||||
return
|
||||
if kind == "sqlite":
|
||||
try:
|
||||
yield await _make_repo(tmp_path)
|
||||
finally:
|
||||
await _cleanup()
|
||||
return
|
||||
uri = os.environ.get("TEST_POSTGRES_URI") or os.environ.get("DEERFLOW_TEST_POSTGRES_URL")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI or DEERFLOW_TEST_POSTGRES_URL")
|
||||
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.persistence.engine import get_engine, get_session_factory, init_engine_from_config
|
||||
|
||||
schema = f"owned_metadata_{uuid4().hex}"
|
||||
await init_engine_from_config(DatabaseConfig(backend="postgres", postgres_url=asyncpg_test_url(uri), postgres_schema=schema))
|
||||
try:
|
||||
yield RunRepository(get_session_factory())
|
||||
finally:
|
||||
async with get_engine().begin() as connection:
|
||||
await connection.execute(DropSchema(schema, cascade=True))
|
||||
await _cleanup()
|
||||
|
||||
@staticmethod
|
||||
async def _put_owned(store, *, status="running", lease_delta=timedelta(seconds=30)):
|
||||
lease = (datetime.now(UTC) + lease_delta).isoformat()
|
||||
await store.put(
|
||||
"run-1",
|
||||
thread_id="t1",
|
||||
status=status,
|
||||
owner_worker_id="worker-a",
|
||||
lease_expires_at=lease,
|
||||
error="kept",
|
||||
metadata={"existing": 1},
|
||||
)
|
||||
# MemoryRunStore.get returns its live row; compare against a copy.
|
||||
return copy.deepcopy(await store.get("run-1"))
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("kind", ["memory", "sqlite", "postgresql"])
|
||||
@pytest.mark.parametrize("status", ["pending", "running"])
|
||||
async def test_live_owner_merges_only_metadata(self, kind, status, tmp_path):
|
||||
async with self._store(kind, tmp_path) as store:
|
||||
before = await self._put_owned(store, status=status)
|
||||
|
||||
assert await store.merge_owned_metadata("run-1", {"evidence": {"v": 1}}, owner_worker_id="worker-a") is True
|
||||
|
||||
after = await store.get("run-1")
|
||||
assert after["metadata"] == {"existing": 1, "evidence": {"v": 1}}
|
||||
for key in ("status", "owner_worker_id", "lease_expires_at", "error", "stop_reason"):
|
||||
assert after[key] == before[key]
|
||||
assert after["change_seq"] > before["change_seq"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("kind", ["memory", "sqlite", "postgresql"])
|
||||
@pytest.mark.parametrize("case", ["foreign-owner", "expired-lease", "terminal-row", "peer-takeover"])
|
||||
async def test_rejects_write_without_live_ownership(self, kind, case, tmp_path):
|
||||
async with self._store(kind, tmp_path) as store:
|
||||
owner = "worker-a"
|
||||
if case == "foreign-owner":
|
||||
await self._put_owned(store)
|
||||
owner = "worker-b"
|
||||
elif case == "expired-lease":
|
||||
await self._put_owned(store, lease_delta=timedelta(seconds=-5))
|
||||
elif case == "terminal-row":
|
||||
await self._put_owned(store, status="success")
|
||||
else:
|
||||
await self._put_owned(store, lease_delta=timedelta(seconds=-60))
|
||||
assert await store.claim_for_takeover("run-1", grace_seconds=10, error="reclaimed", stop_reason="orphan_recovered") is True
|
||||
before = copy.deepcopy(await store.get("run-1"))
|
||||
|
||||
assert await store.merge_owned_metadata("run-1", {"evidence": {"v": 1}}, owner_worker_id=owner) is False
|
||||
|
||||
after = await store.get("run-1")
|
||||
for key in ("status", "metadata", "owner_worker_id", "lease_expires_at", "error", "stop_reason"):
|
||||
assert after[key] == before[key]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_missing_row_is_rejected(self, tmp_path):
|
||||
repo = await _make_repo(tmp_path)
|
||||
try:
|
||||
assert await repo.merge_owned_metadata("missing", {"evidence": {}}, owner_worker_id="worker-a") is False
|
||||
finally:
|
||||
await _cleanup()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user