diff --git a/backend/packages/harness/deerflow/domain/schedule/__init__.py b/backend/packages/harness/deerflow/domain/schedule/__init__.py index 38497879a..8c89fcd89 100644 --- a/backend/packages/harness/deerflow/domain/schedule/__init__.py +++ b/backend/packages/harness/deerflow/domain/schedule/__init__.py @@ -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 ( diff --git a/backend/packages/harness/deerflow/domain/schedule/ports.py b/backend/packages/harness/deerflow/domain/schedule/ports.py new file mode 100644 index 000000000..f491f8f35 --- /dev/null +++ b/backend/packages/harness/deerflow/domain/schedule/ports.py @@ -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. + """ + ... diff --git a/backend/tests/schedule_fakes.py b/backend/tests/schedule_fakes.py new file mode 100644 index 000000000..3d1ffdbb9 --- /dev/null +++ b/backend/tests/schedule_fakes.py @@ -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 diff --git a/backend/tests/test_schedule_fakes.py b/backend/tests/test_schedule_fakes.py new file mode 100644 index 000000000..72a14b2c8 --- /dev/null +++ b/backend/tests/test_schedule_fakes.py @@ -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