mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
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:
@@ -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)
|
||||
Reference in New Issue
Block a user