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,
|
ge=60,
|
||||||
description="Claim heartbeat staleness threshold (seconds)",
|
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(
|
spawn_cooldown_seconds: int = Field(
|
||||||
default=60,
|
default=60,
|
||||||
ge=1,
|
ge=1,
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ import httpx
|
|||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from roboco.services.llm import AgentRoute
|
from roboco.services.llm import AgentRoute
|
||||||
|
from roboco.services.task import TaskService
|
||||||
import structlog
|
import structlog
|
||||||
from fastapi import status as http_status
|
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
|
# is in a loop — without this gate the orchestrator re-spawns every
|
||||||
# tick forever (seen in production on 2026-04-22).
|
# tick forever (seen in production on 2026-04-22).
|
||||||
self._pm_respawn_tracker: dict[tuple[str, str], dict[str, Any]] = {}
|
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
|
# LIFECYCLE
|
||||||
@@ -3523,6 +3530,63 @@ Start now: evidence(task_id="{task_id}")
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error("Dispatcher loop error", error=str(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:
|
async def _dispatch_all_work(self) -> None:
|
||||||
"""Run all dispatchers to check for and assign work.
|
"""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
|
`_dispatch_qa_work` claimed for QA and the next dispatcher
|
||||||
re-spawned the dev on the same claimed row) are defanged by
|
re-spawned the dev on the same claimed row) are defanged by
|
||||||
early dispatchers marking the task handled.
|
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()
|
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
|
# Orchestrator uses SYSTEM role for internal API calls
|
||||||
# Using a well-known UUID for the orchestrator identity
|
# Using a well-known UUID for the orchestrator identity
|
||||||
headers = {
|
headers = {
|
||||||
|
|||||||
@@ -1763,6 +1763,41 @@ class TaskService(BaseService):
|
|||||||
.values(last_heartbeat_at=datetime.now(UTC))
|
.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(
|
async def block(
|
||||||
self,
|
self,
|
||||||
task_id: UUID,
|
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