From b30102039867e7f1fb07760ec934b69f4973a661 Mon Sep 17 00:00:00 2001 From: Renn F Date: Sun, 3 May 2026 04:47:51 +0200 Subject: [PATCH] feat(orchestrator): reap stale claims via last_heartbeat_at Dispatch loop now releases tasks whose holder has gone silent past ROBOCO_CLAIM_HEARTBEAT_TTL_SECONDS (default 300s). Closes the 'dead container squats task forever' failure mode that the schema hinted at but no code enforced. --- roboco/config.py | 8 ++ roboco/runtime/orchestrator.py | 78 +++++++++++++++ roboco/services/task.py | 35 +++++++ tests/unit/runtime/test_stale_claim_reaper.py | 94 +++++++++++++++++++ 4 files changed, 215 insertions(+) create mode 100644 tests/unit/runtime/test_stale_claim_reaper.py diff --git a/roboco/config.py b/roboco/config.py index 69e0209c..452b06bb 100644 --- a/roboco/config.py +++ b/roboco/config.py @@ -305,6 +305,14 @@ class Settings(BaseSettings): ge=60, description="Claim heartbeat staleness threshold (seconds)", ) + claim_heartbeat_ttl_seconds: int = Field( + default=300, + ge=60, + description=( + "If an agent's last_heartbeat_at is older than this, the dispatcher" + " considers the claim dead and releases the task back to pending." + ), + ) spawn_cooldown_seconds: int = Field( default=60, ge=1, diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index cd81119b..39484cc1 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -27,6 +27,7 @@ import httpx if TYPE_CHECKING: from roboco.services.llm import AgentRoute + from roboco.services.task import TaskService import structlog from fastapi import status as http_status @@ -450,6 +451,12 @@ class AgentOrchestrator: # is in a loop — without this gate the orchestrator re-spawns every # tick forever (seen in production on 2026-04-22). self._pm_respawn_tracker: dict[tuple[str, str], dict[str, Any]] = {} + # Stale-claim reaper config + injectable service slot. Production + # code leaves `_task_svc` None so the reaper opens its own per-tick + # session; tests can pre-set a mock on an instance built via + # `__new__` to bypass this dance. + self._claim_heartbeat_ttl: int = settings.claim_heartbeat_ttl_seconds + self._task_svc: TaskService | None = None # ========================================================================= # LIFECYCLE @@ -3523,6 +3530,63 @@ Start now: evidence(task_id="{task_id}") except Exception as e: logger.error("Dispatcher loop error", error=str(e)) + async def _reap_stale_claims(self) -> None: + """Release claimed/in_progress tasks whose holder hasn't heart-beat in TTL. + + Closes the "dead container squats task forever" failure mode that + the schema hinted at (``last_heartbeat_at`` since migration 006) but + no code enforced. The runtime decision (cutoff, iteration) lives + here in the orchestrator; the actual UPDATE statements live in + ``TaskService.unclaim_for_reaper``. + + If an instance has a pre-bound ``_task_svc`` (used by tests and + conceivable future caching), use it directly. Otherwise open a + per-tick session — short-lived because the reaper runs on every + dispatch cycle and the work is cheap (one SELECT plus N UPDATEs + for the typically-empty stale set). + """ + if self._task_svc is not None: + await self._reap_with_service(self._task_svc) + return + + from roboco.db.base import get_session_factory + from roboco.services.task import TaskService + + factory = get_session_factory() + async with factory() as db: + svc = TaskService(db) + await self._reap_with_service(svc) + await db.commit() + + async def _reap_with_service(self, svc: "TaskService") -> None: + """Inner reap loop, parameterized by the TaskService to use. + + Wraps each ``unclaim_for_reaper`` in try/except so a single bad row + doesn't abort the dispatch tick — the reaper must keep ticking even + if one task's release somehow fails. + """ + from roboco.utils.converters import require_uuid + + cutoff = datetime.now(UTC) - timedelta(seconds=self._claim_heartbeat_ttl) + candidates = await svc.list_in_progress_or_claimed() + for t in candidates: + ts = t.last_heartbeat_at + if ts is None or ts < cutoff: + task_id = require_uuid(t.id) + try: + await svc.unclaim_for_reaper(task_id) + logger.warning( + "stale claim reaped", + task_id=str(task_id), + last_heartbeat=ts.isoformat() if ts else None, + ) + except Exception as exc: + logger.error( + "stale-claim reap failed; continuing", + task_id=str(task_id), + error=str(exc), + ) + async def _dispatch_all_work(self) -> None: """Run all dispatchers to check for and assign work. @@ -3536,8 +3600,22 @@ Start now: evidence(task_id="{task_id}") `_dispatch_qa_work` claimed for QA and the next dispatcher re-spawned the dev on the same claimed row) are defanged by early dispatchers marking the task handled. + + The stale-claim reaper runs first, before any dispatcher tries to + spawn an agent for a task whose previous holder is dead. Without + this ordering, the spawn pass could race against a stale claim and + skip work the reaper would have freed in the same tick. """ self._tick_handled_tasks = set() + + # Free any tasks whose claim went stale before the spawn pass runs. + # Wrapped because a reaper failure must not block dispatch — the + # next tick will retry. + try: + await self._reap_stale_claims() + except Exception as e: + logger.error("Stale-claim reaper failed; continuing tick", error=str(e)) + # Orchestrator uses SYSTEM role for internal API calls # Using a well-known UUID for the orchestrator identity headers = { diff --git a/roboco/services/task.py b/roboco/services/task.py index 7a5b77c6..910d0348 100644 --- a/roboco/services/task.py +++ b/roboco/services/task.py @@ -1763,6 +1763,41 @@ class TaskService(BaseService): .values(last_heartbeat_at=datetime.now(UTC)) ) + async def list_in_progress_or_claimed(self) -> list[TaskTable]: + """All tasks currently in claimed or in_progress state. + + Used by the orchestrator's stale-claim reaper to find rows whose + holder may have gone silent. Returns the bare row set; the reaper + applies the heartbeat-TTL filter in Python because the cutoff is a + runtime decision tied to settings, not a column. + """ + result = await self.session.execute( + select(TaskTable).where( + TaskTable.status.in_([TaskStatus.CLAIMED, TaskStatus.IN_PROGRESS]) + ) + ) + return list(result.scalars().all()) + + async def unclaim_for_reaper(self, task_id: UUID) -> None: + """Reaper-only unclaim: skip role checks, force the row back to pending. + + Bypasses the normal claim guards because the holder is provably dead + (no heartbeat past TTL) — the operation is named with ``_for_reaper`` + so callers cannot accidentally use it as a regular unclaim path. + Clears ``assigned_to`` and ``last_heartbeat_at`` so the next claim + starts fresh. + """ + await self.session.execute( + update(TaskTable) + .where(TaskTable.id == task_id) + .values( + status=TaskStatus.PENDING, + assigned_to=None, + last_heartbeat_at=None, + ) + ) + await self.session.flush() + async def block( self, task_id: UUID, diff --git a/tests/unit/runtime/test_stale_claim_reaper.py b/tests/unit/runtime/test_stale_claim_reaper.py new file mode 100644 index 00000000..32652156 --- /dev/null +++ b/tests/unit/runtime/test_stale_claim_reaper.py @@ -0,0 +1,94 @@ +"""Reaper releases tasks whose last_heartbeat_at exceeds the TTL. + +The orchestrator dispatcher periodically calls `_reap_stale_claims` to +release tasks whose holder has gone silent past the heartbeat TTL. The +schema-level column `last_heartbeat_at` (DateTime(timezone=True)) has +existed since migration 006; this test covers the runtime decision that +turns a stale heartbeat into a freed claim. + +Datetimes used here are timezone-aware UTC because the underlying column +is tz-aware — comparing naive vs aware would raise TypeError in production +even though it would silently work against an in-memory mock. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from unittest.mock import AsyncMock +from uuid import uuid4 + +import pytest +from roboco.runtime.orchestrator import AgentOrchestrator + + +@pytest.mark.asyncio +async def test_reap_stale_claims_releases_dead_holders() -> None: + """A task past TTL is unclaimed; a fresh one is left alone.""" + stale_id = uuid4() + fresh_id = uuid4() + now = datetime.now(UTC) + stale_task = type( + "T", + (), + {"id": stale_id, "last_heartbeat_at": now - timedelta(seconds=600)}, + )() + fresh_task = type( + "T", + (), + {"id": fresh_id, "last_heartbeat_at": now - timedelta(seconds=10)}, + )() + + orch = AgentOrchestrator.__new__(AgentOrchestrator) # bypass __init__ + orch._claim_heartbeat_ttl = 300 + orch._task_svc = AsyncMock() + orch._task_svc.list_in_progress_or_claimed.return_value = [stale_task, fresh_task] + orch._task_svc.unclaim_for_reaper = AsyncMock() + + await orch._reap_stale_claims() + + orch._task_svc.unclaim_for_reaper.assert_awaited_once_with(stale_id) + + +@pytest.mark.asyncio +async def test_reap_stale_claims_releases_holders_with_null_heartbeat() -> None: + """A claimed task that never heartbeated (NULL column) is treated as stale.""" + null_id = uuid4() + null_task = type("T", (), {"id": null_id, "last_heartbeat_at": None})() + + orch = AgentOrchestrator.__new__(AgentOrchestrator) + orch._claim_heartbeat_ttl = 300 + orch._task_svc = AsyncMock() + orch._task_svc.list_in_progress_or_claimed.return_value = [null_task] + orch._task_svc.unclaim_for_reaper = AsyncMock() + + await orch._reap_stale_claims() + + orch._task_svc.unclaim_for_reaper.assert_awaited_once_with(null_id) + + +@pytest.mark.asyncio +async def test_reap_stale_claims_swallows_unclaim_errors() -> None: + """An unclaim_for_reaper failure must not abort the reap loop.""" + stale_a = uuid4() + stale_b = uuid4() + now = datetime.now(UTC) + task_a = type( + "T", (), {"id": stale_a, "last_heartbeat_at": now - timedelta(seconds=600)} + )() + task_b = type( + "T", (), {"id": stale_b, "last_heartbeat_at": now - timedelta(seconds=900)} + )() + + orch = AgentOrchestrator.__new__(AgentOrchestrator) + orch._claim_heartbeat_ttl = 300 + orch._task_svc = AsyncMock() + orch._task_svc.list_in_progress_or_claimed.return_value = [task_a, task_b] + orch._task_svc.unclaim_for_reaper = AsyncMock( + side_effect=[RuntimeError("transient"), None] + ) + + await orch._reap_stale_claims() + + # Both stale tasks attempted; second succeeded despite first raising. + expected_attempts = 2 + assert orch._task_svc.unclaim_for_reaper.await_count == expected_attempts