diff --git a/api/services/billing_service.py b/api/services/billing_service.py index 2ee7179f432..4a829590f16 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -50,9 +50,26 @@ class QuotaReleaseResult(TypedDict): released: int +class QuotaBalanceResult(TypedDict): + available: int + reserved: int + quota: int + usage: int + + +class QuotaConsumeCappedResult(TypedDict): + deducted: int + available: int + reserved: int + quota: int + usage: int + + _quota_reserve_adapter = TypeAdapter(QuotaReserveResult) _quota_commit_adapter = TypeAdapter(QuotaCommitResult) _quota_release_adapter = TypeAdapter(QuotaReleaseResult) +_quota_balance_adapter = TypeAdapter(QuotaBalanceResult) +_quota_consume_capped_adapter = TypeAdapter(QuotaConsumeCappedResult) class _TenantFeatureQuota(TypedDict): @@ -176,6 +193,7 @@ class DismissNotificationDict(TypedDict): class BillingService: base_url = os.environ.get("BILLING_API_URL", "BILLING_API_URL") + quota_base_url = os.environ.get("BILLING_QUOTA_API_URL") or base_url secret_key = os.environ.get("BILLING_API_SECRET_KEY", "BILLING_API_SECRET_KEY") compliance_download_rate_limiter = RateLimiter("compliance_download_rate_limiter", 4, 60) @@ -215,12 +233,18 @@ class BillingService: def get_quota_info(cls, tenant_id: str) -> TenantFeatureQuotaInfo: params = {"tenant_id": tenant_id} return _tenant_feature_quota_info_adapter.validate_python( - cls._send_request("GET", "/quota/info", params=params) + cls._send_quota_request("GET", "/quota/info", params=params) ) @classmethod def quota_reserve( - cls, tenant_id: str, feature_key: str, request_id: str, amount: int = 1, meta: dict | None = None + cls, + tenant_id: str, + feature_key: str, + request_id: str, + amount: int = 1, + meta: dict | None = None, + bucket: str = "", ) -> QuotaReserveResult: """Reserve quota before task execution.""" payload: dict = { @@ -229,13 +253,21 @@ class BillingService: "request_id": request_id, "amount": amount, } + if bucket: + payload["bucket"] = bucket if meta: payload["meta"] = meta - return _quota_reserve_adapter.validate_python(cls._send_request("POST", "/quota/reserve", json=payload)) + return _quota_reserve_adapter.validate_python(cls._send_quota_request("POST", "/quota/reserve", json=payload)) @classmethod def quota_commit( - cls, tenant_id: str, feature_key: str, reservation_id: str, actual_amount: int, meta: dict | None = None + cls, + tenant_id: str, + feature_key: str, + reservation_id: str, + actual_amount: int, + meta: dict | None = None, + bucket: str = "", ) -> QuotaCommitResult: """Commit a reservation with actual consumption.""" payload: dict = { @@ -244,23 +276,57 @@ class BillingService: "reservation_id": reservation_id, "actual_amount": actual_amount, } + if bucket: + payload["bucket"] = bucket if meta: payload["meta"] = meta - return _quota_commit_adapter.validate_python(cls._send_request("POST", "/quota/commit", json=payload)) + return _quota_commit_adapter.validate_python(cls._send_quota_request("POST", "/quota/commit", json=payload)) @classmethod - def quota_release(cls, tenant_id: str, feature_key: str, reservation_id: str) -> QuotaReleaseResult: + def quota_release( + cls, tenant_id: str, feature_key: str, reservation_id: str, bucket: str = "" + ) -> QuotaReleaseResult: """Release a reservation (cancel, return frozen quota).""" - return _quota_release_adapter.validate_python( - cls._send_request( - "POST", - "/quota/release", - json={ - "tenant_id": tenant_id, - "feature_key": feature_key, - "reservation_id": reservation_id, - }, - ) + payload = { + "tenant_id": tenant_id, + "feature_key": feature_key, + "reservation_id": reservation_id, + } + if bucket: + payload["bucket"] = bucket + return _quota_release_adapter.validate_python(cls._send_quota_request("POST", "/quota/release", json=payload)) + + @classmethod + def quota_get_balance(cls, tenant_id: str, feature_key: str, bucket: str = "") -> QuotaBalanceResult: + """Get quota balance for a feature bucket.""" + params = {"tenant_id": tenant_id, "feature_key": feature_key} + if bucket: + params["bucket"] = bucket + return _quota_balance_adapter.validate_python(cls._send_quota_request("GET", "/quota/balance", params=params)) + + @classmethod + def quota_consume_capped( + cls, + tenant_id: str, + feature_key: str, + request_id: str, + amount: int, + meta: dict | None = None, + bucket: str = "", + ) -> QuotaConsumeCappedResult: + """Consume up to the available quota and return the actual deducted amount.""" + payload: dict = { + "tenant_id": tenant_id, + "feature_key": feature_key, + "request_id": request_id, + "amount": amount, + } + if bucket: + payload["bucket"] = bucket + if meta: + payload["meta"] = meta + return _quota_consume_capped_adapter.validate_python( + cls._send_quota_request("POST", "/quota/consume-capped", json=payload) ) @classmethod @@ -334,6 +400,12 @@ class BillingService: params = {"tenant_id": tenant_id, "feature_key": feature_key} return cls._send_request("GET", "/billing/tenant_feature_plan/usage", params=params) + @classmethod + def _send_quota_request( + cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None + ): + return cls._send_request(method, endpoint, json=json, params=params, base_url=cls.quota_base_url) + @classmethod @retry( wait=wait_fixed(2), @@ -341,10 +413,17 @@ class BillingService: retry=retry_if_exception_type(httpx.RequestError), reraise=True, ) - def _send_request(cls, method: Literal["GET", "POST", "DELETE", "PUT"], endpoint: str, json=None, params=None): + def _send_request( + cls, + method: Literal["GET", "POST", "DELETE", "PUT"], + endpoint: str, + json=None, + params=None, + base_url: str | None = None, + ): headers = {"Content-Type": "application/json", "Billing-Api-Secret-Key": cls.secret_key} - url = f"{cls.base_url}{endpoint}" + url = f"{base_url or cls.base_url}{endpoint}" response = _http_client.request(method, url, json=json, params=params, headers=headers, follow_redirects=True) if method == "GET" and response.status_code != httpx.codes.OK: raise ValueError("Unable to retrieve billing information. Please try again later or contact support.") diff --git a/api/services/credit_pool_service.py b/api/services/credit_pool_service.py index afc49181185..837bf52c082 100644 --- a/api/services/credit_pool_service.py +++ b/api/services/credit_pool_service.py @@ -7,6 +7,8 @@ from piling up database transactions while preserving cross-tenant concurrency. import logging from collections.abc import Callable +from dataclasses import dataclass +from uuid import uuid4 from sqlalchemy import select from sqlalchemy.orm import Session @@ -19,11 +21,43 @@ from models.enums import ProviderQuotaType logger = logging.getLogger(__name__) +FEATURE_KEY_CREDIT_POOL = "credit_pool" CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS = 10 CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS = 5 +@dataclass(frozen=True) +class CreditPoolBalance: + tenant_id: str + pool_type: str + quota_limit: int + quota_used: int + + @property + def remaining_credits(self) -> int: + if self.quota_limit == -1: + return -1 + return max(0, self.quota_limit - self.quota_used) + + def has_sufficient_credits(self, required_credits: int) -> bool: + return self.quota_limit == -1 or self.remaining_credits >= required_credits + + class CreditPoolService: + @staticmethod + def _normalize_pool_type(pool_type: str | ProviderQuotaType) -> str: + return pool_type.value if isinstance(pool_type, ProviderQuotaType) else str(pool_type) + + @staticmethod + def _use_billing_quota() -> bool: + return bool(dify_config.BILLING_ENABLED) + + @staticmethod + def _require_session(session: Session | None) -> Session: + if session is None: + raise ValueError("session is required when billing quota is disabled") + return session + @staticmethod def _get_tenant_lock_key(tenant_id: str) -> str: return f"credit_pool:tenant:{tenant_id}:deduct_lock" @@ -77,13 +111,36 @@ class CreditPoolService: return credit_pool @classmethod - def get_pool(cls, tenant_id: str, pool_type: str = "trial", *, session: Session) -> TenantCreditPool | None: + def get_pool( + cls, + tenant_id: str, + pool_type: str | ProviderQuotaType = "trial", + *, + session: Session | None = None, + ) -> TenantCreditPool | CreditPoolBalance | None: """get tenant credit pool""" + normalized_pool_type = cls._normalize_pool_type(pool_type) + if cls._use_billing_quota(): + from services.billing_service import BillingService + + balance = BillingService.quota_get_balance( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + ) + return CreditPoolBalance( + tenant_id=tenant_id, + pool_type=normalized_pool_type, + quota_limit=balance["quota"], + quota_used=balance["usage"], + ) + + session = cls._require_session(session) return session.scalar( select(TenantCreditPool) .where( TenantCreditPool.tenant_id == tenant_id, - TenantCreditPool.pool_type == pool_type, + TenantCreditPool.pool_type == normalized_pool_type, ) .limit(1) ) @@ -93,31 +150,77 @@ class CreditPoolService: cls, tenant_id: str, credits_required: int, - pool_type: str = "trial", + pool_type: str | ProviderQuotaType = "trial", *, - session: Session, + session: Session | None = None, ) -> bool: """check if credits are available without deducting""" pool = cls.get_pool(tenant_id, pool_type, session=session) if not pool: return False - return pool.remaining_credits >= credits_required + return pool.has_sufficient_credits(credits_required) @classmethod def check_and_deduct_credits( cls, tenant_id: str, credits_required: int, - pool_type: str = "trial", + pool_type: str | ProviderQuotaType = "trial", *, - session: Session, + session: Session | None = None, ) -> int: """Deduct exactly the requested credits or raise without mutating the pool.""" if credits_required <= 0: return 0 + normalized_pool_type = cls._normalize_pool_type(pool_type) + + if cls._use_billing_quota(): + from services.billing_service import BillingService + + request_id = str(uuid4()) + result = BillingService.quota_reserve( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + request_id=request_id, + amount=credits_required, + meta={"source": "credit_pool.check_and_deduct"}, + ) + reservation_id = result.get("reservation_id", "") + if not reservation_id: + raise QuotaExceededError("Insufficient credits remaining") + try: + BillingService.quota_commit( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + reservation_id=reservation_id, + actual_amount=credits_required, + meta={"source": "credit_pool.check_and_deduct"}, + ) + except Exception: + try: + BillingService.quota_release( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + reservation_id=reservation_id, + ) + except Exception: + logger.warning( + "Failed to release reserved credit pool quota, tenant_id=%s, pool_type=%s, reservation_id=%s", + tenant_id, + normalized_pool_type, + reservation_id, + exc_info=True, + ) + raise + return credits_required + + session = cls._require_session(session) def deduct() -> int: - pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type) + pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type) if not pool: raise QuotaExceededError("Credit pool not found") @@ -144,18 +247,34 @@ class CreditPoolService: cls, tenant_id: str, credits_required: int, - pool_type: str = "trial", + pool_type: str | ProviderQuotaType = "trial", *, - session: Session, + session: Session | None = None, ) -> int: """Deduct up to the available balance and return the actual deducted credits.""" if credits_required <= 0: return 0 + normalized_pool_type = cls._normalize_pool_type(pool_type) + + if cls._use_billing_quota(): + from services.billing_service import BillingService + + result = BillingService.quota_consume_capped( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket=normalized_pool_type, + request_id=str(uuid4()), + amount=credits_required, + meta={"source": "credit_pool.deduct_capped"}, + ) + return result["deducted"] + + session = cls._require_session(session) def deduct() -> int: - pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=pool_type) + pool = cls._get_locked_pool(session=session, tenant_id=tenant_id, pool_type=normalized_pool_type) if not pool: - logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, pool_type) + logger.warning("Credit pool not found, tenant_id=%s, pool_type=%s", tenant_id, normalized_pool_type) return 0 deducted_credits = min(credits_required, pool.remaining_credits) diff --git a/api/tests/unit_tests/services/test_billing_service.py b/api/tests/unit_tests/services/test_billing_service.py index dc691176114..2b1bdb5c5c4 100644 --- a/api/tests/unit_tests/services/test_billing_service.py +++ b/api/tests/unit_tests/services/test_billing_service.py @@ -73,6 +73,23 @@ class TestBillingServiceSendRequest: assert call_args[1]["headers"]["Billing-Api-Secret-Key"] == "test-secret-key" assert call_args[1]["headers"]["Content-Type"] == "application/json" + def test_send_request_with_base_url_override(self, mock_httpx_request, mock_billing_config): + """Quota APIs can use the new billing service without changing legacy billing calls.""" + # Arrange + expected_response = {"result": "success"} + mock_response = MagicMock() + mock_response.status_code = httpx.codes.OK + mock_response.json.return_value = expected_response + mock_httpx_request.return_value = mock_response + + # Act + result = BillingService._send_request("GET", "/quota/balance", base_url="https://quota.example.com") + + # Assert + assert result == expected_response + call_args = mock_httpx_request.call_args + assert call_args[0][1] == "https://quota.example.com/quota/balance" + @pytest.mark.parametrize( "status_code", [httpx.codes.NOT_FOUND, httpx.codes.INTERNAL_SERVER_ERROR, httpx.codes.BAD_REQUEST] ) @@ -393,6 +410,20 @@ class TestBillingServiceSubscriptionInfo: params={"tenant_id": tenant_id}, ) + def test_quota_get_balance_uses_quota_request(self): + tenant_id = "tenant-123" + with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request: + mock_send_quota_request.return_value = {"quota": "200", "usage": "6", "available": "194", "reserved": "0"} + + result = BillingService.quota_get_balance(tenant_id, "credit_pool", bucket="trial") + + assert result == {"quota": 200, "usage": 6, "available": 194, "reserved": 0} + mock_send_quota_request.assert_called_once_with( + "GET", + "/quota/balance", + params={"tenant_id": tenant_id, "feature_key": "credit_pool", "bucket": "trial"}, + ) + def test_get_knowledge_rate_limit_with_defaults(self, mock_send_request): """Test knowledge rate limit retrieval with default values.""" # Arrange @@ -518,19 +549,20 @@ class TestBillingServiceUsageCalculation: assert result == expected_response mock_send_request.assert_called_once_with("GET", "/tenant-feature-usage/info", params={"tenant_id": tenant_id}) - def test_get_quota_info(self, mock_send_request): + def test_get_quota_info(self): """Test retrieval of quota info from new endpoint.""" # Arrange tenant_id = "tenant-123" expected_response = {"trigger_event": {"limit": 100, "usage": 30}, "api_rate_limit": {"limit": -1, "usage": 0}} - mock_send_request.return_value = expected_response + with patch.object(BillingService, "_send_quota_request") as mock_send_quota_request: + mock_send_quota_request.return_value = expected_response - # Act - result = BillingService.get_quota_info(tenant_id) + # Act + result = BillingService.get_quota_info(tenant_id) # Assert assert result == expected_response - mock_send_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id}) + mock_send_quota_request.assert_called_once_with("GET", "/quota/info", params={"tenant_id": tenant_id}) def test_update_tenant_feature_plan_usage_positive_delta(self, mock_send_request): """Test updating tenant feature usage with positive delta (adding credits).""" @@ -614,7 +646,7 @@ class TestBillingServiceQuotaOperations: @pytest.fixture def mock_send_request(self): - with patch.object(BillingService, "_send_request") as mock: + with patch.object(BillingService, "_send_quota_request") as mock: yield mock def test_quota_reserve_success(self, mock_send_request): @@ -652,6 +684,16 @@ class TestBillingServiceQuotaOperations: call_json = mock_send_request.call_args[1]["json"] assert call_json["meta"] == {"source": "webhook"} + def test_quota_reserve_with_bucket(self, mock_send_request): + mock_send_request.return_value = {"reservation_id": "rid-2", "available": 98, "reserved": 1} + + BillingService.quota_reserve( + tenant_id="t1", feature_key="credit_pool", request_id="req-2", amount=1, bucket="trial" + ) + + call_json = mock_send_request.call_args[1]["json"] + assert call_json["bucket"] == "trial" + def test_quota_commit_success(self, mock_send_request): expected = {"available": 98, "reserved": 0, "refunded": 0} mock_send_request.return_value = expected @@ -696,6 +738,20 @@ class TestBillingServiceQuotaOperations: call_json = mock_send_request.call_args[1]["json"] assert call_json["meta"] == {"reason": "partial"} + def test_quota_commit_with_bucket(self, mock_send_request): + mock_send_request.return_value = {"available": 97, "reserved": 0, "refunded": 0} + + BillingService.quota_commit( + tenant_id="t1", + feature_key="credit_pool", + reservation_id="rid-1", + actual_amount=1, + bucket="paid", + ) + + call_json = mock_send_request.call_args[1]["json"] + assert call_json["bucket"] == "paid" + def test_quota_release_success(self, mock_send_request): expected = {"available": 100, "reserved": 0, "released": 1} mock_send_request.return_value = expected @@ -720,6 +776,64 @@ class TestBillingServiceQuotaOperations: assert result["released"] == 1 assert isinstance(result["released"], int) + def test_quota_release_with_bucket(self, mock_send_request): + mock_send_request.return_value = {"available": 100, "reserved": 0, "released": 1} + + BillingService.quota_release(tenant_id="t1", feature_key="credit_pool", reservation_id="rid-1", bucket="trial") + + call_json = mock_send_request.call_args[1]["json"] + assert call_json["bucket"] == "trial" + + def test_quota_consume_capped_success(self, mock_send_request): + mock_send_request.return_value = { + "deducted": "2", + "available": "8", + "reserved": "0", + "quota": "10", + "usage": "2", + } + + result = BillingService.quota_consume_capped( + tenant_id="t1", + feature_key="credit_pool", + request_id="req-1", + amount=5, + bucket="paid", + meta={"source": "test"}, + ) + + assert result == {"deducted": 2, "available": 8, "reserved": 0, "quota": 10, "usage": 2} + mock_send_request.assert_called_once_with( + "POST", + "/quota/consume-capped", + json={ + "tenant_id": "t1", + "feature_key": "credit_pool", + "request_id": "req-1", + "amount": 5, + "bucket": "paid", + "meta": {"source": "test"}, + }, + ) + + def test_send_quota_request_uses_quota_base_url(self): + with ( + patch.object(BillingService, "quota_base_url", "https://quota.example.com/v1"), + patch.object(BillingService, "_send_request") as mock_send_request, + ): + mock_send_request.return_value = {"ok": True} + + result = BillingService._send_quota_request("GET", "/quota/info", params={"tenant_id": "t1"}) + + assert result == {"ok": True} + mock_send_request.assert_called_once_with( + "GET", + "/quota/info", + json=None, + params={"tenant_id": "t1"}, + base_url="https://quota.example.com/v1", + ) + def test_get_quota_info_coerces_string_to_int(self, mock_send_request): """Test that TypeAdapter coerces string values to int for get_quota_info.""" mock_send_request.return_value = { diff --git a/api/tests/unit_tests/services/test_credit_pool_service.py b/api/tests/unit_tests/services/test_credit_pool_service.py index f31d067525a..8cafd3af590 100644 --- a/api/tests/unit_tests/services/test_credit_pool_service.py +++ b/api/tests/unit_tests/services/test_credit_pool_service.py @@ -1,5 +1,6 @@ +from collections.abc import Generator from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch from uuid import uuid4 import pytest @@ -13,6 +14,8 @@ from models.enums import ProviderQuotaType from services.credit_pool_service import ( CREDIT_POOL_TENANT_LOCK_BLOCKING_TIMEOUT_SECONDS, CREDIT_POOL_TENANT_LOCK_TIMEOUT_SECONDS, + FEATURE_KEY_CREDIT_POOL, + CreditPoolBalance, CreditPoolService, ) @@ -51,6 +54,12 @@ def _make_redis_lock() -> MagicMock: return lock +@pytest.fixture(autouse=True) +def _disable_billing_quota_by_default() -> Generator[None, None, None]: + with patch("services.credit_pool_service.dify_config.BILLING_ENABLED", False): + yield + + def test_get_pool_uses_provided_session() -> None: engine, tenant_id, _ = _create_engine_with_pool(quota_limit=10, quota_used=2) @@ -62,6 +71,13 @@ def test_get_pool_uses_provided_session() -> None: assert pool.quota_used == 2 +def test_credit_pool_balance_unlimited_remaining_and_sufficiency() -> None: + pool = CreditPoolBalance(tenant_id="tenant-1", pool_type="paid", quota_limit=-1, quota_used=999) + + assert pool.remaining_credits == -1 + assert pool.has_sufficient_credits(10_000) + + def test_check_and_deduct_credits_deducts_exact_amount_when_sufficient() -> None: engine, tenant_id, pool_id = _create_engine_with_pool(quota_limit=10, quota_used=2) @@ -211,7 +227,7 @@ def test_check_and_deduct_credits_uses_tenant_redis_lock_before_db_deduction() - ) redis_lock.acquire.assert_called_once_with(blocking=True) redis_lock.release.assert_called_once_with() - get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.TRIAL) + get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="trial") def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> None: @@ -240,7 +256,152 @@ def test_deduct_credits_capped_uses_tenant_redis_lock_before_db_deduction() -> N ) redis_lock.acquire.assert_called_once_with(blocking=True) redis_lock.release.assert_called_once_with() - get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID) + get_locked_pool.assert_called_once_with(session=session, tenant_id=tenant_id, pool_type="paid") + + +def test_get_pool_uses_billing_quota_balance_when_enabled() -> None: + tenant_id = "tenant-1" + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_get_balance") as quota_get_balance, + ): + quota_get_balance.return_value = {"quota": 1000, "usage": 250, "available": 750, "reserved": 0} + + pool = CreditPoolService.get_pool(tenant_id=tenant_id, pool_type=ProviderQuotaType.PAID) + + assert isinstance(pool, CreditPoolBalance) + assert pool.quota_limit == 1000 + assert pool.quota_used == 250 + assert pool.remaining_credits == 750 + quota_get_balance.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="paid", + ) + + +def test_check_and_deduct_credits_uses_billing_reserve_and_commit_when_enabled() -> None: + tenant_id = "tenant-1" + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + patch("services.billing_service.BillingService.quota_commit") as quota_commit, + patch("services.billing_service.BillingService.quota_release") as quota_release, + ): + quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3} + + result = CreditPoolService.check_and_deduct_credits( + tenant_id=tenant_id, + credits_required=3, + pool_type=ProviderQuotaType.TRIAL, + ) + + assert result == 3 + quota_reserve.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + request_id=ANY, + amount=3, + meta={"source": "credit_pool.check_and_deduct"}, + ) + quota_commit.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + reservation_id="reservation-1", + actual_amount=3, + meta={"source": "credit_pool.check_and_deduct"}, + ) + quota_release.assert_not_called() + + +def test_check_and_deduct_credits_raises_when_billing_reserve_is_insufficient() -> None: + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + ): + quota_reserve.return_value = {"reservation_id": "", "available": 1, "reserved": 0} + + with pytest.raises(QuotaExceededError, match="Insufficient credits remaining"): + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3) + + +def test_check_and_deduct_credits_releases_billing_reservation_when_commit_fails() -> None: + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")), + patch("services.billing_service.BillingService.quota_release") as quota_release, + ): + quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3} + + with pytest.raises(RuntimeError, match="commit failed"): + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3) + + quota_release.assert_called_once_with( + tenant_id="tenant-1", + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + reservation_id="reservation-1", + ) + + +def test_check_and_deduct_credits_logs_when_billing_release_fails() -> None: + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_reserve") as quota_reserve, + patch("services.billing_service.BillingService.quota_commit", side_effect=RuntimeError("commit failed")), + patch( + "services.billing_service.BillingService.quota_release", side_effect=RuntimeError("release failed") + ) as quota_release, + patch("services.credit_pool_service.logger.warning") as logger_warning, + ): + quota_reserve.return_value = {"reservation_id": "reservation-1", "available": 7, "reserved": 3} + + with pytest.raises(RuntimeError, match="commit failed"): + CreditPoolService.check_and_deduct_credits(tenant_id="tenant-1", credits_required=3) + + quota_release.assert_called_once_with( + tenant_id="tenant-1", + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="trial", + reservation_id="reservation-1", + ) + logger_warning.assert_called_once() + assert logger_warning.call_args.args[3] == "reservation-1" + assert logger_warning.call_args.kwargs["exc_info"] is True + + +def test_deduct_credits_capped_uses_billing_consume_capped_when_enabled() -> None: + tenant_id = "tenant-1" + with ( + patch("services.credit_pool_service.dify_config.BILLING_ENABLED", True), + patch("services.billing_service.BillingService.quota_consume_capped") as quota_consume_capped, + ): + quota_consume_capped.return_value = { + "deducted": 2, + "available": 0, + "reserved": 0, + "quota": 10, + "usage": 10, + } + + result = CreditPoolService.deduct_credits_capped( + tenant_id=tenant_id, + credits_required=5, + pool_type=ProviderQuotaType.PAID, + ) + + assert result == 2 + quota_consume_capped.assert_called_once_with( + tenant_id=tenant_id, + feature_key=FEATURE_KEY_CREDIT_POOL, + bucket="paid", + request_id=ANY, + amount=5, + meta={"source": "credit_pool.deduct_capped"}, + ) @pytest.mark.parametrize(