feat: trigger billing (#28335)
Signed-off-by: lyzno1 <yuanyouhuilyz@gmail.com> Co-authored-by: lyzno1 <yuanyouhuilyz@gmail.com> Co-authored-by: lyzno1 <92089059+lyzno1@users.noreply.github.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
@@ -10,19 +10,14 @@ from core.app.apps.completion.app_generator import CompletionAppGenerator
|
||||
from core.app.apps.workflow.app_generator import WorkflowAppGenerator
|
||||
from core.app.entities.app_invoke_entities import InvokeFrom
|
||||
from core.app.features.rate_limiting import RateLimit
|
||||
from enums.cloud_plan import CloudPlan
|
||||
from libs.helper import RateLimiter
|
||||
from enums.quota_type import QuotaType, unlimited
|
||||
from models.model import Account, App, AppMode, EndUser
|
||||
from models.workflow import Workflow
|
||||
from services.billing_service import BillingService
|
||||
from services.errors.app import WorkflowIdFormatError, WorkflowNotFoundError
|
||||
from services.errors.llm import InvokeRateLimitError
|
||||
from services.errors.app import InvokeRateLimitError, QuotaExceededError, WorkflowIdFormatError, WorkflowNotFoundError
|
||||
from services.workflow_service import WorkflowService
|
||||
|
||||
|
||||
class AppGenerateService:
|
||||
system_rate_limiter = RateLimiter("app_daily_rate_limiter", dify_config.APP_DAILY_RATE_LIMIT, 86400)
|
||||
|
||||
@classmethod
|
||||
def generate(
|
||||
cls,
|
||||
@@ -42,17 +37,12 @@ class AppGenerateService:
|
||||
:param streaming: streaming
|
||||
:return:
|
||||
"""
|
||||
# system level rate limiter
|
||||
quota_charge = unlimited()
|
||||
if dify_config.BILLING_ENABLED:
|
||||
# check if it's free plan
|
||||
limit_info = BillingService.get_info(app_model.tenant_id)
|
||||
if limit_info["subscription"]["plan"] == CloudPlan.SANDBOX:
|
||||
if cls.system_rate_limiter.is_rate_limited(app_model.tenant_id):
|
||||
raise InvokeRateLimitError(
|
||||
"Rate limit exceeded, please upgrade your plan "
|
||||
f"or your RPD was {dify_config.APP_DAILY_RATE_LIMIT} requests/day"
|
||||
)
|
||||
cls.system_rate_limiter.increment_rate_limit(app_model.tenant_id)
|
||||
try:
|
||||
quota_charge = QuotaType.WORKFLOW.consume(app_model.tenant_id)
|
||||
except QuotaExceededError:
|
||||
raise InvokeRateLimitError(f"Workflow execution quota limit reached for tenant {app_model.tenant_id}")
|
||||
|
||||
# app level rate limiter
|
||||
max_active_request = cls._get_max_active_requests(app_model)
|
||||
@@ -124,6 +114,7 @@ class AppGenerateService:
|
||||
else:
|
||||
raise ValueError(f"Invalid app mode {app_model.mode}")
|
||||
except Exception:
|
||||
quota_charge.refund()
|
||||
rate_limit.exit(request_id)
|
||||
raise
|
||||
finally:
|
||||
|
||||
@@ -13,18 +13,17 @@ from celery.result import AsyncResult
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from enums.quota_type import QuotaType
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from models.account import Account
|
||||
from models.enums import CreatorUserRole, WorkflowTriggerStatus
|
||||
from models.model import App, EndUser
|
||||
from models.trigger import WorkflowTriggerLog
|
||||
from models.workflow import Workflow
|
||||
from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository
|
||||
from services.errors.app import InvokeDailyRateLimitError, WorkflowNotFoundError
|
||||
from services.errors.app import InvokeRateLimitError, QuotaExceededError, WorkflowNotFoundError
|
||||
from services.workflow.entities import AsyncTriggerResponse, TriggerData, WorkflowTaskData
|
||||
from services.workflow.queue_dispatcher import QueueDispatcherManager, QueuePriority
|
||||
from services.workflow.rate_limiter import TenantDailyRateLimiter
|
||||
from services.workflow_service import WorkflowService
|
||||
from tasks.async_workflow_tasks import (
|
||||
execute_workflow_professional,
|
||||
@@ -82,7 +81,6 @@ class AsyncWorkflowService:
|
||||
trigger_log_repo = SQLAlchemyWorkflowTriggerLogRepository(session)
|
||||
dispatcher_manager = QueueDispatcherManager()
|
||||
workflow_service = WorkflowService()
|
||||
rate_limiter = TenantDailyRateLimiter(redis_client)
|
||||
|
||||
# 1. Validate app exists
|
||||
app_model = session.scalar(select(App).where(App.id == trigger_data.app_id))
|
||||
@@ -127,25 +125,19 @@ class AsyncWorkflowService:
|
||||
trigger_log = trigger_log_repo.create(trigger_log)
|
||||
session.commit()
|
||||
|
||||
# 7. Check and consume daily quota
|
||||
if not dispatcher.consume_quota(trigger_data.tenant_id):
|
||||
# 7. Check and consume quota
|
||||
try:
|
||||
QuotaType.WORKFLOW.consume(trigger_data.tenant_id)
|
||||
except QuotaExceededError as e:
|
||||
# Update trigger log status
|
||||
trigger_log.status = WorkflowTriggerStatus.RATE_LIMITED
|
||||
trigger_log.error = f"Daily limit reached for {dispatcher.get_queue_name()}"
|
||||
trigger_log.error = f"Quota limit reached: {e}"
|
||||
trigger_log_repo.update(trigger_log)
|
||||
session.commit()
|
||||
|
||||
tenant_owner_tz = rate_limiter.get_tenant_owner_timezone(trigger_data.tenant_id)
|
||||
|
||||
remaining = rate_limiter.get_remaining_quota(trigger_data.tenant_id, dispatcher.get_daily_limit())
|
||||
|
||||
reset_time = rate_limiter.get_quota_reset_time(trigger_data.tenant_id, tenant_owner_tz)
|
||||
|
||||
raise InvokeDailyRateLimitError(
|
||||
f"Daily workflow execution limit reached. "
|
||||
f"Limit resets at {reset_time.strftime('%Y-%m-%d %H:%M:%S %Z')}. "
|
||||
f"Remaining quota: {remaining}"
|
||||
)
|
||||
raise InvokeRateLimitError(
|
||||
f"Workflow execution quota limit reached for tenant {trigger_data.tenant_id}"
|
||||
) from e
|
||||
|
||||
# 8. Create task data
|
||||
queue_name = dispatcher.get_queue_name()
|
||||
|
||||
@@ -24,6 +24,13 @@ class BillingService:
|
||||
billing_info = cls._send_request("GET", "/subscription/info", params=params)
|
||||
return billing_info
|
||||
|
||||
@classmethod
|
||||
def get_tenant_feature_plan_usage_info(cls, tenant_id: str):
|
||||
params = {"tenant_id": tenant_id}
|
||||
|
||||
usage_info = cls._send_request("GET", "/tenant-feature-usage/info", params=params)
|
||||
return usage_info
|
||||
|
||||
@classmethod
|
||||
def get_knowledge_rate_limit(cls, tenant_id: str):
|
||||
params = {"tenant_id": tenant_id}
|
||||
@@ -55,6 +62,44 @@ class BillingService:
|
||||
params = {"prefilled_email": prefilled_email, "tenant_id": tenant_id}
|
||||
return cls._send_request("GET", "/invoices", params=params)
|
||||
|
||||
@classmethod
|
||||
def update_tenant_feature_plan_usage(cls, tenant_id: str, feature_key: str, delta: int) -> dict:
|
||||
"""
|
||||
Update tenant feature plan usage.
|
||||
|
||||
Args:
|
||||
tenant_id: Tenant identifier
|
||||
feature_key: Feature key (e.g., 'trigger', 'workflow')
|
||||
delta: Usage delta (positive to add, negative to consume)
|
||||
|
||||
Returns:
|
||||
Response dict with 'result' and 'history_id'
|
||||
Example: {"result": "success", "history_id": "uuid"}
|
||||
"""
|
||||
return cls._send_request(
|
||||
"POST",
|
||||
"/tenant-feature-usage/usage",
|
||||
params={"tenant_id": tenant_id, "feature_key": feature_key, "delta": delta},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def refund_tenant_feature_plan_usage(cls, history_id: str) -> dict:
|
||||
"""
|
||||
Refund a previous usage charge.
|
||||
|
||||
Args:
|
||||
history_id: The history_id returned from update_tenant_feature_plan_usage
|
||||
|
||||
Returns:
|
||||
Response dict with 'result' and 'history_id'
|
||||
"""
|
||||
return cls._send_request("POST", "/tenant-feature-usage/refund", params={"quota_usage_history_id": history_id})
|
||||
|
||||
@classmethod
|
||||
def get_tenant_feature_plan_usage(cls, tenant_id: str, feature_key: str):
|
||||
params = {"tenant_id": tenant_id, "feature_key": feature_key}
|
||||
return cls._send_request("GET", "/billing/tenant_feature_plan/usage", params=params)
|
||||
|
||||
@classmethod
|
||||
@retry(
|
||||
wait=wait_fixed(2),
|
||||
@@ -69,6 +114,8 @@ class BillingService:
|
||||
response = httpx.request(method, url, json=json, params=params, headers=headers)
|
||||
if method == "GET" and response.status_code != httpx.codes.OK:
|
||||
raise ValueError("Unable to retrieve billing information. Please try again later or contact support.")
|
||||
if method == "POST" and response.status_code != httpx.codes.OK:
|
||||
raise ValueError(f"Unable to send request to {url}. Please try again later or contact support.")
|
||||
return response.json()
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -18,7 +18,29 @@ class WorkflowIdFormatError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class InvokeDailyRateLimitError(Exception):
|
||||
"""Raised when daily rate limit is exceeded for workflow invocations."""
|
||||
class InvokeRateLimitError(Exception):
|
||||
"""Raised when rate limit is exceeded for workflow invocations."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class QuotaExceededError(ValueError):
|
||||
"""Raised when billing quota is exceeded for a feature."""
|
||||
|
||||
def __init__(self, feature: str, tenant_id: str, required: int):
|
||||
self.feature = feature
|
||||
self.tenant_id = tenant_id
|
||||
self.required = required
|
||||
super().__init__(f"Quota exceeded for feature '{feature}' (tenant: {tenant_id}). Required: {required}")
|
||||
|
||||
|
||||
class TriggerNodeLimitExceededError(ValueError):
|
||||
"""Raised when trigger node count exceeds the plan limit."""
|
||||
|
||||
def __init__(self, count: int, limit: int):
|
||||
self.count = count
|
||||
self.limit = limit
|
||||
super().__init__(
|
||||
f"Trigger node count ({count}) exceeds the limit ({limit}) for your subscription plan. "
|
||||
f"Please upgrade your plan or reduce the number of trigger nodes."
|
||||
)
|
||||
|
||||
@@ -54,6 +54,12 @@ class LicenseLimitationModel(BaseModel):
|
||||
return (self.limit - self.size) >= required
|
||||
|
||||
|
||||
class Quota(BaseModel):
|
||||
usage: int = 0
|
||||
limit: int = 0
|
||||
reset_date: int = -1
|
||||
|
||||
|
||||
class LicenseStatus(StrEnum):
|
||||
NONE = "none"
|
||||
INACTIVE = "inactive"
|
||||
@@ -129,6 +135,8 @@ class FeatureModel(BaseModel):
|
||||
webapp_copyright_enabled: bool = False
|
||||
workspace_members: LicenseLimitationModel = LicenseLimitationModel(enabled=False, size=0, limit=0)
|
||||
is_allow_transfer_workspace: bool = True
|
||||
trigger_event: Quota = Quota(usage=0, limit=3000, reset_date=0)
|
||||
api_rate_limit: Quota = Quota(usage=0, limit=5000, reset_date=0)
|
||||
# pydantic configs
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
knowledge_pipeline: KnowledgePipeline = KnowledgePipeline()
|
||||
@@ -236,6 +244,8 @@ class FeatureService:
|
||||
def _fulfill_params_from_billing_api(cls, features: FeatureModel, tenant_id: str):
|
||||
billing_info = BillingService.get_info(tenant_id)
|
||||
|
||||
features_usage_info = BillingService.get_tenant_feature_plan_usage_info(tenant_id)
|
||||
|
||||
features.billing.enabled = billing_info["enabled"]
|
||||
features.billing.subscription.plan = billing_info["subscription"]["plan"]
|
||||
features.billing.subscription.interval = billing_info["subscription"]["interval"]
|
||||
@@ -246,6 +256,16 @@ class FeatureService:
|
||||
else:
|
||||
features.is_allow_transfer_workspace = False
|
||||
|
||||
if "trigger_event" in features_usage_info:
|
||||
features.trigger_event.usage = features_usage_info["trigger_event"]["usage"]
|
||||
features.trigger_event.limit = features_usage_info["trigger_event"]["limit"]
|
||||
features.trigger_event.reset_date = features_usage_info["trigger_event"].get("reset_date", -1)
|
||||
|
||||
if "api_rate_limit" in features_usage_info:
|
||||
features.api_rate_limit.usage = features_usage_info["api_rate_limit"]["usage"]
|
||||
features.api_rate_limit.limit = features_usage_info["api_rate_limit"]["limit"]
|
||||
features.api_rate_limit.reset_date = features_usage_info["api_rate_limit"].get("reset_date", -1)
|
||||
|
||||
if "members" in billing_info:
|
||||
features.members.size = billing_info["members"]["size"]
|
||||
features.members.limit = billing_info["members"]["limit"]
|
||||
|
||||
46
api/services/trigger/app_trigger_service.py
Normal file
46
api/services/trigger/app_trigger_service.py
Normal file
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
AppTrigger management service.
|
||||
|
||||
Handles AppTrigger model CRUD operations and status management.
|
||||
This service centralizes all AppTrigger-related business logic.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from sqlalchemy import update
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from extensions.ext_database import db
|
||||
from models.enums import AppTriggerStatus
|
||||
from models.trigger import AppTrigger
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AppTriggerService:
|
||||
"""Service for managing AppTrigger lifecycle and status."""
|
||||
|
||||
@staticmethod
|
||||
def mark_tenant_triggers_rate_limited(tenant_id: str) -> None:
|
||||
"""
|
||||
Mark all enabled triggers for a tenant as rate limited due to quota exceeded.
|
||||
|
||||
This method is called when a tenant's quota is exhausted. It updates all
|
||||
enabled triggers to RATE_LIMITED status to prevent further executions until
|
||||
quota is restored.
|
||||
|
||||
Args:
|
||||
tenant_id: Tenant ID whose triggers should be marked as rate limited
|
||||
|
||||
"""
|
||||
try:
|
||||
with Session(db.engine) as session:
|
||||
session.execute(
|
||||
update(AppTrigger)
|
||||
.where(AppTrigger.tenant_id == tenant_id, AppTrigger.status == AppTriggerStatus.ENABLED)
|
||||
.values(status=AppTriggerStatus.RATE_LIMITED)
|
||||
)
|
||||
session.commit()
|
||||
logger.info("Marked all enabled triggers as rate limited for tenant %s", tenant_id)
|
||||
except Exception:
|
||||
logger.exception("Failed to mark all enabled triggers as rate limited for tenant %s", tenant_id)
|
||||
@@ -18,6 +18,7 @@ from core.file.models import FileTransferMethod
|
||||
from core.tools.tool_file_manager import ToolFileManager
|
||||
from core.variables.types import SegmentType
|
||||
from core.workflow.enums import NodeType
|
||||
from enums.quota_type import QuotaType
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
from factories import file_factory
|
||||
@@ -27,6 +28,8 @@ from models.trigger import AppTrigger, WorkflowWebhookTrigger
|
||||
from models.workflow import Workflow
|
||||
from services.async_workflow_service import AsyncWorkflowService
|
||||
from services.end_user_service import EndUserService
|
||||
from services.errors.app import QuotaExceededError
|
||||
from services.trigger.app_trigger_service import AppTriggerService
|
||||
from services.workflow.entities import WebhookTriggerData
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -98,6 +101,12 @@ class WebhookService:
|
||||
raise ValueError(f"App trigger not found for webhook {webhook_id}")
|
||||
|
||||
# Only check enabled status if not in debug mode
|
||||
|
||||
if app_trigger.status == AppTriggerStatus.RATE_LIMITED:
|
||||
raise ValueError(
|
||||
f"Webhook trigger is rate limited for webhook {webhook_id}, please upgrade your plan."
|
||||
)
|
||||
|
||||
if app_trigger.status != AppTriggerStatus.ENABLED:
|
||||
raise ValueError(f"Webhook trigger is disabled for webhook {webhook_id}")
|
||||
|
||||
@@ -729,6 +738,18 @@ class WebhookService:
|
||||
user_id=None,
|
||||
)
|
||||
|
||||
# consume quota before triggering workflow execution
|
||||
try:
|
||||
QuotaType.TRIGGER.consume(webhook_trigger.tenant_id)
|
||||
except QuotaExceededError:
|
||||
AppTriggerService.mark_tenant_triggers_rate_limited(webhook_trigger.tenant_id)
|
||||
logger.info(
|
||||
"Tenant %s rate limited, skipping webhook trigger %s",
|
||||
webhook_trigger.tenant_id,
|
||||
webhook_trigger.webhook_id,
|
||||
)
|
||||
raise
|
||||
|
||||
# Trigger workflow execution asynchronously
|
||||
AsyncWorkflowService.trigger_workflow_async(
|
||||
session,
|
||||
|
||||
@@ -2,16 +2,14 @@
|
||||
Queue dispatcher system for async workflow execution.
|
||||
|
||||
Implements an ABC-based pattern for handling different subscription tiers
|
||||
with appropriate queue routing and rate limiting.
|
||||
with appropriate queue routing and priority assignment.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import StrEnum
|
||||
|
||||
from configs import dify_config
|
||||
from extensions.ext_redis import redis_client
|
||||
from services.billing_service import BillingService
|
||||
from services.workflow.rate_limiter import TenantDailyRateLimiter
|
||||
|
||||
|
||||
class QueuePriority(StrEnum):
|
||||
@@ -25,50 +23,16 @@ class QueuePriority(StrEnum):
|
||||
class BaseQueueDispatcher(ABC):
|
||||
"""Abstract base class for queue dispatchers"""
|
||||
|
||||
def __init__(self):
|
||||
self.rate_limiter = TenantDailyRateLimiter(redis_client)
|
||||
|
||||
@abstractmethod
|
||||
def get_queue_name(self) -> str:
|
||||
"""Get the queue name for this dispatcher"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_daily_limit(self) -> int:
|
||||
"""Get daily execution limit"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_priority(self) -> int:
|
||||
"""Get task priority level"""
|
||||
pass
|
||||
|
||||
def check_daily_quota(self, tenant_id: str) -> bool:
|
||||
"""
|
||||
Check if tenant has remaining daily quota
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
|
||||
Returns:
|
||||
True if quota available, False otherwise
|
||||
"""
|
||||
# Check without consuming
|
||||
remaining = self.rate_limiter.get_remaining_quota(tenant_id=tenant_id, max_daily_limit=self.get_daily_limit())
|
||||
return remaining > 0
|
||||
|
||||
def consume_quota(self, tenant_id: str) -> bool:
|
||||
"""
|
||||
Consume one execution from daily quota
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
|
||||
Returns:
|
||||
True if quota consumed successfully, False if limit reached
|
||||
"""
|
||||
return self.rate_limiter.check_and_consume(tenant_id=tenant_id, max_daily_limit=self.get_daily_limit())
|
||||
|
||||
|
||||
class ProfessionalQueueDispatcher(BaseQueueDispatcher):
|
||||
"""Dispatcher for professional tier"""
|
||||
@@ -76,9 +40,6 @@ class ProfessionalQueueDispatcher(BaseQueueDispatcher):
|
||||
def get_queue_name(self) -> str:
|
||||
return QueuePriority.PROFESSIONAL
|
||||
|
||||
def get_daily_limit(self) -> int:
|
||||
return int(1e9)
|
||||
|
||||
def get_priority(self) -> int:
|
||||
return 100
|
||||
|
||||
@@ -89,9 +50,6 @@ class TeamQueueDispatcher(BaseQueueDispatcher):
|
||||
def get_queue_name(self) -> str:
|
||||
return QueuePriority.TEAM
|
||||
|
||||
def get_daily_limit(self) -> int:
|
||||
return int(1e9)
|
||||
|
||||
def get_priority(self) -> int:
|
||||
return 50
|
||||
|
||||
@@ -102,9 +60,6 @@ class SandboxQueueDispatcher(BaseQueueDispatcher):
|
||||
def get_queue_name(self) -> str:
|
||||
return QueuePriority.SANDBOX
|
||||
|
||||
def get_daily_limit(self) -> int:
|
||||
return dify_config.APP_DAILY_RATE_LIMIT
|
||||
|
||||
def get_priority(self) -> int:
|
||||
return 10
|
||||
|
||||
|
||||
@@ -1,183 +0,0 @@
|
||||
"""
|
||||
Day-based rate limiter for workflow executions.
|
||||
|
||||
Implements UTC-based daily quotas that reset at midnight UTC for consistent rate limiting.
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime, time, timedelta
|
||||
from typing import Union
|
||||
|
||||
import pytz
|
||||
from redis import Redis
|
||||
from sqlalchemy import select
|
||||
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import RedisClientWrapper
|
||||
from models.account import Account, TenantAccountJoin, TenantAccountRole
|
||||
|
||||
|
||||
class TenantDailyRateLimiter:
|
||||
"""
|
||||
Day-based rate limiter that resets at midnight UTC
|
||||
|
||||
This class provides Redis-based rate limiting with the following features:
|
||||
- Daily quotas that reset at midnight UTC for consistency
|
||||
- Atomic check-and-consume operations
|
||||
- Automatic cleanup of stale counters
|
||||
- Timezone-aware error messages for better UX
|
||||
"""
|
||||
|
||||
def __init__(self, redis_client: Union[Redis, RedisClientWrapper]):
|
||||
self.redis = redis_client
|
||||
|
||||
def get_tenant_owner_timezone(self, tenant_id: str) -> str:
|
||||
"""
|
||||
Get timezone of tenant owner
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
|
||||
Returns:
|
||||
Timezone string (e.g., 'America/New_York', 'UTC')
|
||||
"""
|
||||
# Query to get tenant owner's timezone using scalar and select
|
||||
owner = db.session.scalar(
|
||||
select(Account)
|
||||
.join(TenantAccountJoin, TenantAccountJoin.account_id == Account.id)
|
||||
.where(TenantAccountJoin.tenant_id == tenant_id, TenantAccountJoin.role == TenantAccountRole.OWNER)
|
||||
)
|
||||
|
||||
if not owner:
|
||||
return "UTC"
|
||||
|
||||
return owner.timezone or "UTC"
|
||||
|
||||
def _get_day_key(self, tenant_id: str) -> str:
|
||||
"""
|
||||
Get Redis key for current UTC day
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
|
||||
Returns:
|
||||
Redis key for the current UTC day
|
||||
"""
|
||||
utc_now = datetime.now(UTC)
|
||||
date_str = utc_now.strftime("%Y-%m-%d")
|
||||
return f"workflow:daily_limit:{tenant_id}:{date_str}"
|
||||
|
||||
def _get_ttl_seconds(self) -> int:
|
||||
"""
|
||||
Calculate seconds until UTC midnight
|
||||
|
||||
Returns:
|
||||
Number of seconds until UTC midnight
|
||||
"""
|
||||
utc_now = datetime.now(UTC)
|
||||
|
||||
# Get next midnight in UTC
|
||||
next_midnight = datetime.combine(utc_now.date() + timedelta(days=1), time.min)
|
||||
next_midnight = next_midnight.replace(tzinfo=UTC)
|
||||
|
||||
return int((next_midnight - utc_now).total_seconds())
|
||||
|
||||
def check_and_consume(self, tenant_id: str, max_daily_limit: int) -> bool:
|
||||
"""
|
||||
Check if quota available and consume one execution
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
max_daily_limit: Maximum daily limit
|
||||
|
||||
Returns:
|
||||
True if quota consumed successfully, False if limit reached
|
||||
"""
|
||||
key = self._get_day_key(tenant_id)
|
||||
ttl = self._get_ttl_seconds()
|
||||
|
||||
# Check current usage
|
||||
current = self.redis.get(key)
|
||||
|
||||
if current is None:
|
||||
# First execution of the day - set to 1
|
||||
self.redis.setex(key, ttl, 1)
|
||||
return True
|
||||
|
||||
current_count = int(current)
|
||||
if current_count < max_daily_limit:
|
||||
# Within limit, increment
|
||||
new_count = self.redis.incr(key)
|
||||
# Update TTL
|
||||
self.redis.expire(key, ttl)
|
||||
|
||||
# Double-check in case of race condition
|
||||
if new_count <= max_daily_limit:
|
||||
return True
|
||||
else:
|
||||
# Race condition occurred, decrement back
|
||||
self.redis.decr(key)
|
||||
return False
|
||||
else:
|
||||
# Limit exceeded
|
||||
return False
|
||||
|
||||
def get_remaining_quota(self, tenant_id: str, max_daily_limit: int) -> int:
|
||||
"""
|
||||
Get remaining quota for the day
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
max_daily_limit: Maximum daily limit
|
||||
|
||||
Returns:
|
||||
Number of remaining executions for the day
|
||||
"""
|
||||
key = self._get_day_key(tenant_id)
|
||||
used = int(self.redis.get(key) or 0)
|
||||
return max(0, max_daily_limit - used)
|
||||
|
||||
def get_current_usage(self, tenant_id: str) -> int:
|
||||
"""
|
||||
Get current usage for the day
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
|
||||
Returns:
|
||||
Number of executions used today
|
||||
"""
|
||||
key = self._get_day_key(tenant_id)
|
||||
return int(self.redis.get(key) or 0)
|
||||
|
||||
def reset_quota(self, tenant_id: str) -> bool:
|
||||
"""
|
||||
Reset quota for testing purposes
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
|
||||
Returns:
|
||||
True if key was deleted, False if key didn't exist
|
||||
"""
|
||||
key = self._get_day_key(tenant_id)
|
||||
return bool(self.redis.delete(key))
|
||||
|
||||
def get_quota_reset_time(self, tenant_id: str, timezone_str: str) -> datetime:
|
||||
"""
|
||||
Get the time when quota will reset (next UTC midnight in tenant's timezone)
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant identifier
|
||||
timezone_str: Tenant's timezone for display purposes
|
||||
|
||||
Returns:
|
||||
Datetime when quota resets (next UTC midnight in tenant's timezone)
|
||||
"""
|
||||
tz = pytz.timezone(timezone_str)
|
||||
utc_now = datetime.now(UTC)
|
||||
|
||||
# Get next midnight in UTC, then convert to tenant's timezone
|
||||
next_utc_midnight = datetime.combine(utc_now.date() + timedelta(days=1), time.min)
|
||||
next_utc_midnight = pytz.UTC.localize(next_utc_midnight)
|
||||
|
||||
return next_utc_midnight.astimezone(tz)
|
||||
@@ -7,6 +7,7 @@ from typing import Any, cast
|
||||
from sqlalchemy import exists, select
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
from configs import dify_config
|
||||
from core.app.app_config.entities import VariableEntityType
|
||||
from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfigManager
|
||||
from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager
|
||||
@@ -25,6 +26,7 @@ from core.workflow.nodes.node_mapping import LATEST_VERSION, NODE_TYPE_CLASSES_M
|
||||
from core.workflow.nodes.start.entities import StartNodeData
|
||||
from core.workflow.system_variable import SystemVariable
|
||||
from core.workflow.workflow_entry import WorkflowEntry
|
||||
from enums.cloud_plan import CloudPlan
|
||||
from events.app_event import app_draft_workflow_was_synced, app_published_workflow_was_updated
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_storage import storage
|
||||
@@ -35,8 +37,9 @@ from models.model import App, AppMode
|
||||
from models.tools import WorkflowToolProvider
|
||||
from models.workflow import Workflow, WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom, WorkflowType
|
||||
from repositories.factory import DifyAPIRepositoryFactory
|
||||
from services.billing_service import BillingService
|
||||
from services.enterprise.plugin_manager_service import PluginCredentialType
|
||||
from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError
|
||||
from services.errors.app import IsDraftWorkflowError, TriggerNodeLimitExceededError, WorkflowHashNotEqualError
|
||||
from services.workflow.workflow_converter import WorkflowConverter
|
||||
|
||||
from .errors.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError
|
||||
@@ -272,6 +275,21 @@ class WorkflowService:
|
||||
# validate graph structure
|
||||
self.validate_graph_structure(graph=draft_workflow.graph_dict)
|
||||
|
||||
# billing check
|
||||
if dify_config.BILLING_ENABLED:
|
||||
limit_info = BillingService.get_info(app_model.tenant_id)
|
||||
if limit_info["subscription"]["plan"] == CloudPlan.SANDBOX:
|
||||
# Check trigger node count limit for SANDBOX plan
|
||||
trigger_node_count = sum(
|
||||
1
|
||||
for _, node_data in draft_workflow.walk_nodes()
|
||||
if (node_type_str := node_data.get("type"))
|
||||
and isinstance(node_type_str, str)
|
||||
and NodeType(node_type_str).is_trigger_node
|
||||
)
|
||||
if trigger_node_count > 2:
|
||||
raise TriggerNodeLimitExceededError(count=trigger_node_count, limit=2)
|
||||
|
||||
# create new workflow
|
||||
workflow = Workflow.new(
|
||||
tenant_id=app_model.tenant_id,
|
||||
|
||||
Reference in New Issue
Block a user