From 6643b4c3750e0d854418fb7dd57291d80b187210 Mon Sep 17 00:00:00 2001 From: Renn F Date: Sun, 3 May 2026 05:01:50 +0200 Subject: [PATCH] feat(gateway): heartbeat on every hot verb i_will_work_on, i_have_committed, i_am_done, i_am_blocked, i_will_plan, and pm_give_me_work now call task.heartbeat() so the reaper has fresh data. Closes the loop opened by tasks 2a + 2b. --- roboco/services/gateway/choreographer.py | 31 +-- tests/unit/gateway/test_heartbeat_wired.py | 228 +++++++++++++++++++++ 2 files changed, 247 insertions(+), 12 deletions(-) create mode 100644 tests/unit/gateway/test_heartbeat_wired.py diff --git a/roboco/services/gateway/choreographer.py b/roboco/services/gateway/choreographer.py index 592a0fab..63dfc8bd 100644 --- a/roboco/services/gateway/choreographer.py +++ b/roboco/services/gateway/choreographer.py @@ -123,6 +123,11 @@ class Choreographer: def evidence_repo(self) -> Any: return self._deps.evidence_repo + async def _touch(self, task_id: UUID | None) -> None: + """Best-effort heartbeat write; silent on missing task.""" + if task_id is not None: + await self.task.heartbeat(task_id) + # --- Phase 1 (developer) verbs --- async def give_me_work(self, agent_id: UUID) -> Envelope: @@ -198,9 +203,7 @@ class Choreographer: guard := pm_cannot_execute_code_guard(role, task_type) ): return guard - if not skip_role_typed and ( - guard := role_typed_claim_guard(role, task_type) - ): + if not skip_role_typed and (guard := role_typed_claim_guard(role, task_type)): return guard in_progress = await self.task.list_in_progress_for_agent(agent_id) if guard := already_active_guard(in_progress, task.id): @@ -292,6 +295,7 @@ class Choreographer: context_briefing=briefing, ) + await self._touch(task_id) return Envelope.ok( status=str(t.status), task_id=str(task_id), @@ -328,6 +332,7 @@ class Choreographer: context_briefing=await self._briefing_for(agent_id, t.id), ) await self.task.add_progress(t.id, agent_id, message) + await self._touch(t.id) return Envelope.ok( status=str(t.status), task_id=str(t.id), @@ -363,9 +368,7 @@ class Choreographer: return rejection # 2. Field-level gates (Gate Set E) — strict. - if rejection := await self._check_submit_qa_field_gates( - agent_id, task_id, t - ): + if rejection := await self._check_submit_qa_field_gates(agent_id, task_id, t): return rejection # 3. Submit (no catch-up). @@ -373,6 +376,7 @@ class Choreographer: if submitted is not None: t = submitted await self._notify_qa(agent_id, task_id, t) + await self._touch(task_id) return await self._build_i_am_done_ok(agent_id, task_id, t) async def i_am_done_with_catchup( @@ -575,6 +579,7 @@ class Choreographer: agent_id=agent_id, task_id=task_id, content=reason ) t = await self.task.escalate(agent_id, task_id, reason) + await self._touch(task_id) return Envelope.ok( status=str(t.status), task_id=str(task_id), @@ -636,8 +641,10 @@ class Choreographer: return None first = pending[0] agent = await self.task.agent_for(agent_id) - verb = "i_will_plan" if agent and agent.role in ("cell_pm", "main_pm") else ( - "i_will_work_on" + verb = ( + "i_will_plan" + if agent and agent.role in ("cell_pm", "main_pm") + else ("i_will_work_on") ) return Envelope.invalid_state( message=( @@ -1054,6 +1061,7 @@ class Choreographer: ) await self.task.set_plan(task_id, plan) t = await self.task.start(task_id, pm_agent_id) + await self._touch(task_id) return Envelope.ok( status=str(t.status) if t else "in_progress", task_id=str(task_id), @@ -1402,8 +1410,7 @@ class Choreographer: missing=["subtasks not all terminal"], remediate=( "all subtasks must be in completed/cancelled before" - " bubbling up. Non-terminal subtasks: " - + non_terminal + " bubbling up. Non-terminal subtasks: " + non_terminal ), context_briefing=await self._briefing_for(pm_agent_id, task_id), ) @@ -1440,6 +1447,7 @@ class Choreographer: assigned = await self.task.list_assigned_for_agent(pm_agent_id) if assigned: t = assigned[0] + await self._touch(t.id) return Envelope.ok( status=str(t.status), task_id=str(t.id), @@ -1604,8 +1612,7 @@ class Choreographer: missing=["subtasks not all terminal"], remediate=( "all subtasks must be in completed/cancelled before" - " completing parent. Non-terminal subtasks: " - + non_terminal + " completing parent. Non-terminal subtasks: " + non_terminal ), context_briefing=await self._briefing_for(pm_agent_id, task_id), ) diff --git a/tests/unit/gateway/test_heartbeat_wired.py b/tests/unit/gateway/test_heartbeat_wired.py new file mode 100644 index 00000000..09bd7e3b --- /dev/null +++ b/tests/unit/gateway/test_heartbeat_wired.py @@ -0,0 +1,228 @@ +"""Choreographer hot verbs must touch last_heartbeat_at via task.heartbeat().""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps + + +def _make_deps(**overrides: AsyncMock) -> ChoreographerDeps: + task = overrides.get("task", AsyncMock()) + work_session = overrides.get("work_session", AsyncMock()) + git = overrides.get("git", AsyncMock()) + a2a = overrides.get("a2a", AsyncMock()) + journal = overrides.get("journal", AsyncMock()) + audit = overrides.get("audit", AsyncMock()) + evidence_repo = overrides.get("evidence_repo", AsyncMock()) + for method in ( + "list_unread_a2a", + "list_unread_mentions", + "list_pending_notifications", + "task_metadata_gaps", + "recent_team_activity", + "blockers_in_lane", + ): + getattr(evidence_repo, method).return_value = [] + return ChoreographerDeps( + task=task, + work_session=work_session, + git=git, + a2a=a2a, + journal=journal, + audit=audit, + evidence_repo=evidence_repo, + ) + + +@pytest.mark.asyncio +async def test_i_will_work_on_calls_heartbeat() -> None: + aid = uuid4() + tid = uuid4() + pending = MagicMock( + id=tid, + status="pending", + plan=None, + assigned_to=None, + parent_task_id=None, + sequence=0, + task_type="code", + ) + in_progress = MagicMock( + id=tid, status="in_progress", plan={"text": "go"}, assigned_to=aid + ) + task_svc = AsyncMock() + task_svc.get.return_value = pending + task_svc.agent_for.return_value = MagicMock(role="developer", team="backend") + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.get_subtasks.return_value = [] + task_svc.claim.return_value = MagicMock( + id=tid, status="claimed", plan=None, assigned_to=aid + ) + task_svc.set_plan.return_value = MagicMock( + id=tid, status="claimed", plan={"text": "go"}, assigned_to=aid + ) + task_svc.start.return_value = in_progress + deps = _make_deps(task=task_svc) + c = Choreographer(deps) + + await c.i_will_work_on(aid, tid, plan="go") + + task_svc.heartbeat.assert_awaited_with(tid) + + +@pytest.mark.asyncio +async def test_i_have_committed_calls_heartbeat() -> None: + aid = uuid4() + tid = uuid4() + t = MagicMock(id=tid, status="in_progress", assigned_to=aid, plan="x") + task_svc = AsyncMock() + task_svc.get_active_task_for_agent.return_value = t + deps = _make_deps(task=task_svc) + c = Choreographer(deps) + + await c.i_have_committed(aid, "did the thing") + + task_svc.heartbeat.assert_awaited_with(tid) + + +@pytest.mark.asyncio +async def test_i_am_done_calls_heartbeat() -> None: + aid = uuid4() + tid = uuid4() + t = MagicMock( + id=tid, + status="verifying", + assigned_to=aid, + plan="x", + self_verified=True, + commits=[MagicMock()], + pr_number=42, + pr_url="http://example/pr/42", + team="backend", + work_session_id=None, + ) + submitted = MagicMock( + id=tid, + status="awaiting_qa", + assigned_to=None, + plan="x", + team="backend", + work_session_id=None, + pr_url="http://example/pr/42", + ) + task_svc = AsyncMock() + task_svc.get.return_value = t + task_svc.submit_qa.return_value = submitted + task_svc.qa_agent_for_team.return_value = None + journal_svc = AsyncMock() + journal_svc.has_reflect_for_task.return_value = True + evidence_repo = AsyncMock() + for method in ( + "list_unread_a2a", + "list_unread_mentions", + "list_pending_notifications", + "task_metadata_gaps", + "recent_team_activity", + "blockers_in_lane", + ): + getattr(evidence_repo, method).return_value = [] + evidence_repo.journal_highlights_for_task.return_value = [] + deps = _make_deps(task=task_svc, journal=journal_svc, evidence_repo=evidence_repo) + c = Choreographer(deps) + + # Patch tracing_gate to pass — _check_tracing_gates uses + # check_requirements, which inspects the Mock's progress_updates etc. + # By configuring t with all-required attrs we let the real gate pass. + t.progress_updates = [MagicMock()] + t.acceptance_criteria = [] + + await c.i_am_done(aid, tid, "done") + + task_svc.heartbeat.assert_awaited_with(tid) + + +@pytest.mark.asyncio +async def test_i_am_blocked_calls_heartbeat() -> None: + aid = uuid4() + tid = uuid4() + t = MagicMock(id=tid, status="in_progress", assigned_to=aid) + blocked = MagicMock(id=tid, status="blocked", assigned_to=aid) + task_svc = AsyncMock() + task_svc.get.return_value = t + task_svc.escalate.return_value = blocked + journal_svc = AsyncMock() + deps = _make_deps(task=task_svc, journal=journal_svc) + c = Choreographer(deps) + + await c.i_am_blocked(aid, tid, "stuck on X") + + task_svc.heartbeat.assert_awaited_with(tid) + + +@pytest.mark.asyncio +async def test_pm_give_me_work_calls_heartbeat_when_returning_task() -> None: + pm_id = uuid4() + tid = uuid4() + assigned = MagicMock(id=tid, status="pending") + task_svc = AsyncMock() + task_svc.list_assigned_for_agent.return_value = [assigned] + deps = _make_deps(task=task_svc) + c = Choreographer(deps) + + await c.pm_give_me_work(pm_id) + + task_svc.heartbeat.assert_awaited_with(tid) + + +@pytest.mark.asyncio +async def test_pm_give_me_work_does_not_heartbeat_on_idle() -> None: + pm_id = uuid4() + task_svc = AsyncMock() + task_svc.list_assigned_for_agent.return_value = [] + deps = _make_deps(task=task_svc) + c = Choreographer(deps) + + await c.pm_give_me_work(pm_id) + + task_svc.heartbeat.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_i_will_plan_calls_heartbeat() -> None: + pm_id = uuid4() + tid = uuid4() + pending = MagicMock( + id=tid, + status="pending", + plan=None, + assigned_to=None, + parent_task_id=None, + sequence=0, + task_type="planning", + ) + in_progress = MagicMock( + id=tid, status="in_progress", plan="plan-text", assigned_to=pm_id + ) + task_svc = AsyncMock() + task_svc.get.return_value = pending + task_svc.agent_for.return_value = MagicMock(role="cell_pm", team="backend") + task_svc.list_in_progress_for_agent.return_value = [] + task_svc.list_paused_for_agent.return_value = [] + task_svc.get_subtasks.return_value = [] + task_svc.claim.return_value = MagicMock( + id=tid, status="claimed", plan=None, assigned_to=pm_id + ) + task_svc.set_plan.return_value = MagicMock( + id=tid, status="claimed", plan="plan-text", assigned_to=pm_id + ) + task_svc.start.return_value = in_progress + deps = _make_deps(task=task_svc) + c = Choreographer(deps) + + await c.i_will_plan(pm_id, tid, plan="plan-text") + + task_svc.heartbeat.assert_awaited_with(tid)