feat(budgets): per-task and per-project cost budgets (flag-gated) (#654)

* fix(notifications): exponential backoff + CAS claim for expired-unacked re-escalation The sweep re-escalated every expired unacked ack-required notification on every ~60s tick, forever — the live incident: 3 fresh blocker escalations + Telegram DMs per minute from a static stale pile. Now each notification carries reescalation_count / last_reescalated_at / reescalation_delivered_count (migration 079): first fire at expiry, then doubling intervals from 1h capped at 24h, hard stop after ROBOCO_NOTIFICATION_MAX_REESCALATIONS (default 5) with one permanent log carrying attempts-vs-delivered so 'seen and ignored' is distinguishable from 'route never worked'. The due/wait/capped decision is a pure function in foundation/policy/communications.py. Per adversarial review, the attempt slot is claimed by compare-and-set (UPDATE ... WHERE reescalation_count = :n) BEFORE delivery — the previous draft leaned on the 60s dedup window, which never engages for BLOCKER_ESCALATION (_LOOP_PRONE_TYPES excludes it), so concurrent sweeps would have double-delivered. A lost claim skips delivery outright. Legacy rows read as count=0 and keep today's first-fire semantics. 61 tests incl. a two-session CAS race and a real alembic upgrade/downgrade round trip. * feat(budgets): per-task and per-project cost budgets (flag-gated) tasks.budget_usd + projects.monthly_budget_usd (migration 080, chained on 079; adds ix_agent_spawn_sessions_task_id since both enforcement seams filter on bare task_id). Behind ROBOCO_TASK_BUDGETS_ENABLED (default off, feature-flags card) — verifiably inert when off. Claim-time: a project-month-spend guard applies to WORK-STARTING claims only (i_will_work_on / i_will_plan) — per adversarial review, review/ doc/gate/inbound-PR claims are exempt so in-flight work can always finish reviewing and merging at cap. Spend counts closed sessions' estimated_cost_usd PLUS open sessions priced live from token snapshots (the original closed-only sum read parallel long sessions as $0). Sweep-side: the existing budget sweep also prices the active task's spend vs budget_usd (TaskType defaults when null); on breach the task is BLOCKED (HUMAN resolver, budget marker) BEFORE the graceful stop so the unclaim no-ops and the dispatcher never respawns onto it, and the CEO notification names both recovery steps. unblock on a budget-blocked task re-checks live spend and refuses while still over — no silent re-breach loop. Panel: budget inputs in both dialogs (0 rejected — a zero budget silently blocks everything), spend logic consolidated in TaskService.task_spend_usd. 42 new tests incl. a real-DB spend-query suite and a two-tick non-refire sweep test. --------- Co-authored-by: Renn F <rennf93@users.noreply.github.com>
2026-08-03 07:23:24 +02:00 · 2026-07-23 00:06:14 +02:00
parent 1d5a8e846f
commit 7c8453e210
31 changed files with 2288 additions and 65 deletions
@@ -0,0 +1,340 @@
+"""Task-level $ budget sweep (ROBOCO_TASK_BUDGETS_ENABLED).
+
+`_sweep_budget_exceeded` gains a second trigger alongside the existing
+tool-call halt: when the flag is on and an active task's own $ budget is
+breached, the task is BLOCKED + the CEO notified (`_handle_task_budget_breach`)
+BEFORE the agent is gracefully stopped — never a mid-verb kill, and never left
+to bounce through `pending` for an instant re-claim.
+"""
+
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from roboco.config import settings
+from roboco.models.base import BlockerResolverType, TaskStatus, TaskType
+from roboco.runtime.orchestrator import AgentOrchestrator, AgentState
+
+_MOCK_TASK_SPEND_USD = 3.0
+
+
+def _make_orchestrator() -> AgentOrchestrator:
+    with patch.object(AgentOrchestrator, "__init__", return_value=None):
+        orch = AgentOrchestrator.__new__(AgentOrchestrator)
+    orch._instances = {}
+    orch._lock = MagicMock()
+    return orch
+
+
+def _instance(task_id: str | None) -> MagicMock:
+    inst = MagicMock()
+    inst.state = AgentState.ACTIVE
+    inst.container_id = "deadbeef1234"
+    inst.current_task_id = task_id
+    inst.error_count = 0
+    inst.config = MagicMock(git_context=None)
+    return inst
+
+
+def _db_ctx(db: Any) -> Any:
+    @asynccontextmanager
+    async def _ctx() -> Any:
+        yield db
+
+    return _ctx
+
+
+@pytest.mark.asyncio
+async def test_task_budget_breach_blocks_before_graceful_stop(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    orch = _make_orchestrator()
+    task_id = "11111111-1111-1111-1111-111111111111"
+    orch._instances = {"be-dev-1": _instance(task_id)}
+    monkeypatch.setattr(settings, "task_budgets_enabled", True)
+
+    with (
+        patch.object(
+            AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None)
+        ),
+        patch.object(orch, "_task_budget_breach", AsyncMock(return_value=(5.0, 7.5))),
+        patch.object(orch, "_handle_task_budget_breach", AsyncMock()) as handle_mock,
+        patch.object(orch, "stop_agent", AsyncMock()) as stop_mock,
+    ):
+        await orch._sweep_budget_exceeded()
+
+    # Block + notify runs, and runs BEFORE stop_agent (never a mid-verb kill —
+    # graceful=True, and the task is already blocked by the time the agent dies).
+    handle_mock.assert_awaited_once_with(task_id, cap_usd=5.0, spend_usd=7.5)
+    stop_mock.assert_awaited_once_with(
+        "be-dev-1",
+        graceful=True,
+        release_claim=True,
+        stop_reason="budget_exceeded_task",
+    )
+
+
+@pytest.mark.asyncio
+async def test_flag_off_never_checks_task_budget(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    orch = _make_orchestrator()
+    task_id = "11111111-1111-1111-1111-111111111111"
+    orch._instances = {"be-dev-1": _instance(task_id)}
+    monkeypatch.setattr(settings, "task_budgets_enabled", False)
+
+    with (
+        patch.object(
+            AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None)
+        ),
+        patch.object(orch, "_task_budget_breach", AsyncMock()) as breach_mock,
+        patch.object(orch, "stop_agent", AsyncMock()) as stop_mock,
+    ):
+        await orch._sweep_budget_exceeded()
+
+    breach_mock.assert_not_awaited()
+    stop_mock.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_under_budget_is_a_no_op(monkeypatch: pytest.MonkeyPatch) -> None:
+    orch = _make_orchestrator()
+    task_id = "11111111-1111-1111-1111-111111111111"
+    orch._instances = {"be-dev-1": _instance(task_id)}
+    monkeypatch.setattr(settings, "task_budgets_enabled", True)
+
+    with (
+        patch.object(
+            AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None)
+        ),
+        patch.object(orch, "_task_budget_breach", AsyncMock(return_value=None)),
+        patch.object(orch, "_handle_task_budget_breach", AsyncMock()) as handle_mock,
+        patch.object(orch, "stop_agent", AsyncMock()) as stop_mock,
+    ):
+        await orch._sweep_budget_exceeded()
+
+    handle_mock.assert_not_awaited()
+    stop_mock.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_tool_call_halt_path_is_unchanged(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """The pre-existing tool-call halt trigger still fires with its own
+    stop_reason, independent of the new $ budget path."""
+    orch = _make_orchestrator()
+    orch._instances = {"be-dev-1": _instance(None)}
+    monkeypatch.setattr(settings, "task_budgets_enabled", False)
+
+    halt_status = {"halt": True, "total": 301, "halt_threshold": 300}
+    with (
+        patch.object(
+            AgentOrchestrator,
+            "_fetch_budget_status",
+            AsyncMock(return_value=halt_status),
+        ),
+        patch.object(orch, "stop_agent", AsyncMock()) as stop_mock,
+    ):
+        await orch._sweep_budget_exceeded()
+
+    stop_mock.assert_awaited_once_with(
+        "be-dev-1", graceful=True, release_claim=True, stop_reason="budget_sweep"
+    )
+
+
+@pytest.mark.asyncio
+async def test_no_task_id_skips_task_budget_check(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A taskless spawn (current_task_id=None) never reaches the $ budget
+    check even with the flag on."""
+    orch = _make_orchestrator()
+    orch._instances = {"be-dev-1": _instance(None)}
+    monkeypatch.setattr(settings, "task_budgets_enabled", True)
+
+    with (
+        patch.object(
+            AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None)
+        ),
+        patch.object(orch, "_task_budget_breach", AsyncMock()) as breach_mock,
+        patch.object(orch, "stop_agent", AsyncMock()) as stop_mock,
+    ):
+        await orch._sweep_budget_exceeded()
+
+    breach_mock.assert_not_awaited()
+    stop_mock.assert_not_awaited()
+
+
+# ---------------------------------------------------------------------------
+# _handle_task_budget_breach: the block + notify write itself
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_handle_breach_blocks_task_and_notifies_ceo() -> None:
+    orch = _make_orchestrator()
+    task_id = "22222222-2222-2222-2222-222222222222"
+    task = MagicMock(status=TaskStatus.IN_PROGRESS)
+    db = MagicMock()
+
+    task_svc = MagicMock()
+    task_svc.get = AsyncMock(return_value=task)
+    task_svc.admin_set_status = AsyncMock()
+    delivery = MagicMock()
+    delivery.notify_ceo_of_budget_breach = AsyncMock()
+
+    with (
+        patch("roboco.db.base.get_db_context", _db_ctx(db)),
+        patch("roboco.services.task.TaskService", return_value=task_svc),
+        patch(
+            "roboco.services.notification_delivery.get_notification_delivery_service",
+            return_value=delivery,
+        ),
+    ):
+        await orch._handle_task_budget_breach(task_id, cap_usd=5.0, spend_usd=8.0)
+
+    assert task.blocker_resolver_type == BlockerResolverType.HUMAN
+    task_svc.admin_set_status.assert_awaited_once()
+    args, _kwargs = task_svc.admin_set_status.call_args
+    assert args[1] == TaskStatus.BLOCKED
+    delivery.notify_ceo_of_budget_breach.assert_awaited_once_with(
+        task=task, task_id=args[0], cap_usd=5.0, spend_usd=8.0
+    )
+
+
+@pytest.mark.asyncio
+async def test_handle_breach_skips_a_task_that_already_moved_on() -> None:
+    """A stale re-check racing the task's own progress (e.g. it completed
+    between the read and the write) must not block/notify."""
+    orch = _make_orchestrator()
+    task_id = "33333333-3333-3333-3333-333333333333"
+    task = MagicMock(status=TaskStatus.COMPLETED)
+    db = MagicMock()
+
+    task_svc = MagicMock()
+    task_svc.get = AsyncMock(return_value=task)
+    task_svc.admin_set_status = AsyncMock()
+    delivery = MagicMock()
+    delivery.notify_ceo_of_budget_breach = AsyncMock()
+
+    with (
+        patch("roboco.db.base.get_db_context", _db_ctx(db)),
+        patch("roboco.services.task.TaskService", return_value=task_svc),
+        patch(
+            "roboco.services.notification_delivery.get_notification_delivery_service",
+            return_value=delivery,
+        ),
+    ):
+        await orch._handle_task_budget_breach(task_id, cap_usd=5.0, spend_usd=8.0)
+
+    task_svc.admin_set_status.assert_not_awaited()
+    delivery.notify_ceo_of_budget_breach.assert_not_awaited()
+
+
+# ---------------------------------------------------------------------------
+# _task_budget_breach: cap resolution (null -> TaskType default) + spend sum
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_breach_falls_back_to_tasktype_default_when_budget_null() -> None:
+    """Cap resolution (task.budget_usd null -> TaskType default) and spend
+    both delegate to TaskService now (task_spend_usd's own open-session
+    pricing is covered directly by its shared implementation — see
+    test_project_month_spend_usd_db.py's real-DB open-session case)."""
+    orch = _make_orchestrator()
+    task_id = "44444444-4444-4444-4444-444444444444"
+    task = MagicMock(
+        status=TaskStatus.IN_PROGRESS, task_type=TaskType.DOCUMENTATION, budget_usd=None
+    )
+    task_svc = MagicMock()
+    task_svc.get = AsyncMock(return_value=task)
+    task_svc.task_spend_usd = AsyncMock(return_value=_MOCK_TASK_SPEND_USD)
+    db = MagicMock()
+
+    with (
+        patch("roboco.db.base.get_db_context", _db_ctx(db)),
+        patch("roboco.services.task.TaskService", return_value=task_svc),
+    ):
+        breach = await orch._task_budget_breach(task_id)
+
+    assert breach is not None
+    cap_usd, spend_usd = breach
+    assert cap_usd == 1.0  # TASK_TYPE_DEFAULT_BUDGET_USD[DOCUMENTATION]
+    assert spend_usd == _MOCK_TASK_SPEND_USD
+
+
+@pytest.mark.asyncio
+async def test_breach_none_when_task_left_claimed_or_in_progress() -> None:
+    """A stale re-check (the task already reached e.g. awaiting_qa) is not a
+    breach — the spend query is never even issued."""
+    orch = _make_orchestrator()
+    task_id = "55555555-5555-5555-5555-555555555555"
+    task = MagicMock(
+        status=TaskStatus.AWAITING_QA, task_type=TaskType.CODE, budget_usd=1.0
+    )
+    task_svc = MagicMock()
+    task_svc.get = AsyncMock(return_value=task)
+    task_svc.task_spend_usd = AsyncMock()
+    db = MagicMock()
+
+    with (
+        patch("roboco.db.base.get_db_context", _db_ctx(db)),
+        patch("roboco.services.task.TaskService", return_value=task_svc),
+    ):
+        breach = await orch._task_budget_breach(task_id)
+
+    assert breach is None
+    task_svc.task_spend_usd.assert_not_awaited()
+
+
+# ---------------------------------------------------------------------------
+# Repeated ticks: a blocked task must not re-fire the breach handling.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_repeated_ticks_do_not_refire_once_blocked(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Two consecutive _sweep_budget_exceeded ticks against the same still-
+    registered instance: tick 1 detects the breach and blocks the task; by
+    tick 2 the task is BLOCKED (no longer CLAIMED/IN_PROGRESS), so
+    _task_budget_breach's own status guard returns None — the sweep never
+    re-blocks / re-notifies / re-stops a task that's already been handled."""
+    orch = _make_orchestrator()
+    task_id = "77777777-7777-7777-7777-777777777777"
+    orch._instances = {"be-dev-1": _instance(task_id)}
+    monkeypatch.setattr(settings, "task_budgets_enabled", True)
+
+    in_progress_task = MagicMock(
+        status=TaskStatus.IN_PROGRESS, task_type=TaskType.CODE, budget_usd=5.0
+    )
+    # Simulates the task having been transitioned to BLOCKED by tick 1's
+    # (mocked-out) _handle_task_budget_breach before tick 2 re-checks it.
+    blocked_task = MagicMock(
+        status=TaskStatus.BLOCKED, task_type=TaskType.CODE, budget_usd=5.0
+    )
+    task_svc = MagicMock()
+    task_svc.get = AsyncMock(side_effect=[in_progress_task, blocked_task])
+    task_svc.task_spend_usd = AsyncMock(return_value=7.0)
+    db = MagicMock()
+
+    with (
+        patch.object(
+            AgentOrchestrator, "_fetch_budget_status", AsyncMock(return_value=None)
+        ),
+        patch("roboco.db.base.get_db_context", _db_ctx(db)),
+        patch("roboco.services.task.TaskService", return_value=task_svc),
+        patch.object(orch, "_handle_task_budget_breach", AsyncMock()) as handle_mock,
+        patch.object(orch, "stop_agent", AsyncMock()) as stop_mock,
+    ):
+        await orch._sweep_budget_exceeded()  # tick 1: breach detected
+        await orch._sweep_budget_exceeded()  # tick 2: already blocked
+
+    handle_mock.assert_awaited_once()
+    stop_mock.assert_awaited_once()