github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-08-01 19:06:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(schedule): declare the schedule output ports

Browse Source 
Four Protocols the domain declares and the outer ring will implement,
plus the two DTOs that keep infrastructure types out of the inner ring:

- ScheduledTaskRepository / ScheduledRunRepository, exchanging domain
  objects rather than the bare dicts the current repositories return
- RunLauncher, whose contract is that only ThreadBusyError or
  LaunchFailedError may escape -- that translation is what keeps the run
  runtime and the web framework out of the domain
- ThreadLookup, one method rather than the whole thread store
- LaunchedRun and RunOutcome, so the completion path stops taking a
  runtime record the purity test would reject

Two deliberate departures from the earlier sketch. There is no Clock
port: `now` is already an explicit parameter throughout, so the domain
never reads a clock and the tests are already deterministic -- adding
one would only create a second source of truth for the same value. And
`record_launch` is not expressed as `save(task)`, because
`protect_terminal` makes it a compare-and-set against a concurrently
finalizing run; a read-modify-write through the aggregate would
reintroduce the race the flag exists to close.

The in-memory doubles model the active-slot rule rather than skipping
it: a double that never refuses a second active run would let the
service's conflict collapse go untested. Their semantics are pinned by
test_schedule_fakes.py, which becomes the contract suite once the SQL
adapters land and both tiers run the same cases.

Concurrency is out of scope for the doubles and says so in their module
docstring -- it stays covered against a real database in
test_scheduled_task_dispatch_race.py.
This commit is contained in:
rayhpeng 2026-07-28 11:27:11 +08:00
parent 45626f6383
commit ab5166ef07
4 changed files with 952 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
backend/packages/harness/deerflow/domain/schedule/__init__.py
View File

@ -1,10 +1,10 @@
"""Schedule bounded context: standing instructions to run a prompt on time.
Public API of the context. Import domain objects from here; ports will live in
Public API of the context. Import domain objects from here; import ports from
`deerflow.domain.schedule.ports` -- they are contracts consumed by adapters and
tests, not everyday call-site symbols.
`ScheduleService` is not exported yet: this commit lands the model only.
`ScheduleService` is not exported yet: the application service has not landed.
"""
from deerflow.domain.schedule.model import (

297
backend/packages/harness/deerflow/domain/schedule/ports.py Normal file
View File

@ -0,0 +1,297 @@
"""Output ports of the schedule context.
Contracts the domain declares and the outer ring implements. Signatures are
technology-neutral: no SQL, table names, HTTP status codes, or run-runtime
types appear here, and every method exchanges domain objects.
The docstrings are longer than the code on purpose -- they are the semantic
contract the adapter must honour and the contract tests are written from.
"""
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from typing import Protocol, runtime_checkable
from deerflow.domain.schedule.model import (
RunStatus,
ScheduledRun,
ScheduledTask,
TaskStatus,
)
@dataclass(frozen=True)
class LaunchedRun:
"""What the launcher reports back once a run is admitted.
`thread_id` is echoed rather than assumed: the launcher is free to return a
different thread than the one requested, and the task's bookkeeping records
what actually ran.
"""
run_id: str
thread_id: str
@dataclass(frozen=True)
class RunOutcome:
"""A launched run reaching a terminal state, in domain vocabulary.
The app layer converts its own run record into this before calling the
service, so the domain never imports the run runtime. That converter also
owns the filtering the completion hook used to do inline -- a run that
carries no scheduled-task metadata, or has not reached a terminal state,
simply produces no RunOutcome and the service is never called.
`status` is narrowed to the three terminal outcomes a scheduled run can
report: SUCCESS, FAILED, INTERRUPTED. INTERRUPTED is deliberately distinct
from FAILED -- a cancel or same-thread takeover is not an execution
failure, and the task ends CANCELLED rather than FAILED.
"""
task_id: str
record_id: str
run_id: str
user_id: str
status: RunStatus
error: str | None
@runtime_checkable
class ScheduledTaskRepository(Protocol):
"""Persistence port for the task aggregate.
Every read is scoped by `user_id`: a task belonging to someone else is
reported as absent (None / False / omitted from a list), never as a
permission error -- the caller must not be able to distinguish "not yours"
from "does not exist".
"""
async def add(self, task: ScheduledTask) -> ScheduledTask:
"""Insert a new task and return the stored state."""
...
async def get(self, task_id: str, *, user_id: str) -> ScheduledTask | None:
"""Return the task, or None when it is absent or owned by someone else."""
...
async def list_by_user(self, user_id: str) -> list[ScheduledTask]:
"""Every task owned by the user, newest first."""
...
async def list_by_user_and_thread(self, user_id: str, thread_id: str) -> list[ScheduledTask]:
"""The user's tasks bound to one thread, newest first.
Only `reuse_thread` tasks can match: a `fresh_thread_per_run` task
carries no thread.
"""
...
async def save(self, task: ScheduledTask) -> ScheduledTask | None:
"""Persist a whole aggregate, keyed by its own id and owner.
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.
Not to be used for the post-dispatch write -- see `record_launch`.
"""
...
async def delete(self, task_id: str, *, user_id: str) -> bool:
"""Remove the task. False when it was absent or owned by someone else."""
...
async def claim_due(self, *, now: datetime, lease_owner: str, lease_seconds: int, limit: int) -> list[ScheduledTask]:
"""Atomically take ownership of up to `limit` tasks that are due.
A task is due when its next fire time has passed AND either it is
claimable (enabled, with no live claim), or it is stuck mid-dispatch
with an expired claim -- the process that took it died between claiming
and dispatching, and it must not stay unreachable forever.
Claiming marks the tasks as running and stamps the claim, so a
concurrent claimer cannot take the same ones. Returns them in the state
they were left in *after* the claim.
`lease_owner` is recorded for diagnostics only and is never read back:
expiry alone decides whether a claim can be taken over. Do not build
ownership enforcement on it without also adding the read side.
Atomicity is the implementation's responsibility. An in-memory double
can satisfy every rule above under single-threaded use while providing
no concurrency guarantee at all; that difference is out of scope for
the contract tests and is covered separately against a real database.
"""
...
async def record_launch(
self,
task_id: str,
*,
status: TaskStatus,
next_run_at: datetime | None,
last_run_at: datetime | None,
last_run_id: str | None,
last_thread_id: str | None,
last_error: str | None,
increment_run_count: bool,
protect_terminal: bool = False,
) -> None:
"""Write the outcome of a dispatch and release the claim.
Deliberately NOT expressed as `save(task)`: `protect_terminal` makes
this a compare-and-set against a run that may be finalizing
concurrently. When it is set and the stored task has already reached a
terminal status, the status and error are left alone and only the
scheduling bookkeeping is written -- a read-modify-write through the
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.
"""
...
async def cancel_stuck_once_tasks(self, *, error: str) -> int:
"""Reconcile `once` tasks orphaned mid-flight by a process crash.
A launched `once` task waits in running for its completion hook, and
its claim was released at launch -- so the expired-claim branch of
`claim_due` can never see it, and after a crash the hook is gone. This
marks those tasks cancelled and returns how many were reconciled.
Tasks still holding a claim are left alone: they were claimed but not
launched, and expired-claim reclaim recovers them safely.
"""
...
@runtime_checkable
class ScheduledRunRepository(Protocol):
"""Persistence port for execution records.
Unlike tasks, runs are not read by owner: they are only ever reached
through a task the caller already proved it owns.
"""
async def add(self, run: ScheduledRun) -> ScheduledRun:
"""Insert an execution record.
Raises:
ActiveRunConflictError: the task already holds its single active
slot. Translating the storage-level rejection into this domain
error is the adapter's job and is load-bearing -- the service
collapses it to exactly the same outcome as the `has_active`
fast path, and those two paths must stay indistinguishable.
A terminal record (a skip tombstone) is outside the active-slot rule
and must never raise it.
"""
...
async def list_by_task(self, task_id: str, *, limit: int, offset: int) -> list[ScheduledRun]:
"""One task's execution history, newest first."""
...
async def count_active(self) -> int:
"""How many executions are active across ALL tasks.
Global on purpose: it bounds total scheduled concurrency, not per-task
overlap. Long runs accumulate across polls, so each poll may only claim
into whatever budget is left.
"""
...
async def has_active(self, task_id: str) -> bool:
"""Whether this task currently holds its active slot.
A non-atomic fast path by nature -- it cannot be relied on to exclude a
concurrent dispatch. `add` is the arbiter; this exists to avoid the
common case of doing pointless work.
"""
...
async def update_status(
self,
record_id: str,
*,
status: RunStatus,
run_id: str | None = None,
error: str | None = None,
started_at: datetime | None = None,
finished_at: datetime | None = None,
protect_terminal: bool = False,
) -> None:
"""Advance an execution record.
With `protect_terminal`, a record that already reached a terminal state
keeps its status and error; only bookkeeping the terminal write could
not have known (the run id, the start time) is backfilled. A run that
fails fast can reach its completion hook before the launch path's own
write lands, and the completion is the authoritative one.
Unknown record ids are ignored rather than raising: the caller is
writing bookkeeping, not asserting existence.
"""
...
async def mark_stale_active(self, *, error: str) -> int:
"""Terminalize executions orphaned by a process crash, returning the count.
Runs execute in-process, so any record still active at startup belongs
to a process that is gone. This is only sound while a single scheduler
instance owns the table.
"""
...
class RunLauncher(Protocol):
"""Output port for actually starting the work.
The contract the adapter MUST honour, because it is what keeps the run
runtime and the web framework out of the inner ring:
- the execution thread is already busy -> ThreadBusyError
- anything else goes wrong -> LaunchFailedError
Nothing else may escape. The domain distinguishes those two because they
lead to different outcomes -- a busy thread on a scheduled dispatch is a
skipped occurrence, while a genuine failure is recorded as one.
"""
async def launch(
self,
*,
thread_id: str,
assistant_id: str | None,
prompt: str,
owner_user_id: str | None,
metadata: dict[str, str],
) -> LaunchedRun:
"""Start one execution and return its identity.
`metadata` is opaque correlation data the domain attaches so the
eventual outcome can be traced back to this task and record; the
adapter must carry it through untouched.
"""
...
class ThreadLookup(Protocol):
"""Narrow port: the only question the schedule context asks about threads.
Deliberately not the full thread store -- depending on this one method
keeps the context decoupled from the wider conversation model.
"""
async def exists_for_user(self, thread_id: str, user_id: str) -> bool:
"""Whether this thread exists AND the user may use it.
Both halves matter: binding a task to a thread that does not exist yet
is as invalid as binding it to someone else's. The two are deliberately
not distinguished in the result -- reporting them differently would let
a caller probe for the existence of threads they cannot see.
"""
...

275
backend/tests/schedule_fakes.py Normal file
View File

@ -0,0 +1,275 @@
"""Shared zero-IO doubles for the schedule context.
Lives in its own module rather than inside a test file so the port contract
suite and the service tests can both import it as an ordinary module, without
relying on the tests directory happening to be on ``sys.path``.
**Concurrency is explicitly out of scope.** These doubles model the
single-threaded semantics of each port -- which rows `claim_due` selects, that
`add` refuses a second active run -- and provide no atomicity whatsoever. A
green run here says nothing about two dispatchers racing; that is covered
against a real database in ``test_scheduled_task_dispatch_race.py``. Do not
read a passing contract suite as licence to run more than one scheduler.
"""
from __future__ import annotations
from dataclasses import dataclass, field, replace
from datetime import datetime, timedelta
from deerflow.domain.schedule.model import (
ACTIVE_RUN_STATUSES,
TERMINAL_RUN_STATUSES,
TERMINAL_TASK_STATUSES,
ActiveRunConflictError,
RunStatus,
ScheduledRun,
ScheduledTask,
TaskStatus,
)
from deerflow.domain.schedule.ports import LaunchedRun
@dataclass
class _TaskRow:
"""One stored task: the aggregate plus the columns the domain has no
business knowing about.
The lease lives here rather than on ``ScheduledTask`` for the same reason
it lives in the table and not in the domain -- it is a storage-level
concurrency control, not a business fact about the task.
"""
task: ScheduledTask
lease_owner: str | None = None
lease_expires_at: datetime | None = None
class InMemoryScheduledTaskRepository:
"""ScheduledTaskRepository double backed by a dict."""
def __init__(self) -> None:
self._rows: dict[str, _TaskRow] = {}
# -- inspection helpers (tests only, not part of the port) -----------
def lease_of(self, task_id: str) -> tuple[str | None, datetime | None]:
row = self._rows[task_id]
return row.lease_owner, row.lease_expires_at
def seed(self, task: ScheduledTask, *, lease_owner: str | None = None, lease_expires_at: datetime | None = None) -> ScheduledTask:
"""Install a task directly, optionally already claimed."""
self._rows[task.task_id] = _TaskRow(task=task, lease_owner=lease_owner, lease_expires_at=lease_expires_at)
return task
# -- port ------------------------------------------------------------
async def add(self, task: ScheduledTask) -> ScheduledTask:
self._rows[task.task_id] = _TaskRow(task=task)
return task
async def get(self, task_id: str, *, user_id: str) -> ScheduledTask | None:
row = self._rows.get(task_id)
if row is None or row.task.user_id != user_id:
return None
return row.task
async def list_by_user(self, user_id: str) -> list[ScheduledTask]:
rows = [row.task for row in self._rows.values() if row.task.user_id == user_id]
return sorted(rows, key=lambda t: (t.created_at, t.task_id), reverse=True)
async def list_by_user_and_thread(self, user_id: str, thread_id: str) -> list[ScheduledTask]:
tasks = await self.list_by_user(user_id)
return [task for task in tasks if task.thread_id == thread_id]
async def save(self, task: ScheduledTask) -> ScheduledTask | None:
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
async def delete(self, task_id: str, *, user_id: str) -> bool:
row = self._rows.get(task_id)
if row is None or row.task.user_id != user_id:
return False
del self._rows[task_id]
return True
async def claim_due(self, *, now: datetime, lease_owner: str, lease_seconds: int, limit: int) -> list[ScheduledTask]:
def claimable(row: _TaskRow) -> bool:
task = row.task
if task.next_run_at is None or task.next_run_at > now:
return False
expired = row.lease_expires_at is not None and row.lease_expires_at < now
if task.status is TaskStatus.ENABLED:
return row.lease_expires_at is None or expired
# Stuck mid-dispatch: the claimer died between claim and launch.
return task.status is TaskStatus.RUNNING and expired
due = sorted(
(row for row in self._rows.values() if claimable(row)),
key=lambda row: (row.task.next_run_at, row.task.task_id),
)[:limit]
claimed = []
for row in due:
row.lease_owner = lease_owner
row.lease_expires_at = now + timedelta(seconds=lease_seconds)
row.task = replace(row.task, status=TaskStatus.RUNNING)
claimed.append(row.task)
return claimed
async def record_launch(
self,
task_id: str,
*,
status: TaskStatus,
next_run_at: datetime | None,
last_run_at: datetime | None,
last_run_id: str | None,
last_thread_id: str | None,
last_error: str | None,
increment_run_count: bool,
protect_terminal: bool = False,
) -> None:
row = self._rows.get(task_id)
if row is None:
return
task = row.task
if not (protect_terminal and task.status in TERMINAL_TASK_STATUSES):
task = replace(task, status=status, last_error=last_error)
task = replace(
task,
next_run_at=next_run_at,
last_run_at=last_run_at,
last_run_id=last_run_id,
last_thread_id=last_thread_id,
)
if increment_run_count:
task = replace(task, run_count=task.run_count + 1)
row.task = task
row.lease_owner = None
row.lease_expires_at = None
async def cancel_stuck_once_tasks(self, *, error: str) -> int:
cancelled = 0
for row in self._rows.values():
stuck = row.task.status is TaskStatus.RUNNING and row.task.schedule.schedule_type.value == "once" and row.lease_expires_at is None
if stuck:
row.task = replace(row.task, status=TaskStatus.CANCELLED, last_error=error)
cancelled += 1
return cancelled
class InMemoryScheduledRunRepository:
"""ScheduledRunRepository double that also models the active-slot rule.
``add`` refusing a second active record for one task is not an
implementation detail to be skipped here: the service collapses that
rejection to the same outcome as its own fast path, and a double that never
raises would let that collapse go untested.
"""
def __init__(self) -> None:
self._rows: dict[str, ScheduledRun] = {}
def all_runs(self) -> list[ScheduledRun]:
"""Inspection helper (tests only, not part of the port)."""
return list(self._rows.values())
async def add(self, run: ScheduledRun) -> ScheduledRun:
if run.is_active and any(other.task_id == run.task_id and other.is_active for other in self._rows.values()):
raise ActiveRunConflictError(f"scheduled task {run.task_id!r} already has an active run")
self._rows[run.record_id] = run
return run
async def list_by_task(self, task_id: str, *, limit: int = 50, offset: int = 0) -> list[ScheduledRun]:
rows = [run for run in self._rows.values() if run.task_id == task_id]
rows.sort(key=lambda run: (run.created_at, run.record_id), reverse=True)
return rows[offset : offset + limit]
async def count_active(self) -> int:
return sum(1 for run in self._rows.values() if run.is_active)
async def has_active(self, task_id: str) -> bool:
return any(run.task_id == task_id and run.is_active for run in self._rows.values())
async def update_status(
self,
record_id: str,
*,
status: RunStatus,
run_id: str | None = None,
error: str | None = None,
started_at: datetime | None = None,
finished_at: datetime | None = None,
protect_terminal: bool = False,
) -> None:
run = self._rows.get(record_id)
if run is None:
return
if protect_terminal and run.status in TERMINAL_RUN_STATUSES:
# Keep the terminal verdict; only backfill what it could not know.
if run.run_id is None and run_id is not None:
run = replace(run, run_id=run_id)
if run.started_at is None and started_at is not None:
run = replace(run, started_at=started_at)
self._rows[record_id] = run
return
run = replace(run, status=status, run_id=run_id, error=error)
if started_at is not None:
run = replace(run, started_at=started_at)
if finished_at is not None:
run = replace(run, finished_at=finished_at)
self._rows[record_id] = run
async def mark_stale_active(self, *, error: str) -> int:
stale = [run for run in self._rows.values() if run.status in ACTIVE_RUN_STATUSES]
for run in stale:
self._rows[run.record_id] = replace(run, status=RunStatus.INTERRUPTED, error=error)
return len(stale)
@dataclass
class FakeRunLauncher:
"""RunLauncher double.
`fail_with` makes every launch raise instead -- set it to a ThreadBusyError
or a LaunchFailedError to drive the two branches the domain distinguishes.
"""
fail_with: Exception | None = None
calls: list[dict] = field(default_factory=list)
async def launch(
self,
*,
thread_id: str,
assistant_id: str | None,
prompt: str,
owner_user_id: str | None,
metadata: dict[str, str],
) -> LaunchedRun:
self.calls.append(
{
"thread_id": thread_id,
"assistant_id": assistant_id,
"prompt": prompt,
"owner_user_id": owner_user_id,
"metadata": metadata,
}
)
if self.fail_with is not None:
raise self.fail_with
return LaunchedRun(run_id=f"run-{len(self.calls)}", thread_id=thread_id)
class FakeThreadLookup:
"""ThreadLookup double backed by a thread_id -> owner mapping."""
def __init__(self, threads: dict[str, str] | None = None) -> None:
self._threads = dict(threads or {})
async def exists_for_user(self, thread_id: str, user_id: str) -> bool:
return self._threads.get(thread_id) == user_id

378
backend/tests/test_schedule_fakes.py Normal file
View File

@ -0,0 +1,378 @@
"""Semantics of the in-memory schedule doubles.
These are the rules ``docs`` calls the port contract: what `claim_due` selects,
when `add` refuses, what `protect_terminal` preserves. They are asserted here
against the fakes so the doubles cannot drift from the ports they stand in for
while the service tests lean on them.
This file is the seed of the contract suite: when the SQL adapters land, the
same cases run against both implementations and any divergence between them
becomes a failure rather than a surprise in production. Concurrency stays out
of scope in both tiers -- see the note in ``schedule_fakes``.
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
import pytest
from schedule_fakes import (
FakeRunLauncher,
FakeThreadLookup,
InMemoryScheduledRunRepository,
InMemoryScheduledTaskRepository,
)
from deerflow.domain.schedule.model import (
ActiveRunConflictError,
ContextMode,
RunStatus,
ScheduledRun,
ScheduledTask,
ScheduleSpec,
TaskStatus,
TriggerKind,
)
from deerflow.domain.schedule.ports import ScheduledRunRepository, ScheduledTaskRepository
pytestmark = pytest.mark.asyncio
NOW = datetime(2026, 7, 27, 12, 0, tzinfo=UTC)
CRON = ScheduleSpec.cron_schedule("0 9 * * *", "UTC")
ONCE = ScheduleSpec.once_at(datetime(2026, 8, 1, 9, 0, tzinfo=UTC), "UTC")
def make_task(
task_id: str = "task-1",
*,
user_id: str = "user-1",
status: TaskStatus = TaskStatus.ENABLED,
next_run_at: datetime | None = None,
schedule: ScheduleSpec = CRON,
thread_id: str | None = None,
context_mode: ContextMode = ContextMode.FRESH_THREAD_PER_RUN,
) -> ScheduledTask:
return ScheduledTask(
task_id=task_id,
user_id=user_id,
title=task_id,
prompt="do the thing",
schedule=schedule,
status=status,
next_run_at=next_run_at,
thread_id=thread_id,
context_mode=context_mode,
)
class TestProtocolConformance:
"""`runtime_checkable` only checks that the methods exist, not their
signatures -- enough to catch a rename that updates one side only.
Declared async purely to match this module's asyncio mark; there is
nothing to await.
"""
async def test_task_repository_satisfies_the_port(self):
assert isinstance(InMemoryScheduledTaskRepository(), ScheduledTaskRepository)
async def test_run_repository_satisfies_the_port(self):
assert isinstance(InMemoryScheduledRunRepository(), ScheduledRunRepository)
class TestTaskOwnershipIsolation:
"""Another user's task must read as absent, never as forbidden."""
async def test_get_hides_another_users_task(self):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task(user_id="owner"))
assert await repo.get("task-1", user_id="owner") is not None
assert await repo.get("task-1", user_id="intruder") is None
async def test_save_refuses_another_users_task(self):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task(user_id="owner"))
assert await repo.save(make_task(user_id="intruder")) is None
assert (await repo.get("task-1", user_id="owner")).user_id == "owner"
async def test_delete_refuses_another_users_task(self):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task(user_id="owner"))
assert await repo.delete("task-1", user_id="intruder") is False
assert await repo.get("task-1", user_id="owner") is not None
async def test_list_by_thread_only_matches_bound_tasks(self):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task("bound", context_mode=ContextMode.REUSE_THREAD, thread_id="thread-1"))
await repo.add(make_task("unbound"))
listed = await repo.list_by_user_and_thread("user-1", "thread-1")
assert [task.task_id for task in listed] == ["bound"]
class TestClaimDue:
async def test_claims_an_enabled_due_task_and_stamps_the_lease(self):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task(next_run_at=NOW - timedelta(minutes=1)))
claimed = await repo.claim_due(now=NOW, lease_owner="worker-1", lease_seconds=120, limit=10)
assert [task.task_id for task in claimed] == ["task-1"]
assert claimed[0].status is TaskStatus.RUNNING
owner, expires = repo.lease_of("task-1")
assert owner == "worker-1"
assert expires == NOW + timedelta(seconds=120)
@pytest.mark.parametrize(
("label", "task_kwargs"),
[
("not yet due", {"next_run_at": NOW + timedelta(minutes=1)}),
("never scheduled", {"next_run_at": None}),
("paused", {"status": TaskStatus.PAUSED, "next_run_at": NOW - timedelta(minutes=1)}),
("completed", {"status": TaskStatus.COMPLETED, "next_run_at": NOW - timedelta(minutes=1)}),
],
)
async def test_does_not_claim(self, label, task_kwargs):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task(**task_kwargs))
assert await repo.claim_due(now=NOW, lease_owner="w", lease_seconds=120, limit=10) == [], label
async def test_does_not_steal_a_live_claim(self):
repo = InMemoryScheduledTaskRepository()
repo.seed(
make_task(next_run_at=NOW - timedelta(minutes=1)),
lease_owner="worker-1",
lease_expires_at=NOW + timedelta(seconds=60),
)
assert await repo.claim_due(now=NOW, lease_owner="worker-2", lease_seconds=120, limit=10) == []
async def test_reclaims_a_task_stuck_mid_dispatch(self):
"""The claimer died between claiming and launching: status is running,
the lease has expired, and the task must not stay unreachable."""
repo = InMemoryScheduledTaskRepository()
repo.seed(
make_task(status=TaskStatus.RUNNING, next_run_at=NOW - timedelta(minutes=1)),
lease_owner="dead-worker",
lease_expires_at=NOW - timedelta(seconds=1),
)
claimed = await repo.claim_due(now=NOW, lease_owner="worker-2", lease_seconds=120, limit=10)
assert [task.task_id for task in claimed] == ["task-1"]
assert repo.lease_of("task-1")[0] == "worker-2"
async def test_claims_the_most_overdue_first_and_honours_the_limit(self):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task("late", next_run_at=NOW - timedelta(hours=2)))
await repo.add(make_task("later", next_run_at=NOW - timedelta(hours=1)))
claimed = await repo.claim_due(now=NOW, lease_owner="w", lease_seconds=120, limit=1)
assert [task.task_id for task in claimed] == ["late"]
class TestRecordLaunch:
async def test_writes_bookkeeping_and_releases_the_claim(self):
repo = InMemoryScheduledTaskRepository()
repo.seed(make_task(next_run_at=NOW), lease_owner="w", lease_expires_at=NOW + timedelta(seconds=60))
await repo.record_launch(
"task-1",
status=TaskStatus.ENABLED,
next_run_at=NOW + timedelta(days=1),
last_run_at=NOW,
last_run_id="run-1",
last_thread_id="thread-1",
last_error=None,
increment_run_count=True,
)
task = await repo.get("task-1", user_id="user-1")
assert task.status is TaskStatus.ENABLED
assert task.last_run_id == "run-1"
assert task.run_count == 1
assert repo.lease_of("task-1") == (None, None)
async def test_protect_terminal_keeps_a_concurrently_finalized_verdict(self):
"""A fast-failing run's completion hook lands before the launch path's
own write; the completion is authoritative."""
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task(status=TaskStatus.COMPLETED, schedule=ONCE))
await repo.record_launch(
"task-1",
status=TaskStatus.RUNNING,
next_run_at=None,
last_run_at=NOW,
last_run_id="run-1",
last_thread_id="thread-1",
last_error="stale",
increment_run_count=True,
protect_terminal=True,
)
task = await repo.get("task-1", user_id="user-1")
assert task.status is TaskStatus.COMPLETED, "the terminal status must survive"
assert task.last_error is None, "the terminal error must survive"
assert task.last_run_id == "run-1", "bookkeeping is still recorded"
assert task.run_count == 1
async def test_without_protect_terminal_the_write_wins(self):
repo = InMemoryScheduledTaskRepository()
await repo.add(make_task(status=TaskStatus.COMPLETED, schedule=ONCE))
await repo.record_launch(
"task-1",
status=TaskStatus.FAILED,
next_run_at=None,
last_run_at=NOW,
last_run_id=None,
last_thread_id=None,
last_error="boom",
increment_run_count=False,
)
task = await repo.get("task-1", user_id="user-1")
assert task.status is TaskStatus.FAILED
assert task.last_error == "boom"
async def test_unknown_task_is_ignored(self):
repo = InMemoryScheduledTaskRepository()
await repo.record_launch(
"nope",
status=TaskStatus.ENABLED,
next_run_at=None,
last_run_at=None,
last_run_id=None,
last_thread_id=None,
last_error=None,
increment_run_count=False,
)
class TestCancelStuckOnceTasks:
async def test_cancels_a_launched_once_task_with_no_claim(self):
repo = InMemoryScheduledTaskRepository()
repo.seed(make_task(status=TaskStatus.RUNNING, schedule=ONCE))
assert await repo.cancel_stuck_once_tasks(error="restarted") == 1
task = await repo.get("task-1", user_id="user-1")
assert task.status is TaskStatus.CANCELLED
assert task.last_error == "restarted"
async def test_leaves_a_claimed_task_to_lease_expiry(self):
"""Claimed but not launched -- expired-claim reclaim recovers it, and
cancelling here would throw away a dispatch that never happened."""
repo = InMemoryScheduledTaskRepository()
repo.seed(
make_task(status=TaskStatus.RUNNING, schedule=ONCE),
lease_owner="w",
lease_expires_at=NOW + timedelta(seconds=60),
)
assert await repo.cancel_stuck_once_tasks(error="restarted") == 0
async def test_leaves_cron_tasks_alone(self):
repo = InMemoryScheduledTaskRepository()
repo.seed(make_task(status=TaskStatus.RUNNING, schedule=CRON))
assert await repo.cancel_stuck_once_tasks(error="restarted") == 0
class TestActiveSlot:
def _queued(self, task_id: str = "task-1") -> ScheduledRun:
return ScheduledRun.queued(task_id=task_id, thread_id="thread-1", scheduled_for=NOW, trigger=TriggerKind.SCHEDULED)
def _tombstone(self, task_id: str = "task-1") -> ScheduledRun:
return ScheduledRun.skipped_tombstone(task_id=task_id, thread_id="thread-1", scheduled_for=NOW, trigger=TriggerKind.SCHEDULED)
async def test_second_active_record_is_refused(self):
repo = InMemoryScheduledRunRepository()
await repo.add(self._queued())
with pytest.raises(ActiveRunConflictError):
await repo.add(self._queued())
async def test_a_tombstone_never_conflicts(self):
"""Terminal from birth, so it sits outside the active-slot rule -- this
is why the skip path cannot reuse the queued factory."""
repo = InMemoryScheduledRunRepository()
await repo.add(self._queued())
await repo.add(self._tombstone())
assert await repo.count_active() == 1
async def test_another_task_is_unaffected(self):
repo = InMemoryScheduledRunRepository()
await repo.add(self._queued("task-1"))
await repo.add(self._queued("task-2"))
assert await repo.count_active() == 2
async def test_slot_frees_up_once_the_record_terminalizes(self):
repo = InMemoryScheduledRunRepository()
first = await repo.add(self._queued())
await repo.update_status(first.record_id, status=RunStatus.SUCCESS, finished_at=NOW)
await repo.add(self._queued())
assert await repo.count_active() == 1
async def test_has_active_is_scoped_to_one_task_while_count_is_global(self):
repo = InMemoryScheduledRunRepository()
await repo.add(self._queued("task-1"))
await repo.add(self._queued("task-2"))
assert await repo.has_active("task-1") is True
assert await repo.has_active("task-3") is False
assert await repo.count_active() == 2
class TestRunStatusWrites:
async def test_protect_terminal_backfills_without_overwriting(self):
repo = InMemoryScheduledRunRepository()
run = await repo.add(ScheduledRun.queued(task_id="t", thread_id="th", scheduled_for=NOW, trigger=TriggerKind.SCHEDULED))
await repo.update_status(run.record_id, status=RunStatus.FAILED, error="boom", finished_at=NOW)
# The launch path's write arrives late.
await repo.update_status(run.record_id, status=RunStatus.RUNNING, run_id="run-1", started_at=NOW, protect_terminal=True)
stored = repo.all_runs()[0]
assert stored.status is RunStatus.FAILED
assert stored.error == "boom"
assert stored.run_id == "run-1", "the id the completion could not know is backfilled"
assert stored.started_at == NOW
async def test_unknown_record_is_ignored(self):
repo = InMemoryScheduledRunRepository()
await repo.update_status("nope", status=RunStatus.SUCCESS)
async def test_mark_stale_active_terminalizes_orphans(self):
repo = InMemoryScheduledRunRepository()
active = await repo.add(ScheduledRun.queued(task_id="t", thread_id="th", scheduled_for=NOW, trigger=TriggerKind.SCHEDULED))
done = await repo.add(ScheduledRun.skipped_tombstone(task_id="t", thread_id="th", scheduled_for=NOW, trigger=TriggerKind.SCHEDULED))
assert await repo.mark_stale_active(error="gateway restarted") == 1
by_id = {run.record_id: run for run in repo.all_runs()}
assert by_id[active.record_id].status is RunStatus.INTERRUPTED
assert by_id[active.record_id].error == "gateway restarted"
assert by_id[done.record_id].status is RunStatus.SKIPPED
class TestLauncherAndThreadLookup:
async def test_launcher_records_the_call_and_echoes_the_thread(self):
launcher = FakeRunLauncher()
launched = await launcher.launch(
thread_id="thread-1",
assistant_id="lead_agent",
prompt="go",
owner_user_id="user-1",
metadata={"scheduled_task_id": "task-1"},
)
assert launched.thread_id == "thread-1"
assert launcher.calls[0]["metadata"] == {"scheduled_task_id": "task-1"}
async def test_launcher_can_be_driven_into_either_failure_branch(self):
boom = RuntimeError("nope")
launcher = FakeRunLauncher(fail_with=boom)
with pytest.raises(RuntimeError):
await launcher.launch(thread_id="t", assistant_id=None, prompt="p", owner_user_id=None, metadata={})
assert len(launcher.calls) == 1, "the attempt is still recorded"
async def test_thread_lookup_requires_both_existence_and_ownership(self):
lookup = FakeThreadLookup({"thread-1": "user-1"})
assert await lookup.exists_for_user("thread-1", "user-1") is True
assert await lookup.exists_for_user("thread-1", "user-2") is False
assert await lookup.exists_for_user("missing", "user-1") is False