diff --git a/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000004_add_sandbox_and_storage_meters.py b/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000004_add_sandbox_and_storage_meters.py index 35c6d6340b..732b0df7cc 100644 --- a/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000004_add_sandbox_and_storage_meters.py +++ b/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000004_add_sandbox_and_storage_meters.py @@ -1,4 +1,4 @@ -"""add sandbox compute + storage_bytes meters to meters_type +"""add sandbox compute + bytes meters to meters_type Revision ID: ee0000000004 Revises: ee0000000003 @@ -29,7 +29,7 @@ def upgrade() -> None: op.execute( "ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_GPU_CORE_SECONDS'" ) - op.execute("ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'STORAGE_BYTES'") + op.execute("ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'BYTES'") def downgrade() -> None: diff --git a/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000005_add_sandbox_credit_meters.py b/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000005_add_sandbox_credit_meters.py new file mode 100644 index 0000000000..9de04039e6 --- /dev/null +++ b/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000005_add_sandbox_credit_meters.py @@ -0,0 +1,40 @@ +"""add sandbox + wallet debit meters to meters_type + +Revision ID: ee0000000005 +Revises: ee0000000004 +Create Date: 2026-07-02 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "ee0000000005" +down_revision: Union[str, None] = "ee0000000004" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute( + "ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_CPU_CORE_DEBITS'" + ) + op.execute( + "ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_RAM_GIBI_DEBITS'" + ) + op.execute( + "ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_SSD_GIBI_DEBITS'" + ) + op.execute( + "ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_GPU_CORE_DEBITS'" + ) + op.execute("ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_DEBITS'") + op.execute("ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'LLM_DEBITS'") + op.execute("ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'GATEWAY_DEBITS'") + op.execute("ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'WALLET_DEBITS'") + + +def downgrade() -> None: + # Postgres cannot drop an enum label; leave the values in place. + pass diff --git a/api/ee/src/core/access/entitlements/types.py b/api/ee/src/core/access/entitlements/types.py index 30e00f4ff8..74b37a76e6 100644 --- a/api/ee/src/core/access/entitlements/types.py +++ b/api/ee/src/core/access/entitlements/types.py @@ -55,15 +55,22 @@ class Counter(str, Enum): CREDITS_CONSUMED = "credits_consumed" EVENTS_INGESTED = "events_ingested" RECORDS_INGESTED = "records_ingested" - SANDBOX_CPU_CORE_SECONDS = "sandbox_cpu_core_seconds" - SANDBOX_RAM_GIBI_SECONDS = "sandbox_ram_gibi_seconds" - SANDBOX_SSD_GIBI_SECONDS = "sandbox_ssd_gibi_seconds" - SANDBOX_GPU_CORE_SECONDS = "sandbox_gpu_core_seconds" + # Billing layer (Track C): per-dimension + family + wallet debit roll-ups. + # Raw *_seconds dimension meters are cost-explainer data (traces/analytics), + # not billing meters, and are intentionally not modeled here. + SANDBOX_CPU_CORE_DEBITS = "sandbox_cpu_core_debits" + SANDBOX_RAM_GIBI_DEBITS = "sandbox_ram_gibi_debits" + SANDBOX_SSD_GIBI_DEBITS = "sandbox_ssd_gibi_debits" + SANDBOX_GPU_CORE_DEBITS = "sandbox_gpu_core_debits" + SANDBOX_DEBITS = "sandbox_debits" + LLM_DEBITS = "llm_debits" + GATEWAY_DEBITS = "gateway_debits" + WALLET_DEBITS = "wallet_debits" class Gauge(str, Enum): USERS = "users" - STORAGE_BYTES = "storage_bytes" + BYTES = "bytes" class Constraint(str, Enum): @@ -242,6 +249,9 @@ class Throttle(BaseModel): }, ], }, + # TODO(pricing): sandbox_debits tiers — measurement-only; wallet + # debits are prepaid (billed at top-up time), never reported to + # Stripe in arrears, so no REPORTS wiring applies here. }, "features": [ "Unlimited prompts", @@ -278,6 +288,9 @@ class Throttle(BaseModel): }, ], }, + # TODO(pricing): sandbox_debits tiers — measurement-only; wallet + # debits are prepaid (billed at top-up time), never reported to + # Stripe in arrears, so no REPORTS wiring applies here. }, "features": [ "Everything in Pro", @@ -362,16 +375,29 @@ class Throttle(BaseModel): retention=Retention.WEEKLY, period=Period.MONTHLY, ), - Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + # TODO(pricing): real free/limit once sandbox billing goes live. + Counter.SANDBOX_CPU_CORE_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + Counter.SANDBOX_RAM_GIBI_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + Counter.SANDBOX_SSD_GIBI_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + Counter.SANDBOX_GPU_CORE_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.LLM_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.GATEWAY_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.WALLET_DEBITS: Quota( period=Period.MONTHLY, ), }, @@ -381,7 +407,7 @@ class Throttle(BaseModel): limit=2, strict=True, ), - Gauge.STORAGE_BYTES: Quota( + Gauge.BYTES: Quota( free=1_073_741_824, limit=1_073_741_824, strict=True, @@ -471,16 +497,29 @@ class Throttle(BaseModel): retention=Retention.MONTHLY, period=Period.MONTHLY, ), - Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + # TODO(pricing): real free/limit once sandbox billing goes live. + Counter.SANDBOX_CPU_CORE_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + Counter.SANDBOX_SSD_GIBI_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + Counter.SANDBOX_GPU_CORE_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + Counter.SANDBOX_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.LLM_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.GATEWAY_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.WALLET_DEBITS: Quota( period=Period.MONTHLY, ), }, @@ -488,7 +527,7 @@ class Throttle(BaseModel): Gauge.USERS: Quota( strict=True, ), - Gauge.STORAGE_BYTES: Quota( + Gauge.BYTES: Quota( free=5_368_709_120, limit=10_737_418_240, strict=True, @@ -578,16 +617,29 @@ class Throttle(BaseModel): retention=Retention.QUARTERLY, period=Period.MONTHLY, ), - Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + # TODO(pricing): real free/limit once sandbox billing goes live. + Counter.SANDBOX_CPU_CORE_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + Counter.SANDBOX_GPU_CORE_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + Counter.SANDBOX_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + Counter.LLM_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.GATEWAY_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.WALLET_DEBITS: Quota( period=Period.MONTHLY, ), }, @@ -595,7 +647,7 @@ class Throttle(BaseModel): Gauge.USERS: Quota( strict=True, ), - Gauge.STORAGE_BYTES: Quota( + Gauge.BYTES: Quota( free=53_687_091_200, strict=True, ), @@ -680,16 +732,29 @@ class Throttle(BaseModel): Counter.RECORDS_INGESTED: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + # TODO(pricing): real free/limit once sandbox billing goes live. + Counter.SANDBOX_CPU_CORE_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_GPU_CORE_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + Counter.SANDBOX_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + Counter.LLM_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + Counter.GATEWAY_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.WALLET_DEBITS: Quota( period=Period.MONTHLY, ), }, @@ -730,16 +795,29 @@ class Throttle(BaseModel): Counter.RECORDS_INGESTED: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + # TODO(pricing): real free/limit once sandbox billing goes live. + Counter.SANDBOX_CPU_CORE_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_GPU_CORE_DEBITS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + Counter.LLM_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + Counter.GATEWAY_DEBITS: Quota( period=Period.MONTHLY, ), - Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + Counter.WALLET_DEBITS: Quota( period=Period.MONTHLY, ), }, @@ -758,6 +836,9 @@ class Throttle(BaseModel): # name to report under (`REPORTS[key]`). REPORTS: dict[str, str] = { Counter.TRACES_INGESTED.value: "traces", + # Wallet/sandbox debit meters are intentionally absent: the wallet total + # is prepaid (money moves at top-up time), so it must never be reported + # to Stripe in arrears. } CONSTRAINTS = { @@ -771,7 +852,7 @@ class Throttle(BaseModel): ], Tracker.GAUGES: [ Gauge.USERS, - Gauge.STORAGE_BYTES, + Gauge.BYTES, ], }, Constraint.READ_ONLY: { @@ -782,10 +863,14 @@ class Throttle(BaseModel): Counter.CREDITS_CONSUMED, Counter.EVENTS_INGESTED, Counter.RECORDS_INGESTED, - Counter.SANDBOX_CPU_CORE_SECONDS, - Counter.SANDBOX_RAM_GIBI_SECONDS, - Counter.SANDBOX_SSD_GIBI_SECONDS, - Counter.SANDBOX_GPU_CORE_SECONDS, + Counter.SANDBOX_CPU_CORE_DEBITS, + Counter.SANDBOX_RAM_GIBI_DEBITS, + Counter.SANDBOX_SSD_GIBI_DEBITS, + Counter.SANDBOX_GPU_CORE_DEBITS, + Counter.SANDBOX_DEBITS, + Counter.LLM_DEBITS, + Counter.GATEWAY_DEBITS, + Counter.WALLET_DEBITS, ], }, } diff --git a/api/ee/src/core/meters/types.py b/api/ee/src/core/meters/types.py index 85108b6311..763bbc8b59 100644 --- a/api/ee/src/core/meters/types.py +++ b/api/ee/src/core/meters/types.py @@ -28,13 +28,17 @@ class Meters(str, Enum): CREDITS_CONSUMED = Counter.CREDITS_CONSUMED.value EVENTS_INGESTED = Counter.EVENTS_INGESTED.value RECORDS_INGESTED = Counter.RECORDS_INGESTED.value - SANDBOX_CPU_CORE_SECONDS = Counter.SANDBOX_CPU_CORE_SECONDS.value - SANDBOX_RAM_GIBI_SECONDS = Counter.SANDBOX_RAM_GIBI_SECONDS.value - SANDBOX_SSD_GIBI_SECONDS = Counter.SANDBOX_SSD_GIBI_SECONDS.value - SANDBOX_GPU_CORE_SECONDS = Counter.SANDBOX_GPU_CORE_SECONDS.value + SANDBOX_CPU_CORE_DEBITS = Counter.SANDBOX_CPU_CORE_DEBITS.value + SANDBOX_RAM_GIBI_DEBITS = Counter.SANDBOX_RAM_GIBI_DEBITS.value + SANDBOX_SSD_GIBI_DEBITS = Counter.SANDBOX_SSD_GIBI_DEBITS.value + SANDBOX_GPU_CORE_DEBITS = Counter.SANDBOX_GPU_CORE_DEBITS.value + SANDBOX_DEBITS = Counter.SANDBOX_DEBITS.value + LLM_DEBITS = Counter.LLM_DEBITS.value + GATEWAY_DEBITS = Counter.GATEWAY_DEBITS.value + WALLET_DEBITS = Counter.WALLET_DEBITS.value # GAUGES USERS = Gauge.USERS.value - STORAGE_BYTES = Gauge.STORAGE_BYTES.value + BYTES = Gauge.BYTES.value class MeterScope(BaseModel): diff --git a/api/ee/src/core/sandboxes/credits.py b/api/ee/src/core/sandboxes/credits.py new file mode 100644 index 0000000000..ec8f2e42e3 --- /dev/null +++ b/api/ee/src/core/sandboxes/credits.py @@ -0,0 +1,149 @@ +"""Sandbox credit conversion: raw resource-seconds -> credits. + +Credits are dimensionless billing units emitted to Stripe. Stripe multiplies by +the per-credit price for the plan at billing time -- no credit<->money math +lives here. Stored as millicredits (credits x 1000, truncated) so the +int-typed MeterDTO.delta field preserves sub-credit precision; Stripe's price +denominator accounts for the x1000 factor. + +to_credits() is the single conversion function; pure, no I/O, Decimal +throughout. Dimension names mirror the locked sandbox meter key scheme +(CPU/RAM/SSD/GPU) -- see docs/designs/sandbox-metering/NAMING.md. +""" + +from __future__ import annotations + +from decimal import Decimal +from enum import Enum + +from pydantic import BaseModel, ConfigDict + +from oss.src.utils.env import env + + +class Dimension(str, Enum): + """Resource dimensions billed per second.""" + + CPU = "cpu" # vCPU-s + RAM = "ram" # GiB-s of RAM + SSD = "ssd" # GiB-s of disk + GPU = "gpu" # GPU-s + + +class ProviderRates(BaseModel): + """Credits-per-unit rates for one provider across all dimensions. + + Each rate is credits per raw unit-second for that dimension. + Stripe owns the credit->money price; these are dimensionless conversion + rates. + """ + + # Reference: a standard 2vCPU/2GiB/10GiB machine for 1 minute is a useful + # cross-check (informational only -- not a billing input). + cpu: Decimal # credits per vCPU-s (e.g. 0.0014) + ram: Decimal # credits per GiB-s of RAM (e.g. 0.00045) + ssd: Decimal # credits per GiB-s of disk (e.g. 0.000003) + gpu: Decimal # credits per GPU-s; 0 = no GPU billing for this provider + + model_config = ConfigDict(extra="forbid") + + +# Default per-provider x per-dimension rate table. +# Informational cross-check: 2 vCPU x 60s x 0.0014 + 2 GiB x 60s x 0.00045 +# + 10 GiB x 60s x 0.000003 ~= 0.2232 credits/min per reference machine. +DEFAULT_PROVIDER_RATES: dict[str, ProviderRates] = { + "e2b": ProviderRates( + cpu=Decimal("0.0014"), # vCPU-s -> credits + ram=Decimal("0.00045"), # GiB-s RAM -> credits + ssd=Decimal("0.000003"), # GiB-s disk -> credits + gpu=Decimal("0"), # E2B has no GPU meter + ), + "daytona": ProviderRates( + cpu=Decimal("0.0014"), # vCPU-s -> credits + ram=Decimal("0.00045"), # GiB-s RAM -> credits + ssd=Decimal("0.000003"), # GiB-s disk -> credits + gpu=Decimal("0"), # GPU rate: no default; set via env override + ), + "local": ProviderRates( + cpu=Decimal("0"), # zero-rated: local sandbox has no billing cost + ram=Decimal("0"), + ssd=Decimal("0"), + gpu=Decimal("0"), + ), +} + +_RATES: dict[str, ProviderRates] | None = None + + +def _build_rates() -> dict[str, ProviderRates]: + raw = env.sandbox.credit_rates + if not raw: + return dict(DEFAULT_PROVIDER_RATES) + + result: dict[str, ProviderRates] = dict(DEFAULT_PROVIDER_RATES) + for provider, dims in raw.items(): + if not isinstance(dims, dict): + continue + base = DEFAULT_PROVIDER_RATES.get(str(provider).lower()) + base_data: dict[str, Decimal] = ( + base.model_dump() + if base + else { + "cpu": Decimal("0"), + "ram": Decimal("0"), + "ssd": Decimal("0"), + "gpu": Decimal("0"), + } + ) + for dim_key, rate_str in dims.items(): + if dim_key in base_data: + try: + base_data[dim_key] = Decimal(str(rate_str)) + except Exception: # pylint: disable=broad-exception-caught + pass + try: + result[str(provider).lower()] = ProviderRates(**base_data) + except Exception: # pylint: disable=broad-exception-caught + pass + return result + + +def _get_rates() -> dict[str, ProviderRates]: + global _RATES + if _RATES is None: + _RATES = _build_rates() + return _RATES + + +def to_credits( + *, + provider: str, + dimension: Dimension | str, + raw_units: Decimal, +) -> Decimal: + """Convert raw resource-seconds to credits. + + Args: + provider: Provider slug ("e2b", "daytona", "local", ...). + dimension: Dimension enum or string ("cpu", "ram", "ssd", "gpu"). + raw_units: Raw resource-seconds as Decimal. Values <= 0 return 0. + + Returns: + Credits as Decimal (>= 0). Returns 0 when provider/dimension has no + rate. + """ + if raw_units <= Decimal("0"): + return Decimal("0") + + rates = _get_rates() + provider_rates = rates.get(str(provider).lower()) + if provider_rates is None: + return Decimal("0") + + dim_str = ( + dimension.value if isinstance(dimension, Dimension) else str(dimension).lower() + ) + rate = getattr(provider_rates, dim_str, None) + if rate is None: + return Decimal("0") + return raw_units * rate diff --git a/api/ee/src/core/sandboxes/gating.py b/api/ee/src/core/sandboxes/gating.py new file mode 100644 index 0000000000..159f975c67 --- /dev/null +++ b/api/ee/src/core/sandboxes/gating.py @@ -0,0 +1,192 @@ +"""Sandbox entitlement gating: create-time quota check (Layer 1) and +post-hoc true-up (Layer 2), both against WALLET_DEBITS -- the cross-family +grand total (LLM_DEBITS + SANDBOX_DEBITS + GATEWAY_DEBITS), since the wallet +gate must block on total prepaid balance, not just the sandbox sub-total. + +Two-layer design mirrors the existing tracing entitlements pattern: + +Layer 1 (create-time soft pre-check) + `check_sandbox_quota()` -- called before launching a sandbox. Uses the + Redis-cached WALLET_DEBITS value (`cache=True`) plus an in-flight + accrual estimate (converted via `credits.to_credits`) so the org does not + exceed its quota the moment a new sandbox starts. Fails open: any + infrastructure error allows the launch and logs a warning. + +Layer 2 (post-hoc true-up) + `check_sandbox_credits_true_up()` -- called after + `SandboxMeteringService.record_usage()` has adjusted the raw *_seconds + and *_debits meters (see sink.py). Re-checks WALLET_DEBITS + (`cache=False`, delta=0 -- a read-only true-up, the write already + happened in the sink) and returns whether the org is within quota. + +Layer 2b (mid-session kill -- stub, deferred) + When `check_sandbox_credits_true_up()` returns `False` the caller should + kill the active sandbox session. The `DELETE /sessions/streams/{id}` + + runner `/kill` endpoint is not yet implemented; until it is, a WARNING + log is the only action taken. See docs/designs/sandbox-metering/tasks.md. + +RBAC (RUN_SESSIONS permission) is strictly separate from entitlement checks. +This module only concerns itself with quota; callers must enforce RBAC +before invoking either function here. +""" + +from decimal import Decimal +from typing import Optional +from uuid import UUID + +from oss.src.utils.env import env +from oss.src.utils.common import is_ee +from oss.src.utils.logging import get_module_logger + +from ee.src.core.sandboxes.credits import Dimension, to_credits + +log = get_module_logger(__name__) + + +async def check_sandbox_quota( + *, + organization_id: UUID, + provider: str = "e2b", +) -> tuple[bool, Optional[str]]: + """Layer 1: create-time soft pre-check. + + Fetches the cached meter value for `WALLET_DEBITS` and adds an + estimated in-flight accrual (CPU-dimension credits for + `env.sandbox.estimated_vcpu` vCPUs over `env.sandbox.estimated_run_seconds` + seconds) before comparing against the plan quota. + + Args: + organization_id: The org that will own the new sandbox. + provider: Provider slug the new sandbox will launch on; selects the + rate table for the accrual estimate. + + Returns: + ``(allowed, reason)`` where ``allowed=False`` means the org is + at or over quota. ``reason`` is a human-readable string for + HTTP 429 responses; ``None`` when allowed. + + Fail-open: any error except `EntitlementsException` returns + ``(True, None)`` and logs a warning so infra issues never block + sandbox creation. + """ + if not is_ee(): + return True, None + + try: + return await _check_sandbox_quota_ee( + organization_id=organization_id, + provider=provider, + ) + except Exception: # pylint: disable=broad-exception-caught + log.warning( + "[sandboxes] check_sandbox_quota failed for org=%s; failing open", + organization_id, + exc_info=True, + ) + return True, None + + +async def _check_sandbox_quota_ee( + *, + organization_id: UUID, + provider: str, +) -> tuple[bool, Optional[str]]: + from ee.src.core.access.entitlements.types import Counter + from ee.src.core.access.entitlements.service import ( + check_entitlements, + EntitlementsException, + ) + from ee.src.core.meters.types import MeterScope + + sandbox_cfg = env.sandbox + estimated_accrual_seconds = Decimal( + sandbox_cfg.estimated_vcpu * sandbox_cfg.estimated_run_seconds + ) + estimated_credits = to_credits( + provider=provider, + dimension=Dimension.CPU, + raw_units=estimated_accrual_seconds, + ) + estimated_millicredits = int(estimated_credits * Decimal("1000")) + + meter_scope = MeterScope(organization_id=organization_id) + + try: + allowed, _, _ = await check_entitlements( + key=Counter.WALLET_DEBITS, + delta=estimated_millicredits, + cache=True, + scope=meter_scope, + ) + except EntitlementsException: + # Config / programming bug -- propagate so it surfaces clearly. + raise + + if not allowed: + return False, ( + "You have reached your sandbox usage quota for this billing period. " + "Please upgrade your plan or wait for the quota to reset." + ) + + return True, None + + +async def check_sandbox_credits_true_up( + *, + organization_id: UUID, +) -> bool: + """Layer 2: post-hoc true-up against WALLET_DEBITS. + + Call after `SandboxMeteringService.record_usage()` (which already wrote + the meters via sink.py) to check whether the org is now over quota. + `delta=0` -- this is a read/check, not a write; the authoritative + adjust() already happened in the sink. + + Returns ``True`` when within quota, ``False`` when over. A ``False`` + return should trigger a sandbox kill (Layer 2b), but the kill endpoint + is not yet implemented -- a WARNING is logged and the caller should + treat this as a known gap. + + Fails open on infrastructure errors (non-`EntitlementsException` + exceptions). + """ + if not is_ee(): + return True + + try: + return await _check_sandbox_credits_true_up_ee(organization_id=organization_id) + except Exception: # pylint: disable=broad-exception-caught + log.warning( + "[sandboxes] check_sandbox_credits_true_up failed for org=%s; failing open", + organization_id, + exc_info=True, + ) + return True + + +async def _check_sandbox_credits_true_up_ee( + *, + organization_id: UUID, +) -> bool: + from ee.src.core.access.entitlements.types import Counter + from ee.src.core.access.entitlements.service import check_entitlements + from ee.src.core.meters.types import MeterScope + + meter_scope = MeterScope(organization_id=organization_id) + + allowed, _, _ = await check_entitlements( + key=Counter.WALLET_DEBITS, + delta=0, + cache=False, + scope=meter_scope, + ) + + if not allowed: + log.warning( + "[sandboxes] over-quota org=%s counter=%s " + "-- kill endpoint not yet implemented; skipping sandbox termination", + organization_id, + Counter.WALLET_DEBITS.value, + ) + + return allowed diff --git a/api/ee/src/core/sandboxes/service.py b/api/ee/src/core/sandboxes/service.py index 072346046d..8f2aab7dc9 100644 --- a/api/ee/src/core/sandboxes/service.py +++ b/api/ee/src/core/sandboxes/service.py @@ -15,12 +15,10 @@ ) from oss.src.utils.logging import get_module_logger -from ee.src.core.access.entitlements.service import check_entitlements -from ee.src.core.access.entitlements.types import Counter from ee.src.core.meters.service import MetersService -from ee.src.core.meters.types import MeterScope from ee.src.core.sandboxes.dtos import SandboxUsageDTO, SandboxUsageResult from ee.src.core.sandboxes.exceptions import SandboxWebhookSignatureError +from ee.src.core.sandboxes.sink import record_usage_credits log = get_module_logger(__name__) @@ -53,11 +51,7 @@ def __init__(self, *, meters_service: MetersService): # ------------------------------------------------------------------ async def record_usage(self, usage: SandboxUsageDTO) -> SandboxUsageResult: - """Persist sandbox resource-second usage into the meters layer. - - Calls check_entitlements(cache=False) per meter so the Layer-2 - atomic adjust() runs, giving an authoritative quota check. - The call is NON-BLOCKING in Phase 1 (quotas are soft). + """Persist sandbox resource-second usage into the wallet debit meters. Deduped on usage.delivery_id via Redis SET NX (webhook redelivery double-counting guard). Missing delivery_id skips dedup (best-effort). @@ -78,34 +72,18 @@ async def record_usage(self, usage: SandboxUsageDTO) -> SandboxUsageResult: ) org_id = usage.organization_id - scope = MeterScope(organization_id=org_id) - - meter_deltas: list[tuple[Counter, int]] = [ - (Counter.SANDBOX_CPU_CORE_SECONDS, usage.vcpu_seconds), - (Counter.SANDBOX_RAM_GIBI_SECONDS, usage.ram_gib_seconds), - (Counter.SANDBOX_SSD_GIBI_SECONDS, usage.disk_gib_seconds), - (Counter.SANDBOX_GPU_CORE_SECONDS, usage.gpu_seconds), - ] - - for counter, delta in meter_deltas: - if delta <= 0: - continue - try: - # cache=False → Layer-2 hard check (atomic DB adjust). - # Fails open on error per check_entitlements contract. - await check_entitlements( - key=counter, - delta=delta, - cache=False, - scope=scope, - ) - except Exception: - log.warning( - "[sandboxes] check_entitlements failed for %s/%s", - org_id, - counter, - exc_info=True, - ) + + # Billing layer: per-dimension + sandbox + wallet debits, derived from + # the raw seconds (see sink.py). Raw *_seconds are cost-explainer data + # (traces/analytics), not billing meters, so they are not adjusted here. + await record_usage_credits( + provider=usage.provider, + organization_id=org_id, + cpu_seconds=Decimal(usage.vcpu_seconds), + ram_seconds=Decimal(usage.ram_gib_seconds), + ssd_seconds=Decimal(usage.disk_gib_seconds), + gpu_seconds=Decimal(usage.gpu_seconds) if usage.gpu_seconds else None, + ) log.info( "[sandboxes] recorded provider=%s sandbox=%s org=%s " diff --git a/api/ee/src/core/sandboxes/sink.py b/api/ee/src/core/sandboxes/sink.py new file mode 100644 index 0000000000..3be207b776 --- /dev/null +++ b/api/ee/src/core/sandboxes/sink.py @@ -0,0 +1,137 @@ +"""record_usage_credits() -- wallet debit sink for sandbox resource-second events. + +Called by `SandboxMeteringService.record_usage()` with the raw resource +seconds for one usage event (the raw *_seconds dimensions themselves are +cost-explainer data, not billing meters, so they are not adjusted here). +Per event it: + + 1. Computes per-dimension *_debits from the raw seconds via `to_credits()`. + 2. Adjusts each per-dimension *_debits meter. + 3. Sums per-dimension debits into sandbox_debits and adjusts that meter. + 4. Adds the same sandbox_debits total into wallet_debits, the cross-family + grand total (LLM_DEBITS + SANDBOX_DEBITS + GATEWAY_DEBITS). + +Credit deltas are stored as millicredits (credits x 1000, truncated) to +preserve sub-credit precision without changing the int-typed MeterDTO.delta +field. Stripe's per-millicredit price accounts for the x1000 factor. + +All adjustments are org-scoped. Gate: is_ee() (caller gates on provider.enabled). +""" + +from __future__ import annotations + +from decimal import Decimal +from typing import Optional +from uuid import UUID + +from oss.src.utils.logging import get_module_logger +from oss.src.utils.common import is_ee + +from ee.src.core.access.entitlements.types import Counter +from ee.src.core.access.entitlements.service import check_entitlements +from ee.src.core.meters.types import MeterScope +from ee.src.core.sandboxes.credits import Dimension, to_credits + + +log = get_module_logger(__name__) + +# Millicredits scale factor: credits x 1000 stored as int in MeterDTO.delta. +_MILLICREDITS = Decimal("1000") + + +async def record_usage_credits( + *, + provider: str, + organization_id: UUID, + # + cpu_seconds: Decimal = Decimal("0"), + ram_seconds: Decimal = Decimal("0"), + ssd_seconds: Decimal = Decimal("0"), + gpu_seconds: Optional[Decimal] = None, +) -> None: + """Record per-dimension + sandbox + wallet debits for one usage event. + + Writes SANDBOX_{CPU,RAM,SSD,GPU}_DEBITS, SANDBOX_DEBITS, and accumulates + the sandbox total into WALLET_DEBITS in one call. No-ops silently when + is_ee() is False. + + Args: + provider: Provider slug ("e2b", "daytona", "local", ...). + organization_id: Org to meter under. + cpu_seconds: vCPU-s consumed (Decimal >= 0). + ram_seconds: GiB-s of RAM consumed (Decimal >= 0). + ssd_seconds: GiB-s of disk consumed (Decimal >= 0). + gpu_seconds: GPU-s consumed; None means provider has no GPU meter. + """ + if not is_ee(): + return + + meter_scope = MeterScope(organization_id=organization_id) + + dimension_pairs: list[tuple[Counter, Dimension, Decimal]] = [ + (Counter.SANDBOX_CPU_CORE_DEBITS, Dimension.CPU, cpu_seconds), + (Counter.SANDBOX_RAM_GIBI_DEBITS, Dimension.RAM, ram_seconds), + (Counter.SANDBOX_SSD_GIBI_DEBITS, Dimension.SSD, ssd_seconds), + ] + if gpu_seconds is not None: + dimension_pairs.append( + (Counter.SANDBOX_GPU_CORE_DEBITS, Dimension.GPU, gpu_seconds) + ) + + # The total is the sum of the per-dimension millicredits actually written, + # not a re-truncation of the exact credit sum. This keeps SANDBOX_DEBITS + # exactly reconcilable against the per-dimension breakdown meters (a dim that + # rounds to 0 millicredits contributes 0 to both the meter and the total). + total_millicredits = 0 + for counter, dimension, units in dimension_pairs: + credits = to_credits(provider=provider, dimension=dimension, raw_units=units) + # Store as millicredits to preserve sub-credit precision in the int field. + millicredits = int(credits * _MILLICREDITS) + if millicredits <= 0: + continue + total_millicredits += millicredits + try: + await check_entitlements( + key=counter, + delta=millicredits, + cache=False, + scope=meter_scope, + ) + except Exception: # pylint: disable=broad-exception-caught + log.warning( + "[sandboxes] failed to record %s for org=%s", + counter.value, + organization_id, + exc_info=True, + ) + + if total_millicredits <= 0: + return + + try: + await check_entitlements( + key=Counter.SANDBOX_DEBITS, + delta=total_millicredits, + cache=False, + scope=meter_scope, + ) + except Exception: # pylint: disable=broad-exception-caught + log.warning( + "[sandboxes] failed to record sandbox_debits for org=%s", + organization_id, + exc_info=True, + ) + + try: + await check_entitlements( + key=Counter.WALLET_DEBITS, + delta=total_millicredits, + cache=False, + scope=meter_scope, + ) + except Exception: # pylint: disable=broad-exception-caught + log.warning( + "[sandboxes] failed to record wallet_debits for org=%s", + organization_id, + exc_info=True, + ) diff --git a/api/ee/src/core/storage/service.py b/api/ee/src/core/storage/service.py index 3e9846827b..b3eff42e3b 100644 --- a/api/ee/src/core/storage/service.py +++ b/api/ee/src/core/storage/service.py @@ -28,7 +28,7 @@ async def record_storage_delta( scope = MeterScope(organization_id=org_id) allowed, _, _ = await check_entitlements( - key=Gauge.STORAGE_BYTES, + key=Gauge.BYTES, delta=delta_bytes, scope=scope, ) @@ -66,7 +66,7 @@ async def reconcile_org_storage( meters = await _meters_service().fetch( scope=scope, - key=Meters.STORAGE_BYTES, + key=Meters.BYTES, period=period, ) current = (meters[0].value if meters else 0) or 0 @@ -76,7 +76,7 @@ async def reconcile_org_storage( return await check_entitlements( - key=Gauge.STORAGE_BYTES, + key=Gauge.BYTES, delta=delta, scope=scope, period=period, diff --git a/api/ee/tests/pytest/unit/test_sandbox_credits.py b/api/ee/tests/pytest/unit/test_sandbox_credits.py new file mode 100644 index 0000000000..fda1d38efc --- /dev/null +++ b/api/ee/tests/pytest/unit/test_sandbox_credits.py @@ -0,0 +1,314 @@ +"""Unit tests for sandbox credit conversion (to_credits). + +Tests the pure conversion function against the rate table defined in +credits.DEFAULT_PROVIDER_RATES. No dollar/cents values are asserted; +credit->money is Stripe's concern. The cross-check scenario at the end is +informational only -- it verifies the rate table is self-consistent, not +that any dollar amount is correct. +""" + +from decimal import Decimal + + +from ee.src.core.sandboxes.credits import ( + DEFAULT_PROVIDER_RATES, + Dimension, + ProviderRates, + to_credits, +) + + +class TestToCreditsDefaultRates: + """Per-dimension conversion against the code-default rate table.""" + + # E2B — each dimension on its own rate + + def test_cpu_e2b_1000s(self): + result = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("1000") + ) + expected = Decimal("1000") * DEFAULT_PROVIDER_RATES["e2b"].cpu + assert result == expected + + def test_ram_e2b_1000s(self): + result = to_credits( + provider="e2b", dimension=Dimension.RAM, raw_units=Decimal("1000") + ) + expected = Decimal("1000") * DEFAULT_PROVIDER_RATES["e2b"].ram + assert result == expected + + def test_ssd_e2b_1000s(self): + result = to_credits( + provider="e2b", dimension=Dimension.SSD, raw_units=Decimal("1000") + ) + expected = Decimal("1000") * DEFAULT_PROVIDER_RATES["e2b"].ssd + assert result == expected + + def test_gpu_e2b_is_zero(self): + # E2B GPU rate is 0 by default. + result = to_credits( + provider="e2b", dimension=Dimension.GPU, raw_units=Decimal("100") + ) + assert result == Decimal("0") + + # Daytona — same default rates as E2B + + def test_cpu_daytona_1000s(self): + result = to_credits( + provider="daytona", dimension=Dimension.CPU, raw_units=Decimal("1000") + ) + expected = Decimal("1000") * DEFAULT_PROVIDER_RATES["daytona"].cpu + assert result == expected + + def test_ram_daytona_1000s(self): + result = to_credits( + provider="daytona", dimension=Dimension.RAM, raw_units=Decimal("1000") + ) + expected = Decimal("1000") * DEFAULT_PROVIDER_RATES["daytona"].ram + assert result == expected + + def test_ssd_daytona_1000s(self): + result = to_credits( + provider="daytona", dimension=Dimension.SSD, raw_units=Decimal("1000") + ) + expected = Decimal("1000") * DEFAULT_PROVIDER_RATES["daytona"].ssd + assert result == expected + + def test_gpu_daytona_is_zero(self): + # Daytona GPU rate is 0 by default; set via env override. + result = to_credits( + provider="daytona", dimension=Dimension.GPU, raw_units=Decimal("100") + ) + assert result == Decimal("0") + + # Local — zero-rated + + def test_local_cpu_is_zero(self): + result = to_credits( + provider="local", dimension=Dimension.CPU, raw_units=Decimal("1000") + ) + assert result == Decimal("0") + + def test_local_ram_is_zero(self): + result = to_credits( + provider="local", dimension=Dimension.RAM, raw_units=Decimal("1000") + ) + assert result == Decimal("0") + + def test_local_ssd_is_zero(self): + result = to_credits( + provider="local", dimension=Dimension.SSD, raw_units=Decimal("1000") + ) + assert result == Decimal("0") + + +class TestDimensionIndependence: + """Each dimension converts on its own rate; changing one doesn't affect others.""" + + def test_each_dimension_uses_own_rate(self): + cpu_rate = DEFAULT_PROVIDER_RATES["e2b"].cpu + ram_rate = DEFAULT_PROVIDER_RATES["e2b"].ram + ssd_rate = DEFAULT_PROVIDER_RATES["e2b"].ssd + + units = Decimal("3600") + cpu_credits = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=units + ) + ram_credits = to_credits( + provider="e2b", dimension=Dimension.RAM, raw_units=units + ) + ssd_credits = to_credits( + provider="e2b", dimension=Dimension.SSD, raw_units=units + ) + + assert cpu_credits == units * cpu_rate + assert ram_credits == units * ram_rate + assert ssd_credits == units * ssd_rate + # All three must be distinct (rates differ). + assert cpu_credits != ram_credits + assert ram_credits != ssd_credits + + def test_rates_are_distinct_per_dimension(self): + rates = DEFAULT_PROVIDER_RATES["e2b"] + assert rates.cpu != rates.ram + assert rates.ram != rates.ssd + assert rates.cpu != rates.ssd + + +class TestToCreditsEdgeCases: + """Edge and boundary cases.""" + + def test_zero_input_returns_zero(self): + result = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("0") + ) + assert result == Decimal("0") + + def test_negative_input_returns_zero(self): + result = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("-500") + ) + assert result == Decimal("0") + + def test_unknown_provider_returns_zero(self): + result = to_credits( + provider="unknown_cloud_xyz", + dimension=Dimension.CPU, + raw_units=Decimal("1000"), + ) + assert result == Decimal("0") + + def test_unknown_dimension_string_returns_zero(self): + result = to_credits( + provider="e2b", dimension="quantum_processor", raw_units=Decimal("1000") + ) + assert result == Decimal("0") + + def test_provider_case_insensitive(self): + lower = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("1000") + ) + upper = to_credits( + provider="E2B", dimension=Dimension.CPU, raw_units=Decimal("1000") + ) + assert lower == upper + + def test_dimension_as_string(self): + enum_result = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("1000") + ) + str_result = to_credits( + provider="e2b", dimension="cpu", raw_units=Decimal("1000") + ) + assert enum_result == str_result + + def test_dimension_string_case_insensitive(self): + lower = to_credits(provider="e2b", dimension="cpu", raw_units=Decimal("1000")) + upper = to_credits(provider="e2b", dimension="CPU", raw_units=Decimal("1000")) + assert lower == upper + + +class TestToCreditsDecimalPrecision: + """Decimal arithmetic must be exact; no float drift.""" + + def test_no_float_drift_cpu(self): + # 10800 vCPU-s at rate 0.0014 = 15.12 exactly + result = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("10800") + ) + assert result == Decimal("10800") * DEFAULT_PROVIDER_RATES["e2b"].cpu + + def test_no_float_drift_ram(self): + result = to_credits( + provider="e2b", dimension=Dimension.RAM, raw_units=Decimal("14400") + ) + assert result == Decimal("14400") * DEFAULT_PROVIDER_RATES["e2b"].ram + + def test_no_float_drift_ssd(self): + result = to_credits( + provider="e2b", dimension=Dimension.SSD, raw_units=Decimal("360000") + ) + assert result == Decimal("360000") * DEFAULT_PROVIDER_RATES["e2b"].ssd + + def test_result_type_is_decimal(self): + result = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("1") + ) + assert isinstance(result, Decimal) + + def test_sum_of_dimensions_is_exact(self): + """Full event: cpu + ram + ssd credits sum without float drift.""" + cpu = to_credits( + provider="e2b", dimension=Dimension.CPU, raw_units=Decimal("3600") + ) + ram = to_credits( + provider="e2b", dimension=Dimension.RAM, raw_units=Decimal("7200") + ) + ssd = to_credits( + provider="e2b", dimension=Dimension.SSD, raw_units=Decimal("7200") + ) + + rates = DEFAULT_PROVIDER_RATES["e2b"] + assert cpu == Decimal("3600") * rates.cpu + assert ram == Decimal("7200") * rates.ram + assert ssd == Decimal("7200") * rates.ssd + + total = cpu + ram + ssd + expected = ( + Decimal("3600") * rates.cpu + + Decimal("7200") * rates.ram + + Decimal("7200") * rates.ssd + ) + assert total == expected + + +class TestProviderRatesModel: + """ProviderRates is a typed Pydantic model with named fields.""" + + def test_provider_rates_has_named_fields(self): + rates = DEFAULT_PROVIDER_RATES["e2b"] + assert isinstance(rates, ProviderRates) + assert isinstance(rates.cpu, Decimal) + assert isinstance(rates.ram, Decimal) + assert isinstance(rates.ssd, Decimal) + assert isinstance(rates.gpu, Decimal) + + def test_all_default_providers_present(self): + assert "e2b" in DEFAULT_PROVIDER_RATES + assert "daytona" in DEFAULT_PROVIDER_RATES + assert "local" in DEFAULT_PROVIDER_RATES + + def test_local_all_zero(self): + rates = DEFAULT_PROVIDER_RATES["local"] + assert rates.cpu == Decimal("0") + assert rates.ram == Decimal("0") + assert rates.ssd == Decimal("0") + assert rates.gpu == Decimal("0") + + +class TestReferenceScenarioCrossCheck: + """Informational cross-check: rate table applied to a reference machine-minute. + + This is a sanity check on the rate table, NOT a billing input. + A 2vCPU / 2GiB-RAM / 10GiB-disk machine running for 1 minute (60s) with the + default E2B rates should produce a credits figure consistent with the table. + This test asserts against the rate table itself — it does NOT assert a dollar + amount (credit->money is Stripe's job). + """ + + def test_reference_machine_minute_e2b(self): + cpu_s = Decimal("2") * Decimal("60") # 2 vCPU × 60s = 120 vCPU·s + ram_s = Decimal("2") * Decimal("60") # 2 GiB × 60s = 120 GiB·s + ssd_s = Decimal("10") * Decimal("60") # 10 GiB × 60s = 600 GiB·s + + rates = DEFAULT_PROVIDER_RATES["e2b"] + cpu_credits = cpu_s * rates.cpu + ram_credits = ram_s * rates.ram + ssd_credits = ssd_s * rates.ssd + total_credits = cpu_credits + ram_credits + ssd_credits + + # Each dimension contributes independently. + assert cpu_credits == Decimal("120") * rates.cpu + assert ram_credits == Decimal("120") * rates.ram + assert ssd_credits == Decimal("600") * rates.ssd + + # Total must equal the per-dimension sum. + assert total_credits == cpu_credits + ram_credits + ssd_credits + + # Sanity bound: reference machine should produce a positive credit amount. + assert total_credits > Decimal("0") + + def test_reference_machine_minute_daytona(self): + cpu_s = Decimal("2") * Decimal("60") + ram_s = Decimal("2") * Decimal("60") + ssd_s = Decimal("10") * Decimal("60") + + rates = DEFAULT_PROVIDER_RATES["daytona"] + total_credits = cpu_s * rates.cpu + ram_s * rates.ram + ssd_s * rates.ssd + assert total_credits > Decimal("0") + # Daytona and E2B share default rates; cross-check they agree. + e2b_rates = DEFAULT_PROVIDER_RATES["e2b"] + e2b_total = ( + cpu_s * e2b_rates.cpu + ram_s * e2b_rates.ram + ssd_s * e2b_rates.ssd + ) + assert total_credits == e2b_total diff --git a/api/ee/tests/pytest/unit/test_sandbox_gating.py b/api/ee/tests/pytest/unit/test_sandbox_gating.py new file mode 100644 index 0000000000..c9481b0ca3 --- /dev/null +++ b/api/ee/tests/pytest/unit/test_sandbox_gating.py @@ -0,0 +1,239 @@ +"""Unit tests for ee.src.core.sandboxes.gating. + +These tests mock `check_entitlements` / `is_ee` so no live DB or Redis is +needed. The goal is to pin the logic that: + +- `check_sandbox_quota` returns (False, reason) when the entitlements check + is denied, and (True, None) when allowed. +- `check_sandbox_credits_true_up` returns False when WALLET_DEBITS is + over-quota, True when within quota. +- Both functions fail open (return allowed=True) on unexpected exceptions. +- No sandbox/wallet debit meter is reported to Stripe (wallet debits are + prepaid, never reported in arrears). +""" + +import pytest +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + + +_ORG_ID = uuid4() + + +@pytest.mark.asyncio +async def test_check_sandbox_quota_allowed(): + """Layer 1: entitlements returns True -> quota check passes.""" + with ( + patch("ee.src.core.sandboxes.gating.is_ee", return_value=True), + patch( + "ee.src.core.sandboxes.gating._check_sandbox_quota_ee", + new_callable=AsyncMock, + return_value=(True, None), + ), + ): + from ee.src.core.sandboxes.gating import check_sandbox_quota + + allowed, reason = await check_sandbox_quota(organization_id=_ORG_ID) + + assert allowed is True + assert reason is None + + +@pytest.mark.asyncio +async def test_check_sandbox_quota_denied(): + """Layer 1: entitlements returns False -> quota check blocks.""" + msg = "You have reached your sandbox usage quota for this billing period." + + with ( + patch("ee.src.core.sandboxes.gating.is_ee", return_value=True), + patch( + "ee.src.core.sandboxes.gating._check_sandbox_quota_ee", + new_callable=AsyncMock, + return_value=(False, msg), + ), + ): + from ee.src.core.sandboxes.gating import check_sandbox_quota + + allowed, reason = await check_sandbox_quota(organization_id=_ORG_ID) + + assert allowed is False + assert reason == msg + + +@pytest.mark.asyncio +async def test_check_sandbox_quota_fails_open(): + """Layer 1: unexpected error -> fail open (True, None).""" + with ( + patch("ee.src.core.sandboxes.gating.is_ee", return_value=True), + patch( + "ee.src.core.sandboxes.gating._check_sandbox_quota_ee", + new_callable=AsyncMock, + side_effect=RuntimeError("redis down"), + ), + ): + from ee.src.core.sandboxes.gating import check_sandbox_quota + + allowed, reason = await check_sandbox_quota(organization_id=_ORG_ID) + + assert allowed is True + assert reason is None + + +@pytest.mark.asyncio +async def test_check_sandbox_quota_oss_passthrough(): + """Non-EE deployments always get (True, None).""" + with patch("ee.src.core.sandboxes.gating.is_ee", return_value=False): + from ee.src.core.sandboxes.gating import check_sandbox_quota + + allowed, reason = await check_sandbox_quota(organization_id=_ORG_ID) + + assert allowed is True + assert reason is None + + +@pytest.mark.asyncio +async def test_true_up_within_quota(): + """Layer 2: within quota -> returns True.""" + with ( + patch("ee.src.core.sandboxes.gating.is_ee", return_value=True), + patch( + "ee.src.core.sandboxes.gating._check_sandbox_credits_true_up_ee", + new_callable=AsyncMock, + return_value=True, + ), + ): + from ee.src.core.sandboxes.gating import check_sandbox_credits_true_up + + result = await check_sandbox_credits_true_up(organization_id=_ORG_ID) + + assert result is True + + +@pytest.mark.asyncio +async def test_true_up_over_quota(): + """Layer 2: over quota -> returns False.""" + with ( + patch("ee.src.core.sandboxes.gating.is_ee", return_value=True), + patch( + "ee.src.core.sandboxes.gating._check_sandbox_credits_true_up_ee", + new_callable=AsyncMock, + return_value=False, + ), + ): + from ee.src.core.sandboxes.gating import check_sandbox_credits_true_up + + result = await check_sandbox_credits_true_up(organization_id=_ORG_ID) + + assert result is False + + +@pytest.mark.asyncio +async def test_true_up_fails_open(): + """Layer 2: unexpected error -> fail open (True).""" + with ( + patch("ee.src.core.sandboxes.gating.is_ee", return_value=True), + patch( + "ee.src.core.sandboxes.gating._check_sandbox_credits_true_up_ee", + new_callable=AsyncMock, + side_effect=ConnectionError("db unreachable"), + ), + ): + from ee.src.core.sandboxes.gating import check_sandbox_credits_true_up + + result = await check_sandbox_credits_true_up(organization_id=_ORG_ID) + + assert result is True + + +@pytest.mark.asyncio +async def test_true_up_oss_passthrough(): + """Non-EE deployments always get True.""" + with patch("ee.src.core.sandboxes.gating.is_ee", return_value=False): + from ee.src.core.sandboxes.gating import check_sandbox_credits_true_up + + result = await check_sandbox_credits_true_up(organization_id=_ORG_ID) + + assert result is True + + +def test_sandbox_debit_counter_values(): + """Confirm the debit Counter slugs match the Meters enum values.""" + from ee.src.core.access.entitlements.types import Counter + from ee.src.core.meters.types import Meters + + assert ( + Meters["SANDBOX_CPU_CORE_DEBITS"].value == Counter.SANDBOX_CPU_CORE_DEBITS.value + ) + assert ( + Meters["SANDBOX_RAM_GIBI_DEBITS"].value == Counter.SANDBOX_RAM_GIBI_DEBITS.value + ) + assert ( + Meters["SANDBOX_SSD_GIBI_DEBITS"].value == Counter.SANDBOX_SSD_GIBI_DEBITS.value + ) + assert ( + Meters["SANDBOX_GPU_CORE_DEBITS"].value == Counter.SANDBOX_GPU_CORE_DEBITS.value + ) + assert Meters["SANDBOX_DEBITS"].value == Counter.SANDBOX_DEBITS.value + assert Meters["LLM_DEBITS"].value == Counter.LLM_DEBITS.value + assert Meters["GATEWAY_DEBITS"].value == Counter.GATEWAY_DEBITS.value + assert Meters["WALLET_DEBITS"].value == Counter.WALLET_DEBITS.value + + +def test_no_sandbox_or_wallet_debit_meter_in_reports(): + """Wallet debits are prepaid: no *_DEBITS meter is ever reported to Stripe.""" + from ee.src.core.access.entitlements.types import Counter, REPORTS + + assert Counter.SANDBOX_CPU_CORE_DEBITS.value not in REPORTS + assert Counter.SANDBOX_RAM_GIBI_DEBITS.value not in REPORTS + assert Counter.SANDBOX_SSD_GIBI_DEBITS.value not in REPORTS + assert Counter.SANDBOX_GPU_CORE_DEBITS.value not in REPORTS + assert Counter.SANDBOX_DEBITS.value not in REPORTS + assert Counter.LLM_DEBITS.value not in REPORTS + assert Counter.GATEWAY_DEBITS.value not in REPORTS + assert Counter.WALLET_DEBITS.value not in REPORTS + + +def test_sandbox_debit_counters_in_read_only_constraint(): + """All sandbox/wallet debit counters must be in CONSTRAINTS[READ_ONLY][COUNTERS].""" + from ee.src.core.access.entitlements.types import ( + Counter, + CONSTRAINTS, + Constraint, + Tracker, + ) + + read_only_counters = CONSTRAINTS[Constraint.READ_ONLY][Tracker.COUNTERS] + for counter in ( + Counter.SANDBOX_CPU_CORE_DEBITS, + Counter.SANDBOX_RAM_GIBI_DEBITS, + Counter.SANDBOX_SSD_GIBI_DEBITS, + Counter.SANDBOX_GPU_CORE_DEBITS, + Counter.SANDBOX_DEBITS, + Counter.LLM_DEBITS, + Counter.GATEWAY_DEBITS, + Counter.WALLET_DEBITS, + ): + assert counter in read_only_counters, f"{counter} missing from READ_ONLY" + + +def test_all_plans_have_sandbox_and_wallet_debit_quotas(): + """Every default plan must carry a Quota for all 8 debit counters.""" + from ee.src.core.access.entitlements.types import ( + Counter, + DEFAULT_ENTITLEMENTS, + Tracker, + ) + + for plan, entitlements in DEFAULT_ENTITLEMENTS.items(): + counters = entitlements[Tracker.COUNTERS] + for counter in ( + Counter.SANDBOX_CPU_CORE_DEBITS, + Counter.SANDBOX_RAM_GIBI_DEBITS, + Counter.SANDBOX_SSD_GIBI_DEBITS, + Counter.SANDBOX_GPU_CORE_DEBITS, + Counter.SANDBOX_DEBITS, + Counter.LLM_DEBITS, + Counter.GATEWAY_DEBITS, + Counter.WALLET_DEBITS, + ): + assert counter in counters, f"{plan}: missing {counter}" diff --git a/api/oss/src/apis/fastapi/sessions/router.py b/api/oss/src/apis/fastapi/sessions/router.py index 4953c0069d..9a1c2cff8b 100644 --- a/api/oss/src/apis/fastapi/sessions/router.py +++ b/api/oss/src/apis/fastapi/sessions/router.py @@ -15,6 +15,7 @@ from fastapi.responses import JSONResponse from typing import Any, Optional, Union +from oss.src.utils.common import is_ee from oss.src.utils.exceptions import intercept_exceptions from oss.src.utils.logging import get_module_logger @@ -234,6 +235,19 @@ async def set_session_stream( if not has_permission: raise FORBIDDEN_EXCEPTION + # Entitlement gate (EE): may-run (RBAC above) is separate from has-quota. + # Function-local guarded import so OSS never imports ee.* at module top. + if is_ee(): + from ee.src.core.sandboxes.gating import check_sandbox_quota + + allowed, reason = await check_sandbox_quota( + organization_id=UUID(request.state.organization_id), + ) + if not allowed: + raise HTTPException( + status_code=429, detail=reason or "Sandbox quota exceeded." + ) + await self._service.check_concurrency_cap(project_id=project_id) result = await self._service.command( diff --git a/api/oss/src/utils/env.py b/api/oss/src/utils/env.py index 901f88ef0d..a948ed7928 100644 --- a/api/oss/src/utils/env.py +++ b/api/oss/src/utils/env.py @@ -1075,6 +1075,34 @@ def enabled(self) -> bool: return bool(self.uri_volatile or self.uri_durable) +# --------------------------------------------------------------------------- +# sandbox credit rates (billing conversion for sandbox compute meters) +# --------------------------------------------------------------------------- + + +class SandboxConfig(BaseModel): + """Sandbox credit-rate overrides + create-time quota estimate. + + `credit_rates` is a JSON object keyed by provider slug ("e2b", "daytona", + "local"), each value a partial `{"cpu": "0.0014", "ram": "...", ...}` rate + override merged onto the code-default table in + `ee.src.core.sandboxes.credits.DEFAULT_PROVIDER_RATES`. Values are strings + (parsed as Decimal) to avoid float precision loss in env JSON. + + `estimated_vcpu` / `estimated_run_seconds` size the in-flight accrual + estimate `check_sandbox_quota` (Layer 1) adds atop the cached meter value + before a new sandbox is allowed to launch. + """ + + credit_rates: dict | None = _load_json_env_dict("AGENTA_SANDBOX_CREDIT_RATES") + estimated_vcpu: int = int(os.getenv("AGENTA_SANDBOX_ESTIMATED_VCPU") or 2) + estimated_run_seconds: int = int( + os.getenv("AGENTA_SANDBOX_ESTIMATED_RUN_SECONDS") or 300 + ) + + model_config = ConfigDict(extra="ignore") + + # --------------------------------------------------------------------------- # email delivery # --------------------------------------------------------------------------- @@ -1280,6 +1308,7 @@ class EnvironSettings(BaseModel): postgres: PostgresConfig = PostgresConfig() posthog: PostHogConfig = PostHogConfig() redis: RedisConfig = RedisConfig() + sandbox: SandboxConfig = SandboxConfig() smtp: SmtpConfig = SmtpConfig() sendgrid: SendgridConfig = SendgridConfig() store: StoreConfig = StoreConfig() diff --git a/docs/designs/sandbox-metering/NAMING.md b/docs/designs/sandbox-metering/NAMING.md index fd1ac132dd..33dbcf6144 100644 --- a/docs/designs/sandbox-metering/NAMING.md +++ b/docs/designs/sandbox-metering/NAMING.md @@ -12,7 +12,7 @@ Scheme: `SANDBOX___SECONDS` — plain resource token + unit toke - Plain hardware resource names (CPU/RAM/SSD/GPU). Unit token: `CORE` = per-core-second (compute), `GIBI` = per-GiB-second (SI gibi = 2^30) for memory/disk. `SSD` = sandbox disk compute-time (allocated disk x time). -- Storage GAUGE is separate: `Gauge.STORAGE_BYTES` (`storage_bytes`) — persisted +- Storage GAUGE is separate: `Gauge.BYTES` (`bytes`) — persisted bytes at rest, distinct from `SANDBOX_SSD_GIBI_SECONDS`. - Track C per-dimension credit meters mirror the resource+unit tokens (`SANDBOX_CPU_CORE_CREDITS`, `SANDBOX_RAM_GIBI_CREDITS`, diff --git a/docs/designs/sandbox-metering/TRACK_B_FINDINGS.md b/docs/designs/sandbox-metering/TRACK_B_FINDINGS.md index c0b166b065..5e5ddfe4a1 100644 --- a/docs/designs/sandbox-metering/TRACK_B_FINDINGS.md +++ b/docs/designs/sandbox-metering/TRACK_B_FINDINGS.md @@ -46,7 +46,7 @@ simplest scheme — plain 3-letter resource tokens, no unit token: | `SANDBOX_SSD_SECONDS` | `sandbox_ssd_seconds` | | `SANDBOX_GPU_SECONDS` | `sandbox_gpu_seconds` | -Plus `Gauge.STORAGE_BYTES` (`storage_bytes`) — the storage-size gauge, distinct from +Plus `Gauge.BYTES` (`bytes`) — the storage-size gauge, distinct from `SANDBOX_SSD_SECONDS` (sandbox disk *compute-time*, not stored bytes). Applied consistently to: `Counter` enum, `Meters` mirror, `DEFAULT_ENTITLEMENTS` quotas, `CONSTRAINTS`, the `ee0000000004` migration's enum labels, and the sandboxes service's @@ -57,7 +57,7 @@ list; the code is the source of truth. ## Entitlements (measurement only) -- `Counter.SANDBOX_{CPU,RAM,SSD,GPU}_SECONDS` and `Gauge.STORAGE_BYTES` added. +- `Counter.SANDBOX_{CPU,RAM,SSD,GPU}_SECONDS` and `Gauge.BYTES` added. - Every plan (`HOBBY`, `PRO`, `BUSINESS`, `AGENTA_AI`, `SELF_HOSTED_ENTERPRISE`) gets a non-blocking `Quota(period=Period.MONTHLY)` for each sandbox counter — no `free`/`limit`/`strict`, so `check_entitlements` records but never blocks. @@ -68,7 +68,7 @@ list; the code is the source of truth. PRO 5 GiB free / 10 GiB limit (strict), BUSINESS 50 GiB free only (strict, no hard limit). AGENTA_AI and SELF_HOSTED_ENTERPRISE get no storage cap (unlimited, matching their existing unlimited-everything pattern). -- `CONSTRAINTS[BLOCKED][GAUGES]` gained `Gauge.STORAGE_BYTES`; `CONSTRAINTS[READ_ONLY][COUNTERS]` +- `CONSTRAINTS[BLOCKED][GAUGES]` gained `Gauge.BYTES`; `CONSTRAINTS[READ_ONLY][COUNTERS]` gained the 4 sandbox counters (same treatment as every other counter). - **`REPORTS` is untouched** — still `{Counter.TRACES_INGESTED.value: "traces"}`. No Stripe line items, no billing wiring for sandboxes or storage. @@ -93,7 +93,7 @@ existing S3-compatible client: - `env.py`: added one field, `StoreConfig.reconcile_enabled` (`AGENTA_STORE_RECONCILE_ENABLED`, default `false`). No new top-level config class. - `storage/service.py` (delta tracking + `reconcile_org_storage`) needed no changes — - it already only touched `Gauge.STORAGE_BYTES` / `Meters.STORAGE_BYTES` via + it already only touched `Gauge.BYTES` / `Meters.BYTES` via `check_entitlements`, no direct env access. ## Migration @@ -102,7 +102,7 @@ existing S3-compatible client: `down_revision = "ee0000000003"` (the current head — `add_records_ingested_meter`). Appends 5 enum labels to `meters_type` via `ALTER TYPE ... ADD VALUE IF NOT EXISTS`: `SANDBOX_CPU_SECONDS`, `SANDBOX_RAM_SECONDS`, `SANDBOX_SSD_SECONDS`, -`SANDBOX_GPU_SECONDS`, `STORAGE_BYTES` (uppercase Python-enum-member-name labels, matching +`SANDBOX_GPU_SECONDS`, `BYTES` (uppercase Python-enum-member-name labels, matching the existing `SQLEnum(Meters, name="meters_type")` convention — verified against `ee0000000002`'s `CREATE TYPE` and `ee0000000003`). `downgrade()` is a no-op (Postgres can't drop enum labels), matching `ee0000000003`. Chain diff --git a/docs/designs/sandbox-metering/TRACK_C_FINDINGS.md b/docs/designs/sandbox-metering/TRACK_C_FINDINGS.md new file mode 100644 index 0000000000..8a06194fec --- /dev/null +++ b/docs/designs/sandbox-metering/TRACK_C_FINDINGS.md @@ -0,0 +1,225 @@ +# Track C — new billing (credits + gating) + +Adds the billing layer on top of Track B's measurement-only sandbox meters: +the credits unit that rolls up the raw per-dimension sandbox meters into one +billable meter, plus create-time/mid-session entitlement gating. Built on +`feat/metering-track-b` (sandbox compute meters + storage gauge, nothing in +`REPORTS`). + +## What was brought in + +From `feat/metering-credits-layer` (`api/ee/src/core/sandbox_metering/`): +- `credits.py` -> `api/ee/src/core/sandboxes/credits.py` -- renamed + `Dimension.{VCPU,RAM,DISK,GPU}` to `{CPU,RAM,SSD,GPU}` and `ProviderRates` + fields `vcpu/ram/disk/gpu` to `cpu/ram/ssd/gpu`, matching the locked naming + in `NAMING.md`. Rewired `env.sandbox.credit_rates.rates` (old, nonexistent + `env.sandbox_metering`-style path) to the new `env.sandbox.credit_rates` + dict added to `api/oss/src/utils/env.py`. +- `sink.py` -> `api/ee/src/core/sandboxes/sink.py` -- renamed + `record_usage()` to `record_usage_credits()` (the pre-existing + `SandboxMeteringService.record_usage()` already owns that name for the raw + seconds); dropped the raw `*_seconds` adjust loop (Track B's service + already does that) so this module is credits-only. Wired as the last step + of `SandboxMeteringService.record_usage()` in `core/sandboxes/service.py`. +- `test_sandbox_credits.py` -> `api/ee/tests/pytest/unit/test_sandbox_credits.py`, + same rename treatment (42 assertions, all renamed VCPU/DISK -> CPU/SSD). + +From `feat/sandbox-metering-phase-2` (`api/ee/src/core/sandbox/metering.py`, +singular -- a different, unmerged branch from Track B's `sandboxes`): +- `metering.py` -> `api/ee/src/core/sandboxes/gating.py`. Renamed + `check_sandbox_quota()` (kept) and `record_sandbox_usage()` -> + `check_sandbox_credits_true_up()` -- Track B's service already performs + the authoritative meter writes (raw seconds + credits via sink.py), so + Layer 2 here is now a **read-only recheck** (`delta=0`) of + `SANDBOX_CREDITS` rather than a second write path. Deleted the inline + `_derive_credits()` (ad hoc `vcpu_seconds + ram_seconds` sum) entirely -- + credits now come exclusively from `credits.to_credits()`. +- `test_sandbox_metering.py` -> `api/ee/tests/pytest/unit/test_sandbox_gating.py`, + adapted to the new function names/module path, plus the REPORTS/ + CONSTRAINTS/DEFAULT_ENTITLEMENTS assertions from that file kept (still + valid against the new key set). + +The 4 junk files (`.agents/skills/agenta-package-practices/SKILL.md`, +`web/AGENTS.md`, `web/packages/agenta-entities/src/loadable/controller.ts`, +`web/packages/agenta-entities/tests/unit/trace-run-error.test.ts`) were +never touched — confirmed neither scratch branch's relevant file list nor +this branch's diff includes them. + +## Credit keys (locked naming) + +| Counter key | value | +|---------------------------|---------------------------| +| `SANDBOX_CPU_CREDITS` | `sandbox_cpu_credits` | +| `SANDBOX_RAM_CREDITS` | `sandbox_ram_credits` | +| `SANDBOX_SSD_CREDITS` | `sandbox_ssd_credits` | +| `SANDBOX_GPU_CREDITS` | `sandbox_gpu_credits` | +| `SANDBOX_CREDITS` | `sandbox_credits` | + +Added to `Counter` (`entitlements/types.py`) and mirrored into `Meters` +(`meters/types.py`). Every default plan (`HOBBY`, `PRO`, `BUSINESS`, +`AGENTA_AI`, `SELF_HOSTED_ENTERPRISE`) gets a non-blocking +`Quota(period=Period.MONTHLY)` for all 5 — no `free`/`limit`/`strict`, same +pattern as Track B's raw-second quotas, with a `# TODO(pricing)` marker +since real free/limit numbers aren't decided yet. Added to +`CONSTRAINTS[READ_ONLY][COUNTERS]`. + +## Rate table (`core/sandboxes/credits.py`) + +`Dimension` enum: `CPU | RAM | SSD | GPU` (mirrors the meter key tokens). +`ProviderRates` is a typed Pydantic model (`cpu`/`ram`/`ssd`/`gpu`, all +`Decimal`) — one instance per provider in `DEFAULT_PROVIDER_RATES` (`e2b`, +`daytona`, `local`). `local` is zero-rated (no billing cost). `to_credits()` +is pure: `raw_units: Decimal -> Decimal`, no I/O, no money math — Stripe +owns credit->money conversion, same division of responsibility as +`traces_ingested`. Values <= 0, unknown providers, and unknown dimensions +all return `Decimal("0")`. + +Env override: `AGENTA_SANDBOX_CREDIT_RATES` (JSON, `env.sandbox.credit_rates` +in `api/oss/src/utils/env.py`), a dict keyed by provider slug with partial +per-dimension string overrides merged onto the code defaults, parsed lazily +and cached in a module-level `_RATES` dict (first call wins — same lazy-init +pattern as the rest of the entitlements layer). + +## Sink (`core/sandboxes/sink.py`) + +`record_usage_credits()` is called from `SandboxMeteringService.record_usage()` +(in `core/sandboxes/service.py`) immediately after the existing raw +`SANDBOX_{CPU,RAM,SSD,GPU}_SECONDS` adjust loop. Per event: converts each +dimension's raw seconds to credits via `to_credits()`, adjusts the 4 +per-dimension `*_CREDITS` meters, sums into `total_credits`, and adjusts +`SANDBOX_CREDITS`. All deltas are stored as **millicredits** (`credits x +1000`, truncated to `int`) to keep the int-typed `MeterDTO.delta` field +precise without a schema change — a Stripe per-millicredit price accounts +for the factor at billing time. One `check_entitlements(cache=False)` call +per meter, org-scoped via `MeterScope(organization_id=...)`; each call is +independently try/excepted and fails open with a warning log (mirrors +Track B's existing per-meter error handling in the same service). + +## REPORTS wiring + +```python +REPORTS: dict[str, str] = { + Counter.TRACES_INGESTED.value: "traces", + Counter.SANDBOX_CREDITS.value: "sandbox_credits", +} +``` + +Only `sandbox_credits` (2 entries total). The 4 raw `*_seconds` meters and +the 4 per-dimension `*_credits` meters are recorded (adjusted, cached, +queryable) but **not** in `REPORTS` — so nothing per-dimension is billed to +Stripe. Flipping per-dimension billing on later is a one-line `REPORTS` +addition, no other code change, since `MetersService`'s report path already +resolves any `REPORTS`-listed key generically. + +Confirmed at runtime: `REPORTS == {'traces_ingested': 'traces', +'sandbox_credits': 'sandbox_credits'}`. + +Storage (`Gauge.BYTES`) stays out of `REPORTS`, unchanged from +Track B — deferred per that track's findings. + +## Pricing slots + +`AGENTA_BILLING_PRICING` (`ee/src/core/subscriptions/settings.py`) has no +code default (`_default_pricing()` returns `{}`) and is fully generic: any +top-level key present in `REPORTS`'s *values* becomes a valid Stripe meter +slot name once an operator sets a `{"price": "price_..."}` entry for it. No +code change was needed there — `sandbox_credits` already works as a slot +name because `REPORTS` now maps `Counter.SANDBOX_CREDITS.value -> +"sandbox_credits"`. + +`DEFAULT_CATALOG` (the user-facing pricing-modal display metadata, separate +from Stripe wiring) got `# TODO(pricing)` comment placeholders in the Pro +and Business tiers' `price` blocks, next to the existing `traces` tiered +entry — no numbers, since real sandbox pricing isn't decided. Hobby/ +Enterprise/Agenta entries were left untouched (Hobby has no `price.traces` +either; Enterprise/Agenta have no `price` block at all). + +## Gating (`core/sandboxes/gating.py`) + +Two layers, both gated on `Counter.SANDBOX_CREDITS` only (not the raw or +per-dimension meters): + +- **Layer 1** `check_sandbox_quota(organization_id, provider="e2b")` -- + create-time soft pre-check before a sandbox launches. Uses + `check_entitlements(cache=True)` (Redis-cached read) plus an estimated + in-flight accrual: `env.sandbox.estimated_vcpu * + env.sandbox.estimated_run_seconds` seconds of CPU-dimension usage, + converted through `credits.to_credits()` for the target provider (so the + estimate honors the same rate table as real usage, unlike phase-2's + hardcoded `1 credit = 1 vCPU-second` constant). Returns + `(allowed, reason)`; fails open (`True, None`) on any non- + `EntitlementsException` error. +- **Layer 2** `check_sandbox_credits_true_up(organization_id)` -- called + after `SandboxMeteringService.record_usage()` has already written the + authoritative meters via the sink. This is a **read-only recheck** + (`delta=0`, `cache=False`) of `SANDBOX_CREDITS`, not a second write -- + the old phase-2 design re-adjusted all 4 counters a second time here, + which Track B's service already owns. Returns `True`/`False`; `False` + logs a warning that session-kill (Layer 2b) isn't wired yet + (`DELETE /sessions/streams/{id}` + runner `/kill`, tracked in + `tasks.md`). + +RBAC (`RUN_SESSIONS` permission) is untouched by this module — callers must +check it separately before calling either gating function; `gating.py` only +ever looks at `Counter.SANDBOX_CREDITS`. + +New env knobs (`api/oss/src/utils/env.py`, `SandboxConfig` / +`env.sandbox`): `AGENTA_SANDBOX_CREDIT_RATES` (credit-rate overrides, used +by `credits.py`), `AGENTA_SANDBOX_ESTIMATED_VCPU` (default `2`), +`AGENTA_SANDBOX_ESTIMATED_RUN_SECONDS` (default `300`) — the Layer-1 +accrual estimate inputs. + +## Migration + +`api/ee/databases/postgres/migrations/core_ee/versions/ee0000000005_add_sandbox_credit_meters.py`, +`down_revision = "ee0000000004"` (Track B's sandbox + storage meters +migration, the current head). Appends 5 enum labels to `meters_type`: +`SANDBOX_CPU_CREDITS`, `SANDBOX_RAM_CREDITS`, `SANDBOX_SSD_CREDITS`, +`SANDBOX_GPU_CREDITS`, `SANDBOX_CREDITS` via `ALTER TYPE ... ADD VALUE IF +NOT EXISTS` (uppercase Python-enum-member-name labels, matching +`ee0000000003`/`ee0000000004`). `downgrade()` is a no-op (Postgres can't +drop enum labels). Verified `revision="ee0000000005"`, +`down_revision="ee0000000004"` by direct module import — chain is linear, +single head. + +## Tests + +- `api/ee/tests/pytest/unit/test_sandbox_credits.py` -- 30 tests on + `to_credits()`/`DEFAULT_PROVIDER_RATES`/`ProviderRates`: per-dimension + conversion for e2b/daytona/local, dimension independence, edge cases + (zero/negative/unknown provider/unknown dimension/case-insensitivity), + Decimal-exactness (no float drift), and the informational reference- + machine-minute cross-check. +- `api/ee/tests/pytest/unit/test_sandbox_gating.py` -- 12 tests on + `check_sandbox_quota()` / `check_sandbox_credits_true_up()` (allowed / + denied / fails-open / OSS-passthrough, mocked `check_entitlements` and + `is_ee`), plus 4 wiring assertions: `Meters` mirrors `Counter` for all 5 + credit keys, `REPORTS` contains only `sandbox_credits` (not the raw or + per-dimension meters), all 5 credit counters are in + `CONSTRAINTS[READ_ONLY][COUNTERS]`, and every `DEFAULT_ENTITLEMENTS` plan + carries a `Quota` for all 5. +- `cd api && ruff format . && ruff check --fix .` -- clean, no errors + (`1263 files left unchanged`, `All checks passed!`). +- `uv run pytest ee/tests/pytest/unit/test_sandbox_credits.py + ee/tests/pytest/unit/test_sandbox_gating.py -q` -- 42 passed. +- `uv run pytest ee/tests/pytest/unit -q -k "entitle or sandbox or meter"` + -- 98 passed (no regressions in adjacent entitlements/meters/sandboxes + unit tests). +- Import-tested `ee.src.core.sandboxes.{credits,sink,gating,service}`, + `ee.src.core.access.entitlements.types`, `ee.src.core.meters.types`, and + the full `ee.src.main` composition root (exercises service + router + wiring end-to-end) -- all clean. + +## Commits + +1. `feat(billing): sandbox credits unit + per-dimension rate table` -- + `credits.py` + its test, renamed to the CPU/RAM/SSD/GPU scheme. +2. `feat(billing): wire sandbox credits sink into the meters + REPORTS layer` + -- `sink.py`, wired into `service.py`; `entitlements/types.py` (Counter, + quotas, REPORTS, CONSTRAINTS, DEFAULT_CATALOG TODO placeholders); + `meters/types.py` mirror; `env.py` `SandboxConfig`. +3. `feat(billing): create-time + true-up gating on sandbox_credits` -- + `gating.py` + its test. +4. `chore(db): credit meters enum migration` -- `ee0000000005`. + +Not pushed, per instructions. diff --git a/docs/designs/sandbox-metering/specs.md b/docs/designs/sandbox-metering/specs.md index cd575458fc..ca3e3252d8 100644 --- a/docs/designs/sandbox-metering/specs.md +++ b/docs/designs/sandbox-metering/specs.md @@ -453,11 +453,11 @@ Counters accumulate monotonically over a billing period and reset; a gauge is a **down on delete**. The only existing gauge is `Gauge.USERS`, and it is the template: -- **New gauge:** `Gauge.STORAGE_BYTES = "storage_bytes"` (mirror into `Meters`). +- **New gauge:** `Gauge.BYTES = "bytes"` (mirror into `Meters`). Consider a second gauge keyed at project scope if per-project caps are needed (`MeterScope` already supports `project_id` under `organization_id`). - **Delta semantics:** adjust the gauge by **signed delta** the same way USERS - does (`check_entitlements(key=Gauge.STORAGE_BYTES, delta=+bytes)` on write, + does (`check_entitlements(key=Gauge.BYTES, delta=+bytes)` on write, `delta=-bytes` on delete). The meter `value` is the live total. Stripe sync (`report()`) already treats gauges as **absolute quantity** via `Subscription.modify` — correct for "current GB stored." @@ -481,7 +481,7 @@ template: gate here (unlike sandbox compute's post-hoc cost): check the gauge before accepting an upload. - **Billing:** negligible per-GiB early; defer `REPORTS`/Stripe until volume - justifies it. Add `Gauge.STORAGE_BYTES.value: "storage"` to `REPORTS` + a + justifies it. Add `Gauge.BYTES.value: "storage"` to `REPORTS` + a per-plan `"storage"` price when ready — no mechanism change. - **Retention:** later. A retention sweep that deletes old mount data simply emits negative deltas (and the periodic reconcile corrects any drift). diff --git a/docs/designs/sandbox-metering/tasks.md b/docs/designs/sandbox-metering/tasks.md index 59d72b49b5..d2ac69380b 100644 --- a/docs/designs/sandbox-metering/tasks.md +++ b/docs/designs/sandbox-metering/tasks.md @@ -222,12 +222,12 @@ for typical size. No provider integration, no new mechanism. Unit = **count** Structurally different: a **gauge** (level), not a counter. `Gauge.USERS` is the template. Caps now; billing/retention later. -- [ ] Add `Gauge.STORAGE_BYTES = "storage_bytes"` to `entitlements/types.py`, +- [ ] Add `Gauge.BYTES = "bytes"` to `entitlements/types.py`, mirror into `Meters`. Decide org-scope only vs also project-scope gauge. - [ ] Mount-layout prerequisite: ensure the S3/SeaweedFS path convention encodes `org/project` prefix so stored size is attributable to a `MeterScope` without per-file bookkeeping. -- [ ] Incremental deltas: on mount write `check_entitlements(key=Gauge.STORAGE_BYTES, +- [ ] Incremental deltas: on mount write `check_entitlements(key=Gauge.BYTES, delta=+bytes, scope=org[/project])`; on delete `delta=-bytes`. Gauge `value` = live total. - [ ] Periodic reconcile job (mirror the Daytona poll/lock pattern): read @@ -241,7 +241,7 @@ template. Caps now; billing/retention later. - [ ] Provider/scope toggle + `is_ee()` gating on the reconcile job, mirroring Phase 1. - [ ] Billing (defer until volume justifies): add - `Gauge.STORAGE_BYTES.value: "storage"` to `REPORTS` + per-plan `"storage"` + `Gauge.BYTES.value: "storage"` to `REPORTS` + per-plan `"storage"` price. `report()` syncs gauges as absolute quantity (`Subscription.modify`) — no change. - [ ] Retention (later): a sweep deleting old mount data emits negative deltas; diff --git a/docs/designs/track-c-billing/specs.md b/docs/designs/track-c-billing/specs.md new file mode 100644 index 0000000000..d069320d60 --- /dev/null +++ b/docs/designs/track-c-billing/specs.md @@ -0,0 +1,182 @@ +# Track C — billing: debit population, gating, and the wallet + +Track C owns **billing**: converting measured usage into wallet debits +(populate), checking and enforcing balances (charge/gate), and the wallet +itself — funding ledger, allowance cron, tier, top-ups, auto-recharge. Track B +(metering) defines the meters and measures; Track D (BYOS) adds +bring-your-own-secrets zero-rating on top of this track. + +Companion design docs (big-agents-audit collection): `tiers-and-unified-wallet.md` +(the model: billing boundary §0, grant kinds + amounts §1.5/§3.5, promotion §2, +Stripe §3, data model §3.5, entitlement chain §4, tier-change §5, meter taxonomy +§6) and `wallet-enforcement-matrix.md` (credit/debit/measure/enforce per +resource family; the preventive rule §4). This spec is the build cut of those +docs; where they conflict, they win and this file gets fixed. + +Stacked on `feat/metering-track-b` — rebase onto its re-partitioned tip first +(see tasks C1). + +## Part 1 — sandbox debit population + gating (reconcile existing work) + +End state of `api/ee/src/core/sandboxes/` on this track: + +- **`debits.py`** (renamed from `credits.py`): the provider × dimension rate + table and `to_debits()` (renamed from `to_credits()`). Millicredit int + convention unchanged (1 credit ≈ $0.01 list price; millicredits = credits × + 1000, truncated). A `local` provider row is all-zero (local sandboxes are + free). +- **`sink.py`**: `record_usage_debits()` (renamed from + `record_usage_credits()`). Per usage event: per-dimension + `SANDBOX_{CPU_CORE,RAM_GIBI,SSD_GIBI,GPU_CORE}_DEBITS`, their sum into + `SANDBOX_DEBITS`, and the same total into `WALLET_DEBITS` (the cross-family + grand total). The total stays reconcilable with the per-dimension writes + (sum of written millicredits, not a re-truncation). +- **`gating.py`**: both layers read `WALLET_DEBITS` (the wallet gates on the + cross-family total, not the sandbox sub-total). Layer 1 + `check_sandbox_quota()` at kickoff; Layer 2 `check_sandbox_credits_true_up()` + post-debit. Keep the limit an injected/quota-supplied parameter — Part 2 + swaps its source from the static plan quota to the wallet balance. +- **`service.py`**: the one-line populate patch on Track B's measurement-only + `record_usage()` — call `record_usage_debits(...)` after dedup/parsing. This + is the only Track C edit to a Track B file. +- **tests**: `test_sandbox_gating.py` (and any sink tests) against the + `*_DEBITS` names. + +The `Counter`/`Meters` enum members and the `ee0000000005` migration belong to +Track B; after the rebase, any enum edits still in this track are redundant and +must be dropped. + +## Part 2 — the wallet + +### 2.1 `wallet_credits` table — the funding ledger + +Append-only ledger of credits INTO the wallet. New EE table, reusing the house +mixins (`IdentifierDBA` uuid7 `id`, `OrganizationScopeDBA`, `LifecycleDBA`): + +| Column | Type | Notes | +|---|---|---| +| `id` | UUID (uuid7) | PK. Mint the id explicitly in DAO insert mappings (ORM `default=` does not fire on `insert().values()`). | +| `organization_id` | UUID FK, indexed | scope | +| `kind` | Enum `credit_kind` | `signup_gift`, `plan_allowance`, `card_topup`, `admin_promotion`, `support_adjustment` — exactly five, flat, full words | +| `amount_millicredits` | BigInteger | list-price credits added, millicredit int convention | +| `granted_at` | TIMESTAMP | when spendable | +| `expires_at` | TIMESTAMP nullable | null = never. Set for `plan_allowance` (month-end) and optionally time-boxed `admin_promotion`; null for `signup_gift`, `card_topup`, `support_adjustment` | +| `source_reference` | String nullable | Stripe object id or period key; **UNIQUE(organization_id, source_reference) WHERE source_reference IS NOT NULL** — the idempotency key | +| `metadata` | JSONB nullable | e.g. `{"tier_at_grant": 2}` | + +**Balance** (the only hot-path read): + +``` +balance(organization) = + Σ amount_millicredits WHERE expires_at IS NULL OR expires_at > now() + − WALLET_DEBITS meter value +``` + +One filtered sum — no FIFO, no per-grant allocation, no stored mutable balance. +Only `plan_allowance` reliably expires, so expiry is just the `WHERE` predicate +ceasing to count a row. Cache the balance in the entitlements Redis namespace; +invalidate on any credit mint and on debit writes. Balance CAN go negative +(post-flight debits land after pre-flight estimates); that is expected and +bounded — see 2.6. + +No new "payments"/"transactions" table: cumulative income (the tier promotion +signal) is `Σ amount_millicredits` over kinds +(`plan_allowance`, `card_topup`, `admin_promotion`) on this same table. + +### 2.2 `subscriptions` extension — tier + auto-recharge + +``` +tier = Column(SmallInteger, nullable=False, default=0) +autorecharge_threshold_millicredits = Column(BigInteger, nullable=True) # null = off +autorecharge_target_millicredits = Column(BigInteger, nullable=True) +``` + +`tier` is a materialized derivation (recompute writes it; nothing else does), +stored here because the gate already loads the subscription row per request. +Extend `SubscriptionDTO` + DAO read/update mechanically. + +### 2.3 Minting — daily cron + signup hook + webhooks + +All amounts are config dials (env-configured like the existing billing +settings); zero = no-op mint. Candidates: signup gift $5; allowance free $1 / +Pro $10 / Business $100 / Enterprise $1000 per month. + +- **`signup_gift`** — minted once, at first-organization provisioning (the + user's signup org only; later orgs get nothing). Never expires. +- **`plan_allowance`** — minted by a **daily cron for everyone** (reuse the + meters cron host): per org, "is the period due, and (for paid plans) is it + paid?" → mint once, idempotently, `source_reference = + allowance::`; `expires_at` = period end. Payment + gates the mint (the cron reads subscription state the existing webhooks + maintain); the cron is only the trigger. Amount = the allowance dial for the + plan actually invoiced (Stripe proration means mid-cycle plan changes come + out right by construction). The mint is **never suspended** — with a negative + balance it applies against the debt. +- **`card_topup`** — one-time Checkout `mode="payment"` ("credits" product) → + `checkout.session.completed` webhook (a **new** case in the closed 4-event + switch) mints with `source_reference` = the session id. Auto-recharge mints + arrive via `payment_intent.succeeded`; handle `payment_intent.payment_failed` + (no mint, log/notify — the org simply drifts toward the gate). Mandate the + webhook secret in production while touching the switch. +- **`admin_promotion` / `support_adjustment`** — admin endpoints (Access-header + gated, with an audit log line). + +### 2.4 Auto-recharge + +Scheduled check (same cron host): `balance < autorecharge_threshold` → create +an off-session PaymentIntent (`off_session=true, confirm=true`) for +`target − balance` against the stored payment method. Success mints via the +webhook path above; decline mints nothing. + +### 2.5 Tier recompute — plain sum, monotonic + +``` +computed = highest threshold T in {25, 250, 2500} credits such that + Σ mints(plan_allowance, card_topup, admin_promotion) ≥ T +tier = max(computed, manual_override) # + card-on-file required above the floor +``` + +No weights, no rolling windows, no demotion of any kind: nothing pulls a tier +down except an admin action. Wallet exhaustion **blocks** (gate fails) but +never demotes. Recompute triggers: any credit mint, subscription change +webhooks, card attach/detach, admin override. Each recompute writes +`subscriptions.tier` and bumps the entitlements cache. + +Tier gates a four-axis vector per family (concurrency / power / duration / +wallet ceiling) — ship the tier→limits table as config; sandbox first, the +power axis dormant until resource-class selection exists. + +### 2.6 Gating — the preventive rule, soft at launch + +Shared gate logic (callable from API, services, runner), per invocation site: + +``` +allow ⇔ balance ≥ f × max_run_cost(family, tier, request) +``` + +`max_run_cost` is computable because the tier caps resources: sandbox = +duration_cap × class rate; LLM = max token windows × model rates; gateway = +flat per-call. `f` is a per-family config fraction (1.0 = fully preventive). +**Launch posture: all three families non-strict (soft)** via the existing +`Quota.strict` semantics — compute and warn, don't block; flip per family by +config. Sandbox is wired now (sessions router + runner kickoff already call +`check_sandbox_quota`); LLM (services level — the agent path passes through the +services handler via `AgentResult.usage`) and gateway sinks/gates land after +this track against the same interface. + +The gate's ceiling source becomes the wallet balance (2.1) instead of the +static plan quota; `check_entitlements`/`adjust(delta, limit)` stays the atom. + +### 2.7 True-up → kill caller + +After each debit lands, recompute balance; if ≤ 0, terminate the org's +in-flight sessions. The targeted-kill primitive (parameterized runner +`POST /kill` with a session/stream/sandbox identifier; the bare form stays the +orphan hatch) is mainline sessions/runner work outside this track — C ships the +true-up caller with the kill invocation behind a feature flag/log until that +lands (today's behavior: warn). + +## Out of scope + +LLM/gateway sinks (later, same interface), BYOS `secret_origin` zero-rating +(Track D), the runner-side per-turn LLM gate (optional tightening), wallet UI. diff --git a/docs/designs/track-c-billing/tasks.md b/docs/designs/track-c-billing/tasks.md new file mode 100644 index 0000000000..9d9274c95d --- /dev/null +++ b/docs/designs/track-c-billing/tasks.md @@ -0,0 +1,64 @@ +# Track C — tasks + +Execution order for `feat/metering-track-c` (stacked on +`feat/metering-track-b`). Safety tag before re-partition: +`safety-track-c-pre-repartition`. Do C1 only after Track B's re-partition +commit exists. + +## C1 — rebase + re-partition delta +- [ ] Commit or stash-inventory the current uncommitted work; rebase + `feat/metering-track-c` onto the new `feat/metering-track-b` tip (or + re-cut: cherry-pick/squash C-owned content onto B — pick whichever gives + clean, reviewable commits; content end-state wins over history). +- [ ] Drop everything B now owns: enum members, quota entries, + `ee0000000005`, `ee0000000004` edits, the seconds-block removal in + `service.py`. After C1, `git diff feat/metering-track-b...HEAD` touches + ONLY: `sandboxes/{debits,sink,gating}.py`, the one-line sink call in + `sandboxes/service.py`, tests, and this track's docs. +- [ ] Rename `credits.py` → `debits.py`; `to_credits()` → `to_debits()`; + `record_usage_credits()` → `record_usage_debits()`; update imports, + docstrings, and comments (consumption side says debits; "credits" only + for the funding ledger / millicredit unit). +- [ ] Verify sink writes per-dimension + `SANDBOX_DEBITS` + `WALLET_DEBITS`; + gating reads `WALLET_DEBITS`; tests green; ruff clean. Commit. + +## C2 — wallet schema +- [ ] `wallet_credits` DBE/DBA/DTO/DAO (columns per specs 2.1; mint `id` in the + insert mapping; partial UNIQUE(organization_id, source_reference)). +- [ ] Migration for `wallet_credits` + the `credit_kind` enum. +- [ ] `subscriptions`: `tier` + auto-recharge columns, DTO/DAO extension, + migration. +- [ ] `balance(organization)` query + Redis caching + invalidation on + mint/debit. Unit tests: expiry predicate, negative balance, idempotent + re-mint rejected. + +## C3 — minting +- [ ] Config dials: signup-gift amount, per-plan allowance amounts (zero = + no-op), thresholds `{25, 250, 2500}`. +- [ ] Signup hook: mint `signup_gift` for the signup organization only. +- [ ] Daily allowance cron (meters cron host): due+paid check, idempotent + period-keyed mint, `expires_at` = period end. +- [ ] Webhook switch: add `checkout.session.completed`, + `payment_intent.succeeded`, `payment_intent.payment_failed`; mandate the + webhook secret in production. +- [ ] Top-up Checkout session creation endpoint. +- [ ] Admin endpoints for `admin_promotion` / `support_adjustment` with an + audit log line. + +## C4 — tier +- [ ] Recompute function (plain sum ≥ threshold, card-on-file above floor, + `max(computed, manual)`), triggered on mint / subscription webhook / card + events; writes `subscriptions.tier`; bumps entitlements cache. +- [ ] Tier→limits config table (concurrency / power / duration / ceiling per + family; sandbox instantiated, others declared). + +## C5 — gating + auto-recharge + true-up +- [ ] Gate reads wallet balance as the ceiling (preventive rule, per-family + `f`, soft/`strict=False` launch default) — swap the source in + `check_sandbox_quota`, keep the signature. +- [ ] Auto-recharge scheduled check → off-session PaymentIntent. +- [ ] True-up caller after debit writes: balance ≤ 0 → targeted kill invocation + behind a flag (until the parameterized `/kill` lands mainline), else + warn. +- [ ] Acceptance tests: mint→balance→gate→debit→true-up round trip; EE-only + guards via `is_ee()`.