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.
This commit is contained in:
Renn F
2026-05-03 05:01:50 +02:00
parent c12aad3005
commit 6643b4c375
2 changed files with 247 additions and 12 deletions
+19 -12
View File
@@ -123,6 +123,11 @@ class Choreographer:
def evidence_repo(self) -> Any: def evidence_repo(self) -> Any:
return self._deps.evidence_repo 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 --- # --- Phase 1 (developer) verbs ---
async def give_me_work(self, agent_id: UUID) -> Envelope: 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) guard := pm_cannot_execute_code_guard(role, task_type)
): ):
return guard return guard
if not skip_role_typed and ( if not skip_role_typed and (guard := role_typed_claim_guard(role, task_type)):
guard := role_typed_claim_guard(role, task_type)
):
return guard return guard
in_progress = await self.task.list_in_progress_for_agent(agent_id) in_progress = await self.task.list_in_progress_for_agent(agent_id)
if guard := already_active_guard(in_progress, task.id): if guard := already_active_guard(in_progress, task.id):
@@ -292,6 +295,7 @@ class Choreographer:
context_briefing=briefing, context_briefing=briefing,
) )
await self._touch(task_id)
return Envelope.ok( return Envelope.ok(
status=str(t.status), status=str(t.status),
task_id=str(task_id), task_id=str(task_id),
@@ -328,6 +332,7 @@ class Choreographer:
context_briefing=await self._briefing_for(agent_id, t.id), context_briefing=await self._briefing_for(agent_id, t.id),
) )
await self.task.add_progress(t.id, agent_id, message) await self.task.add_progress(t.id, agent_id, message)
await self._touch(t.id)
return Envelope.ok( return Envelope.ok(
status=str(t.status), status=str(t.status),
task_id=str(t.id), task_id=str(t.id),
@@ -363,9 +368,7 @@ class Choreographer:
return rejection return rejection
# 2. Field-level gates (Gate Set E) — strict. # 2. Field-level gates (Gate Set E) — strict.
if rejection := await self._check_submit_qa_field_gates( if rejection := await self._check_submit_qa_field_gates(agent_id, task_id, t):
agent_id, task_id, t
):
return rejection return rejection
# 3. Submit (no catch-up). # 3. Submit (no catch-up).
@@ -373,6 +376,7 @@ class Choreographer:
if submitted is not None: if submitted is not None:
t = submitted t = submitted
await self._notify_qa(agent_id, task_id, t) 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) return await self._build_i_am_done_ok(agent_id, task_id, t)
async def i_am_done_with_catchup( async def i_am_done_with_catchup(
@@ -575,6 +579,7 @@ class Choreographer:
agent_id=agent_id, task_id=task_id, content=reason agent_id=agent_id, task_id=task_id, content=reason
) )
t = await self.task.escalate(agent_id, task_id, reason) t = await self.task.escalate(agent_id, task_id, reason)
await self._touch(task_id)
return Envelope.ok( return Envelope.ok(
status=str(t.status), status=str(t.status),
task_id=str(task_id), task_id=str(task_id),
@@ -636,8 +641,10 @@ class Choreographer:
return None return None
first = pending[0] first = pending[0]
agent = await self.task.agent_for(agent_id) agent = await self.task.agent_for(agent_id)
verb = "i_will_plan" if agent and agent.role in ("cell_pm", "main_pm") else ( verb = (
"i_will_work_on" "i_will_plan"
if agent and agent.role in ("cell_pm", "main_pm")
else ("i_will_work_on")
) )
return Envelope.invalid_state( return Envelope.invalid_state(
message=( message=(
@@ -1054,6 +1061,7 @@ class Choreographer:
) )
await self.task.set_plan(task_id, plan) await self.task.set_plan(task_id, plan)
t = await self.task.start(task_id, pm_agent_id) t = await self.task.start(task_id, pm_agent_id)
await self._touch(task_id)
return Envelope.ok( return Envelope.ok(
status=str(t.status) if t else "in_progress", status=str(t.status) if t else "in_progress",
task_id=str(task_id), task_id=str(task_id),
@@ -1402,8 +1410,7 @@ class Choreographer:
missing=["subtasks not all terminal"], missing=["subtasks not all terminal"],
remediate=( remediate=(
"all subtasks must be in completed/cancelled before" "all subtasks must be in completed/cancelled before"
" bubbling up. Non-terminal subtasks: " " bubbling up. Non-terminal subtasks: " + non_terminal
+ non_terminal
), ),
context_briefing=await self._briefing_for(pm_agent_id, task_id), 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) assigned = await self.task.list_assigned_for_agent(pm_agent_id)
if assigned: if assigned:
t = assigned[0] t = assigned[0]
await self._touch(t.id)
return Envelope.ok( return Envelope.ok(
status=str(t.status), status=str(t.status),
task_id=str(t.id), task_id=str(t.id),
@@ -1604,8 +1612,7 @@ class Choreographer:
missing=["subtasks not all terminal"], missing=["subtasks not all terminal"],
remediate=( remediate=(
"all subtasks must be in completed/cancelled before" "all subtasks must be in completed/cancelled before"
" completing parent. Non-terminal subtasks: " " completing parent. Non-terminal subtasks: " + non_terminal
+ non_terminal
), ),
context_briefing=await self._briefing_for(pm_agent_id, task_id), context_briefing=await self._briefing_for(pm_agent_id, task_id),
) )
+228
View File
@@ -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)