fix(gateway): make i_will_plan / i_will_work_on idempotent on re-entry

Smoke 2026-05-04 captured the cycle the prior 63d0adf fix didn't close:
- Spawn 1: i_will_plan succeeds, task pending → claimed → in_progress.
- Agent goes idle (LLM thinking, container exits, respawned).
- Spawn 2: i_will_plan called again. _i_will_plan_preflight rejects
  'task in in_progress, expected pending'. Agent has no recovery path.
- Heartbeat eventually goes stale, reaper drops claim back to pending,
  spawn 3 fires, loop repeats indefinitely.

Fix: when a respawned PM/dev re-enters the verb on a task they already
own in claimed/in_progress, return OK with current state and refresh
the heartbeat instead of rejecting. The verb is now genuinely
idempotent for the caller, which matches what 'I will plan' should
mean — record intent + advance state, regardless of how many times
the agent says it. Different-caller contention still rejects.

Refactored i_will_work_on's pending branch into _i_will_work_on_pending
helper to satisfy PLR0911 after the new branch raised return count.
This commit is contained in:
Renn F
2026-05-04 02:04:39 +02:00
parent cfdd389c43
commit 0a3c963923
2 changed files with 178 additions and 39 deletions
+85 -39
View File
@@ -292,6 +292,43 @@ class Choreographer:
# Format: "<id> (<status>)"
return ", ".join(f"{s.id} ({s.status})" for s in non_terminal)
async def _i_will_work_on_pending(
self,
agent_id: UUID,
task_id: UUID,
t: Any,
plan: str | None,
briefing: dict[str, Any],
) -> tuple[Envelope | None, Any]:
"""Pending-branch dispatch for i_will_work_on. Extracted to keep
the parent's return count under PLR0911. Returns (rejection|None,
task). Caller emits rejection via _emit_rejection and falls
through to the OK envelope when rejection is None.
"""
if guard := await self._run_claim_guards(agent_id=agent_id, task=t):
return self._with_briefing(guard, briefing), t
# claim() transitions pending → claimed; idempotent for same assignee.
t = await self.task.claim(task_id, agent_id)
if t is None:
return Envelope.invalid_state(
message="claim failed",
remediate="task may already be claimed by another agent",
context_briefing=briefing,
), t
if not t.plan and not plan:
return Envelope.tracing_gap(
missing=["plan"],
remediate=(
f"call i_will_work_on(task_id='{task_id}',"
f" plan='<one-paragraph plan describing what you will do>')"
),
context_briefing=briefing,
), t
if plan:
t = await self.task.set_plan(task_id, plan)
t = await self.task.start(task_id, agent_id)
return None, t
async def i_will_work_on(
self, agent_id: UUID, task_id: UUID, plan: str | None = None
) -> Envelope:
@@ -314,48 +351,16 @@ class Choreographer:
t = await self.task.claim(task_id, agent_id)
t = await self.task.start(task_id, agent_id)
elif status == "pending":
# Fresh claim — run all claim-time gates BEFORE mutating state.
if guard := await self._run_claim_guards(agent_id=agent_id, task=t):
rejection, t = await self._i_will_work_on_pending(
agent_id, task_id, t, plan, briefing
)
if rejection is not None:
return await self._emit_rejection(
self._with_briefing(guard, briefing),
rejection,
agent_id=agent_id,
task_id=task_id,
verb="i_will_work_on",
)
# Always call claim() when status is pending — even if the dev
# is already in assigned_to (parent PM may pre-assign at delegate
# time). claim() transitions pending → claimed; idempotent for
# the same assignee. See i_will_plan for the same fix.
t = await self.task.claim(task_id, agent_id)
if t is None:
return await self._emit_rejection(
Envelope.invalid_state(
message="claim failed",
remediate="task may already be claimed by another agent",
context_briefing=briefing,
),
agent_id=agent_id,
task_id=task_id,
verb="i_will_work_on",
)
if not t.plan and not plan:
remediate = (
f"call i_will_work_on(task_id='{task_id}',"
f" plan='<one-paragraph plan describing what you will do>')"
)
return await self._emit_rejection(
Envelope.tracing_gap(
missing=["plan"],
remediate=remediate,
context_briefing=briefing,
),
agent_id=agent_id,
task_id=task_id,
verb="i_will_work_on",
)
if plan:
t = await self.task.set_plan(task_id, plan)
t = await self.task.start(task_id, agent_id)
elif status == "claimed" and t.assigned_to == agent_id:
# Resumption: skip sibling-sequence (already passed at claim).
# Still enforce already_active/paused so concurrent claims fail.
@@ -370,6 +375,13 @@ class Choreographer:
verb="i_will_work_on",
)
t = await self.task.start(task_id, agent_id)
elif status == "in_progress" and t.assigned_to == agent_id:
# Idempotent re-entry: respawned dev re-calling i_will_work_on
# on a task they already own in_progress. Skip start() (would
# reject — wrong source state) but fall through to the OK
# envelope at the end. Heartbeat fires there too, refreshing
# reaper activity.
pass
else:
return await self._emit_rejection(
Envelope.invalid_state(
@@ -1419,7 +1431,15 @@ class Choreographer:
async def _i_will_plan_preflight(
self, pm_agent_id: UUID, task_id: UUID, t: Any, plan: str
) -> Envelope | None:
"""Run i_will_plan's role / status / plan / claim guards. None = pass."""
"""Run i_will_plan's role / status / plan / claim guards. None = pass.
Idempotent on re-entry: if the caller already owns the task in
claimed/in_progress (their previous spawn moved it forward), the
verb returns OK with current state instead of rejecting. Without
this, a respawned PM hits 'task in in_progress, expected pending'
and loops until the reaper drops the claim back to pending —
producing the cycle smoke 2026-05-04 captured.
"""
agent = await self.task.agent_for(pm_agent_id)
if agent is None or agent.role not in ("cell_pm", "main_pm"):
return Envelope.not_authorized(
@@ -1427,7 +1447,13 @@ class Choreographer:
remediate="this verb is reserved for PMs",
context_briefing=await self._briefing_for(pm_agent_id, task_id),
)
if str(t.status) != "pending":
status = str(t.status)
if status != "pending":
# Idempotent re-entry: caller already owns this task in a
# post-claim state. Don't reject; the i_will_plan body will
# short-circuit on the same condition and return OK.
if status in ("claimed", "in_progress") and t.assigned_to == pm_agent_id:
return None
return Envelope.invalid_state(
message=f"task {task_id} is in {t.status}, expected pending",
remediate="call give_me_work() to find a pending task to plan",
@@ -1481,6 +1507,26 @@ class Choreographer:
verb="i_will_plan",
)
# Idempotent re-entry: respawned PM that already owns this task in
# claimed/in_progress short-circuits with the current state. Touch
# the heartbeat so the reaper sees fresh activity, then return OK
# pointing at delegate as the next call. Without this short-circuit
# the body would reach start() — which rejects because status is
# not 'claimed' on the in_progress branch — and emit a misleading
# invalid_state envelope.
status = str(t.status)
if status in ("claimed", "in_progress") and t.assigned_to == pm_agent_id:
await self._touch(task_id)
return Envelope.ok(
status=status,
task_id=str(task_id),
next=(
"task already claimed; delegate(parent_task_id, ...) for"
" each subtask, then i_am_idle"
),
context_briefing=await self._briefing_for(pm_agent_id, task_id),
)
# Always call claim() when status is pending — even if the PM is
# already in assigned_to (CEO pre-assigns root tasks at creation).
# claim() transitions pending → claimed; without that, start() below
@@ -222,6 +222,99 @@ async def test_i_will_plan_surfaces_start_failure_instead_of_faking_ok() -> None
assert "start failed" in body["message"]
@pytest.mark.asyncio
async def test_i_will_plan_idempotent_when_already_in_progress_for_caller() -> None:
"""Regression: respawned PM re-calling i_will_plan on a task they
already moved to in_progress must NOT be rejected. Returns OK with
current state. Smoke 2026-05-04 captured the cycle the old reject
caused: respawn reject reaper drops claim respawn loop.
"""
pm_id = uuid4()
task_id = uuid4()
in_progress = MagicMock(
id=task_id,
status="in_progress",
plan={"text": "x"},
assigned_to=pm_id,
task_type="planning",
parent_task_id=None,
sequence=0,
)
task_svc = AsyncMock()
task_svc.get.return_value = in_progress
task_svc.agent_for.return_value = MagicMock(role="main_pm", team="main_pm")
task_svc.list_in_progress_for_agent.return_value = [in_progress]
task_svc.list_paused_for_agent.return_value = []
deps = _make_deps(task=task_svc)
c = Choreographer(deps)
env = await c.i_will_plan(pm_id, task_id, plan="re-entry plan")
assert env.error is None
assert env.status == "in_progress"
# Heartbeat refreshed so the reaper sees activity.
task_svc.heartbeat.assert_awaited()
# Did NOT re-call claim or start — already past those.
task_svc.claim.assert_not_awaited()
task_svc.start.assert_not_awaited()
@pytest.mark.asyncio
async def test_i_will_plan_idempotent_when_already_claimed_for_caller() -> None:
"""Same regression but task is in claimed (post-claim, pre-start)."""
pm_id = uuid4()
task_id = uuid4()
claimed = MagicMock(
id=task_id,
status="claimed",
plan=None,
assigned_to=pm_id,
task_type="planning",
parent_task_id=None,
sequence=0,
)
task_svc = AsyncMock()
task_svc.get.return_value = claimed
task_svc.agent_for.return_value = MagicMock(role="main_pm", team="main_pm")
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
deps = _make_deps(task=task_svc)
c = Choreographer(deps)
env = await c.i_will_plan(pm_id, task_id, plan="re-entry plan")
assert env.error is None
assert env.status == "claimed"
task_svc.heartbeat.assert_awaited()
@pytest.mark.asyncio
async def test_i_will_plan_still_rejects_in_progress_for_other_agent() -> None:
"""Idempotency only applies to the caller. Different PM still rejected."""
pm_id = uuid4()
other_pm_id = uuid4()
task_id = uuid4()
in_progress = MagicMock(
id=task_id,
status="in_progress",
plan={"text": "x"},
assigned_to=other_pm_id, # different PM owns it
task_type="planning",
parent_task_id=None,
sequence=0,
)
task_svc = AsyncMock()
task_svc.get.return_value = in_progress
task_svc.agent_for.return_value = MagicMock(role="main_pm", team="main_pm")
deps = _make_deps(task=task_svc)
c = Choreographer(deps)
env = await c.i_will_plan(pm_id, task_id, plan="x")
body = env.as_dict()
assert body["error"] == "invalid_state"
@pytest.mark.asyncio
async def test_i_will_plan_returns_tracing_gap_without_plan() -> None:
pm_id = uuid4()