mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
Persist the PM-respawn counter across orchestrator restarts (#275)
* feat(orchestrator): add respawn_tracker table + migration 051 Durable backing for AgentOrchestrator._pm_respawn_tracker (the PM-respawn loop breaker). Kept only in memory it reset to count=1 on every restart, re-burning the strike threshold against a still-wedged task. RespawnTrackerTable mirrors WaitingRecordTable: composite PK (agent_slug, task_id) matching the in-memory key; task_id is intentionally NOT a FK (the startup loader validates against live tasks so a stale counter can't resurrect). Migration 051 verified with a real alembic upgrade head + downgrade -1 + re-upgrade on Postgres. * feat(orchestrator): persist the PM-respawn counter across restarts The PM-respawn loop breaker (_pm_respawn_tracker) lived only in memory, so an orchestrator restart reset a wedged task's strike count to 1 and re-burned the whole threshold (4 spawns x container cost) before the gate fired again. Write-through each gate mutation to the respawn_tracker table via a fire-and-forget _schedule_respawn_persist (on the existing _bg_tasks strong-ref set; a DB hiccup degrades to in-memory-only, never gates/un-gates a spawn), and restore_respawn_tracker() repopulates the counter at startup, validating each row against live tasks (drops terminal/missing) so a stale counter can't resurrect against a fixed task. Best-effort + inert when the table is empty. Cannot manufacture a spawn — the counter only ever suppresses one. (_instances reconcile, the spec's other goal, already shipped as _readopt_running_agents.) * fix(types): cast Mapped[UUID] columns in project routes + self_heal A clean `mypy roboco/ tests/` run surfaces 7 pre-existing errors in files this branch doesn't touch: project-route handlers and self_heal_engine pass a ProjectTable.id (declared Mapped[UUID] against SQLAlchemy's dialect UUID, so mypy infers sqlalchemy.sql.sqltypes.UUID[Any]) where a uuid.UUID is expected. An incremental .mypy_cache had hidden them. Apply the same targeted cast unblock used for the prior batch; the deeper fix (migrating the ~88 Mapped[UUID] columns to Mapped[uuid.UUID]) remains a separate dedicated task. * docs(orchestrator): document respawn_tracker durability Add the orchestrator runtime-state durability note to CLAUDE.md (respawn_tracker write-through + restore; _instances reconciled-from-Docker) + the migration-051 narrative, and a CHANGELOG [Unreleased] Fixed entry. Also type-clean the new respawn_tracker table test (cast __table__ to Table under TYPE_CHECKING). --------- Co-authored-by: Renn F <rennf93@users.noreply.github.com>
This commit is contained in:
Renzo F
@@ -36,6 +36,10 @@ def _new_orchestrator() -> AgentOrchestrator:
|
||||
"""Bypass __init__ so tests don't need a full DI graph."""
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
cast("Any", orch)._pm_respawn_tracker = {}
|
||||
# The gate now write-throughs each mutation to the respawn_tracker table via
|
||||
# _schedule_respawn_persist; these tests cover gate LOGIC only, so stub the
|
||||
# scheduler to a no-op (persistence is covered in test_respawn_persistence).
|
||||
cast("Any", orch)._schedule_respawn_persist = lambda *_a, **_k: None
|
||||
return orch
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,315 @@
|
||||
"""The PM-respawn counter survives an orchestrator restart.
|
||||
|
||||
`_pm_respawn_tracker` is the loop breaker against respawning the same PM on the
|
||||
same task forever. Kept only in memory it reset to count=1 on every restart,
|
||||
re-burning the whole strike threshold against a still-wedged task. These tests
|
||||
cover the write-through persist on each gate mutation, the startup loader (which
|
||||
validates against live tasks and drops terminal/missing rows), and the safety
|
||||
property that a restored counter trips at the persisted threshold — never
|
||||
manufacturing a spawn.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from datetime import UTC, datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from roboco.runtime.orchestrator import AgentOrchestrator
|
||||
|
||||
_SEEDED_COUNT = 3 # a persisted strike count, one below the trip threshold
|
||||
_STRIKE_COUNT = 2
|
||||
_MIN_PERSISTS = 2
|
||||
_TRIP_COUNT = 4 # count > _PM_RESPAWN_MAX_UNPRODUCTIVE (3) fires the gate
|
||||
|
||||
|
||||
def _new_orchestrator() -> AgentOrchestrator:
|
||||
orch = AgentOrchestrator.__new__(AgentOrchestrator)
|
||||
cast("Any", orch)._pm_respawn_tracker = {}
|
||||
cast("Any", orch)._bg_tasks = set()
|
||||
return orch
|
||||
|
||||
|
||||
def _row(task_id: Any, **over: Any) -> SimpleNamespace:
|
||||
base = {
|
||||
"agent_slug": "be-pm",
|
||||
"task_id": task_id,
|
||||
"count": 2,
|
||||
"last_status": "blocked",
|
||||
"last_check": datetime(2026, 6, 26, tzinfo=UTC),
|
||||
"tracing_resets": 0,
|
||||
"notified": False,
|
||||
}
|
||||
base.update(over)
|
||||
return SimpleNamespace(**base)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Pure partition helper
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_partition_keeps_live_nonterminal_rows() -> None:
|
||||
tid = uuid4()
|
||||
rows = [_row(tid, count=3)]
|
||||
restored, stale = AgentOrchestrator._partition_respawn_rows(
|
||||
rows, {tid: "in_progress"}
|
||||
)
|
||||
assert stale == []
|
||||
assert restored[("be-pm", str(tid))]["count"] == _SEEDED_COUNT
|
||||
assert restored[("be-pm", str(tid))]["last_status"] == "blocked"
|
||||
|
||||
|
||||
def test_partition_drops_terminal_and_missing_rows() -> None:
|
||||
done, cancelled, gone = uuid4(), uuid4(), uuid4()
|
||||
rows = [_row(done), _row(cancelled), _row(gone)]
|
||||
restored, stale = AgentOrchestrator._partition_respawn_rows(
|
||||
rows,
|
||||
{done: "completed", cancelled: "cancelled"}, # gone absent entirely
|
||||
)
|
||||
assert restored == {}
|
||||
assert {(s, t) for s, t in stale} == {
|
||||
("be-pm", done),
|
||||
("be-pm", cancelled),
|
||||
("be-pm", gone),
|
||||
}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Startup loader
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _mock_session_factory(respawn_rows: list[Any], live: list[Any]) -> Any:
|
||||
"""A get_session_factory() stub: 1st execute -> respawn rows, 2nd -> live
|
||||
tasks, any further (deletes) -> a throwaway result."""
|
||||
db = AsyncMock()
|
||||
rows_result = MagicMock()
|
||||
rows_result.scalars.return_value.all.return_value = respawn_rows
|
||||
live_result = MagicMock()
|
||||
live_result.all.return_value = live
|
||||
db.execute = AsyncMock(side_effect=[rows_result, live_result, MagicMock()])
|
||||
db.commit = AsyncMock()
|
||||
ctx = MagicMock()
|
||||
ctx.__aenter__ = AsyncMock(return_value=db)
|
||||
ctx.__aexit__ = AsyncMock(return_value=False)
|
||||
factory = MagicMock(return_value=ctx)
|
||||
return factory, db
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loader_populates_dict_with_str_keys() -> None:
|
||||
orch = _new_orchestrator()
|
||||
tid = uuid4()
|
||||
factory, _db = _mock_session_factory(
|
||||
[_row(tid, count=3)], [SimpleNamespace(id=tid, status="in_progress")]
|
||||
)
|
||||
with patch("roboco.db.base.get_session_factory", return_value=factory):
|
||||
restored = await orch.restore_respawn_tracker()
|
||||
assert restored == 1
|
||||
assert orch._pm_respawn_tracker[("be-pm", str(tid))]["count"] == _SEEDED_COUNT
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loader_skips_terminal_rows_and_deletes_them() -> None:
|
||||
orch = _new_orchestrator()
|
||||
done = uuid4()
|
||||
factory, db = _mock_session_factory(
|
||||
[_row(done)], [SimpleNamespace(id=done, status="completed")]
|
||||
)
|
||||
with patch("roboco.db.base.get_session_factory", return_value=factory):
|
||||
restored = await orch.restore_respawn_tracker()
|
||||
assert restored == 0
|
||||
assert orch._pm_respawn_tracker == {}
|
||||
db.commit.assert_awaited() # stale row deleted + committed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loader_empty_on_exception() -> None:
|
||||
orch = _new_orchestrator()
|
||||
with patch(
|
||||
"roboco.db.base.get_session_factory", side_effect=RuntimeError("db down")
|
||||
):
|
||||
restored = await orch.restore_respawn_tracker()
|
||||
assert restored == 0
|
||||
assert orch._pm_respawn_tracker == {}
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Write-through on each gate mutation
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def _capture_persists(orch: AgentOrchestrator) -> list[tuple[str, str, dict[str, Any]]]:
|
||||
"""Record each _schedule_respawn_persist call, snapshotting the payload.
|
||||
|
||||
The real scheduler copies the record (dict(record)) before the background
|
||||
write, so the test must snapshot too — the gate mutates the same dict object
|
||||
in place, so capturing the reference would show only its final state.
|
||||
"""
|
||||
captured: list[tuple[str, str, dict[str, Any]]] = []
|
||||
|
||||
def _cap(agent_slug: str, task_id: str, record: dict[str, Any]) -> None:
|
||||
captured.append((agent_slug, task_id, dict(record)))
|
||||
|
||||
cast("Any", orch)._schedule_respawn_persist = _cap
|
||||
return captured
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_entry_and_strike_schedule_persist() -> None:
|
||||
orch = _new_orchestrator()
|
||||
captured = _capture_persists(orch)
|
||||
task_id = str(uuid4())
|
||||
task = {"id": task_id, "status": "pending"}
|
||||
fake_audit = AsyncMock()
|
||||
fake_audit.has_recent_tracing_gap = AsyncMock(return_value=False)
|
||||
with (
|
||||
patch("roboco.services.audit.get_audit_service", return_value=fake_audit),
|
||||
patch(
|
||||
"roboco.services.notification.NotificationService",
|
||||
return_value=AsyncMock(),
|
||||
),
|
||||
):
|
||||
await orch._pm_respawn_should_gate("be-pm", task) # new entry
|
||||
await orch._pm_respawn_should_gate("be-pm", task) # strike -> count 2
|
||||
assert len(captured) >= _MIN_PERSISTS
|
||||
assert captured[0][0] == "be-pm" and captured[0][1] == task_id
|
||||
assert captured[0][2]["count"] == 1
|
||||
assert captured[1][2]["count"] == _STRIKE_COUNT
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notified_flip_schedules_persist_with_notified_true() -> None:
|
||||
orch = _new_orchestrator()
|
||||
captured = _capture_persists(orch)
|
||||
task_id = str(uuid4())
|
||||
task = {"id": task_id, "status": "pending"}
|
||||
fake_audit = AsyncMock()
|
||||
fake_audit.has_recent_tracing_gap = AsyncMock(return_value=False)
|
||||
with (
|
||||
patch("roboco.services.audit.get_audit_service", return_value=fake_audit),
|
||||
patch(
|
||||
"roboco.services.notification.NotificationService",
|
||||
return_value=AsyncMock(),
|
||||
),
|
||||
):
|
||||
for _ in range(4): # 4th trips the gate + flips notified
|
||||
await orch._pm_respawn_should_gate("be-pm", task)
|
||||
assert any(c[2].get("notified") for c in captured), (
|
||||
"the notified flip must schedule a persist with notified=True"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tracing_reset_schedules_persist_with_reset_count() -> None:
|
||||
orch = _new_orchestrator()
|
||||
captured = _capture_persists(orch)
|
||||
task_id = str(uuid4())
|
||||
task = {"id": task_id, "status": "blocked"}
|
||||
fake_audit = AsyncMock()
|
||||
fake_audit.has_recent_tracing_gap = AsyncMock(return_value=True)
|
||||
with patch("roboco.services.audit.get_audit_service", return_value=fake_audit):
|
||||
await orch._pm_respawn_should_gate("be-pm", task) # new entry
|
||||
await orch._pm_respawn_should_gate("be-pm", task) # tracing reset
|
||||
assert captured[-1][2]["count"] == 1
|
||||
assert captured[-1][2]["tracing_resets"] == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Persist helper safety
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persist_record_swallows_db_failure() -> None:
|
||||
orch = _new_orchestrator()
|
||||
with patch(
|
||||
"roboco.db.base.get_session_factory", side_effect=RuntimeError("db down")
|
||||
):
|
||||
# Must not raise — a persistence failure can never gate/un-gate a spawn.
|
||||
await orch._persist_respawn_record(
|
||||
"be-pm",
|
||||
str(uuid4()),
|
||||
{"count": 2, "last_check": datetime.now(UTC)},
|
||||
)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Safety regression + transparency
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restored_counter_trips_at_persisted_threshold_not_from_one() -> None:
|
||||
"""A restart mid-loop must NOT reset the strike count to 1."""
|
||||
orch = _new_orchestrator()
|
||||
cast("Any", orch)._schedule_respawn_persist = MagicMock()
|
||||
task_id = str(uuid4())
|
||||
# Simulate restore: count=3 (one below the trip), status-stable, resets spent.
|
||||
orch._pm_respawn_tracker[("be-pm", task_id)] = {
|
||||
"count": 3,
|
||||
"last_status": "blocked",
|
||||
"last_check": datetime(2026, 6, 26, tzinfo=UTC),
|
||||
"tracing_resets": 3,
|
||||
"notified": False,
|
||||
}
|
||||
task = {"id": task_id, "status": "blocked"}
|
||||
fake_audit = AsyncMock()
|
||||
fake_audit.has_recent_tracing_gap = AsyncMock(return_value=False)
|
||||
with (
|
||||
patch("roboco.services.audit.get_audit_service", return_value=fake_audit),
|
||||
patch(
|
||||
"roboco.services.notification.NotificationService",
|
||||
return_value=AsyncMock(),
|
||||
),
|
||||
):
|
||||
gated = await orch._pm_respawn_should_gate("be-pm", task)
|
||||
assert gated is True # fires on the next spawn, not re-counted from 1
|
||||
assert orch._pm_respawn_tracker[("be-pm", task_id)]["count"] == _TRIP_COUNT
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restart_midloop_continues_identically_to_no_restart() -> None:
|
||||
"""Transparency: the gate decision depends only on the dict contents.
|
||||
|
||||
Drive a fresh orchestrator through N spawns; separately drive a second one
|
||||
for K spawns, snapshot its dict (the restart point), load that snapshot into
|
||||
a third orchestrator and continue — the tail must equal the no-restart tail.
|
||||
"""
|
||||
task = {"id": "t1", "status": "pending"}
|
||||
spawns = 5
|
||||
restart_after = 2
|
||||
|
||||
async def _spawn(orch: AgentOrchestrator) -> bool:
|
||||
fake_audit = AsyncMock()
|
||||
fake_audit.has_recent_tracing_gap = AsyncMock(return_value=False)
|
||||
with (
|
||||
patch("roboco.services.audit.get_audit_service", return_value=fake_audit),
|
||||
patch(
|
||||
"roboco.services.notification.NotificationService",
|
||||
return_value=AsyncMock(),
|
||||
),
|
||||
):
|
||||
return await orch._pm_respawn_should_gate("be-pm", task)
|
||||
|
||||
no_restart = _new_orchestrator()
|
||||
cast("Any", no_restart)._schedule_respawn_persist = MagicMock()
|
||||
full = [await _spawn(no_restart) for _ in range(spawns)]
|
||||
|
||||
pre = _new_orchestrator()
|
||||
cast("Any", pre)._schedule_respawn_persist = MagicMock()
|
||||
for _ in range(restart_after):
|
||||
await _spawn(pre)
|
||||
snapshot = copy.deepcopy(pre._pm_respawn_tracker)
|
||||
|
||||
loaded = _new_orchestrator()
|
||||
cast("Any", loaded)._schedule_respawn_persist = MagicMock()
|
||||
loaded._pm_respawn_tracker = snapshot
|
||||
tail = [await _spawn(loaded) for _ in range(spawns - restart_after)]
|
||||
|
||||
assert tail == full[restart_after:]
|
||||
Reference in New Issue
Block a user