roboco/tests/unit/gateway/test_budget_unblock_guard.py

"""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)
    # _ensure_pm_decision's journal write is savepoint-guarded — an
    # unconfigured AsyncMock's begin_nested() call returns a raw unawaited
    # coroutine, which `async with` cannot use.
    base["task"].session.begin_nested = MagicMock(
        return_value=MagicMock(
            __aenter__=AsyncMock(return_value=None),
            __aexit__=AsyncMock(return_value=False),
        )
    )
    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()