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:
Maries
2025-11-20 10:15:23 +08:00
committed by GitHub
parent c0b7ffd5d0
commit a1b735a4c0
61 changed files with 1475 additions and 465 deletions

View File

@@ -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:

View File

@@ -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()

View File

@@ -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

View File

@@ -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."
)

View File

@@ -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"]

View 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)

View File

@@ -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,

View File

@@ -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

View File

@@ -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)

View File

@@ -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,