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>
This commit is contained in:
Renzo F
2026-07-23 00:06:14 +02:00
committed by GitHub
co-authored by Renn F
parent 1d5a8e846f
commit 7c8453e210
31 changed files with 2288 additions and 65 deletions
@@ -0,0 +1,27 @@
"""Per-task and per-project cost budgets are gated by a default-off config
flag, registered as a panel-tunable feature flag (mirrors possibilities_matrix)."""
from __future__ import annotations
import os
from unittest import mock
from roboco.config import Settings
from roboco.services.settings import FEATURE_FLAGS, validate_setting
def test_task_budgets_disabled_by_default() -> None:
assert Settings().task_budgets_enabled is False
def test_task_budgets_reads_env_var() -> None:
with mock.patch.dict(os.environ, {"ROBOCO_TASK_BUDGETS_ENABLED": "true"}):
assert Settings().task_budgets_enabled is True
def test_task_budgets_flag_registered_in_feature_flags() -> None:
assert "task_budgets_enabled" in [key for key, _ in FEATURE_FLAGS]
def test_task_budgets_flag_validates_as_bool() -> None:
validate_setting("task_budgets_enabled", "true")
@@ -0,0 +1,112 @@
"""Per-project monthly-budget claim guard (ROBOCO_TASK_BUDGETS_ENABLED).
Two layers: the pure predicate ``project_budget_exceeded_guard`` (over/under/
null cap) and the Choreographer's ``_project_budget_claim_guard`` wiring
(flag-off inert, no-cap inert, spend query only reached when both are set).
"""
from __future__ import annotations
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4
import pytest
from roboco.config import settings
from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps
from roboco.services.gateway.claim_guards import project_budget_exceeded_guard
# ---------------------------------------------------------------------------
# Pure predicate: project_budget_exceeded_guard
# ---------------------------------------------------------------------------
def test_predicate_null_cap_is_inert() -> None:
task = MagicMock(id=uuid4())
assert project_budget_exceeded_guard(task, None, 999.0) is None
def test_predicate_under_cap_passes() -> None:
task = MagicMock(id=uuid4())
assert project_budget_exceeded_guard(task, 10.0, 5.0) is None
def test_predicate_over_cap_refuses() -> None:
task = MagicMock(id=uuid4())
env = project_budget_exceeded_guard(task, 10.0, 15.0)
assert env is not None
body = env.as_dict()
assert body["error"] == "invalid_state"
assert "10.00" in body["message"]
assert "15.00" in body["message"]
assert "project settings" in body["remediate"]
def test_predicate_at_cap_refuses() -> None:
"""At-or-over refuses — the cap itself is not headroom."""
task = MagicMock(id=uuid4())
assert project_budget_exceeded_guard(task, 10.0, 10.0) is not None
# ---------------------------------------------------------------------------
# Choreographer wiring: _project_budget_claim_guard
# ---------------------------------------------------------------------------
def _make_choreographer(*, project_month_spend_usd: float = 0.0) -> Choreographer:
task_svc = AsyncMock()
task_svc.project_month_spend_usd.return_value = project_month_spend_usd
base = {
"task": task_svc,
"work_session": AsyncMock(),
"git": AsyncMock(),
"a2a": AsyncMock(),
"journal": AsyncMock(),
"audit": AsyncMock(),
"evidence_repo": AsyncMock(),
}
return Choreographer(ChoreographerDeps(**base))
def _task_with_project(monthly_budget_usd: float | None) -> MagicMock:
project = MagicMock(id=uuid4(), monthly_budget_usd=monthly_budget_usd)
return MagicMock(id=uuid4(), project=project)
@pytest.mark.asyncio
async def test_flag_off_is_inert_even_over_cap(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(settings, "task_budgets_enabled", False)
c = _make_choreographer(project_month_spend_usd=999.0)
task = _task_with_project(monthly_budget_usd=10.0)
assert await c._project_budget_claim_guard(task) is None
c.task.project_month_spend_usd.assert_not_awaited()
@pytest.mark.asyncio
async def test_no_project_cap_is_inert(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(settings, "task_budgets_enabled", True)
c = _make_choreographer(project_month_spend_usd=999.0)
task = _task_with_project(monthly_budget_usd=None)
assert await c._project_budget_claim_guard(task) is None
c.task.project_month_spend_usd.assert_not_awaited()
@pytest.mark.asyncio
async def test_under_cap_passes(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(settings, "task_budgets_enabled", True)
c = _make_choreographer(project_month_spend_usd=4.0)
task = _task_with_project(monthly_budget_usd=10.0)
assert await c._project_budget_claim_guard(task) is None
@pytest.mark.asyncio
async def test_over_cap_refuses_naming_the_field(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(settings, "task_budgets_enabled", True)
c = _make_choreographer(project_month_spend_usd=12.0)
task = _task_with_project(monthly_budget_usd=10.0)
env = await c._project_budget_claim_guard(task)
assert env is not None
body = env.as_dict()
assert body["error"] == "invalid_state"
assert "Monthly Budget" in body["remediate"]
@@ -0,0 +1,335 @@
"""The project monthly-budget claim guard is scoped to work-STARTING claims.
A project at (or over) its monthly cap still has non-negotiable in-flight
work to finish: QA's ``claim_review``, the PR gate's ``claim_gate_review``,
``claim_doc_task``, and inbound ``claim_pr_review`` never even reach the
spend query (``check_project_budget`` defaults False at those four call
sites) — reviewing/documenting/merging what's already been paid for must
never wedge behind an exhausted cap. Only ``i_will_work_on`` / ``i_will_plan``
(``check_project_budget=True``) — genuinely starting new spend — refuse.
"""
from __future__ import annotations
from typing import Any
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4
import pytest
from roboco.config import settings
from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps
_STEPS = [
{
"title": "Implement the change",
"description": (
"edit the target file, add tests, run them, and stage the "
"change for commit on the task branch"
),
}
]
@pytest.fixture(autouse=True)
def _budgets_on(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(settings, "task_budgets_enabled", True)
def _over_cap_project() -> MagicMock:
"""A project with a $10 cap — the spend stub below always reports $999."""
return MagicMock(id=uuid4(), monthly_budget_usd=10.0)
def _make_deps(task_svc: AsyncMock, **overrides: Any) -> ChoreographerDeps:
base: dict[str, Any] = {
"task": task_svc,
"work_session": AsyncMock(),
"git": AsyncMock(),
"a2a": AsyncMock(),
"journal": AsyncMock(),
"audit": AsyncMock(),
"evidence_repo": AsyncMock(),
}
base.update(overrides)
return ChoreographerDeps(**base)
# ---------------------------------------------------------------------------
# Review-ish claims: guard skipped — each succeeds despite an over-cap project.
# project_month_spend_usd is asserted NOT awaited: the guard is skipped
# entirely, not merely lucky (e.g. a cache hit).
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_claim_review_succeeds_at_cap() -> None:
task_svc = AsyncMock()
t = MagicMock(
id=uuid4(),
status="awaiting_qa",
assigned_to=uuid4(),
parent_task_id=uuid4(),
task_type="code",
dependency_ids=[],
team="backend",
pr_number=10,
pr_url="https://example/pr/10",
branch_name="feature/backend/abc",
batch_id=None,
project=_over_cap_project(),
)
task_svc.get.return_value = t
task_svc.agent_for.return_value = MagicMock(role="qa", slug="be-qa")
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
task_svc.unmet_dependency_ids = AsyncMock(return_value=[])
task_svc.has_earlier_incomplete_code_sibling.return_value = False
task_svc.qa_claim = AsyncMock(return_value=t)
task_svc.project_month_spend_usd = AsyncMock(return_value=999.0)
c = Choreographer(_make_deps(task_svc))
cc: Any = c
cc._build_qa_review_evidence = AsyncMock(return_value={})
env = await c.claim_review(uuid4(), t.id)
body = env.as_dict()
assert body.get("error") is None, body
task_svc.project_month_spend_usd.assert_not_awaited()
@pytest.mark.asyncio
async def test_claim_gate_review_succeeds_at_cap() -> None:
task_svc = AsyncMock()
t = MagicMock(
id=uuid4(),
status="awaiting_pr_review",
assigned_to=uuid4(),
parent_task_id=uuid4(),
task_type="planning",
dependency_ids=[],
team="main_pm",
pr_number=139,
pr_url="https://example/pr/139",
branch_name="feature/main_pm/root",
batch_id=None,
project=_over_cap_project(),
)
task_svc.get.return_value = t
task_svc.agent_for.return_value = MagicMock(
role="pr_reviewer", slug="be-pr-reviewer"
)
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
task_svc.unmet_dependency_ids = AsyncMock(return_value=[])
task_svc.pr_gate_claim = AsyncMock(return_value=t)
task_svc.project_month_spend_usd = AsyncMock(return_value=999.0)
c = Choreographer(_make_deps(task_svc))
cc: Any = c
cc._build_gate_review_evidence = AsyncMock(return_value={"pr_number": 139})
env = await c.claim_gate_review(uuid4(), t.id)
body = env.as_dict()
assert body.get("error") is None, body
task_svc.project_month_spend_usd.assert_not_awaited()
@pytest.mark.asyncio
async def test_claim_doc_task_succeeds_at_cap() -> None:
task_svc = AsyncMock()
t = MagicMock(
id=uuid4(),
status="awaiting_documentation",
assigned_to=None,
parent_task_id=None,
task_type="documentation",
team="backend",
branch_name="feature/backend/abc",
quick_context=None,
documents=[],
commits=[{"sha": "abc123", "message": "[x] work"}],
pr_number=7,
pr_url="https://github.com/x/y/pull/7",
dev_notes="done",
acceptance_criteria_status=[],
work_session_id=uuid4(),
dependency_ids=[],
batch_id=None,
project=_over_cap_project(),
)
task_svc.get.return_value = t
task_svc.agent_for.return_value = MagicMock(role="documenter", team="backend")
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
task_svc.unmet_dependency_ids = AsyncMock(return_value=[])
task_svc.doc_claim.return_value = t
task_svc.project_month_spend_usd = AsyncMock(return_value=999.0)
git_svc = AsyncMock()
git_svc.diff.return_value = "diff"
git_svc.list_changed_files.return_value = ["README.md"]
c = Choreographer(_make_deps(task_svc, git=git_svc))
env = await c.claim_doc_task(uuid4(), t.id)
body = env.as_dict()
assert body["error"] is None, body
task_svc.project_month_spend_usd.assert_not_awaited()
@pytest.mark.asyncio
async def test_claim_pr_review_succeeds_at_cap() -> None:
task_svc = AsyncMock()
t = MagicMock(
id=uuid4(),
status="pending",
assigned_to=None,
parent_task_id=None,
task_type="code",
dependency_ids=[],
team="system",
batch_id=None,
project=_over_cap_project(),
)
task_svc.get.return_value = t
task_svc.agent_for.return_value = MagicMock(
role="pr_reviewer", slug="be-pr-reviewer"
)
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
task_svc.unmet_dependency_ids = AsyncMock(return_value=[])
task_svc.has_earlier_incomplete_code_sibling.return_value = False
task_svc.pr_review_claim = AsyncMock(return_value=t)
task_svc.project_month_spend_usd = AsyncMock(return_value=999.0)
c = Choreographer(_make_deps(task_svc))
cc: Any = c
cc._build_pr_review_evidence = AsyncMock(return_value={})
env = await c.claim_pr_review(uuid4(), t.id)
body = env.as_dict()
assert body.get("error") is None, body
task_svc.project_month_spend_usd.assert_not_awaited()
# ---------------------------------------------------------------------------
# Work-starting claims: guard ON — both refuse at cap.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_i_will_work_on_refuses_at_cap() -> None:
agent_id = uuid4()
task_id = uuid4()
target = MagicMock(
id=task_id,
status="pending",
plan=None,
assigned_to=None,
parent_task_id=None,
sequence=0,
task_type="code",
team="backend",
dependency_ids=[],
project=_over_cap_project(),
)
task_svc = AsyncMock()
task_svc.get.return_value = target
task_svc.agent_for.return_value = MagicMock(
id=agent_id, role="developer", team="backend", slug=None
)
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
task_svc.unmet_dependency_ids = AsyncMock(return_value=[])
task_svc.project_month_spend_usd = AsyncMock(return_value=999.0)
c = Choreographer(_make_deps(task_svc))
env = await c.i_will_work_on(agent_id, task_id, plan="x", steps=_STEPS)
body = env.as_dict()
assert body["error"] == "invalid_state", body
assert "10.00" in body["message"] and "999.00" in body["message"]
task_svc.claim.assert_not_awaited()
task_svc.project_month_spend_usd.assert_awaited_once()
@pytest.mark.asyncio
async def test_i_will_plan_refuses_at_cap() -> None:
pm_id = uuid4()
task_id = uuid4()
target = MagicMock(
id=task_id,
status="pending",
plan=None,
assigned_to=None,
parent_task_id=None,
sequence=0,
task_type="planning",
team="backend",
dependency_ids=[],
project=_over_cap_project(),
)
task_svc = AsyncMock()
task_svc.get.return_value = target
task_svc.agent_for.return_value = MagicMock(
id=pm_id, role="cell_pm", team="backend", slug=None
)
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
task_svc.unmet_dependency_ids = AsyncMock(return_value=[])
task_svc.project_month_spend_usd = AsyncMock(return_value=999.0)
c = Choreographer(_make_deps(task_svc))
env = await c.i_will_plan(
pm_id,
task_id,
plan="break it down",
rich_plan={
"approach": (
"Single-cell decomposition: backend handles the full scope; "
"no frontend or ux work required for this planning task. "
"be-dev-1 owns the change end to end; QA reviews after the "
"PR opens, documentation follows, then be-pm completes and "
"submits up. Strict sequencing, no cross-cell dependencies."
),
"sub_tasks": [
{
"title": "Backend planning slice",
"description": (
"scope the change, assign be-dev-1, who implements "
"with tests and opens the leaf PR for QA review."
),
}
],
},
)
body = env.as_dict()
assert body["error"] == "invalid_state", body
task_svc.claim.assert_not_awaited()
task_svc.project_month_spend_usd.assert_awaited_once()
# ---------------------------------------------------------------------------
# Plumbing contract on _run_claim_guards itself.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_run_claim_guards_only_checks_budget_when_asked() -> None:
task_svc = AsyncMock()
task_svc.list_in_progress_for_agent.return_value = []
task_svc.list_paused_for_agent.return_value = []
task_svc.unmet_dependency_ids = AsyncMock(return_value=[])
c = Choreographer(_make_deps(task_svc))
task = MagicMock(id=uuid4(), dependency_ids=[])
calls = {"n": 0}
async def _fake_guard(_t: Any) -> None:
calls["n"] += 1
cc: Any = c
cc._project_budget_claim_guard = _fake_guard
await c._run_claim_guards(agent_id=uuid4(), task=task, skip_dev_guards=True)
assert calls["n"] == 0, (
"budget guard must not run without check_project_budget=True"
)
await c._run_claim_guards(
agent_id=uuid4(), task=task, skip_dev_guards=True, check_project_budget=True
)
assert calls["n"] == 1
@@ -0,0 +1,165 @@
"""unblock's budget-breach re-check (ROBOCO_TASK_BUDGETS_ENABLED).
A task the orchestrator's budget sweep BLOCKed carries the BUDGET_BLOCKED
marker (`_handle_task_budget_breach`). `unblock` re-checks spend-vs-cap
before letting it through: still over refuses (naming the budget
remediation), so a PM can't silently re-breach the same cap the next tick;
under (the CEO raised the cap) clears the marker and the unblock proceeds
exactly as before.
"""
from __future__ import annotations
from datetime import UTC, datetime
from typing import Any
from unittest.mock import AsyncMock, MagicMock
from uuid import uuid4
import pytest
from roboco.config import settings
from roboco.foundation.policy.content import markers
from roboco.models.base import TaskType
from roboco.services.gateway.choreographer import Choreographer, ChoreographerDeps
def _make_deps(task_svc: AsyncMock) -> ChoreographerDeps:
base: dict[str, Any] = {
"task": task_svc,
"work_session": AsyncMock(),
"git": AsyncMock(),
"a2a": AsyncMock(),
"journal": AsyncMock(),
"audit": AsyncMock(),
"evidence_repo": AsyncMock(),
}
base["journal"].has_decision_for_task.return_value = True
base["journal"].latest_decision_at.return_value = datetime.now(UTC)
return ChoreographerDeps(**base)
def _budget_blocked_task(*, budget_usd: float | None = 5.0) -> MagicMock:
t = MagicMock(
id=uuid4(),
status="blocked",
pre_block_state="in_progress",
pre_block_assignee=uuid4(),
pre_block_metadata={},
dependency_ids=[],
task_type=TaskType.CODE,
budget_usd=budget_usd,
orchestration_markers=None,
)
markers.mark_budget_blocked(t)
return t
@pytest.fixture(autouse=True)
def _budgets_on(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(settings, "task_budgets_enabled", True)
@pytest.mark.asyncio
async def test_unblock_refuses_while_still_over_cap() -> None:
t = _budget_blocked_task(budget_usd=5.0)
task_svc = AsyncMock()
task_svc.get.return_value = t
task_svc.unmet_dependency_ids.return_value = []
task_svc.task_spend_usd.return_value = 7.0
c = Choreographer(_make_deps(task_svc))
env = await c.unblock(uuid4(), t.id, "attempting to resume")
body = env.as_dict()
assert body["error"] == "invalid_state", body
assert "5.00" in body["message"] and "7.00" in body["message"]
assert "budget" in body["remediate"].lower()
task_svc.unblock_with_restore.assert_not_awaited()
# Marker survives — a retry that hasn't actually cleared the cap must
# still be caught.
assert markers.is_budget_blocked(t) is True
@pytest.mark.asyncio
async def test_unblock_succeeds_when_already_under_cap() -> None:
t = _budget_blocked_task(budget_usd=5.0)
task_svc = AsyncMock()
task_svc.get.return_value = t
task_svc.unmet_dependency_ids.return_value = []
task_svc.task_spend_usd.return_value = 3.0
task_svc.unblock_with_restore.return_value = t
c = Choreographer(_make_deps(task_svc))
env = await c.unblock(uuid4(), t.id, "spend never actually breached")
body = env.as_dict()
assert body.get("error") is None, body
task_svc.unblock_with_restore.assert_awaited_once()
assert markers.is_budget_blocked(t) is False
@pytest.mark.asyncio
async def test_raise_then_unblock_succeeds_after_a_prior_refusal() -> None:
"""First attempt: still over -> refused. CEO raises budget_usd. Second
attempt: now under -> succeeds, marker cleared."""
t = _budget_blocked_task(budget_usd=5.0)
task_svc = AsyncMock()
task_svc.get.return_value = t
task_svc.unmet_dependency_ids.return_value = []
task_svc.task_spend_usd.return_value = 7.0
task_svc.unblock_with_restore.return_value = t
c = Choreographer(_make_deps(task_svc))
pm_id = uuid4()
refused = await c.unblock(pm_id, t.id, "attempting resume")
assert refused.as_dict()["error"] == "invalid_state"
task_svc.unblock_with_restore.assert_not_awaited()
# CEO raises the cap; re-block for a fresh attempt (the same task row, as
# it would be across two real requests).
t.budget_usd = 20.0
t.status = "blocked"
succeeded = await c.unblock(pm_id, t.id, "raised the budget, resuming")
body = succeeded.as_dict()
assert body.get("error") is None, body
task_svc.unblock_with_restore.assert_awaited_once()
assert markers.is_budget_blocked(t) is False
@pytest.mark.asyncio
async def test_flag_off_ignores_the_marker(monkeypatch: pytest.MonkeyPatch) -> None:
"""The marker alone must never gate anything with the flag off."""
monkeypatch.setattr(settings, "task_budgets_enabled", False)
t = _budget_blocked_task(budget_usd=5.0)
task_svc = AsyncMock()
task_svc.get.return_value = t
task_svc.unmet_dependency_ids.return_value = []
task_svc.unblock_with_restore.return_value = t
c = Choreographer(_make_deps(task_svc))
env = await c.unblock(uuid4(), t.id, "flag is off")
assert env.as_dict().get("error") is None, env.as_dict()
task_svc.task_spend_usd.assert_not_awaited()
@pytest.mark.asyncio
async def test_non_budget_block_never_reaches_the_spend_query() -> None:
"""A task blocked for an ordinary reason (no BUDGET_BLOCKED marker) skips
the guard entirely — it's not a budget block at all."""
t = MagicMock(
id=uuid4(),
status="blocked",
pre_block_state="in_progress",
pre_block_assignee=uuid4(),
pre_block_metadata={},
dependency_ids=[],
task_type=TaskType.CODE,
budget_usd=5.0,
orchestration_markers=None,
)
task_svc = AsyncMock()
task_svc.get.return_value = t
task_svc.unmet_dependency_ids.return_value = []
task_svc.unblock_with_restore.return_value = t
c = Choreographer(_make_deps(task_svc))
env = await c.unblock(uuid4(), t.id, "manual escalation resolved")
assert env.as_dict().get("error") is None, env.as_dict()
task_svc.task_spend_usd.assert_not_awaited()
@@ -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()
@@ -0,0 +1,212 @@
"""TaskService.project_month_spend_usd against a real Postgres DB.
The join (agent_spawn_sessions.task_id, a plain String(36), to tasks.id via a
cast) and the month-boundary filter are exactly the kind of thing a mocked
unit test can't prove actually executes as real SQL. This also pins the
open-session live-token pricing fix: a still-open session's
estimated_cost_usd is null until close, so it must be priced from its token
columns via calculate_cost, not silently summed as $0.
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
from typing import TYPE_CHECKING
from uuid import UUID, uuid4
import pytest
from roboco.billing.pricing import calculate_cost
from roboco.db.tables import AgentSpawnSessionTable, AgentTable, ProjectTable, TaskTable
from roboco.models.base import (
AgentRole,
AgentStatus,
Complexity,
TaskNature,
TaskStatus,
TaskType,
Team,
)
from roboco.services.task import TaskService
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
_MODEL = "claude-sonnet-5"
async def _seed_project(session: AsyncSession) -> tuple[UUID, UUID]:
"""Seed a system agent + project. Returns ``(project_id, system_agent_id)``
— the latter doubles as the task-row FK target below."""
system_agent = AgentTable(
id=uuid4(),
name="System",
slug=f"system-{uuid4().hex[:8]}",
role=AgentRole.SYSTEM,
team=None,
status=AgentStatus.ACTIVE,
model_config={},
system_prompt="system",
capabilities=[],
permissions={},
metrics={},
)
session.add(system_agent)
await session.flush()
project = ProjectTable(
id=uuid4(),
name="Budget Spend Test Project",
slug=f"budget-spend-{uuid4().hex[:8]}",
git_url="https://github.com/example/budget-spend.git",
default_branch="main",
protected_branches=["main"],
assigned_cell=Team.BACKEND,
created_by=system_agent.id,
is_active=True,
)
session.add(project)
await session.flush()
return UUID(str(project.id)), UUID(str(system_agent.id))
async def _seed_task(session: AsyncSession, project_id: UUID, created_by: UUID) -> UUID:
task = TaskTable(
id=uuid4(),
title="Budget spend fixture task",
description="A description long enough to satisfy any length floor.",
acceptance_criteria=["it exists"],
status=TaskStatus.IN_PROGRESS,
priority=2,
task_type=TaskType.CODE,
nature=TaskNature.TECHNICAL,
estimated_complexity=Complexity.MEDIUM,
created_by=created_by,
team=Team.BACKEND,
project_id=project_id,
)
session.add(task)
await session.flush()
return UUID(str(task.id))
def _spawn_session(
task_id: UUID,
*,
started_at: datetime,
estimated_cost_usd: float | None,
ended_at: datetime | None,
tokens: tuple[int, int] = (0, 0),
) -> AgentSpawnSessionTable:
tokens_input, tokens_output = tokens
return AgentSpawnSessionTable(
id=uuid4(),
agent_slug="be-dev-1",
team="backend",
role="developer",
model=_MODEL,
task_id=str(task_id),
started_at=started_at,
ended_at=ended_at,
estimated_cost_usd=estimated_cost_usd,
tokens_input=tokens_input,
tokens_output=tokens_output,
)
@pytest.mark.asyncio
async def test_sums_closed_and_prices_open_session(db_session: AsyncSession) -> None:
"""A closed session's estimated_cost_usd + an open session's live-token
price (calculate_cost) — not the open session silently counted as $0."""
project_id, agent_id = await _seed_project(db_session)
task_id = await _seed_task(db_session, project_id, agent_id)
now = datetime.now(UTC)
db_session.add(
_spawn_session(
task_id,
started_at=now - timedelta(hours=2),
estimated_cost_usd=2.5,
ended_at=now - timedelta(hours=1),
)
)
db_session.add(
_spawn_session(
task_id,
started_at=now - timedelta(minutes=30),
estimated_cost_usd=None,
ended_at=None,
tokens=(100_000, 50_000),
)
)
await db_session.flush()
expected_open_cost = calculate_cost(
model=_MODEL, tokens_input=100_000, tokens_output=50_000
)
assert expected_open_cost > 0, (
"fixture model must be priced for this to be a real test"
)
svc = TaskService(db_session)
total = await svc.project_month_spend_usd(project_id)
assert total == pytest.approx(2.5 + expected_open_cost)
@pytest.mark.asyncio
async def test_excludes_last_months_session(db_session: AsyncSession) -> None:
"""A session that started before this calendar month's boundary must not
count, even though its cost is closed and non-zero."""
project_id, agent_id = await _seed_project(db_session)
task_id = await _seed_task(db_session, project_id, agent_id)
month_start = datetime.now(UTC).replace(
day=1, hour=0, minute=0, second=0, microsecond=0
)
last_month = month_start - timedelta(days=1)
db_session.add(
_spawn_session(
task_id,
started_at=last_month,
estimated_cost_usd=50.0,
ended_at=last_month + timedelta(hours=1),
)
)
await db_session.flush()
svc = TaskService(db_session)
total = await svc.project_month_spend_usd(project_id)
assert total == 0.0
@pytest.mark.asyncio
async def test_join_excludes_other_projects_tasks(db_session: AsyncSession) -> None:
"""A session on ANOTHER project's task must never bleed into this
project's sum — proves the join filters by project_id, not just presence
in agent_spawn_sessions."""
project_id, agent_id = await _seed_project(db_session)
other_project_id, other_agent_id = await _seed_project(db_session)
my_task_id = await _seed_task(db_session, project_id, agent_id)
other_task_id = await _seed_task(db_session, other_project_id, other_agent_id)
now = datetime.now(UTC)
db_session.add(
_spawn_session(
my_task_id,
started_at=now - timedelta(hours=1),
estimated_cost_usd=1.0,
ended_at=now,
)
)
db_session.add(
_spawn_session(
other_task_id,
started_at=now - timedelta(hours=1),
estimated_cost_usd=999.0,
ended_at=now,
)
)
await db_session.flush()
svc = TaskService(db_session)
total = await svc.project_month_spend_usd(project_id)
assert total == pytest.approx(1.0)