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.
This commit is contained in:
Renn F
2026-05-03 04:47:51 +02:00
parent 37bc4e58ed
commit b301020398
4 changed files with 215 additions and 0 deletions
+8
View File
@@ -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,
+78
View File
@@ -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 = {
+35
View File
@@ -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,
@@ -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