diff --git a/strix/core/agents.py b/strix/core/agents.py index 3f20629d..d03636e1 100644 --- a/strix/core/agents.py +++ b/strix/core/agents.py @@ -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(): diff --git a/strix/core/execution.py b/strix/core/execution.py index b16041bb..d9f8c0a9 100644 --- a/strix/core/execution.py +++ b/strix/core/execution.py @@ -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", diff --git a/tests/test_execution.py b/tests/test_execution.py index 7fc73782..3b2df11c 100644 --- a/tests/test_execution.py +++ b/tests/test_execution.py @@ -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