feat: improve ai-credits display (#38589)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/controllers/console/workspace/workspace.py

@@ -88,6 +88,7 @@ class TenantInfoResponse(ResponseModel):
     custom_config: WorkspaceCustomConfigResponse | None = None
     trial_credits: int | None = None
     trial_credits_used: int | None = None
+    trial_credits_exhausted_at: int | None = None
     next_credit_reset_date: int | None = None
 
     @field_validator("plan", "status", "trial_end_reason", mode="before")

api/openapi/markdown/console-openapi.md

@@ -22008,6 +22008,7 @@ Tag type
 | role | string |  | No |
 | status | string |  | No |
 | trial_credits | integer |  | No |
+| trial_credits_exhausted_at | integer |  | No |
 | trial_credits_used | integer |  | No |
 | trial_end_reason | string |  | No |
 

api/services/billing_service.py

@@ -55,6 +55,7 @@ class QuotaBalanceResult(TypedDict):
     reserved: int
     quota: int
     usage: int
+    exhausted_at: NotRequired[int]
 
 
 class QuotaConsumeCappedResult(TypedDict):

api/services/credit_pool_service.py

@@ -32,6 +32,7 @@ class CreditPoolBalance:
     pool_type: str
     quota_limit: int
     quota_used: int
+    exhausted_at: int | None = None
 
     @property
     def remaining_credits(self) -> int:
@@ -133,6 +134,7 @@ class CreditPoolService:
                 pool_type=normalized_pool_type,
                 quota_limit=balance["quota"],
                 quota_used=balance["usage"],
+                exhausted_at=balance.get("exhausted_at"),
             )
 
         session = cls._require_session(session)

api/services/workspace_service.py

@@ -9,6 +9,15 @@ from services.account_service import TenantService
 from services.feature_service import FeatureService
 
 
+def _set_credit_pool_info(
+    tenant_info: dict[str, object], *, quota_limit: int, quota_used: int, exhausted_at: int | None = None
+) -> None:
+    tenant_info["trial_credits"] = quota_limit
+    tenant_info["trial_credits_used"] = quota_used
+    if isinstance(exhausted_at, int) and exhausted_at > 0 and quota_limit > 0 and quota_used >= quota_limit:
+        tenant_info["trial_credits_exhausted_at"] = exhausted_at
+
+
 class WorkspaceService:
     @classmethod
     def get_tenant_info(cls, tenant: Tenant, session: Session):
@@ -54,7 +63,7 @@ class WorkspaceService:
         if dify_config.EDITION == "CLOUD":
             tenant_info["next_credit_reset_date"] = feature.next_credit_reset_date
 
-            from services.credit_pool_service import CreditPoolService
+            from services.credit_pool_service import CreditPoolBalance, CreditPoolService
 
             paid_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="paid", session=session)
             # if the tenant is not on the sandbox plan and the paid pool is not full, use the paid pool
@@ -63,12 +72,22 @@ class WorkspaceService:
                 and paid_pool is not None
                 and (paid_pool.quota_limit == -1 or paid_pool.quota_limit > paid_pool.quota_used)
             ):
-                tenant_info["trial_credits"] = paid_pool.quota_limit
-                tenant_info["trial_credits_used"] = paid_pool.quota_used
+                exhausted_at = paid_pool.exhausted_at if isinstance(paid_pool, CreditPoolBalance) else None
+                _set_credit_pool_info(
+                    tenant_info,
+                    quota_limit=paid_pool.quota_limit,
+                    quota_used=paid_pool.quota_used,
+                    exhausted_at=exhausted_at,
+                )
             else:
                 trial_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="trial", session=session)
                 if trial_pool:
-                    tenant_info["trial_credits"] = trial_pool.quota_limit
-                    tenant_info["trial_credits_used"] = trial_pool.quota_used
+                    exhausted_at = trial_pool.exhausted_at if isinstance(trial_pool, CreditPoolBalance) else None
+                    _set_credit_pool_info(
+                        tenant_info,
+                        quota_limit=trial_pool.quota_limit,
+                        quota_used=trial_pool.quota_used,
+                        exhausted_at=exhausted_at,
+                    )
 
         return tenant_info

api/tests/test_containers_integration_tests/services/test_workspace_service.py

@@ -7,6 +7,7 @@ from faker import Faker
 from sqlalchemy.orm import Session
 
 from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
+from services.credit_pool_service import CreditPoolBalance
 from services.workspace_service import WorkspaceService
 
 
@@ -720,7 +721,13 @@ class TestWorkspaceService:
         mock_external_service_dependencies["tenant_service"].has_roles.return_value = False
 
         paid_pool = MagicMock(quota_limit=500, quota_used=500)
-        trial_pool = MagicMock(quota_limit=100, quota_used=10)
+        trial_pool = CreditPoolBalance(
+            tenant_id=tenant.id,
+            pool_type="trial",
+            quota_limit=100,
+            quota_used=100,
+            exhausted_at=1748908800,
+        )
 
         with (
             patch("services.workspace_service.current_user", account),
@@ -730,7 +737,8 @@ class TestWorkspaceService:
 
         assert result is not None
         assert result["trial_credits"] == 100
-        assert result["trial_credits_used"] == 10
+        assert result["trial_credits_used"] == 100
+        assert result["trial_credits_exhausted_at"] == 1748908800
 
     def test_get_tenant_info_cloud_fall_back_to_trial_when_paid_none(
         self, db_session_with_containers: Session, mock_external_service_dependencies

api/tests/unit_tests/services/test_billing_service.py

@@ -465,11 +465,17 @@ class TestBillingServiceSubscriptionInfo:
     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"}
+            mock_send_quota_request.return_value = {
+                "quota": "200",
+                "usage": "6",
+                "available": "194",
+                "reserved": "0",
+                "exhausted_at": "1748908800",
+            }
 
             result = BillingService.quota_get_balance(tenant_id, "credit_pool", bucket="trial")
 
-        assert result == {"quota": 200, "usage": 6, "available": 194, "reserved": 0}
+        assert result == {"quota": 200, "usage": 6, "available": 194, "reserved": 0, "exhausted_at": 1748908800}
         mock_send_quota_request.assert_called_once_with(
             "GET",
             "/quota/balance",

api/tests/unit_tests/services/test_credit_pool_service.py

@@ -265,13 +265,20 @@ def test_get_pool_uses_billing_quota_balance_when_enabled() -> None:
         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}
+        quota_get_balance.return_value = {
+            "quota": 1000,
+            "usage": 250,
+            "available": 750,
+            "reserved": 0,
+            "exhausted_at": 1748908800,
+        }
 
         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.exhausted_at == 1748908800
     assert pool.remaining_credits == 750
     quota_get_balance.assert_called_once_with(
         tenant_id=tenant_id,