From bf44d5aade1b786940b2a5473450d7adf34ceb88 Mon Sep 17 00:00:00 2001 From: Renn F Date: Sun, 3 May 2026 10:03:46 +0200 Subject: [PATCH] fix(orchestrator): skip closure spawn if PM just paused via i_am_idle Tiny race: dispatcher decided to spawn for closure between agent's heartbeat and idle-pause. Spawn would land against an already-paused parent. Gate spawn on (status != PAUSED OR last_heartbeat older than cutoff). --- roboco/api/schemas/tasks.py | 2 + roboco/runtime/orchestrator.py | 50 +++++ tests/unit/runtime/test_idle_vs_respawn.py | 202 +++++++++++++++++++++ 3 files changed, 254 insertions(+) create mode 100644 tests/unit/runtime/test_idle_vs_respawn.py diff --git a/roboco/api/schemas/tasks.py b/roboco/api/schemas/tasks.py index 74f89fde..ffd1fdd2 100644 --- a/roboco/api/schemas/tasks.py +++ b/roboco/api/schemas/tasks.py @@ -264,6 +264,7 @@ class TaskResponse(BaseModel): started_at: datetime | None completed_at: datetime | None target_date: datetime | None + last_heartbeat_at: datetime | None = None # Planning estimated_complexity: Complexity @@ -609,6 +610,7 @@ def task_to_response(task: "TaskTable") -> TaskResponse: started_at=task.started_at, completed_at=task.completed_at, target_date=task.target_date, + last_heartbeat_at=getattr(task, "last_heartbeat_at", None), estimated_complexity=task.estimated_complexity, plan=convert_plan(task.plan), checkpoints=convert_checkpoints(task.checkpoints), diff --git a/roboco/runtime/orchestrator.py b/roboco/runtime/orchestrator.py index 82d9af04..b9b5b1ad 100644 --- a/roboco/runtime/orchestrator.py +++ b/roboco/runtime/orchestrator.py @@ -3924,6 +3924,48 @@ Start now: evidence(task_id="{task_id}") in ("awaiting_pm_review", "awaiting_ceo_approval", "completed") ) + @staticmethod + def _coerce_heartbeat(value: Any) -> datetime | None: + """Normalize ``last_heartbeat_at`` to an aware UTC datetime. + + The dispatcher reads tasks via the HTTP API, which serializes + datetimes as ISO-8601 strings; direct service callers (and tests) + may pass ``datetime`` objects. Anything else is treated as + absent so a malformed value can't accidentally arm the gate. + """ + if value is None: + return None + if isinstance(value, datetime): + return value if value.tzinfo else value.replace(tzinfo=UTC) + if isinstance(value, str): + try: + parsed = datetime.fromisoformat(value) + except ValueError: + return None + return parsed if parsed.tzinfo else parsed.replace(tzinfo=UTC) + return None + + def _is_recently_paused(self, task: dict[str, Any]) -> bool: + """A paused task whose heartbeat is fresher than the stale cutoff. + + Closes the ``i_am_idle`` vs closure-respawn race (audit C12): + ``i_am_idle`` auto-pauses in-flight tasks and then sets the agent + IDLE. If the dispatcher ticks between those two writes it sees a + paused parent and would spawn the closure PM against a session + that is mid-shutdown. A fresh ``last_heartbeat_at`` (newer than + ``settings.claim_stale_seconds``) is the signal that the agent + was alive moments ago and a respawn now would race the existing + session. Genuinely-stale paused tasks (or tasks with no heartbeat + recorded) fall through and follow the regular closure path. + """ + if task.get("status") != "paused": + return False + last_hb = self._coerce_heartbeat(task.get("last_heartbeat_at")) + if last_hb is None: + return False + cutoff = datetime.now(UTC) - timedelta(seconds=self._claim_heartbeat_ttl) + return last_hb > cutoff + def _closure_pm_for_team(self, team: str | None) -> str: """Pick the PM that owns closure for a given team.""" if team in ("backend", "frontend", "ux_ui"): @@ -3938,6 +3980,14 @@ Start now: evidence(task_id="{task_id}") if not task_id: return + if self._is_recently_paused(task): + logger.debug( + "Skipping closure spawn for recently-paused parent", + task_id=task_id, + last_heartbeat_at=task.get("last_heartbeat_at"), + ) + return + descendants = await self._fetch_all_descendants(client, task_id) if not descendants: return diff --git a/tests/unit/runtime/test_idle_vs_respawn.py b/tests/unit/runtime/test_idle_vs_respawn.py new file mode 100644 index 00000000..5730e0f7 --- /dev/null +++ b/tests/unit/runtime/test_idle_vs_respawn.py @@ -0,0 +1,202 @@ +"""Closure dispatcher must skip respawn for a parent task that just paused. + +Race scenario (audit C12) +------------------------- +``i_am_idle`` runs ``auto_pause_paused_tasks`` (transitions the agent's +in-flight tasks to ``paused`` and stamps ``last_heartbeat_at``) and then +flips the agent state to IDLE. The closure dispatcher iterates parent +tasks every tick and, for any whose descendants are all terminal, calls +``spawn_agent`` for the closure PM. + +If ``i_am_idle``'s pause lands one tick before the dispatcher runs, the +parent's status is already ``paused`` (so it's in the closure dispatcher's +``parent_statuses`` list) but its ``last_heartbeat_at`` is fresh — the +agent literally just heartbeated before pausing itself. Spawning a fresh +container for that PM here would race the in-flight session that just +called i_am_idle. + +The gate: skip closure spawn when the task is paused AND +``last_heartbeat_at`` is newer than ``settings.claim_stale_seconds``. +A genuinely-stale paused task (heartbeat older than the cutoff) still +gets the closure spawn — the gate is about *recency*, not paused-ness. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 + +import pytest +from roboco.config import settings +from roboco.runtime.orchestrator import AgentOrchestrator + + +def _make_orch() -> AgentOrchestrator: + """Bypass __init__ — tests don't need a full DI graph.""" + orch = AgentOrchestrator.__new__(AgentOrchestrator) + orch._instances = {} + orch._claim_heartbeat_ttl = settings.claim_stale_seconds + return orch + + +def _paused_parent(*, last_heartbeat_at: datetime | str | None) -> dict[str, Any]: + """A parent task that's paused and has the requested heartbeat freshness.""" + return { + "id": str(uuid4()), + "status": "paused", + "team": "backend", + "last_heartbeat_at": last_heartbeat_at, + } + + +@pytest.mark.asyncio +async def test_skips_spawn_when_paused_and_recently_touched_datetime() -> None: + """Recent heartbeat + status=paused = agent just called i_am_idle. + + Last heartbeat is 1 second ago, claim_stale_seconds default is 180s. + The task is fresh — closure spawn must not fire. + """ + orch = _make_orch() + fresh = datetime.now(UTC) - timedelta(seconds=1) + task = _paused_parent(last_heartbeat_at=fresh) + + client = AsyncMock() + + with ( + patch.object(orch, "_fetch_all_descendants", new=AsyncMock()) as fetch_desc, + patch.object(orch, "spawn_agent", new=AsyncMock()) as spawn, + ): + await orch._maybe_spawn_pm_closure(client, task) + + fetch_desc.assert_not_awaited() + spawn.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_skips_spawn_when_paused_and_recently_touched_iso_string() -> None: + """API serializes datetimes as ISO strings; the gate must handle both.""" + orch = _make_orch() + fresh_iso = (datetime.now(UTC) - timedelta(seconds=1)).isoformat() + task = _paused_parent(last_heartbeat_at=fresh_iso) + + client = AsyncMock() + + with ( + patch.object(orch, "_fetch_all_descendants", new=AsyncMock()) as fetch_desc, + patch.object(orch, "spawn_agent", new=AsyncMock()) as spawn, + ): + await orch._maybe_spawn_pm_closure(client, task) + + fetch_desc.assert_not_awaited() + spawn.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_spawns_when_paused_but_heartbeat_is_stale() -> None: + """Heartbeat older than claim_stale_seconds = genuinely-stale agent. + + The closure dispatcher should still spawn here — the i_am_idle race + window has long since closed. + """ + orch = _make_orch() + stale = datetime.now(UTC) - timedelta(seconds=settings.claim_stale_seconds + 30) + task = _paused_parent(last_heartbeat_at=stale) + descendant = {"id": str(uuid4()), "status": "completed"} + + client = AsyncMock() + + with ( + patch.object( + orch, + "_fetch_all_descendants", + new=AsyncMock(return_value=[descendant]), + ), + patch.object(orch, "_already_promoted_for_closure", return_value=False), + patch.object(orch, "_is_agent_active", return_value=False), + patch.object( + orch, + "_build_pm_closure_prompt", + return_value="prompt", + ), + patch.object(orch, "_task_git_context", return_value=MagicMock()), + patch.object(orch, "spawn_agent", new=AsyncMock()) as spawn, + ): + await orch._maybe_spawn_pm_closure(client, task) + + spawn.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_spawns_when_paused_but_heartbeat_missing() -> None: + """No heartbeat at all means the freshness gate cannot trigger. + + The legitimate stale-paused-parent case (e.g. a PM was paused before + Phase 2 added heartbeats) should not be blocked by the new gate. + """ + orch = _make_orch() + task = _paused_parent(last_heartbeat_at=None) + descendant = {"id": str(uuid4()), "status": "completed"} + + client = AsyncMock() + + with ( + patch.object( + orch, + "_fetch_all_descendants", + new=AsyncMock(return_value=[descendant]), + ), + patch.object(orch, "_already_promoted_for_closure", return_value=False), + patch.object(orch, "_is_agent_active", return_value=False), + patch.object( + orch, + "_build_pm_closure_prompt", + return_value="prompt", + ), + patch.object(orch, "_task_git_context", return_value=MagicMock()), + patch.object(orch, "spawn_agent", new=AsyncMock()) as spawn, + ): + await orch._maybe_spawn_pm_closure(client, task) + + spawn.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_non_paused_status_not_gated_by_heartbeat() -> None: + """Status != paused means the gate is not applicable. + + A claimed/in_progress parent with a fresh heartbeat is normal active + work and should follow the existing closure logic unchanged. + """ + orch = _make_orch() + fresh = datetime.now(UTC) - timedelta(seconds=1) + task = { + "id": str(uuid4()), + "status": "in_progress", + "team": "backend", + "last_heartbeat_at": fresh, + } + descendant = {"id": str(uuid4()), "status": "completed"} + + client = AsyncMock() + + with ( + patch.object( + orch, + "_fetch_all_descendants", + new=AsyncMock(return_value=[descendant]), + ), + patch.object(orch, "_already_promoted_for_closure", return_value=False), + patch.object(orch, "_is_agent_active", return_value=False), + patch.object( + orch, + "_build_pm_closure_prompt", + return_value="prompt", + ), + patch.object(orch, "_task_git_context", return_value=MagicMock()), + patch.object(orch, "spawn_agent", new=AsyncMock()) as spawn, + ): + await orch._maybe_spawn_pm_closure(client, task) + + spawn.assert_awaited_once()