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 new file mode 100644 index 0000000000..732b0df7cc --- /dev/null +++ b/api/ee/databases/postgres/migrations/core_ee/versions/ee0000000004_add_sandbox_and_storage_meters.py @@ -0,0 +1,37 @@ +"""add sandbox compute + bytes meters to meters_type + +Revision ID: ee0000000004 +Revises: ee0000000003 +Create Date: 2026-07-02 00:00:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "ee0000000004" +down_revision: Union[str, None] = "ee0000000003" +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_SECONDS'" + ) + op.execute( + "ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_RAM_GIBI_SECONDS'" + ) + op.execute( + "ALTER TYPE meters_type ADD VALUE IF NOT EXISTS 'SANDBOX_SSD_GIBI_SECONDS'" + ) + 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 'BYTES'") + + +def downgrade() -> None: + # Postgres cannot drop an enum label; leave the values in place. + pass diff --git a/api/ee/src/apis/fastapi/billing/router.py b/api/ee/src/apis/fastapi/billing/router.py index 87a4f04de5..bc8f1b4727 100644 --- a/api/ee/src/apis/fastapi/billing/router.py +++ b/api/ee/src/apis/fastapi/billing/router.py @@ -214,6 +214,18 @@ def __init__( methods=["POST"], ) + self.admin_router.add_api_route( + "/storage/reconcile", + self.reconcile_storage, + methods=["POST"], + ) + + self.admin_router.add_api_route( + "/storage/reconcile/unlock", + self.unlock_reconcile_storage, + methods=["POST"], + ) + async def _reset_organization_flags(self, organization_id: str) -> None: organization = await db_manager.get_organization(organization_id) if not organization: @@ -1156,6 +1168,114 @@ async def unlock_report_usage( content={"status": "noop", "released": False}, ) + @intercept_exceptions() + async def reconcile_storage( + self, + ): + log.info("[storage] [reconcile] [endpoint] Trigger") + + LOCK_TTL = 3600 # 1 hour + + try: + lock_owner = await acquire_lock( + namespace="meters:storage:reconcile", + key={}, + ttl=LOCK_TTL, + strict=True, + ) + + if not lock_owner: + log.info("[storage] [reconcile] [endpoint] Skipped (ongoing)") + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"status": "skipped"}, + ) + + async def _renew_lock(): + return await renew_lock( + namespace="meters:storage:reconcile", + key={}, + ttl=LOCK_TTL, + owner=lock_owner, + ) + + try: + from ee.src.core.storage.reconcile import run_storage_reconcile + + await run_storage_reconcile(renew=_renew_lock) + + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"status": "success"}, + ) + + except Exception: + log.error( + "[storage] [reconcile] [endpoint] Failed", + exc_info=True, + ) + return JSONResponse( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content={"status": "error", "message": "Reconcile failed"}, + ) + + finally: + released = await release_lock( + namespace="meters:storage:reconcile", + key={}, + owner=lock_owner, + ) + if released: + log.info("[storage] [reconcile] [endpoint] Lock released") + else: + log.warn( + "[storage] [reconcile] [endpoint] Lock release skipped (expired/lost)" + ) + + except Exception: + log.error( + "[storage] [reconcile] [endpoint] Fatal error", + exc_info=True, + ) + return JSONResponse( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content={"status": "error", "message": "Fatal error"}, + ) + + @intercept_exceptions() + async def unlock_reconcile_storage( + self, + ): + log.warn("[storage] [reconcile] [unlock] Trigger") + + try: + released = await release_lock( + namespace="meters:storage:reconcile", + key={}, + strict=True, + ) + except Exception: + log.error( + "[storage] [reconcile] [unlock] Failed to release lock", exc_info=True + ) + return JSONResponse( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + content={"status": "error", "message": "Lock backend error"}, + ) + + if released: + log.warn("[storage] [reconcile] [unlock] Lock force-released") + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"status": "success", "released": True}, + ) + + log.info("[storage] [reconcile] [unlock] No lock found") + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"status": "noop", "released": False}, + ) + # ROUTES @intercept_exceptions() diff --git a/api/ee/src/apis/fastapi/sandboxes/__init__.py b/api/ee/src/apis/fastapi/sandboxes/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/api/ee/src/apis/fastapi/sandboxes/models.py b/api/ee/src/apis/fastapi/sandboxes/models.py new file mode 100644 index 0000000000..0ab2b1a1ff --- /dev/null +++ b/api/ee/src/apis/fastapi/sandboxes/models.py @@ -0,0 +1,28 @@ +from typing import Optional +from uuid import UUID + +from pydantic import BaseModel + + +class E2BEventPayload(BaseModel): + """Shape of an E2B webhook event payload (subset of documented fields). + + E2B issue #1103: header/signature mismatch — keep this loose and log + raw payloads on first delivery for verification. + """ + + event: Optional[str] = None + sandbox_id: Optional[str] = None + team_id: Optional[str] = None + # Resource allocation at time of event + vcpu: Optional[int] = None + memory_mb: Optional[int] = None + # Duration of the billing window (ms) + duration_ms: Optional[int] = None + start_timestamp: Optional[str] = None + + +class DaytonaPollRequest(BaseModel): + organization_id: UUID + period_start: str # RFC3339 + period_end: str # RFC3339 diff --git a/api/ee/src/apis/fastapi/sandboxes/router.py b/api/ee/src/apis/fastapi/sandboxes/router.py new file mode 100644 index 0000000000..9107e0a6c7 --- /dev/null +++ b/api/ee/src/apis/fastapi/sandboxes/router.py @@ -0,0 +1,188 @@ +"""Sandbox metering ingestion routes. + +Public (unauthenticated) webhook receiver for E2B events. +Admin-only endpoint to trigger a Daytona usage poll on-demand. +""" + +import math +from typing import Optional + +from fastapi import APIRouter, Request, status +from fastapi.responses import JSONResponse + +from oss.src.utils.env import env +from oss.src.utils.logging import get_module_logger +from oss.src.utils.exceptions import intercept_exceptions +from oss.src.utils.common import is_ee + +from ee.src.core.sandboxes.service import SandboxMeteringService +from ee.src.core.sandboxes.dtos import SandboxUsageDTO +from ee.src.core.sandboxes.exceptions import SandboxWebhookSignatureError +from ee.src.apis.fastapi.sandboxes.models import DaytonaPollRequest + +log = get_module_logger(__name__) + +_GIB = 1024.0 * 1024.0 * 1024.0 + + +def _mb_ms_to_gib_seconds(memory_mb: Optional[int], duration_ms: Optional[int]) -> int: + """Convert memory_mb × duration_ms to GiB-seconds (ceiling).""" + if not memory_mb or not duration_ms: + return 0 + gib = memory_mb / 1024.0 + seconds = duration_ms / 1000.0 + return max(1, math.ceil(gib * seconds)) + + +def _vcpu_ms_to_vcpu_seconds(vcpu: Optional[int], duration_ms: Optional[int]) -> int: + """Convert vCPU count × duration_ms to vCPU-seconds (ceiling).""" + if not vcpu or not duration_ms: + return 0 + return max(1, math.ceil(vcpu * duration_ms / 1000.0)) + + +class SandboxMeteringRouter: + def __init__(self, *, sandboxes_service: SandboxMeteringService): + self.service = sandboxes_service + + # Public webhook receiver (no auth — verified by HMAC). + self.router = APIRouter() + self.router.add_api_route( + "/e2b/events/", + self.receive_e2b_event, + methods=["POST"], + operation_id="sandboxes_e2b_event", + include_in_schema=False, + ) + + # Admin-only routes. + self.admin_router = APIRouter() + self.admin_router.add_api_route( + "/daytona/poll", + self.trigger_daytona_poll, + methods=["POST"], + operation_id="sandboxes_daytona_poll", + ) + + @intercept_exceptions() + async def receive_e2b_event(self, request: Request): + """Receive and verify an E2B sandbox lifecycle webhook event.""" + if not is_ee() or not env.e2b.enabled: + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"status": "disabled"}, + ) + + raw_body = await request.body() + signature = request.headers.get("e2b-signature", "") + delivery_id = request.headers.get("e2b-delivery-id", "") + + # HMAC verification against env.e2b.webhook_secret. + if not self.service.verify_e2b_signature( + raw_body=raw_body, + signature_header=signature, + ): + log.warning( + "[sandboxes] E2B HMAC verification failed delivery_id=%s sig=%r", + delivery_id, + signature[:32], + ) + raise SandboxWebhookSignatureError() + + # Parse payload (loose — log raw on unknown events for issue #1103 debugging). + try: + import json as _json + + payload = _json.loads(raw_body) + except Exception: + return JSONResponse( + status_code=status.HTTP_400_BAD_REQUEST, + content={"status": "invalid_json"}, + ) + + log.debug( + "[sandboxes] E2B event delivery_id=%s payload=%s", + delivery_id, + payload, + ) + + event_type = payload.get("event", "") + sandbox_id = payload.get("sandbox_id") or payload.get("id") or "unknown" + team_id = payload.get("team_id", "") + vcpu = payload.get("vcpu") + memory_mb = payload.get("memory_mb") + duration_ms = payload.get("duration_ms") + + # Only billing-relevant events carry resource usage. + billable_events = {"sandbox.killed", "sandbox.paused", "sandbox.checkpointed"} + if event_type not in billable_events: + return JSONResponse( + status_code=status.HTTP_202_ACCEPTED, + content={"status": "accepted", "metered": False}, + ) + + # Map team_id → organization_id. + # Phase 1: expect team_id to be set to the Agenta organization UUID. + # Phase 2 will add a proper lookup table. + from uuid import UUID as _UUID + + try: + org_id = _UUID(team_id) + except (ValueError, TypeError): + log.warning( + "[sandboxes] E2B team_id is not a valid UUID: %r — skipping metering", + team_id, + ) + return JSONResponse( + status_code=status.HTTP_202_ACCEPTED, + content={ + "status": "accepted", + "metered": False, + "reason": "team_id_not_uuid", + }, + ) + + usage = SandboxUsageDTO( + organization_id=org_id, + provider="e2b", + sandbox_id=sandbox_id, + vcpu_seconds=_vcpu_ms_to_vcpu_seconds(vcpu, duration_ms), + ram_gib_seconds=_mb_ms_to_gib_seconds(memory_mb, duration_ms), + disk_gib_seconds=0, # E2B doesn't report disk per event + gpu_seconds=0, + delivery_id=delivery_id, + ) + + result = await self.service.record_usage(usage) + + return JSONResponse( + status_code=status.HTTP_202_ACCEPTED, + content={ + "status": "accepted", + "metered": result.accepted, + "deduped": result.deduped, + }, + ) + + @intercept_exceptions() + async def trigger_daytona_poll(self, body: DaytonaPollRequest): + """Admin endpoint to trigger an on-demand Daytona usage poll.""" + if not is_ee() or not env.daytona.enabled: + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"status": "disabled"}, + ) + + ran = await self.service.run_daytona_poll( + org_id=body.organization_id, + api_key=env.daytona.api_key, + analytics_url=env.daytona.analytics_url, + daytona_organization_id=env.daytona.organization_id, + period_start=body.period_start, + period_end=body.period_end, + ) + + return JSONResponse( + status_code=status.HTTP_200_OK, + content={"status": "ran" if ran else "skipped"}, + ) diff --git a/api/ee/src/core/access/entitlements/types.py b/api/ee/src/core/access/entitlements/types.py index a1ff6ffebb..9d2e7dc5a5 100644 --- a/api/ee/src/core/access/entitlements/types.py +++ b/api/ee/src/core/access/entitlements/types.py @@ -55,10 +55,15 @@ 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" class Gauge(str, Enum): USERS = "users" + BYTES = "bytes" class Constraint(str, Enum): @@ -357,6 +362,18 @@ class Throttle(BaseModel): retention=Retention.WEEKLY, period=Period.MONTHLY, ), + Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), }, Tracker.GAUGES: { Gauge.USERS: Quota( @@ -364,6 +381,11 @@ class Throttle(BaseModel): limit=2, strict=True, ), + Gauge.BYTES: Quota( + free=1_073_741_824, + limit=1_073_741_824, + strict=True, + ), }, Tracker.THROTTLES: [ Throttle( @@ -449,11 +471,28 @@ class Throttle(BaseModel): retention=Retention.MONTHLY, period=Period.MONTHLY, ), + Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), }, Tracker.GAUGES: { Gauge.USERS: Quota( strict=True, ), + Gauge.BYTES: Quota( + free=5_368_709_120, + limit=10_737_418_240, + strict=True, + ), }, Tracker.THROTTLES: [ Throttle( @@ -539,11 +578,27 @@ class Throttle(BaseModel): retention=Retention.QUARTERLY, period=Period.MONTHLY, ), + Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), }, Tracker.GAUGES: { Gauge.USERS: Quota( strict=True, ), + Gauge.BYTES: Quota( + free=53_687_091_200, + strict=True, + ), }, Tracker.THROTTLES: [ Throttle( @@ -625,6 +680,18 @@ class Throttle(BaseModel): Counter.RECORDS_INGESTED: Quota( period=Period.MONTHLY, ), + Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), }, Tracker.GAUGES: { Gauge.USERS: Quota( @@ -663,6 +730,18 @@ class Throttle(BaseModel): Counter.RECORDS_INGESTED: Quota( period=Period.MONTHLY, ), + Counter.SANDBOX_CPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_RAM_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_SSD_GIBI_SECONDS: Quota( + period=Period.MONTHLY, + ), + Counter.SANDBOX_GPU_CORE_SECONDS: Quota( + period=Period.MONTHLY, + ), }, Tracker.GAUGES: { Gauge.USERS: Quota( @@ -692,6 +771,7 @@ class Throttle(BaseModel): ], Tracker.GAUGES: [ Gauge.USERS, + Gauge.BYTES, ], }, Constraint.READ_ONLY: { @@ -702,6 +782,10 @@ 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, ], }, } diff --git a/api/ee/src/core/meters/types.py b/api/ee/src/core/meters/types.py index 5f9aa92f77..ec2b5d69ce 100644 --- a/api/ee/src/core/meters/types.py +++ b/api/ee/src/core/meters/types.py @@ -28,8 +28,13 @@ 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 # GAUGES USERS = Gauge.USERS.value + BYTES = Gauge.BYTES.value class MeterScope(BaseModel): diff --git a/api/ee/src/core/sandboxes/__init__.py b/api/ee/src/core/sandboxes/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/api/ee/src/core/sandboxes/dtos.py b/api/ee/src/core/sandboxes/dtos.py new file mode 100644 index 0000000000..bf01ec06c2 --- /dev/null +++ b/api/ee/src/core/sandboxes/dtos.py @@ -0,0 +1,27 @@ +from typing import Optional +from uuid import UUID + +from pydantic import BaseModel + + +class SandboxUsageDTO(BaseModel): + """Sandbox resource usage from a single provider billing event or poll.""" + + organization_id: UUID + provider: str # "e2b" | "daytona" + sandbox_id: str + + # Physical resource-second quantities (what we store and meter). + vcpu_seconds: int = 0 + ram_gib_seconds: int = 0 + disk_gib_seconds: int = 0 + gpu_seconds: int = 0 + + # E2B: delivery-id header value for idempotency dedupe. + delivery_id: Optional[str] = None + + +class SandboxUsageResult(BaseModel): + accepted: bool + delivery_id: Optional[str] = None + deduped: bool = False diff --git a/api/ee/src/core/sandboxes/exceptions.py b/api/ee/src/core/sandboxes/exceptions.py new file mode 100644 index 0000000000..e1c8ff9a4c --- /dev/null +++ b/api/ee/src/core/sandboxes/exceptions.py @@ -0,0 +1,8 @@ +class SandboxMeteringError(Exception): + pass + + +class SandboxWebhookSignatureError(SandboxMeteringError): + def __init__(self, message: str = "Webhook signature verification failed."): + self.message = message + super().__init__(message) diff --git a/api/ee/src/core/sandboxes/service.py b/api/ee/src/core/sandboxes/service.py new file mode 100644 index 0000000000..072346046d --- /dev/null +++ b/api/ee/src/core/sandboxes/service.py @@ -0,0 +1,248 @@ +"""Sandbox metering service: record_usage(), E2B signature verification, Daytona poll.""" + +import hashlib +import hmac +import math +from decimal import Decimal +from uuid import UUID + +import httpx + +from oss.src.utils.env import env +from oss.src.utils.locking import ( + acquire_lock, + release_lock, +) +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 + +log = get_module_logger(__name__) + +# Daytona poll lock +_DAYTONA_LOCK_NS = "sandboxes:daytona" +_DAYTONA_LOCK_KEY = "poll" +_DAYTONA_LOCK_TTL = 120 # 2 min — poll should complete well within this + +# Webhook redelivery dedup (E2B `e2b-delivery-id`). +_DELIVERY_DEDUP_NS = "sandboxes:e2b:delivery" +_DELIVERY_DEDUP_TTL = 48 * 60 * 60 # 48h — comfortably beyond E2B's redelivery window + +# Daytona reports decimal GB (10^9 bytes); meters store binary GiB (2^30 bytes). +_GB_TO_GIB = Decimal(10**9) / Decimal(2**30) + + +def _gb_to_gib_seconds(gb_seconds: Decimal) -> int: + """Convert Daytona GB-seconds to GiB-seconds (ceiling, matches E2B rounding).""" + if gb_seconds <= 0: + return 0 + return max(1, math.ceil(gb_seconds * _GB_TO_GIB)) + + +class SandboxMeteringService: + def __init__(self, *, meters_service: MetersService): + self.meters_service = meters_service + + # ------------------------------------------------------------------ + # Core: record usage for one sandbox billing event + # ------------------------------------------------------------------ + + 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). + + Deduped on usage.delivery_id via Redis SET NX (webhook redelivery + double-counting guard). Missing delivery_id skips dedup (best-effort). + """ + if usage.delivery_id: + claimed = await acquire_lock( + namespace=_DELIVERY_DEDUP_NS, + key=usage.delivery_id, + ttl=_DELIVERY_DEDUP_TTL, + ) + if not claimed: + log.info( + "[sandboxes] duplicate delivery_id=%s — skipping meter writes", + usage.delivery_id, + ) + return SandboxUsageResult( + accepted=True, delivery_id=usage.delivery_id, deduped=True + ) + + 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, + ) + + log.info( + "[sandboxes] recorded provider=%s sandbox=%s org=%s " + "vcpu_s=%d ram_s=%d disk_s=%d gpu_s=%d", + usage.provider, + usage.sandbox_id, + org_id, + usage.vcpu_seconds, + usage.ram_gib_seconds, + usage.disk_gib_seconds, + usage.gpu_seconds, + ) + + return SandboxUsageResult(accepted=True, delivery_id=usage.delivery_id) + + # ------------------------------------------------------------------ + # E2B: webhook signature verification + # ------------------------------------------------------------------ + + def verify_e2b_signature(self, *, raw_body: bytes, signature_header: str) -> bool: + """Verify E2B webhook HMAC signature against env.e2b.webhook_secret. + + Mirrors the Stripe pattern: the operator sets E2B_WEBHOOK_SECRET and + registers the webhook with E2B out-of-band; no secret machinery here. + E2B signs: sha256(secret + raw_body) → hex, sent in e2b-signature header. + """ + secret = env.e2b.webhook_secret + if not secret: + raise SandboxWebhookSignatureError("E2B_WEBHOOK_SECRET is not configured.") + + expected = hmac.new( + secret.encode(), + raw_body, + hashlib.sha256, + ).hexdigest() + try: + return hmac.compare_digest(expected, signature_header.strip()) + except Exception: + return False + + # ------------------------------------------------------------------ + # Daytona: periodic poll + # ------------------------------------------------------------------ + + async def daytona_poll( + self, + *, + org_id: UUID, + api_key: str, + analytics_url: str, + daytona_organization_id: str, + period_start: str, + period_end: str, + ) -> None: + """Poll Daytona usage/aggregated and adjust meters. + + Daytona returns cumulative totals for the window → SET absolute + value (delta = total - current) so re-polls are idempotent. + """ + headers = { + "Authorization": f"Bearer {api_key}", + "X-Daytona-Organization-ID": daytona_organization_id, + } + url = f"{analytics_url}/organization/{daytona_organization_id}/usage/aggregated" + params = {"from": period_start, "to": period_end} + + try: + async with httpx.AsyncClient(timeout=15) as client: + resp = await client.get(url, headers=headers, params=params) + resp.raise_for_status() + data = resp.json() + except Exception as exc: + log.error("[sandboxes] Daytona poll failed: %s", exc) + return + + vcpu_seconds = int(data.get("totalCPUSeconds") or 0) + ram_gib_seconds = _gb_to_gib_seconds( + Decimal(str(data.get("totalRAMGBSeconds") or 0)) + ) + disk_gib_seconds = _gb_to_gib_seconds( + Decimal(str(data.get("totalDiskGBSeconds") or 0)) + ) + gpu_seconds = int(data.get("totalGPUSeconds") or 0) + + log.info( + "[sandboxes] Daytona poll org=%s vcpu_s=%d ram_s=%d disk_s=%d gpu_s=%d", + org_id, + vcpu_seconds, + ram_gib_seconds, + disk_gib_seconds, + gpu_seconds, + ) + + usage = SandboxUsageDTO( + organization_id=org_id, + provider="daytona", + sandbox_id="__aggregate__", + vcpu_seconds=vcpu_seconds, + ram_gib_seconds=ram_gib_seconds, + disk_gib_seconds=disk_gib_seconds, + gpu_seconds=gpu_seconds, + ) + await self.record_usage(usage) + + async def run_daytona_poll( + self, + *, + org_id: UUID, + api_key: str, + analytics_url: str, + daytona_organization_id: str, + period_start: str, + period_end: str, + ) -> bool: + """Acquire lock, run poll, release. Returns True if poll ran.""" + lock_owner = await acquire_lock( + namespace=_DAYTONA_LOCK_NS, + key=_DAYTONA_LOCK_KEY, + ttl=_DAYTONA_LOCK_TTL, + ) + if not lock_owner: + log.info("[sandboxes] Daytona poll already in progress, skipping") + return False + + try: + await self.daytona_poll( + org_id=org_id, + api_key=api_key, + analytics_url=analytics_url, + daytona_organization_id=daytona_organization_id, + period_start=period_start, + period_end=period_end, + ) + return True + finally: + await release_lock( + namespace=_DAYTONA_LOCK_NS, + key=_DAYTONA_LOCK_KEY, + owner=lock_owner, + ) diff --git a/api/ee/src/core/storage/__init__.py b/api/ee/src/core/storage/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/api/ee/src/core/storage/adapters.py b/api/ee/src/core/storage/adapters.py new file mode 100644 index 0000000000..ac2caae974 --- /dev/null +++ b/api/ee/src/core/storage/adapters.py @@ -0,0 +1,53 @@ +"""Storage-size adapter: authoritative per-org byte count via the shared object store. + +Mount object keys carry no org_id component (see MountsService._storage_key), so an +org's authoritative size is the sum over that org's projects' `mounts//` +prefixes, not a single org-level ListObjectsV2 scan. +""" + +from uuid import UUID + +from oss.src.utils.logging import get_module_logger +from oss.src.utils.env import env +from oss.src.core.store.storage import ObjectStore +from oss.src.services.db_manager import fetch_projects_by_organization + +from ee.src.core.storage.paths import project_prefix + +log = get_module_logger(__name__) + + +async def get_org_storage_bytes(*, org_id: UUID) -> int: + """Authoritative total bytes for one org: sum of its projects' mount prefixes.""" + if not env.store.enabled: + log.debug("[storage] store not configured; returning 0") + return 0 + + projects = await fetch_projects_by_organization(str(org_id)) + if not projects: + return 0 + + store = ObjectStore( + endpoint_url=env.store.endpoint_url, + access_key=env.store.access_key, + secret_key=env.store.secret_key, + region=env.store.region, + sts_endpoint_url=env.store.sts_endpoint_url, + signing_key=env.store.signing_key, + ) + bucket = env.store.bucket or "" + + total = 0 + for project in projects: + prefix = project_prefix(project.id) + try: + objects = await store.list_objects_v2(bucket=bucket, prefix=prefix) + total += sum(size for _, size in objects) + except Exception: + log.warning( + "[storage] size query failed for org=%s project=%s", + org_id, + project.id, + exc_info=True, + ) + return total diff --git a/api/ee/src/core/storage/paths.py b/api/ee/src/core/storage/paths.py new file mode 100644 index 0000000000..ee1d402ce9 --- /dev/null +++ b/api/ee/src/core/storage/paths.py @@ -0,0 +1,14 @@ +from uuid import UUID + +from oss.src.utils.env import env + + +def project_prefix(project_id: UUID) -> str: + """S3/SeaweedFS prefix for one project's mounts: [/]mounts//. + + Matches MountsService._storage_key (api/oss/src/core/mounts/service.py): + org_id is not a key component, so an org's bytes are the sum over its projects. + """ + ns = (env.store.namespace or "").strip("/") + base = f"{ns}/mounts/{project_id}" if ns else f"mounts/{project_id}" + return f"{base}/" diff --git a/api/ee/src/core/storage/reconcile.py b/api/ee/src/core/storage/reconcile.py new file mode 100644 index 0000000000..b9503f5e53 --- /dev/null +++ b/api/ee/src/core/storage/reconcile.py @@ -0,0 +1,69 @@ +"""Periodic reconcile job: reads authoritative storage size for all orgs, corrects gauge.""" + +from typing import Optional, Callable, Awaitable + +from oss.src.utils.logging import get_module_logger + +log = get_module_logger(__name__) + + +async def run_storage_reconcile( + *, + renew: Optional[Callable[[], Awaitable[bool]]] = None, +) -> None: + """Iterate over all active orgs and reconcile their storage gauge. + + Gated on is_ee() + env.store.reconcile_enabled. + """ + from oss.src.utils.common import is_ee + from oss.src.utils.env import env + + if not is_ee(): + log.debug("[storage] reconcile skipped (not EE)") + return + + if not env.store.reconcile_enabled: + log.debug( + "[storage] reconcile skipped (AGENTA_STORE_RECONCILE_ENABLED not set)" + ) + return + + if not env.store.enabled: + log.debug("[storage] reconcile skipped (no storage provider configured)") + return + + try: + from ee.src.dbs.postgres.subscriptions.dao import SubscriptionsDAO + + dao = SubscriptionsDAO() + subscriptions = await dao.list_active() + except Exception: + log.error("[storage] failed to list active subscriptions", exc_info=True) + return + + from ee.src.core.storage.service import reconcile_org_storage + + for sub in subscriptions: + try: + org_id = sub.organization_id + if org_id is None: + continue + await reconcile_org_storage(org_id=org_id) + except Exception: + log.warning( + "[storage] reconcile failed for org=%s", + getattr(sub, "organization_id", "?"), + exc_info=True, + ) + + if renew: + try: + ok = await renew() + if not ok: + log.error("[storage] lock renewal rejected; stopping reconcile") + return + except Exception: + log.error( + "[storage] lock renewal error; stopping reconcile", exc_info=True + ) + return diff --git a/api/ee/src/core/storage/service.py b/api/ee/src/core/storage/service.py new file mode 100644 index 0000000000..b3eff42e3b --- /dev/null +++ b/api/ee/src/core/storage/service.py @@ -0,0 +1,92 @@ +"""Storage gauge: delta tracking + periodic reconciliation.""" + +from uuid import UUID + +from oss.src.utils.logging import get_module_logger + +log = get_module_logger(__name__) + + +async def record_storage_delta( + *, + org_id: UUID, + delta_bytes: int, +) -> bool: + """Adjust the storage gauge by `delta_bytes` (positive=write, negative=delete). + + Returns True if allowed (under quota), False if capped. Fails open on errors. + """ + from oss.src.utils.common import is_ee + + if not is_ee(): + return True + + try: + from ee.src.core.access.entitlements.types import Gauge + from ee.src.core.access.entitlements.service import check_entitlements + from ee.src.core.meters.types import MeterScope + + scope = MeterScope(organization_id=org_id) + allowed, _, _ = await check_entitlements( + key=Gauge.BYTES, + delta=delta_bytes, + scope=scope, + ) + return allowed + except Exception: + log.warning( + "[storage] record_storage_delta failed; failing open", exc_info=True + ) + return True + + +async def reconcile_org_storage( + *, + org_id: UUID, +) -> None: + """Reconcile storage gauge: read authoritative size, set gauge via delta.""" + from oss.src.utils.common import is_ee + + if not is_ee(): + return + + try: + from ee.src.core.access.entitlements.types import Gauge + from ee.src.core.access.entitlements.service import ( + check_entitlements, + _meters_service, + ) + from ee.src.core.meters.types import MeterScope, MeterPeriod, Meters + from ee.src.core.storage.adapters import get_org_storage_bytes + + authoritative = await get_org_storage_bytes(org_id=org_id) + + scope = MeterScope(organization_id=org_id) + period = MeterPeriod() + + meters = await _meters_service().fetch( + scope=scope, + key=Meters.BYTES, + period=period, + ) + current = (meters[0].value if meters else 0) or 0 + delta = authoritative - current + + if delta == 0: + return + + await check_entitlements( + key=Gauge.BYTES, + delta=delta, + scope=scope, + period=period, + ) + log.info( + "[storage] reconciled org=%s authoritative=%d current=%d delta=%d", + org_id, + authoritative, + current, + delta, + ) + except Exception: + log.warning("[storage] reconcile_org_storage failed", exc_info=True) diff --git a/api/ee/src/core/storage/types.py b/api/ee/src/core/storage/types.py new file mode 100644 index 0000000000..a5679966ba --- /dev/null +++ b/api/ee/src/core/storage/types.py @@ -0,0 +1,12 @@ +class StorageError(Exception): + pass + + +class StorageQuotaExceeded(StorageError): + def __init__(self, org_id, bytes_used: int, bytes_limit: int): + self.org_id = org_id + self.bytes_used = bytes_used + self.bytes_limit = bytes_limit + super().__init__( + f"Storage quota exceeded for org {org_id}: {bytes_used} >= {bytes_limit}" + ) diff --git a/api/ee/src/core/subscriptions/interfaces.py b/api/ee/src/core/subscriptions/interfaces.py index 2c47a9a302..5dd6625a14 100644 --- a/api/ee/src/core/subscriptions/interfaces.py +++ b/api/ee/src/core/subscriptions/interfaces.py @@ -54,3 +54,9 @@ async def update( - Optional[SubscriptionDTO]: The updated subscription if found, else None. """ raise NotImplementedError + + async def list_active( + self, + ) -> list[SubscriptionDTO]: + """Return all active subscriptions (active=True).""" + raise NotImplementedError diff --git a/api/ee/src/dbs/postgres/subscriptions/dao.py b/api/ee/src/dbs/postgres/subscriptions/dao.py index 58f7f26f03..988c65295d 100644 --- a/api/ee/src/dbs/postgres/subscriptions/dao.py +++ b/api/ee/src/dbs/postgres/subscriptions/dao.py @@ -59,6 +59,18 @@ async def read( return subscription_dto + async def list_active( + self, + ) -> list[SubscriptionDTO]: + async with self.engine.session() as session: + result = await session.execute( + select(SubscriptionDBE).where( + SubscriptionDBE.active == True, # noqa: E712 + ) + ) + dbes = result.scalars().all() + return [map_dbe_to_dto(dbe) for dbe in dbes] + async def update( self, *, diff --git a/api/ee/src/main.py b/api/ee/src/main.py index 3261c2ebd0..df93f80711 100644 --- a/api/ee/src/main.py +++ b/api/ee/src/main.py @@ -32,6 +32,8 @@ from ee.src.apis.fastapi.events.router import EventsRouter, EventsRetentionRouter from ee.src.apis.fastapi.sessions.records.router import RecordsRetentionRouter from ee.src.apis.fastapi.organizations.router import router as organization_router +from ee.src.apis.fastapi.sandboxes.router import SandboxMeteringRouter +from ee.src.core.sandboxes.service import SandboxMeteringService from ee.src.core.access.entitlements.service import bootstrap_entitlements_services # DBS -------------------------------------------------------------------------- @@ -93,6 +95,10 @@ subscriptions_dao=subscriptions_dao, ) +sandboxes_service = SandboxMeteringService( + meters_service=meters_service, +) + # Wire entitlements module against the freshly-built services so the # `BillingRouter` and the entitlements helper share one instance each. bootstrap_entitlements_services( @@ -109,6 +115,10 @@ meters_service=meters_service, ) +sandboxes_router = SandboxMeteringRouter( + sandboxes_service=sandboxes_service, +) + spans_retention_router = SpansRetentionRouter( tracing_retention_service=tracing_retention_service, ) @@ -172,6 +182,20 @@ def extend_main(app: FastAPI): include_in_schema=False, ) + app.include_router( + router=sandboxes_router.router, + prefix="/webhooks/sandboxes", + tags=["Sandbox Metering"], + include_in_schema=False, + ) + + app.include_router( + router=sandboxes_router.admin_router, + prefix="/admin/sandboxes", + tags=["Admin"], + include_in_schema=False, + ) + app.include_router( router=events_router.router, prefix="/events", diff --git a/api/oss/src/services/db_manager.py b/api/oss/src/services/db_manager.py index c1f04ac5a9..4c905d3185 100644 --- a/api/oss/src/services/db_manager.py +++ b/api/oss/src/services/db_manager.py @@ -115,6 +115,27 @@ async def fetch_projects_by_workspace( return result.scalars().all() +async def fetch_projects_by_organization( + organization_id: str, +) -> List[ProjectDB]: + """ + Retrieve all projects that belong to an organization ordered by creation date. + Args: + organization_id (str): Organization identifier. + Returns: + List[ProjectDB]: Projects scoped to the organization. + """ + + engine = get_transactions_engine() + async with engine.session() as session: + result = await session.execute( + select(ProjectDB) + .filter(ProjectDB.organization_id == uuid.UUID(organization_id)) + .order_by(ProjectDB.created_at.asc()) + ) + return result.scalars().all() + + async def get_project_by_workspace( workspace_id: str, *, diff --git a/api/oss/src/utils/env.py b/api/oss/src/utils/env.py index 121f0a2ce4..901f88ef0d 100644 --- a/api/oss/src/utils/env.py +++ b/api/oss/src/utils/env.py @@ -555,11 +555,37 @@ def enabled(self) -> bool: class DaytonaConfig(BaseModel): api_key: str | None = os.getenv("DAYTONA_API_KEY") api_url: str | None = os.getenv("DAYTONA_API_URL") + # Usage/aggregated metering poll (sandbox compute meters). + analytics_url: str | None = os.getenv("DAYTONA_ANALYTICS_URL") + organization_id: str | None = os.getenv("DAYTONA_ORGANIZATION_ID") snapshot: str | None = os.getenv("DAYTONA_SNAPSHOT") target: str | None = os.getenv("DAYTONA_TARGET") model_config = ConfigDict(extra="ignore") + @property + def enabled(self) -> bool: + return bool(self.api_key and self.organization_id and self.analytics_url) + + +# --------------------------------------------------------------------------- +# e2b — sandbox compute metering (webhook-fed). +# --------------------------------------------------------------------------- + + +class E2BConfig(BaseModel): + api_key: str | None = os.getenv("E2B_API_KEY") + api_url: str = os.getenv("E2B_API_URL") or "https://api.e2b.app" + webhook_url: str | None = os.getenv("E2B_WEBHOOK_URL") + webhook_secret: str | None = os.getenv("E2B_WEBHOOK_SECRET") + + model_config = ConfigDict(extra="ignore") + + @property + def enabled(self) -> bool: + """E2B webhook receiver enabled only if API key and webhook secret are present""" + return bool(self.api_key and self.webhook_secret) + # --------------------------------------------------------------------------- # docker @@ -883,6 +909,12 @@ class StoreConfig(BaseModel): jwt_private_key: str | None = os.getenv("AGENTA_STORE_JWT_PRIVATE_KEY") jwt_issuer: str = os.getenv("AGENTA_STORE_JWT_ISSUER") or "http://api:8000" + # Gate for the periodic storage-gauge reconcile job (EE): off by default since it + # walks every org's prefix via ListObjectsV2/filer-list. + reconcile_enabled: bool = ( + os.getenv("AGENTA_STORE_RECONCILE_ENABLED") or "false" + ).lower() in _TRUTHY + model_config = ConfigDict(extra="ignore") @property @@ -1239,6 +1271,7 @@ class EnvironSettings(BaseModel): crisp: CrispConfig = CrispConfig() daytona: DaytonaConfig = DaytonaConfig() docker: DockerConfig = DockerConfig() + e2b: E2BConfig = E2BConfig() identity: IdentityConfig = IdentityConfig() llm: LLMConfig = LLMConfig() loops: LoopsConfig = LoopsConfig() diff --git a/docs/designs/sandbox-metering/NAMING.md b/docs/designs/sandbox-metering/NAMING.md new file mode 100644 index 0000000000..33dbcf6144 --- /dev/null +++ b/docs/designs/sandbox-metering/NAMING.md @@ -0,0 +1,20 @@ +# Sandbox meter key naming (locked) + +Scheme: `SANDBOX___SECONDS` — plain resource token + unit token. + +| Resource | Unit | Counter key | value | +|----------|------|-----------------------------|------------------------------| +| CPU | CORE | `SANDBOX_CPU_CORE_SECONDS` | `sandbox_cpu_core_seconds` | +| RAM | GIBI | `SANDBOX_RAM_GIBI_SECONDS` | `sandbox_ram_gibi_seconds` | +| SSD | GIBI | `SANDBOX_SSD_GIBI_SECONDS` | `sandbox_ssd_gibi_seconds` | +| GPU | CORE | `SANDBOX_GPU_CORE_SECONDS` | `sandbox_gpu_core_seconds` | + +- 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.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`, + `SANDBOX_SSD_GIBI_CREDITS`, `SANDBOX_GPU_CORE_CREDITS`) + total + `SANDBOX_CREDITS`. diff --git a/docs/designs/sandbox-metering/TRACK_B_FINDINGS.md b/docs/designs/sandbox-metering/TRACK_B_FINDINGS.md new file mode 100644 index 0000000000..5e5ddfe4a1 --- /dev/null +++ b/docs/designs/sandbox-metering/TRACK_B_FINDINGS.md @@ -0,0 +1,154 @@ +# Track B — new metering, measurement only + +Consolidates `feat/sandbox-metering-phase-1` (sandbox compute providers + sink) and +`feat/sandbox-metering-phase-4` (storage gauge) into `feat/metering-track-b`, on top of +`feat/add-sandbox-metering` (Track A) + big-agents. Pure measurement: nothing added to +`REPORTS`, no Stripe pricing. + +## What was brought in + +From phase-1 (renamed `sandbox_metering` -> `sandboxes` at both core and API level): +- `api/ee/src/core/sandboxes/{__init__,dtos,exceptions,service}.py` +- `api/ee/src/apis/fastapi/sandboxes/{__init__,models,router}.py` +- `main.py` sandboxes service/router wiring; `entrypoints/routers.py` best-effort E2B + webhook registration in `lifespan()`. + +From phase-4 (storage gauge), rewired to the existing shared store config: +- `api/ee/src/core/storage/{__init__,types,paths,adapters,service,reconcile}.py` +- `subscriptions/interfaces.py` + `dbs/postgres/subscriptions/dao.py`: `list_active()` +- `billing/router.py`: `POST /admin/billing/storage/reconcile` + + `.../storage/reconcile/unlock`, mirroring the existing `usage/report` lock pattern. + +The 4 junk files listed in the task (`.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. + +## Renames + +`sandbox_metering` -> `sandboxes` everywhere: directory names, import paths, log-tag +prefixes (`[sandbox_metering]` -> `[sandboxes]`), Redis key namespaces +(`sandbox_metering:e2b`/`:daytona` -> `sandboxes:e2b`/`:daytona`), FastAPI +`operation_id`s, and the router constructor kwarg (`sandbox_metering_service` -> +`sandboxes_service`). Class names `SandboxMeteringService`/`SandboxMeteringRouter` were +kept (descriptive class names, not paths). Mount points: +`/webhooks/sandboxes` (public E2B receiver) and `/admin/sandboxes` (Daytona poll +trigger). `core/sandbox/` (phase-2, singular) is Track C and was not touched. + +## Meter key naming (final) + +Went through two naming revisions mid-task; the committed state uses the final, +simplest scheme — plain 3-letter resource tokens, no unit token: + +| Counter key | value | +|-------------------------|--------------------------| +| `SANDBOX_CPU_SECONDS` | `sandbox_cpu_seconds` | +| `SANDBOX_RAM_SECONDS` | `sandbox_ram_seconds` | +| `SANDBOX_SSD_SECONDS` | `sandbox_ssd_seconds` | +| `SANDBOX_GPU_SECONDS` | `sandbox_gpu_seconds` | + +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 +meter-delta mapping. A stray `docs/designs/sandbox-metering/NAMING.md` on disk +(untracked, authored elsewhere) documents an earlier, superseded 4-letter-token variant +of this scheme — not updated as part of this task since it wasn't in Track B's file +list; the code is the source of truth. + +## Entitlements (measurement only) + +- `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. + `SandboxMeteringService.record_usage()` also calls `check_entitlements` with + `cache=False` per delta (Layer-2 atomic adjust) purely to persist the meter row; the + call fails open on error. +- Storage caps carried over from phase-4's design: HOBBY 1 GiB (free=limit=strict), + 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.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. + +## Storage wiring to `env.store` + +Phase-4 originally introduced its own `env.agenta.storage.*` (`StorageConfig`) with a +provider string, bucket, endpoint, and a boto3/httpx-based size adapter. Per the +reconciliation facts, that duplicate config was **deleted** and the storage gauge now +consumes the existing `env.store` (`StoreConfig` in `api/oss/src/utils/env.py`) and the +existing S3-compatible client: + +- `storage/adapters.py`: `get_org_storage_bytes()` builds an `ObjectStore` (from + `oss.src.core.store.storage`, the same client `mounts` already uses) from + `env.store.{endpoint_url,access_key,secret_key,region,sts_endpoint_url,signing_key}` + and sums `list_objects_v2(bucket=env.store.bucket, prefix=org_prefix(org_id))`. Dropped + the separate boto3/httpx-per-provider code path entirely — SeaweedFS vs. real-S3 + selection already lives in `ObjectStore.is_seaweedfs` (keyed on `signing_key` + presence), so `storage/types.py`'s now-dead `StorageProvider` enum was removed too. +- `storage/reconcile.py`: gate changed from `env.agenta.storage.reconcile_enabled` / + `.enabled` to `env.store.reconcile_enabled` / `env.store.enabled`. +- `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.BYTES` / `Meters.BYTES` via + `check_entitlements`, no direct env access. + +## Migration + +`api/ee/databases/postgres/migrations/core_ee/versions/ee0000000004_add_sandbox_and_storage_meters.py`, +`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`, `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 +`ee0000000000 -> ...0001 -> ...0002 -> ...0003 -> ...0004` is linear, single head. + +## Deliberately left out (per task scope) + +- **Records metering** (phase-3 / `RECORDS_INGESTED`): already fully on big-agents + (Counter member, per-plan quotas, `ee0000000003` migration). Not touched. Size-cap + + two-layer wiring for records is separately deferred pending the transcripts->records + rename. +- **REPORTS / Stripe billing**: nothing added for sandboxes or storage. Storage billing + stays commented/deferred (phase-4's `DEFAULT_CATALOG` had `"storage": {"type": + "tiered", "tiers": []}` TODO(pricing) placeholders per plan — not brought in, since + `DEFAULT_CATALOG` pricing entries are a Track A/billing concern, not measurement). + Track A's unrelated retention/pricing changes (WEEKLY->MONTHLY retention bumps, PRO + price bump, per-seat pricing) that were tangled into both phase branches' diffs of + `entitlements/types.py` were **not** cherry-picked — only the pure sandbox-counter and + storage-gauge additions were extracted. +- **`core/sandbox/` (phase-2, singular)**: Track C, not in scope here. +- **Credits keys**: Track C; no `SANDBOX_*_CREDITS` or aggregate credits meter added. + +## Verification + +- `cd api && ruff format . && ruff check .` — clean, no errors. +- `uv run python3 -c "import ee.src.core.sandboxes.service, ee.src.core.storage.service, ee.src.main, ..."` + — all new modules import cleanly, including `ee.src.main` (exercises the full service + + router instantiation and mount wiring). `import entrypoints.routers` fails on + `alembic.util.exc.CommandError: No 'script_location' key found` — a pre-existing + local-env limitation (needs the docker-compose env loaded per repo `AGENTS.md`), not + caused by this change; unaffected by anything in this branch. +- Confirmed `env.e2b.enabled`, `env.daytona.enabled`, `env.store.reconcile_enabled` + properties evaluate without error, and `Meters`/`Counter`/`Gauge`/`REPORTS` reflect the + expected final state at runtime. + +## Commits + +1. `feat(metering): sandbox compute meters + providers (sandboxes domain)` — phase-1 + core+API files under the `sandboxes` rename, entitlements/meters key additions + (first naming pass), `env.py` E2B/Daytona config, `main.py`/`routers.py` wiring. +2. `fixup(metering): rename sandbox meter keys to SANDBOX_{CPU,RAM,SSD,GPU}_SECONDS` — + applies the final naming decision across the Counter enum, Meters mirror, quotas, + constraints, and the sandboxes service; adds the `ee0000000004` migration (was + deferred to this commit since the migration was written after the naming settled). +3. `feat(metering): storage gauge on env.store` — phase-4 storage domain rewired off + the duplicate `StorageConfig` onto `env.store` + the existing `ObjectStore` client, + `list_active()` on subscriptions, billing router admin reconcile endpoints. + +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-b-metering/specs.md b/docs/designs/track-b-metering/specs.md new file mode 100644 index 0000000000..e48ac762e4 --- /dev/null +++ b/docs/designs/track-b-metering/specs.md @@ -0,0 +1,94 @@ +# Track B — wallet-debit meter definitions (metering) + +Track B owns **metering**: measurement pipelines and the meter **definitions**. +It does not populate billing meters and does not charge. The boundary with the +other tracks: + +- **Track B (this track)** — meters exist as defined things: enum members, + quotas, DB enum migration, and measurement-only collection + (webhook/poll → usage DTO). No credit math, no sink calls, no gating. +- **Track C (billing)** — populates the debit meters from measured usage + (rate table + sink), gates on them, and builds the wallet. +- **Track D (BYOS)** — bring-your-own secrets for sandbox/gateway providers and + `secret_origin` zero-rating on top of C. + +Companion design docs (big-agents-audit collection): `monetization-integration.md` +(audited facts), `tiers-and-unified-wallet.md` (wallet + tier model, meter +taxonomy §6), `wallet-enforcement-matrix.md` (credit/debit/measure/enforce per +family). + +## Meter taxonomy — end state + +Direction is uniform: every meter measures consumption **out** of the wallet, +so every meter is `*_DEBITS`. Money **in** lives in the `wallet_credits` table +(Track C). The raw `*_SECONDS` breakdown ("what did RAM cost me") is +cost-explainer data that belongs in traces/analytics, **not** in billing +meters — it is deleted here, including its migration labels. + +`Counter` members (in `api/ee/src/core/access/entitlements/types.py`), mirrored +in `Meters` (`api/ee/src/core/meters/types.py`): + +| Member | Value | Role | +|---|---|---| +| `SANDBOX_CPU_CORE_DEBITS` | `sandbox_cpu_core_debits` | per-dimension visibility | +| `SANDBOX_RAM_GIBI_DEBITS` | `sandbox_ram_gibi_debits` | per-dimension visibility | +| `SANDBOX_SSD_GIBI_DEBITS` | `sandbox_ssd_gibi_debits` | per-dimension visibility | +| `SANDBOX_GPU_CORE_DEBITS` | `sandbox_gpu_core_debits` | per-dimension visibility | +| `SANDBOX_DEBITS` | `sandbox_debits` | family roll-up | +| `LLM_DEBITS` | `llm_debits` | family roll-up (sink lands later) | +| `GATEWAY_DEBITS` | `gateway_debits` | family roll-up (sink lands later) | +| `WALLET_DEBITS` | `wallet_debits` | cross-family grand total; what the wallet gate reads | + +Removed everywhere: `SANDBOX_CPU_CORE_SECONDS`, `SANDBOX_RAM_GIBI_SECONDS`, +`SANDBOX_SSD_GIBI_SECONDS`, `SANDBOX_GPU_CORE_SECONDS`. + +Unchanged: `CREDITS_CONSUMED` (the legacy hosted-keys-gate counter — a +different mechanism; no migration of it), `TRACES_INGESTED` and the rest of the +our-infra pay-as-you-go set, and `REPORTS` (wallet debits are prepaid — billed +at top-up time — and are **never** reported to Stripe in arrears, so no +`*_DEBITS` key ever enters `REPORTS`). + +## File-by-file end state + +1. `api/ee/src/core/access/entitlements/types.py` + - `Counter`: remove the 4 `*_SECONDS` members; add the 8 `*_DEBITS` members + above. + - `DEFAULT_ENTITLEMENTS`: for every plan, remove the 4 `*_SECONDS` quotas and + add the 8 `*_DEBITS` quotas as bare `Quota(period=Period.MONTHLY)` (no + free/limit numbers — `TODO(pricing)`; the wallet supplies the ceiling in + Track C). + - `DEFAULT_CATALOG` comments: the pricing TODO notes reference + `sandbox_debits` and state the prepaid/no-REPORTS rule. + - Any `CONSTRAINTS`/read-only counter list that enumerates the `*_SECONDS` + members: swap to the `*_DEBITS` set. +2. `api/ee/src/core/meters/types.py` — mirror the same member swap in `Meters`. +3. Migration `ee0000000005` (new file in this track, + `api/ee/databases/postgres/migrations/core_ee/versions/`) — docstring + "add sandbox + wallet debit meters to meters_type"; `ADD VALUE IF NOT EXISTS` + for the 8 uppercase `*_DEBITS` labels; revises `ee0000000004`. +4. Migration `ee0000000004_add_sandbox_and_storage_meters.py` — remove the 4 + `*_SECONDS` `ADD VALUE` lines; keep `BYTES`. (Branch-local migration, not + shipped; editing in place is correct — no follow-up migration.) +5. `api/ee/src/core/sandboxes/service.py` — measurement-only: + - `record_usage()` keeps webhook/poll parsing, delivery-id dedup, and the + usage DTO handling; **remove** the `meter_deltas` block that adjusted the + 4 `*_SECONDS` counters. + - **No sink call** — `record_usage_credits`/`record_usage_debits` and its + import belong to Track C (the one-line populate patch lands there). + - Remove imports left unused (`check_entitlements`, `Counter`, `MeterScope`, + the sink import). + +## Also in scope: Daytona poll idempotency + +The Daytona poller delta-adjusts without diffing provider-cumulative totals +against the current meter value; the Redis lock guards concurrent polls, not +double-counting across overlapping windows (monetization-integration.md F7). +This is a measurement-correctness bug and must be fixed before a wallet makes +every debit balance-affecting: persist/diff the provider's cumulative totals and +adjust only the positive delta. Verify current behavior first — recent Track B +commits touched adjacent conversion code. + +## Out of scope for Track B + +`sink.py`, `gating.py`, `credits.py`/`debits.py`, gating tests, any wallet +table/cron/webhook/tier work (all Track C); vault/secret work (Track D). diff --git a/docs/designs/track-b-metering/tasks.md b/docs/designs/track-b-metering/tasks.md new file mode 100644 index 0000000000..3fa3795a03 --- /dev/null +++ b/docs/designs/track-b-metering/tasks.md @@ -0,0 +1,45 @@ +# Track B — tasks + +Execution order for `feat/metering-track-b`. Safety tag before re-partition: +`safety-track-b-pre-repartition`. + +## B1 — enum + quota swap +- [ ] `api/ee/src/core/access/entitlements/types.py`: remove the 4 `*_SECONDS` + `Counter` members; add the 8 `*_DEBITS` members (values lowercase). +- [ ] Same file: per plan, remove `*_SECONDS` quotas; add the 8 `*_DEBITS` + quotas as `Quota(period=Period.MONTHLY)` with the `TODO(pricing)` note. +- [ ] Same file: update `DEFAULT_CATALOG` pricing-TODO comments + (`sandbox_debits`, prepaid/no-REPORTS). +- [ ] Grep the file for any remaining `*_SECONDS` reference (constraints / + read-only lists) and swap to the `*_DEBITS` set. +- [ ] `api/ee/src/core/meters/types.py`: mirror the member swap in `Meters`. + +## B2 — migrations +- [ ] Add `ee0000000005` migration: 8 uppercase `*_DEBITS` `ADD VALUE IF NOT + EXISTS` labels; docstring "add sandbox + wallet debit meters to + meters_type"; revises `ee0000000004`. +- [ ] `ee0000000004_add_sandbox_and_storage_meters.py`: delete the 4 + `*_SECONDS` `ADD VALUE` lines; keep `BYTES`. + +## B3 — measurement-only service +- [ ] `api/ee/src/core/sandboxes/service.py`: remove the `meter_deltas` + `*_SECONDS` block from `record_usage()`; do NOT add a sink call; drop + now-unused imports. + +## B4 — Daytona poll idempotency +- [ ] Read the current poll path (`service.py` / `router.py`); confirm whether + delta adjustment diffs provider-cumulative totals against the meter. +- [ ] If confirmed: make the poll idempotent (diff cumulative totals, adjust + the positive delta only); unit-test overlapping-window redelivery. + +## B5 — verify + commit +- [ ] `grep -rn "SECONDS" api/ee/src/core api/ee/databases` → no sandbox + seconds meter references remain. +- [ ] `grep -rn "record_usage_credits\|record_usage_debits\|sink" + api/ee/src/core/sandboxes/service.py` → empty. +- [ ] ruff format + check on touched files. +- [ ] Run the existing sandbox metering unit tests (skip gating tests — they + live in Track C and are expected red here only if they import removed + members; they should not, since sink/gating files are absent on B). +- [ ] Commit on `feat/metering-track-b` (conventional message, e.g. + `refactor(metering): define wallet debit meters, drop seconds meters`).