fix(core): persist the tool-call recovery counter across resumes

An exhausted agent parked in 'waiting' got a fresh nudge budget on every
600s auto-resume, so a wedged agent could nudge-park-nudge indefinitely.
Track the count on the coordinator, snapshot it, and reset it only on
real input or an explicit lifecycle tool.
This commit is contained in:
Ahmed Allam
2026-08-02 02:15:51 +03:00
committed by Ahmed Allam
parent f6f9469e00
commit 6eec34df24
3 changed files with 82 additions and 3 deletions
+22
View File
@@ -46,6 +46,7 @@ class AgentCoordinator:
self.metadata: dict[str, dict[str, Any]] = {}
self.pending_counts: dict[str, int] = {}
self.errors: dict[str, str] = {}
self.recovery_counts: dict[str, int] = {}
self.runtimes: dict[str, AgentRuntime] = {}
self._lock = asyncio.Lock()
self._snapshot_path: Path | None = None
@@ -187,6 +188,25 @@ class AgentCoordinator:
async def park_waiting(self, agent_id: str) -> None:
await self.set_status(agent_id, "waiting")
async def record_recovery(self, agent_id: str) -> int:
"""Count a turn that ended without a lifecycle tool call; return the new total.
Persisted so a resumed agent cannot earn a fresh nudge budget on every
auto-resume and loop forever.
"""
async with self._lock:
count = self.recovery_counts.get(agent_id, 0) + 1
self.recovery_counts[agent_id] = count
await self._maybe_snapshot()
return count
async def reset_recovery(self, agent_id: str) -> None:
"""Clear the nudge budget after real progress (new message or a lifecycle tool)."""
async with self._lock:
if self.recovery_counts.pop(agent_id, None) is None:
return
await self._maybe_snapshot()
async def set_status(
self, agent_id: str, status: Status | str, *, error: str | None = None
) -> None:
@@ -397,6 +417,7 @@ class AgentCoordinator:
"names": dict(self.names),
"metadata": {aid: dict(md) for aid, md in self.metadata.items()},
"pending_counts": dict(self.pending_counts),
"recovery_counts": dict(self.recovery_counts),
"mailboxes": {
aid: [dict(m) for m in runtime.mailbox]
for aid, runtime in self.runtimes.items()
@@ -416,6 +437,7 @@ class AgentCoordinator:
self.metadata = {aid: dict(md) for aid, md in snap.get("metadata", {}).items()}
self.pending_counts = dict(snap.get("pending_counts", {}))
self.errors = dict(snap.get("errors", {}))
self.recovery_counts = dict(snap.get("recovery_counts", {}))
mailboxes = snap.get("mailboxes", {})
if isinstance(mailboxes, dict):
for aid, msgs in mailboxes.items():
+7 -3
View File
@@ -241,7 +241,11 @@ async def run_agent_loop(
await coordinator.set_status(agent_id, "stopped")
raise SubagentBudgetReservedError("scan reached the sub-agent budget reserve")
if not woke:
if woke:
# Real input is real progress, so the nudge budget starts over. A bare
# auto-resume is not: it must not hand a wedged agent a fresh budget.
await coordinator.reset_recovery(agent_id)
else:
logger.info("agent %s reached its waiting timeout; auto-resuming", agent_id)
await coordinator.send(
agent_id,
@@ -441,7 +445,6 @@ async def _run_until_lifecycle(
"""
result: RunResultBase | None = None
input_data: Any = initial_input
recoveries = 0
recovery_limit = _INTERACTIVE_TOOL_RECOVERY_LIMIT if interactive else max(1, max_turns)
while True:
@@ -483,9 +486,10 @@ async def _run_until_lifecycle(
status = await _agent_status(coordinator, agent_id)
if status != "running":
await coordinator.reset_recovery(agent_id)
return result
recoveries += 1
recoveries = await coordinator.record_recovery(agent_id)
logger.warning(
"agent %s ended a turn without a lifecycle tool call (interactive=%s); "
"forcing tool continuation (%d/%d): %s",
+53
View File
@@ -931,3 +931,56 @@ async def test_tool_required_message_is_persisted_to_the_session(tmp_path: Any)
assert "finish_scan" in stored[0]["content"]
assert "wait_for_message" in stored[0]["content"]
session.close()
@pytest.mark.asyncio
async def test_recovery_count_survives_a_snapshot_round_trip(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""A resumed agent must not earn a fresh nudge budget and loop forever."""
coordinator = AgentCoordinator()
await coordinator.register("root", "strix", parent_id=None)
calls: list[Any] = []
monkeypatch.setattr(
execution,
"_run_cycle_parked",
_scripted_cycle(coordinator, "root", ["running"], calls),
)
await _drive(coordinator, "root", interactive=True)
assert coordinator.recovery_counts["root"] == execution._INTERACTIVE_TOOL_RECOVERY_LIMIT
restored = AgentCoordinator()
await restored.restore(await coordinator.snapshot())
assert restored.recovery_counts["root"] == execution._INTERACTIVE_TOOL_RECOVERY_LIMIT
# The restored agent is already at its cap, so it parks after a single
# further text-only cycle instead of starting the whole budget over.
resumed_calls: list[Any] = []
monkeypatch.setattr(
execution,
"_run_cycle_parked",
_scripted_cycle(restored, "root", ["running"], resumed_calls),
)
await _drive(restored, "root", interactive=True)
assert len(resumed_calls) == 1
assert restored.statuses["root"] == "waiting"
@pytest.mark.asyncio
async def test_recovery_count_is_cleared_by_a_lifecycle_tool(
monkeypatch: pytest.MonkeyPatch,
) -> None:
coordinator = AgentCoordinator()
await coordinator.register("root", "strix", parent_id=None)
calls: list[Any] = []
monkeypatch.setattr(
execution,
"_run_cycle_parked",
_scripted_cycle(coordinator, "root", ["running", "completed"], calls),
)
await _drive(coordinator, "root", interactive=True)
assert "root" not in coordinator.recovery_counts