diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index 4c3fdb75..0b70f499 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -6481,6 +6481,29 @@ Start by: kind=kind, error=str(exc), ) + # F035: register a WaitingRecord so the probe-resume loop can revive + # this agent when the provider recovers. ``_on_probe_success`` reads + # ``_waiting_records`` filtered by ``waiting_for == "rate_limit_lifted"`` + # + ``context.provider``; without a record here it resumes nobody and + # recovery falls to the 600s stale-claim reaper instead of the + # probe-success path the parking design relies on. Persisted (mirrors + # ``mark_waiting_long``) so a restart still resolves the wait. We do NOT + # call ``mark_waiting_long`` itself — the container is already dead (it + # exited), so there is nothing to stop, and parking keeps OFFLINE (not + # WAITING_LONG) so the reaper's live-skip / health loop ignore it. + task_id = ( + str(instance.current_task_id) if instance.current_task_id else None + ) + record = WaitingRecord( + agent_id=agent_id, + task_id=task_id, + waiting_for="rate_limit_lifted", + waiting_since=datetime.now(UTC), + context={"provider": provider, "kind": kind}, + ) + self._waiting_records[agent_id] = record + with contextlib.suppress(Exception): + await self._persist_waiting_record(record) logger.warning( "Provider unavailable; parked (task retried when it recovers)", provider=provider, @@ -7716,6 +7739,26 @@ Start now: evidence(task_id="{task_id}") instance = instances.get(self._resolve_agent_slug(str(owner))) return instance is not None and instance.state == AgentState.ACTIVE + def _assignee_is_provider_parked(self, task: Any) -> bool: + """True if the task's assignee is parked waiting for a provider to recover. + + A provider-parked agent (session-limit / overload / grok-429) is OFFLINE + with a dead container and a ``rate_limit_lifted`` WaitingRecord; the + probe-resume loop owns its recovery. The stale-claim reaper must skip it + so the claim survives until the probe revives the agent — reaping would + release the claim to pending and probe-success would then respawn the + agent on a task it no longer owns. Defensive on a missing registry. + """ + owner = getattr(task, "assigned_to", None) or getattr(task, "claimed_by", None) + if not owner: + return False + records = getattr(self, "_waiting_records", None) + if not records: + return False + slug = self._resolve_agent_slug(str(owner)) + record = records.get(slug) + return record is not None and record.waiting_for == "rate_limit_lifted" + async def _readopt_running_agents(self) -> int: """Re-adopt still-running agent containers into ``_instances`` at startup. @@ -7960,6 +8003,13 @@ Start now: evidence(task_id="{task_id}") and not await self._maybe_recover_broken_gateway(t) ): continue + # F035: a provider-parked agent (session-limit / overload / + # grok-429) is OFFLINE with a dead container and a + # ``rate_limit_lifted`` WaitingRecord. The probe-resume loop + # owns its recovery — do NOT reap the claim, or probe-success + # would later respawn the agent on a task it no longer owns. + if self._assignee_is_provider_parked(t): + continue task_id = require_uuid(t.id) try: await svc.unclaim_for_reaper(task_id) diff --git a/tests/unit/runtime/test_provider_overload_break.py b/tests/unit/runtime/test_provider_overload_break.py index 96f7e829..7c53c824 100644 --- a/tests/unit/runtime/test_provider_overload_break.py +++ b/tests/unit/runtime/test_provider_overload_break.py @@ -9,11 +9,12 @@ decision points deterministically (logs + tracker + finalize stubbed). from __future__ import annotations +from datetime import UTC, datetime from unittest.mock import AsyncMock import pytest from roboco.config import settings -from roboco.models.runtime import AgentInstance +from roboco.models.runtime import AgentInstance, WaitingRecord from roboco.runtime.orchestrator import ( _OVERLOAD_RETRY_AFTER_S, _RATE_LIMIT_RETRY_AFTER_S, @@ -52,6 +53,9 @@ def orch(monkeypatch: pytest.MonkeyPatch) -> AgentOrchestrator: # transcript fallback empty by default so dev environments with stray # transcripts do not make the tests flaky. monkeypatch.setattr(orch, "_transcript_tail_text", lambda _a, _lines=80: "") + orch._waiting_records = {} + orch._rate_limit_ceo_notified = set() + orch._instances = {} return orch @@ -153,6 +157,64 @@ async def test_park_offlines_and_activates_with_kind( } +@pytest.mark.asyncio +async def test_park_registers_waiting_record_so_probe_can_resume( + orch: AgentOrchestrator, monkeypatch: pytest.MonkeyPatch +) -> None: + # F035: the probe-resume loop reads _waiting_records filtered by + # waiting_for == "rate_limit_lifted" + context.provider. Without a record + # here, _parked_agents_for(provider) returns [] and _on_probe_success + # resumes nobody — recovery falls to the 600s stale-claim reaper instead of + # the probe-success path the parking design relies on. + orch._waiting_records = {} + inst = _instance() + inst.current_task_id = "task-1" + tracker = _FakeTracker() + monkeypatch.setattr(orch, "_make_tracker", lambda _p: tracker) + monkeypatch.setattr(orch, "_finalize_spawn_session", AsyncMock()) + monkeypatch.setattr(orch, "_persist_waiting_record", AsyncMock()) + + await orch._park_provider_unavailable( + "be-dev-1", inst, provider="anthropic", retry_after=45.0, kind="rate_limited" + ) + + assert "be-dev-1" in orch._waiting_records + rec = orch._waiting_records["be-dev-1"] + assert rec.waiting_for == "rate_limit_lifted" + assert rec.context.get("provider") == "anthropic" + assert rec.task_id == "task-1" + assert orch._parked_agents_for("anthropic") == ["be-dev-1"] + + +@pytest.mark.asyncio +async def test_probe_success_respawns_parked_agent( + orch: AgentOrchestrator, monkeypatch: pytest.MonkeyPatch +) -> None: + # F035: once the probe succeeds, the parked agent must be respawned via + # resolve_wait — not left stranded for the 600s reaper. + orch._waiting_records = { + "be-dev-1": WaitingRecord( + agent_id="be-dev-1", + task_id="task-1", + waiting_for="rate_limit_lifted", + waiting_since=datetime.now(UTC), + context={"provider": "anthropic"}, + ) + } + tracker = _FakeTracker() + tracker.clear = AsyncMock() # type: ignore[method-assign] + monkeypatch.setattr(orch, "_make_tracker", lambda _p: tracker) + monkeypatch.setattr(orch, "_delete_waiting_record", AsyncMock()) + monkeypatch.setattr(orch, "_generate_resume_prompt", lambda _r, _res: "resume") + spawn = AsyncMock() + monkeypatch.setattr(orch, "spawn_agent", spawn) + + await orch._on_probe_success("anthropic", tracker) + + spawn.assert_awaited_once() # parked agent respawned, not stranded + assert "be-dev-1" not in orch._waiting_records + + # --------------------------------------------------------------------------- # _handle_stopped_container — overload short-circuits the crash-retry path # --------------------------------------------------------------------------- diff --git a/tests/unit/runtime/test_stale_claim_reaper.py b/tests/unit/runtime/test_stale_claim_reaper.py index 481c458d..98549954 100644 --- a/tests/unit/runtime/test_stale_claim_reaper.py +++ b/tests/unit/runtime/test_stale_claim_reaper.py @@ -18,7 +18,7 @@ from unittest.mock import AsyncMock from uuid import uuid4 import pytest -from roboco.models.runtime import AgentInstance +from roboco.models.runtime import AgentInstance, WaitingRecord from roboco.runtime.orchestrator import AgentOrchestrator, AgentState from roboco.seeds.initial_data import AGENT_UUIDS @@ -376,6 +376,57 @@ async def test_reap_releases_on_registry_miss_when_container_gone( 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,