diff --git a/backend/packages/harness/deerflow/domain/schedule/__init__.py b/backend/packages/harness/deerflow/domain/schedule/__init__.py index f673e972e..b8e6268f9 100644 --- a/backend/packages/harness/deerflow/domain/schedule/__init__.py +++ b/backend/packages/harness/deerflow/domain/schedule/__init__.py @@ -18,10 +18,12 @@ from deerflow.domain.schedule.commands import ( ) from deerflow.domain.schedule.exceptions import ( ActiveRunConflictError, + ConcurrentUpdateError, CorruptStoredScheduleError, InvalidContextModeError, InvalidScheduleError, LaunchFailedError, + LaunchIndeterminateError, ScheduleError, TaskNotFoundError, TaskNotMutableError, @@ -51,6 +53,7 @@ __all__ = [ "TERMINAL_TASK_STATUSES", "UNSET", "ActiveRunConflictError", + "ConcurrentUpdateError", "ContextChange", "ContextMode", "CorruptStoredScheduleError", @@ -61,6 +64,7 @@ __all__ = [ "InvalidContextModeError", "InvalidScheduleError", "LaunchFailedError", + "LaunchIndeterminateError", "PauseTask", "ResumeTask", "RunStatus", diff --git a/backend/packages/harness/deerflow/domain/schedule/exceptions.py b/backend/packages/harness/deerflow/domain/schedule/exceptions.py index 4743cc504..7db1886f4 100644 --- a/backend/packages/harness/deerflow/domain/schedule/exceptions.py +++ b/backend/packages/harness/deerflow/domain/schedule/exceptions.py @@ -44,6 +44,16 @@ class CorruptStoredScheduleError(ScheduleError): """ +class ConcurrentUpdateError(ScheduleError): + """The aggregate changed between this caller's read and its write. + + Raised by the task repository's `save` when the stored version no longer + matches the aggregate's -- a dispatch or completion committed in between. + The service retries the read-modify-write a bounded number of times and + then lets this surface; the router maps it to a retryable conflict. + """ + + class ActiveRunConflictError(ScheduleError): """The task already holds its single active run slot. diff --git a/backend/packages/harness/deerflow/domain/schedule/model/task.py b/backend/packages/harness/deerflow/domain/schedule/model/task.py index bf5b6e848..534552552 100644 --- a/backend/packages/harness/deerflow/domain/schedule/model/task.py +++ b/backend/packages/harness/deerflow/domain/schedule/model/task.py @@ -82,6 +82,15 @@ class ScheduledTask: last_thread_id: str | None = None last_error: str | None = None run_count: int = 0 + version: int = 0 + """Optimistic-concurrency token, owned by the storage write path. + + Every committed write -- `save`'s compare-and-set, `record_launch`, + `record_completion` -- increments the stored value. The aggregate's own + transitions deliberately leave it alone: it identifies the snapshot this + aggregate was read from, which is exactly what `save` compares against + to refuse a write racing the run lifecycle (`ConcurrentUpdateError`). + """ created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) updated_at: datetime = field(default_factory=lambda: datetime.now(UTC)) diff --git a/backend/packages/harness/deerflow/domain/schedule/ports.py b/backend/packages/harness/deerflow/domain/schedule/ports.py index 0ad531d11..c4b721aba 100644 --- a/backend/packages/harness/deerflow/domain/schedule/ports.py +++ b/backend/packages/harness/deerflow/domain/schedule/ports.py @@ -90,13 +90,25 @@ class ScheduledTaskRepository(Protocol): ... async def save(self, task: ScheduledTask) -> ScheduledTask | None: - """Persist a whole aggregate, keyed by its own id and owner. + """Persist a whole aggregate -- a compare-and-set on its version. + + The write commits only when the stored version still equals + `task.version`; the stored value is then incremented and the stored + state returned. A mismatch raises ConcurrentUpdateError: the aggregate + was read before a dispatch/completion committed, and replacing the row + would roll back the fields those writes own (`next_run_at`, + `run_count`, `last_run_id`), re-arming an already-executed occurrence. Whole-aggregate replacement rather than a field patch: the aggregate is immutable, so a caller that changed anything is holding a complete new - value. Returns None when the row is absent or owned by someone else. + value -- the version check is what makes that shape safe. Returns None + when the row is absent or owned by someone else. Not to be used for the post-dispatch write -- see `record_launch`. + + Raises: + ConcurrentUpdateError: the stored version moved past + `task.version` since the caller's read. """ ... @@ -151,7 +163,9 @@ class ScheduledTaskRepository(Protocol): aggregate would reintroduce the very race the flag exists to close. Every field is assigned unconditionally, so a caller preserving a value - must pass the current one back. The claim is always released. + must pass the current one back. The claim is always released. Like + every committed write, this increments the stored version so a stale + `save` racing it is refused. """ ... @@ -185,7 +199,9 @@ class ScheduledTaskRepository(Protocol): Scoped by `user_id` like every other read: a task belonging to someone else is left alone rather than reported. An unknown task is ignored -- - one deleted mid-flight simply has nothing to update. + one deleted mid-flight simply has nothing to update. Like every + committed write, this increments the stored version so a stale `save` + racing it is refused. """ ... diff --git a/backend/packages/harness/deerflow/domain/schedule/service.py b/backend/packages/harness/deerflow/domain/schedule/service.py index 2959a4139..4d87976c9 100644 --- a/backend/packages/harness/deerflow/domain/schedule/service.py +++ b/backend/packages/harness/deerflow/domain/schedule/service.py @@ -30,6 +30,7 @@ from deerflow.domain.schedule.commands import ( ) from deerflow.domain.schedule.exceptions import ( ActiveRunConflictError, + ConcurrentUpdateError, LaunchFailedError, LaunchIndeterminateError, TaskNotFoundError, @@ -56,6 +57,11 @@ from deerflow.domain.schedule.ports import ( logger = logging.getLogger(__name__) +# How many times a read-modify-write is re-attempted when its version CAS +# loses to a concurrent dispatch/completion write before the conflict is +# surfaced to the caller. +_SAVE_ATTEMPTS = 3 + # Shared so the has_active fast path and the active-slot race path produce # byte-identical outcomes for the same "task already has an active run" # condition. Two callers must not be able to tell which one rejected them. @@ -167,33 +173,40 @@ class ScheduleService: transitions, so the re-arm rule and the running-task gate cannot be bypassed by patching fields directly. """ - task = await self.get_task(cmd.task_id, user_id=cmd.user_id) - task.ensure_mutable() - if cmd.context is not UNSET: - task = task.with_context(cmd.context.context_mode, cmd.context.thread_id) - await self._require_thread(task) - if cmd.schedule is not UNSET: - task = task.with_schedule(cmd.schedule, now=now, policy=self._policy) - if cmd.title is not UNSET: - task = replace(task, title=cmd.title) - if cmd.prompt is not UNSET: - task = replace(task, prompt=cmd.prompt) + async def apply(task: ScheduledTask) -> ScheduledTask: + task.ensure_mutable() + if cmd.context is not UNSET: + task = task.with_context(cmd.context.context_mode, cmd.context.thread_id) + await self._require_thread(task) + if cmd.schedule is not UNSET: + task = task.with_schedule(cmd.schedule, now=now, policy=self._policy) + if cmd.title is not UNSET: + task = replace(task, title=cmd.title) + if cmd.prompt is not UNSET: + task = replace(task, prompt=cmd.prompt) + return task - return await self._save(task) + return await self._mutate_task(cmd.task_id, cmd.user_id, apply) async def pause_task(self, cmd: PauseTask) -> ScheduledTask: """Stop claiming this task until it is resumed. Refused while the task is being dispatched -- see `ensure_mutable`. """ - task = await self.get_task(cmd.task_id, user_id=cmd.user_id) - return await self._save(task.paused()) + + async def apply(task: ScheduledTask) -> ScheduledTask: + return task.paused() + + return await self._mutate_task(cmd.task_id, cmd.user_id, apply) async def resume_task(self, cmd: ResumeTask) -> ScheduledTask: """Re-admit this task to claiming. Same gate as `pause_task`.""" - task = await self.get_task(cmd.task_id, user_id=cmd.user_id) - return await self._save(task.resumed()) + + async def apply(task: ScheduledTask) -> ScheduledTask: + return task.resumed() + + return await self._mutate_task(cmd.task_id, cmd.user_id, apply) async def delete_task(self, cmd: DeleteTask) -> None: """Deleting is deliberately not gated on the task being idle: the @@ -361,14 +374,36 @@ class ScheduleService: if not task.thread_id or not await self._threads.exists_for_user(task.thread_id, task.user_id): raise ThreadNotFoundError("Thread not found") - async def _save(self, task: ScheduledTask) -> ScheduledTask: - """Persist an already-validated aggregate. + async def _mutate_task(self, task_id: str, user_id: str, apply) -> ScheduledTask: + """Read-modify-write with optimistic retry. - Read-modify-write rather than a conditional UPDATE: pushing the - "not while running" rule into a storage predicate would put it beyond - the reach of a zero-IO test and give it a second home. The pre-existing - code had the same shape, and closing the window properly needs - optimistic locking, which needs a schema change. + `save` is a version compare-and-set, so a dispatch or completion + committing between this read and this write is detected rather than + overwritten -- a stale aggregate would roll back `next_run_at`, + `run_count`, and `last_run_id`, and the next poll would re-launch an + already-executed occurrence. On conflict the intent is re-applied to a + fresh read (re-validating every aggregate rule against current state); + a write that keeps losing surfaces ConcurrentUpdateError after + `_SAVE_ATTEMPTS` rounds instead of spinning. + + The business rules stay in `apply` (aggregate transitions), not in a + storage predicate: the version comparison is the only thing the + repository decides. + """ + for attempt in range(_SAVE_ATTEMPTS): + task = await self.get_task(task_id, user_id=user_id) + try: + return await self._save(await apply(task)) + except ConcurrentUpdateError: + if attempt == _SAVE_ATTEMPTS - 1: + raise + raise AssertionError("unreachable: the loop either returns or re-raises") + + async def _save(self, task: ScheduledTask) -> ScheduledTask: + """Persist an already-validated aggregate through the version CAS. + + Raises TaskNotFoundError when the row is absent or owned by someone + else; lets ConcurrentUpdateError propagate to the retry loop above. """ saved = await self._tasks.save(task) if saved is None: diff --git a/backend/tests/schedule_fakes.py b/backend/tests/schedule_fakes.py index 4e7fae5f3..d4b036cf1 100644 --- a/backend/tests/schedule_fakes.py +++ b/backend/tests/schedule_fakes.py @@ -17,7 +17,7 @@ from __future__ import annotations from dataclasses import dataclass, field, replace from datetime import datetime, timedelta -from deerflow.domain.schedule.exceptions import ActiveRunConflictError +from deerflow.domain.schedule.exceptions import ActiveRunConflictError, ConcurrentUpdateError from deerflow.domain.schedule.model import ( ACTIVE_RUN_STATUSES, TERMINAL_RUN_STATUSES, @@ -87,8 +87,10 @@ class InMemoryScheduledTaskRepository: row = self._rows.get(task.task_id) if row is None or row.task.user_id != task.user_id: return None - row.task = task - return task + if row.task.version != task.version: + raise ConcurrentUpdateError(f"scheduled task {task.task_id!r} was modified concurrently (expected version {task.version}, stored {row.task.version})") + row.task = replace(task, version=task.version + 1) + return row.task async def delete(self, task_id: str, *, user_id: str) -> bool: row = self._rows.get(task_id) @@ -118,7 +120,7 @@ class InMemoryScheduledTaskRepository: # Stands in for whatever identity a real adapter records. row.lease_owner = "fake-worker" row.lease_expires_at = now + timedelta(seconds=lease_seconds) - row.task = replace(row.task, status=TaskStatus.RUNNING) + row.task = replace(row.task, status=TaskStatus.RUNNING, version=row.task.version + 1) claimed.append(row.task) return claimed @@ -150,7 +152,7 @@ class InMemoryScheduledTaskRepository: ) if increment_run_count: task = replace(task, run_count=task.run_count + 1) - row.task = task + row.task = replace(task, version=task.version + 1) row.lease_owner = None row.lease_expires_at = None @@ -166,7 +168,7 @@ class InMemoryScheduledTaskRepository: if row is None or row.task.user_id != user_id: return # Only the verdict; every scheduling field belongs to record_launch. - row.task = replace(row.task, last_error=error) + row.task = replace(row.task, last_error=error, version=row.task.version + 1) if status is not None: row.task = replace(row.task, status=status) @@ -175,7 +177,7 @@ class InMemoryScheduledTaskRepository: for row in self._rows.values(): stuck = row.task.status is TaskStatus.RUNNING and row.task.schedule.schedule_type is ScheduleType.ONCE and row.lease_expires_at is None if stuck: - row.task = replace(row.task, status=TaskStatus.CANCELLED, last_error=error) + row.task = replace(row.task, status=TaskStatus.CANCELLED, last_error=error, version=row.task.version + 1) cancelled += 1 return cancelled diff --git a/backend/tests/test_schedule_service.py b/backend/tests/test_schedule_service.py index 2cf8fc9d5..94ef3548c 100644 --- a/backend/tests/test_schedule_service.py +++ b/backend/tests/test_schedule_service.py @@ -35,6 +35,7 @@ from deerflow.domain.schedule.commands import ( UpdateScheduledTask, ) from deerflow.domain.schedule.exceptions import ( + ConcurrentUpdateError, InvalidScheduleError, LaunchFailedError, LaunchIndeterminateError, @@ -915,7 +916,10 @@ class TestTaskManagement: updated = await service.update_scheduled_task(UpdateScheduledTask(task_id=task.task_id, user_id="user-1"), now=NOW) - assert updated == task + # The write still commits (and bumps the version -- every committed + # write does), but no business field moves. + assert updated == replace(task, version=updated.version) + assert updated.version == task.version + 1 async def test_update_can_change_the_prompt(self): service = make_service() @@ -1045,3 +1049,79 @@ class TestTaskManagement: assert await service.list_tasks("someone-else") == [] by_thread = await service.list_tasks_by_thread("user-1", "thread-1") assert [t.task_id for t in by_thread] == [bound.task_id] + + +class _AlwaysStaleTaskRepo(_LaunchWritesMidReadTaskRepo): + """Every read races a `record_launch` commit, so no save can ever win.""" + + async def get(self, task_id: str, *, user_id: str): + self.arm() + return await super().get(task_id, user_id=user_id) + + +class TestOptimisticLocking: + """A stale aggregate must not overwrite fields the run lifecycle owns. + + `update/pause/resume` read the aggregate and later persist it whole; a + dispatch or completion can commit between those two operations. The save + is therefore a version compare-and-set: a stale write is refused, the + service re-reads and re-applies its intent, and only the requested change + lands -- never a rollback of `next_run_at` / `run_count` / `last_run_id`. + """ + + async def test_a_stale_update_does_not_roll_back_run_lifecycle_fields(self): + launch_next = NOW + timedelta(hours=1) + tasks = _LaunchWritesMidReadTaskRepo(launch_next_run_at=launch_next) + service = make_service(tasks=tasks) + task = await create_cron_task(service) + tasks.arm() + + updated = await service.update_scheduled_task( + UpdateScheduledTask(task_id=task.task_id, user_id="user-1", title="Renamed"), + now=NOW, + ) + + assert updated.title == "Renamed" + # The launch bookkeeping that committed mid-update survives; on the + # whole-aggregate write shape these roll back and the next poll + # re-launches an already-executed occurrence. + assert updated.run_count == 1 + assert updated.last_run_id == "run-1" + assert updated.next_run_at == launch_next + + async def test_a_stale_pause_does_not_roll_back_run_lifecycle_fields(self): + launch_next = NOW + timedelta(hours=1) + tasks = _LaunchWritesMidReadTaskRepo(launch_next_run_at=launch_next) + service = make_service(tasks=tasks) + task = await create_cron_task(service) + tasks.arm() + + updated = await service.pause_task(PauseTask(task_id=task.task_id, user_id="user-1")) + + assert updated.status is TaskStatus.PAUSED + assert updated.run_count == 1 + assert updated.next_run_at == launch_next + + async def test_conflicts_exhaust_bounded_retries_and_surface(self): + """The retry loop is bounded: a write that keeps losing must surface + the conflict rather than spin or silently overwrite.""" + tasks = _AlwaysStaleTaskRepo(launch_next_run_at=NOW + timedelta(hours=1)) + service = make_service(tasks=tasks) + task = await create_cron_task(service) + + with pytest.raises(ConcurrentUpdateError): + await service.update_scheduled_task( + UpdateScheduledTask(task_id=task.task_id, user_id="user-1", title="Renamed"), + now=NOW, + ) + + async def test_a_clean_save_bumps_the_version(self): + service = make_service() + task = await create_cron_task(service) + + renamed = await service.update_scheduled_task( + UpdateScheduledTask(task_id=task.task_id, user_id="user-1", title="Renamed"), + now=NOW, + ) + + assert renamed.version == task.version + 1