diff --git a/ai_engine/client.py b/ai_engine/client.py index 7f0c45a..ffc9d16 100644 --- a/ai_engine/client.py +++ b/ai_engine/client.py @@ -1037,6 +1037,16 @@ async def _record_invocation( cost_cents = int(round(estimate_call_cost(model, in_tok, out_tok) * 100)) except Exception: cost_cents = 0 + # P0-4 / m12-pr08: auto-attach tenant from the current context + # when callers don't pass one explicitly. This is what wires + # cost attribution + cap enforcement together at the same + # contextvar boundary set by ``cost_breaker.tenant_scope``. + if tenant_id is None: + try: + from ai_engine.cost_breaker import get_tenant_id + tenant_id = get_tenant_id() + except Exception: + tenant_id = None await get_recorder().record( model=model, prompt_text=prompt_text, @@ -1083,6 +1093,20 @@ def _check_budget(self) -> None: "Requests are paused until the next calendar day." ) + async def _enforce_cost_cap(self) -> None: + """P0-4 / m12-pr08: refuse the call when the org is over its daily $ cap. + + Raises :class:`ai_engine.cost_breaker.OrgDailyCostCapExceeded` when + the tenant bound to the current context has spent past its + configured cap. No-op when no tenant is bound or no cap is set. + Fails open on cost-service errors. + """ + try: + from ai_engine.cost_breaker import enforce_org_daily_cost_cap + except Exception: # pragma: no cover — defensive import guard + return + await enforce_org_daily_cost_cap() + def _resolve_model(self, task_type: Optional[str], model: Optional[str]) -> Optional[str]: """Resolve model name from task_type or explicit model param.""" if model: @@ -1121,6 +1145,7 @@ async def complete(self, prompt: str, system: Optional[str] = None, task_type: Optional[str] = None, model: Optional[str] = None) -> str: self._check_budget() + await self._enforce_cost_cap() prompt = _sanitize_prompt_input(prompt) prompt = self._truncate_input(prompt) @@ -1212,6 +1237,7 @@ async def complete_json(self, prompt: str, system: Optional[str] = None, task_type: Optional[str] = None, model: Optional[str] = None) -> Dict[str, Any]: self._check_budget() + await self._enforce_cost_cap() prompt = _sanitize_prompt_input(prompt) prompt = self._truncate_input(prompt) @@ -1455,6 +1481,7 @@ async def stream_completion( * Budget check + prompt sanitisation + truncation still applied. """ self._check_budget() + await self._enforce_cost_cap() prompt = _sanitize_prompt_input(prompt) prompt = self._truncate_input(prompt) candidate_model = self._resolve_model(task_type, model) or self.model @@ -1482,6 +1509,7 @@ async def chat(self, messages: List[Dict[str, str]], system: Optional[str] = Non task_type: Optional[str] = None, model: Optional[str] = None) -> str: self._check_budget() + await self._enforce_cost_cap() # Sanitize user-role messages only (system messages are trusted) messages = [ {**m, "content": _sanitize_prompt_input(m.get("content", ""))} if m.get("role") == "user" else m diff --git a/ai_engine/cost_breaker.py b/ai_engine/cost_breaker.py new file mode 100644 index 0000000..ac3babb --- /dev/null +++ b/ai_engine/cost_breaker.py @@ -0,0 +1,209 @@ +"""Per-org daily cost cap + cascade breaker. + +P0-4 / m12-pr08. Consumes ``backend.app.services.cost_attribution`` to +short-circuit LLM calls when an org has spent past its configured daily +USD ceiling. Tripped at the entry to the cascade — never inside it — +so a tripped cap costs zero LLM tokens. + +Tenant identity is propagated via a :class:`contextvars.ContextVar` +(``set_tenant_id`` / :func:`tenant_scope`). The cap is read from env: + +* ``ORG_DAILY_COST_CAP_USD`` — global default (e.g. ``"50"`` = $50/day). + Empty/unset → cap disabled (open mode). +* ``ORG_DAILY_COST_CAP_OVERRIDES`` — JSON ``{"": }`` for + per-tenant caps. ``0`` disables the cap for that tenant. + +Failure modes are conservative: if the cost service can't be reached or +the cap can't be parsed, the cap is **not** enforced (fail-open). This +matches blueprint guidance — a billing telemetry outage must not cause +a service outage. +""" + +from __future__ import annotations + +import contextvars +import json +import logging +import os +from contextlib import asynccontextmanager +from typing import AsyncIterator, Optional + +logger = logging.getLogger(__name__) + + +_current_tenant_id: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar( + "ai_current_tenant_id", default=None +) + + +# ─── Tenant context ───────────────────────────────────────────────────── + + +def set_tenant_id(tenant_id: Optional[str]) -> contextvars.Token: + """Bind ``tenant_id`` to the current async context. Returns reset token.""" + return _current_tenant_id.set(tenant_id) + + +def reset_tenant_id(token: contextvars.Token) -> None: + """Undo a previous :func:`set_tenant_id`.""" + _current_tenant_id.reset(token) + + +def get_tenant_id() -> Optional[str]: + """Return the tenant bound to the current context (or ``None``).""" + return _current_tenant_id.get() + + +@asynccontextmanager +async def tenant_scope(tenant_id: Optional[str]) -> AsyncIterator[None]: + """Async context manager that binds a tenant for the enclosed block. + + Use at request entry (HTTP handler, worker job pickup, SSE start) + so every nested LLM call inherits the tenant without parameter + drilling. + """ + token = set_tenant_id(tenant_id) + try: + yield + finally: + reset_tenant_id(token) + + +# ─── Cap configuration ────────────────────────────────────────────────── + + +def _parse_global_cap_cents() -> Optional[int]: + raw = (os.getenv("ORG_DAILY_COST_CAP_USD", "") or "").strip() + if not raw: + return None + try: + usd = float(raw) + except ValueError: + logger.warning("invalid_org_daily_cost_cap_usd: raw=%r", raw) + return None + if usd <= 0: + return None + return int(round(usd * 100)) + + +def _parse_overrides_cents() -> dict[str, int]: + raw = (os.getenv("ORG_DAILY_COST_CAP_OVERRIDES", "") or "").strip() + if not raw: + return {} + try: + decoded = json.loads(raw) + except json.JSONDecodeError: + logger.warning("invalid_org_daily_cost_cap_overrides: not json") + return {} + if not isinstance(decoded, dict): + return {} + out: dict[str, int] = {} + for tid, usd in decoded.items(): + try: + cents = int(round(float(usd) * 100)) + except (TypeError, ValueError): + continue + out[str(tid)] = max(0, cents) + return out + + +def cap_cents_for(tenant_id: str) -> Optional[int]: + """Return the cap (in cents) for ``tenant_id``, or ``None`` if disabled. + + Override semantics: + + * ``tenant_id`` explicitly mapped to ``0`` → cap disabled. + * ``tenant_id`` mapped to positive cents → that cap applies. + * Otherwise → fall back to the global env cap. + """ + overrides = _parse_overrides_cents() + if tenant_id in overrides: + cents = overrides[tenant_id] + return cents if cents > 0 else None + return _parse_global_cap_cents() + + +# ─── Exception ────────────────────────────────────────────────────────── + + +class OrgDailyCostCapExceeded(Exception): + """Raised when an org's daily LLM spend would exceed its configured cap. + + Carries the numbers so callers / metrics / logs can render a + meaningful message without re-querying. + """ + + def __init__(self, tenant_id: str, spent_cents: int, cap_cents: int) -> None: + self.tenant_id = tenant_id + self.spent_cents = int(spent_cents) + self.cap_cents = int(cap_cents) + super().__init__( + f"org_daily_cost_cap_exceeded: tenant={tenant_id} " + f"spent_cents={self.spent_cents} cap_cents={self.cap_cents}" + ) + + +# ─── Enforcement ──────────────────────────────────────────────────────── + + +async def _emit_cap_tripped(tenant_id: str, spent_cents: int, cap_cents: int) -> None: + """Best-effort emit of ``cost.cap.tripped`` to the bound event emitter.""" + try: + from ai_engine.agent_events import get_event_emitter + emitter = get_event_emitter() + if emitter is None: + return + await emitter("cost.cap.tripped", { + "tenant_id": tenant_id, + "spent_cents": spent_cents, + "cap_cents": cap_cents, + }) + except Exception as exc: + logger.debug("cost_cap_event_emit_failed: %s", exc) + + +async def enforce_org_daily_cost_cap( + tenant_id: Optional[str] = None, + *, + emit_event: bool = True, +) -> None: + """Raise :class:`OrgDailyCostCapExceeded` if the org is over its daily cap. + + No-op when: + + * No tenant is bound and none is passed. + * No cap is configured for the tenant. + * The cost service is unavailable (fail-open — telemetry outage + must not become an availability outage). + + Reads today's spend via + :class:`backend.app.services.cost_attribution.CostAttributionService`, + which queries the base ``ai_invocations`` table since UTC midnight + so the answer is fresh (not bounded by the 60s MV refresh). + """ + tid = tenant_id or get_tenant_id() + if not tid: + return + cap = cap_cents_for(tid) + if cap is None: + return + + try: + from backend.app.services.cost_attribution import ( + get_cost_attribution_service, + ) + spent = await get_cost_attribution_service().get_org_cost_today_cents(tid) + except Exception as exc: + # Fail-open: telemetry outage must not break the request path. + logger.warning("cost_cap_check_skipped_due_to_error: tenant=%s err=%s", + tid, str(exc)[:200]) + return + + if int(spent) >= int(cap): + if emit_event: + await _emit_cap_tripped(tid, int(spent), int(cap)) + logger.warning( + "org_daily_cost_cap_exceeded: tenant=%s spent_cents=%s cap_cents=%s", + tid, spent, cap, + ) + raise OrgDailyCostCapExceeded(tid, int(spent), int(cap)) diff --git a/backend/tests/unit/test_cost_breaker.py b/backend/tests/unit/test_cost_breaker.py new file mode 100644 index 0000000..49be6aa --- /dev/null +++ b/backend/tests/unit/test_cost_breaker.py @@ -0,0 +1,249 @@ +"""Tests for ai_engine.cost_breaker (P0-4 / m12-pr08).""" + +from __future__ import annotations + +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from ai_engine import cost_breaker as cb + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + monkeypatch.delenv("ORG_DAILY_COST_CAP_USD", raising=False) + monkeypatch.delenv("ORG_DAILY_COST_CAP_OVERRIDES", raising=False) + yield + + +# ─── ContextVar ───────────────────────────────────────────────────────── + + +def test_tenant_context_default_none(): + assert cb.get_tenant_id() is None + + +def test_tenant_context_set_reset(): + token = cb.set_tenant_id("org-1") + try: + assert cb.get_tenant_id() == "org-1" + finally: + cb.reset_tenant_id(token) + assert cb.get_tenant_id() is None + + +def test_tenant_scope_async(): + async def run(): + async with cb.tenant_scope("org-2"): + assert cb.get_tenant_id() == "org-2" + assert cb.get_tenant_id() is None + asyncio.run(run()) + + +# ─── Cap parsing ──────────────────────────────────────────────────────── + + +def test_cap_disabled_when_unset(): + assert cb.cap_cents_for("org-1") is None + + +def test_global_cap_parsed(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "12.50") + assert cb.cap_cents_for("org-1") == 1250 + + +def test_global_cap_invalid_string(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "abc") + assert cb.cap_cents_for("org-1") is None + + +def test_global_cap_zero_disables(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "0") + assert cb.cap_cents_for("org-1") is None + + +def test_per_tenant_override_wins(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "10") + monkeypatch.setenv( + "ORG_DAILY_COST_CAP_OVERRIDES", + json.dumps({"org-special": 99.0}), + ) + assert cb.cap_cents_for("org-special") == 9900 + assert cb.cap_cents_for("org-other") == 1000 + + +def test_per_tenant_override_zero_disables(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "10") + monkeypatch.setenv( + "ORG_DAILY_COST_CAP_OVERRIDES", + json.dumps({"org-vip": 0}), + ) + assert cb.cap_cents_for("org-vip") is None + assert cb.cap_cents_for("org-other") == 1000 + + +def test_overrides_invalid_json(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "5") + monkeypatch.setenv("ORG_DAILY_COST_CAP_OVERRIDES", "not-json") + assert cb.cap_cents_for("org-1") == 500 + + +# ─── Enforcement ──────────────────────────────────────────────────────── + + +def _patch_cost_service(spent_cents: int) -> MagicMock: + svc = MagicMock() + svc.get_org_cost_today_cents = AsyncMock(return_value=spent_cents) + return svc + + +@pytest.mark.asyncio +async def test_enforce_no_tenant_is_noop(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "1") + # No tenant in context, no tenant param → no DB call, no raise. + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service" + ) as get_svc: + await cb.enforce_org_daily_cost_cap() + assert not get_svc.called + + +@pytest.mark.asyncio +async def test_enforce_no_cap_is_noop(monkeypatch): + # Tenant bound but no cap configured → no DB call, no raise. + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service" + ) as get_svc: + await cb.enforce_org_daily_cost_cap(tenant_id="org-1") + assert not get_svc.called + + +@pytest.mark.asyncio +async def test_enforce_under_cap_no_raise(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "10") # 1000 cents + svc = _patch_cost_service(spent_cents=500) + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service", + return_value=svc, + ): + await cb.enforce_org_daily_cost_cap(tenant_id="org-1") + svc.get_org_cost_today_cents.assert_awaited_once_with("org-1") + + +@pytest.mark.asyncio +async def test_enforce_over_cap_raises(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "10") # 1000 cents + svc = _patch_cost_service(spent_cents=1500) + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service", + return_value=svc, + ): + with pytest.raises(cb.OrgDailyCostCapExceeded) as excinfo: + await cb.enforce_org_daily_cost_cap(tenant_id="org-1") + assert excinfo.value.tenant_id == "org-1" + assert excinfo.value.spent_cents == 1500 + assert excinfo.value.cap_cents == 1000 + + +@pytest.mark.asyncio +async def test_enforce_at_exact_cap_raises(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "10") + svc = _patch_cost_service(spent_cents=1000) + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service", + return_value=svc, + ): + with pytest.raises(cb.OrgDailyCostCapExceeded): + await cb.enforce_org_daily_cost_cap(tenant_id="org-1") + + +@pytest.mark.asyncio +async def test_enforce_uses_context_tenant(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "10") + svc = _patch_cost_service(spent_cents=2000) + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service", + return_value=svc, + ): + async with cb.tenant_scope("org-from-ctx"): + with pytest.raises(cb.OrgDailyCostCapExceeded) as excinfo: + await cb.enforce_org_daily_cost_cap() + assert excinfo.value.tenant_id == "org-from-ctx" + svc.get_org_cost_today_cents.assert_awaited_once_with("org-from-ctx") + + +@pytest.mark.asyncio +async def test_enforce_fail_open_on_db_error(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "10") + svc = MagicMock() + svc.get_org_cost_today_cents = AsyncMock( + side_effect=RuntimeError("supabase down"), + ) + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service", + return_value=svc, + ): + # Must not raise — fail-open behaviour. + await cb.enforce_org_daily_cost_cap(tenant_id="org-1") + + +@pytest.mark.asyncio +async def test_enforce_emits_cap_tripped_event(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "5") + svc = _patch_cost_service(spent_cents=900) + + emitted: list[tuple[str, dict]] = [] + + async def emitter(name, payload): + emitted.append((name, payload)) + + from ai_engine.agent_events import event_emitter_scope + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service", + return_value=svc, + ): + with event_emitter_scope(emitter): + with pytest.raises(cb.OrgDailyCostCapExceeded): + await cb.enforce_org_daily_cost_cap(tenant_id="org-1") + + assert len(emitted) == 1 + name, payload = emitted[0] + assert name == "cost.cap.tripped" + assert payload == { + "tenant_id": "org-1", + "spent_cents": 900, + "cap_cents": 500, + } + + +@pytest.mark.asyncio +async def test_enforce_emit_event_disabled(monkeypatch): + monkeypatch.setenv("ORG_DAILY_COST_CAP_USD", "5") + svc = _patch_cost_service(spent_cents=900) + + emitted: list[tuple[str, dict]] = [] + + async def emitter(name, payload): + emitted.append((name, payload)) + + from ai_engine.agent_events import event_emitter_scope + with patch( + "backend.app.services.cost_attribution.get_cost_attribution_service", + return_value=svc, + ): + with event_emitter_scope(emitter): + with pytest.raises(cb.OrgDailyCostCapExceeded): + await cb.enforce_org_daily_cost_cap( + tenant_id="org-1", emit_event=False, + ) + assert emitted == [] + + +def test_exception_message_format(): + exc = cb.OrgDailyCostCapExceeded("org-x", 1234, 1000) + msg = str(exc) + assert "org-x" in msg + assert "1234" in msg + assert "1000" in msg diff --git a/docs/architecture/WORLD_CLASS_ARCHITECTURE_BLUEPRINT.md b/docs/architecture/WORLD_CLASS_ARCHITECTURE_BLUEPRINT.md index 19b675c..a7ebf33 100644 --- a/docs/architecture/WORLD_CLASS_ARCHITECTURE_BLUEPRINT.md +++ b/docs/architecture/WORLD_CLASS_ARCHITECTURE_BLUEPRINT.md @@ -981,7 +981,7 @@ Every fix has a tracking ID. Status updated in PRs. Closing the last entry in th | P0-1 | Partition rotation: `pg_partman` + `PartitionMaintenanceWorkflow` | W1 | platform | **SHIPPED** — m7-pr27a (`supabase/migrations/20260508120000_outbox_partition_rotation.sql`) | | P0-2 | Eliminate in-process generation fallback; fail-fast 503 | W2 | platform | **SHIPPED** — m7-pr27b | | P0-3 | Queue ACK on success only; DLQ + idempotent handlers | W3 | platform | **SHIPPED** — m7-pr27c | -| P0-4 | Per-org daily $ cap + cascade-failure breaker | W8 | platform | TODO | +| P0-4 | Per-org daily $ cap + cascade-failure breaker | W8 | platform | SHIPPED (m12-pr08) — `ai_engine.cost_breaker` reads `ORG_DAILY_COST_CAP_USD` (+ per-tenant overrides), pulls today's spend via `CostAttributionService.get_org_cost_today_cents`, raises `OrgDailyCostCapExceeded` at the entry of every cascade method, emits `cost.cap.tripped` event, fail-open on telemetry errors | | P0-5 | Tool registry: `jsonschema` + capability tokens + sandbox tiers | W5 | ai-team + security | **SHIPPED** — m7-pr29 | | P0-6 | Idempotency middleware ON in production | W9 | platform | **SHIPPED** — `backend/main.py` middleware stack | | P0-7 | SSE `Last-Event-ID` resumption end-to-end | W7 | frontend + platform | **SHIPPED** (m12-pr05) — backend exposes `GET /pipeline/agentic-stream/{session_id}/replay?after_sequence=N` backed by `AgenticEventEmitter.get_events_after()`; `agentic_stream.py` returns `X-Session-ID` header so clients can resume; in-memory store is single-process — durable cross-pod replay still requires Redis/DB-backed session registry (tracked separately) |