Merge remote-tracking branch 'origin/main' into feat/support-agent-sandbox

api/services/account_service.py

@@ -999,6 +999,11 @@ class TenantService:
 
         tenant.encrypt_public_key = generate_key_pair(tenant.id)
         db.session.commit()
+
+        from services.credit_pool_service import CreditPoolService
+
+        CreditPoolService.create_default_pool(tenant.id)
+
         return tenant
 
     @staticmethod

api/services/credit_pool_service.py

@@ -0,0 +1,85 @@
+import logging
+
+from sqlalchemy import update
+from sqlalchemy.orm import Session
+
+from configs import dify_config
+from core.errors.error import QuotaExceededError
+from extensions.ext_database import db
+from models import TenantCreditPool
+
+logger = logging.getLogger(__name__)
+
+
+class CreditPoolService:
+    @classmethod
+    def create_default_pool(cls, tenant_id: str) -> TenantCreditPool:
+        """create default credit pool for new tenant"""
+        credit_pool = TenantCreditPool(
+            tenant_id=tenant_id, quota_limit=dify_config.HOSTED_POOL_CREDITS, quota_used=0, pool_type="trial"
+        )
+        db.session.add(credit_pool)
+        db.session.commit()
+        return credit_pool
+
+    @classmethod
+    def get_pool(cls, tenant_id: str, pool_type: str = "trial") -> TenantCreditPool | None:
+        """get tenant credit pool"""
+        return (
+            db.session.query(TenantCreditPool)
+            .filter_by(
+                tenant_id=tenant_id,
+                pool_type=pool_type,
+            )
+            .first()
+        )
+
+    @classmethod
+    def check_credits_available(
+        cls,
+        tenant_id: str,
+        credits_required: int,
+        pool_type: str = "trial",
+    ) -> bool:
+        """check if credits are available without deducting"""
+        pool = cls.get_pool(tenant_id, pool_type)
+        if not pool:
+            return False
+        return pool.remaining_credits >= credits_required
+
+    @classmethod
+    def check_and_deduct_credits(
+        cls,
+        tenant_id: str,
+        credits_required: int,
+        pool_type: str = "trial",
+    ) -> int:
+        """check and deduct credits, returns actual credits deducted"""
+
+        pool = cls.get_pool(tenant_id, pool_type)
+        if not pool:
+            raise QuotaExceededError("Credit pool not found")
+
+        if pool.remaining_credits <= 0:
+            raise QuotaExceededError("No credits remaining")
+
+        # deduct all remaining credits if less than required
+        actual_credits = min(credits_required, pool.remaining_credits)
+
+        try:
+            with Session(db.engine) as session:
+                stmt = (
+                    update(TenantCreditPool)
+                    .where(
+                        TenantCreditPool.tenant_id == tenant_id,
+                        TenantCreditPool.pool_type == pool_type,
+                    )
+                    .values(quota_used=TenantCreditPool.quota_used + actual_credits)
+                )
+                session.execute(stmt)
+                session.commit()
+        except Exception:
+            logger.exception("Failed to deduct credits for tenant %s", tenant_id)
+            raise QuotaExceededError("Failed to deduct credits")
+
+        return actual_credits

api/services/feature_service.py

@@ -140,6 +140,7 @@ class FeatureModel(BaseModel):
     # pydantic configs
     model_config = ConfigDict(protected_namespaces=())
     knowledge_pipeline: KnowledgePipeline = KnowledgePipeline()
+    next_credit_reset_date: int = 0
 
 
 class KnowledgeRateLimitModel(BaseModel):
@@ -301,6 +302,9 @@ class FeatureService:
         if "knowledge_pipeline_publish_enabled" in billing_info:
             features.knowledge_pipeline.publish_enabled = billing_info["knowledge_pipeline_publish_enabled"]
 
+        if "next_credit_reset_date" in billing_info:
+            features.next_credit_reset_date = billing_info["next_credit_reset_date"]
+
     @classmethod
     def _fulfill_params_from_enterprise(cls, features: SystemFeatureModel):
         enterprise_info = EnterpriseService.get_info()

api/services/trigger/trigger_provider_service.py

@@ -799,7 +799,7 @@ class TriggerProviderService:
         user_id: str,
         provider_id: TriggerProviderID,
         subscription_id: str,
-        credentials: Mapping[str, Any],
+        credentials: dict[str, Any],
     ) -> dict[str, Any]:
         """
         Verify credentials for an existing subscription without updating it.

api/services/workspace_service.py

@@ -31,7 +31,8 @@ class WorkspaceService:
         assert tenant_account_join is not None, "TenantAccountJoin not found"
         tenant_info["role"] = tenant_account_join.role
 
-        can_replace_logo = FeatureService.get_features(tenant.id).can_replace_logo
+        feature = FeatureService.get_features(tenant.id)
+        can_replace_logo = feature.can_replace_logo
 
         if can_replace_logo and TenantService.has_roles(tenant, [TenantAccountRole.OWNER, TenantAccountRole.ADMIN]):
             base_url = dify_config.FILES_URL
@@ -46,5 +47,19 @@ class WorkspaceService:
                 "remove_webapp_brand": remove_webapp_brand,
                 "replace_webapp_logo": replace_webapp_logo,
             }
+        if dify_config.EDITION == "CLOUD":
+            tenant_info["next_credit_reset_date"] = feature.next_credit_reset_date
+
+            from services.credit_pool_service import CreditPoolService
+
+            paid_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="paid")
+            if paid_pool:
+                tenant_info["trial_credits"] = paid_pool.quota_limit
+                tenant_info["trial_credits_used"] = paid_pool.quota_used
+            else:
+                trial_pool = CreditPoolService.get_pool(tenant_id=tenant.id, pool_type="trial")
+                if trial_pool:
+                    tenant_info["trial_credits"] = trial_pool.quota_limit
+                    tenant_info["trial_credits_used"] = trial_pool.quota_used
 
         return tenant_info