mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
_park_provider_unavailable parked the provider + offlined the instance but never registered a WaitingRecord, so _on_probe_success -> _parked_agents_for (always filtered on waiting_for=='rate_limit_lifted') returned [] and resolve_wait revived nobody — recovery fell to the 600s stale-claim reaper instead of the probe-success path the parking design relied on. Register + persist a rate_limit_lifted WaitingRecord at park time (mirrors mark_waiting_long, minus stop_agent — the container is already dead). Companion reaper guard: _reap_with_service now skips provider-parked assignees (_assignee_is_provider_parked) so the claim survives until the probe revives the agent — otherwise the reaper releases the claim to pending and probe-success respawns on a task the agent no longer owns.
464 lines
15 KiB
Python
464 lines
15 KiB
Python
"""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.models.runtime import AgentInstance, WaitingRecord
|
|
from roboco.runtime.orchestrator import AgentOrchestrator, AgentState
|
|
from roboco.seeds.initial_data import AGENT_UUIDS
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reap_stale_claims_releases_dead_holders(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> 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__
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [stale_task, fresh_task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
svc.unclaim_for_reaper.assert_awaited_once_with(stale_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reap_stale_claims_releases_holders_with_null_heartbeat(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> 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)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [null_task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
svc.unclaim_for_reaper.assert_awaited_once_with(null_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reap_stale_claims_swallows_unclaim_errors(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> 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)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [task_a, task_b]
|
|
svc.unclaim_for_reaper = AsyncMock(side_effect=[RuntimeError("transient"), None])
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
# Both stale tasks attempted; second succeeded despite first raising.
|
|
expected_attempts = 2
|
|
assert svc.unclaim_for_reaper.await_count == expected_attempts
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reap_spares_claims_whose_assignee_container_is_alive(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A stale-heartbeat task is NOT reaped while its assignee container lives.
|
|
|
|
A developer deep in a long edit/test cycle outruns the heartbeat TTL; the
|
|
running container is the ground truth, so the claim survives rather than
|
|
being churned out from under live work. A peer task whose assignee has no
|
|
live instance is still reaped.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
live_id = uuid4()
|
|
dead_id = uuid4()
|
|
live_task = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": live_id,
|
|
"last_heartbeat_at": now - timedelta(seconds=600),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
dead_task = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": dead_id,
|
|
"last_heartbeat_at": now - timedelta(seconds=600),
|
|
"assigned_to": AGENT_UUIDS["be-dev-2"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
orch._instances = {
|
|
"be-dev-1": AgentInstance(agent_id="be-dev-1", state=AgentState.ACTIVE)
|
|
}
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [live_task, dead_task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
# The live-assignee task is spared; only the dead one is reaped.
|
|
svc.unclaim_for_reaper.assert_awaited_once_with(dead_id)
|
|
|
|
|
|
def _grok_instance() -> AgentInstance:
|
|
cfg = type("C", (), {"provider_type": "grok"})()
|
|
return AgentInstance(agent_id="be-dev-1", state=AgentState.ACTIVE, config=cfg)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reaper_kills_and_releases_wedged_grok_container(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A GROK container idle past the kill TTL is killed, evicted, and released.
|
|
|
|
Unlike a Claude agent, a wedged grok container is ACTIVE yet fires no
|
|
verb, so the live-instance skip would shield it forever. Past the longer
|
|
grok-idle TTL the watchdog removes the container and drops it from
|
|
`_instances`, so the same reap pass then unclaims the task.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
task_id = uuid4()
|
|
wedged = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": task_id,
|
|
"last_heartbeat_at": now - timedelta(seconds=1200),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
orch._grok_idle_kill_ttl = 900
|
|
orch._instances = {"be-dev-1": _grok_instance()}
|
|
remove_mock = AsyncMock()
|
|
monkeypatch.setattr(orch, "_remove_container", remove_mock)
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [wedged]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
remove_mock.assert_awaited_once_with("roboco-agent-be-dev-1")
|
|
assert "be-dev-1" not in orch._instances # evicted
|
|
svc.unclaim_for_reaper.assert_awaited_once_with(task_id) # released
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reaper_spares_grok_container_within_kill_ttl(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A GROK container stale past the claim TTL but within the kill TTL lives.
|
|
|
|
Only a truly-dead run (idle past the longer grok-idle TTL) is killed — a
|
|
slow-but-working agent is left alone.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
recent = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": uuid4(),
|
|
"last_heartbeat_at": now - timedelta(seconds=600),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
orch._grok_idle_kill_ttl = 900
|
|
orch._instances = {"be-dev-1": _grok_instance()}
|
|
remove_mock = AsyncMock()
|
|
monkeypatch.setattr(orch, "_remove_container", remove_mock)
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [recent]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
remove_mock.assert_not_awaited()
|
|
assert "be-dev-1" in orch._instances
|
|
svc.unclaim_for_reaper.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reaper_never_kills_non_grok_container(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""A non-GROK (Claude) container idle past the kill TTL is still spared.
|
|
|
|
The watchdog only kills GROK runtimes; a quiet Claude agent keeps the
|
|
heartbeat-skip protection regardless of how long it has been silent.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
claude_task = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": uuid4(),
|
|
"last_heartbeat_at": now - timedelta(seconds=1200),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
claude_cfg = type("C", (), {"provider_type": "anthropic"})()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
orch._grok_idle_kill_ttl = 900
|
|
orch._instances = {
|
|
"be-dev-1": AgentInstance(
|
|
agent_id="be-dev-1", state=AgentState.ACTIVE, config=claude_cfg
|
|
)
|
|
}
|
|
remove_mock = AsyncMock()
|
|
monkeypatch.setattr(orch, "_remove_container", remove_mock)
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [claude_task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
remove_mock.assert_not_awaited()
|
|
assert "be-dev-1" in orch._instances
|
|
svc.unclaim_for_reaper.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reap_spares_live_container_on_registry_miss(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Registry amnesia: the orchestrator forgot the instance (empty `_instances`,
|
|
e.g. after a restart) but the container is still running per Docker — the
|
|
task must NOT be reaped. This closes the be-dev-1 over-reap: a live agent the
|
|
orchestrator merely lost track of is no longer churned out from under work.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
task_id = uuid4()
|
|
task = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": task_id,
|
|
"last_heartbeat_at": now - timedelta(seconds=600),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
orch._grok_idle_kill_ttl = 900
|
|
orch._instances = {} # registry lost; container still up
|
|
monkeypatch.setattr(
|
|
orch, "_inspect_container_state", AsyncMock(return_value=(True, None))
|
|
)
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
svc.unclaim_for_reaper.assert_not_awaited() # spared — Docker says it's alive
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reap_releases_on_registry_miss_when_container_gone(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Registry miss AND Docker says the container is gone → genuinely dead → reaped."""
|
|
now = datetime.now(UTC)
|
|
task_id = uuid4()
|
|
task = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": task_id,
|
|
"last_heartbeat_at": now - timedelta(seconds=600),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
orch._grok_idle_kill_ttl = 900
|
|
orch._instances = {}
|
|
monkeypatch.setattr(
|
|
orch, "_inspect_container_state", AsyncMock(return_value=(False, 0))
|
|
)
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
svc.unclaim_for_reaper.assert_awaited_once_with(task_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reap_spares_provider_parked_agent_for_probe_resume(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""F035: a provider-parked agent (dead container, OFFLINE, with a
|
|
``rate_limit_lifted`` WaitingRecord) must NOT be reaped by the stale-claim
|
|
reaper. The probe-resume loop owns its recovery and respawns it when the
|
|
provider recovers; reaping would release the claim to pending, and then
|
|
probe-success would respawn the agent on a task it no longer owns.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
task_id = uuid4()
|
|
task = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": task_id,
|
|
"last_heartbeat_at": now - timedelta(seconds=1200),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
orch._grok_idle_kill_ttl = 900
|
|
orch._instances = {} # parked agent is OFFLINE / not in the registry
|
|
orch._waiting_records = {
|
|
"be-dev-1": WaitingRecord(
|
|
agent_id="be-dev-1",
|
|
task_id=str(task_id),
|
|
waiting_for="rate_limit_lifted",
|
|
waiting_since=now,
|
|
context={"provider": "anthropic"},
|
|
)
|
|
}
|
|
monkeypatch.setattr(
|
|
orch, "_inspect_container_state", AsyncMock(return_value=(False, 0))
|
|
)
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
svc.unclaim_for_reaper.assert_not_awaited() # spared for the probe loop
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_registry_uninitialised_skips_docker_fallback(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""With `_instances` never initialised (None — the __new__ unit harness), the
|
|
Docker fallback is skipped and the stale task reaps as before; no accidental
|
|
Docker probing where there's no registry to be amnesiac about.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
task_id = uuid4()
|
|
task = type(
|
|
"T",
|
|
(),
|
|
{
|
|
"id": task_id,
|
|
"last_heartbeat_at": now - timedelta(seconds=600),
|
|
"assigned_to": AGENT_UUIDS["be-dev-1"],
|
|
"claimed_by": None,
|
|
},
|
|
)()
|
|
|
|
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
|
monkeypatch.setattr(
|
|
orch, "_maybe_recover_broken_gateway", AsyncMock(return_value=False)
|
|
)
|
|
orch._claim_heartbeat_ttl = 300
|
|
# _instances intentionally NOT set -> getattr yields None -> no fallback.
|
|
svc = AsyncMock()
|
|
svc.list_in_progress_or_claimed.return_value = [task]
|
|
svc.unclaim_for_reaper = AsyncMock()
|
|
|
|
await orch._reap_with_service(svc)
|
|
|
|
svc.unclaim_for_reaper.assert_awaited_once_with(task_id)
|