mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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:
@@ -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,
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user