roboco/tests/unit/runtime/test_stale_claim_reaper.py

"""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", stop_reason="reaper_wedged_grok"
    )
    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_spares_claude_container_within_stuck_ttl(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A non-GROK (Claude) container silent 1200s is still spared — it is within
    the much longer claude-stuck TTL, so a slow-but-working agent keeps the
    heartbeat-skip protection. Only a truly-stuck run (silent past the stuck
    TTL) is killed (see test_reaper_kills_stuck_claude_past_stuck_ttl)."""
    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._claude_stuck_kill_ttl = 3600
    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_reaper_kills_stuck_claude_past_stuck_ttl(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """#73: a non-GROK container stuck in a non-verb loop (alive, no heartbeat
    advance past the claude-stuck TTL) is killed, evicted, and released — not
    shielded by the live-container skip forever."""
    now = datetime.now(UTC)
    stuck = type(
        "T",
        (),
        {
            "id": uuid4(),
            "last_heartbeat_at": now - timedelta(seconds=4000),
            "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._claude_stuck_kill_ttl = 3600
    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 = [stuck]
    svc.unclaim_for_reaper = AsyncMock()

    await orch._reap_with_service(svc)

    remove_mock.assert_awaited_once_with(
        "roboco-agent-be-dev-1", stop_reason="reaper_stuck_claude"
    )
    assert "be-dev-1" not in orch._instances  # evicted
    svc.unclaim_for_reaper.assert_awaited_once_with(stuck.id)  # released


def test_stuck_claude_slug_excludes_grok_and_non_active() -> None:
    """#73: the stuck-slug predicate fires only for a non-GROK ACTIVE instance
    past the TTL — a GROK provider (its own kill path), a non-ACTIVE instance, a
    recent heartbeat, and no owner all yield None."""
    now = datetime.now(UTC)
    past = now - timedelta(seconds=4000)
    task = type(
        "T",
        (),
        {
            "id": uuid4(),
            "last_heartbeat_at": past,
            "assigned_to": AGENT_UUIDS["be-dev-1"],
            "claimed_by": None,
        },
    )()
    orch = AgentOrchestrator.__new__(AgentOrchestrator)
    orch._claude_stuck_kill_ttl = 3600

    claude_cfg = type("C", (), {"provider_type": "anthropic"})()
    grok_cfg = type("C", (), {"provider_type": "grok"})()

    # Non-GROK ACTIVE past TTL → slug returned.
    orch._instances = {
        "be-dev-1": AgentInstance(
            agent_id="be-dev-1", state=AgentState.ACTIVE, config=claude_cfg
        )
    }
    assert orch._stuck_claude_slug(task, past) == "be-dev-1"

    # GROK provider → None (the wedged-grok path owns it).
    orch._instances = {
        "be-dev-1": AgentInstance(
            agent_id="be-dev-1", state=AgentState.ACTIVE, config=grok_cfg
        )
    }
    assert orch._stuck_claude_slug(task, past) is None

    # Non-ACTIVE → None.
    orch._instances = {
        "be-dev-1": AgentInstance(
            agent_id="be-dev-1", state=AgentState.OFFLINE, config=claude_cfg
        )
    }
    assert orch._stuck_claude_slug(task, past) is None

    # Recent heartbeat → None.
    orch._instances = {
        "be-dev-1": AgentInstance(
            agent_id="be-dev-1", state=AgentState.ACTIVE, config=claude_cfg
        )
    }
    assert orch._stuck_claude_slug(task, now) is None


@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:
    """A provider-parked agent (dead container, OFFLINE, with a
    ``rate_limit_lifted`` WaitingRecord) must NOT be reaped by the stale-claim
    reaper — reaping would release the claim to pending and 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)