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>

api/tasks/trigger_processing_tasks.py

@@ -26,14 +26,22 @@ from core.trigger.provider import PluginTriggerProviderController
 from core.trigger.trigger_manager import TriggerManager
 from core.workflow.enums import NodeType, WorkflowExecutionStatus
 from core.workflow.nodes.trigger_plugin.entities import TriggerEventNodeData
+from enums.quota_type import QuotaType, unlimited
 from extensions.ext_database import db
-from models.enums import AppTriggerType, CreatorUserRole, WorkflowRunTriggeredFrom, WorkflowTriggerStatus
+from models.enums import (
+    AppTriggerType,
+    CreatorUserRole,
+    WorkflowRunTriggeredFrom,
+    WorkflowTriggerStatus,
+)
 from models.model import EndUser
 from models.provider_ids import TriggerProviderID
 from models.trigger import TriggerSubscription, WorkflowPluginTrigger, WorkflowTriggerLog
 from models.workflow import Workflow, WorkflowAppLog, WorkflowAppLogCreatedFrom, WorkflowRun
 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.trigger.trigger_provider_service import TriggerProviderService
 from services.trigger.trigger_request_service import TriggerHttpRequestCachingService
 from services.trigger.trigger_subscription_operator_service import TriggerSubscriptionOperatorService
@@ -287,6 +295,17 @@ def dispatch_triggered_workflow(
                 icon_dark_filename=trigger_entity.identity.icon_dark or "",
             )
 
+            # consume quota before invoking trigger
+            quota_charge = unlimited()
+            try:
+                quota_charge = QuotaType.TRIGGER.consume(subscription.tenant_id)
+            except QuotaExceededError:
+                AppTriggerService.mark_tenant_triggers_rate_limited(subscription.tenant_id)
+                logger.info(
+                    "Tenant %s rate limited, skipping plugin trigger %s", subscription.tenant_id, plugin_trigger.id
+                )
+                return 0
+
             node_data: TriggerEventNodeData = TriggerEventNodeData.model_validate(event_node)
             invoke_response: TriggerInvokeEventResponse | None = None
             try:
@@ -305,6 +324,8 @@ def dispatch_triggered_workflow(
                     payload=payload,
                 )
             except PluginInvokeError as e:
+                quota_charge.refund()
+
                 error_message = e.to_user_friendly_error(plugin_name=trigger_entity.identity.name)
                 try:
                     end_user = end_users.get(plugin_trigger.app_id)
@@ -326,6 +347,8 @@ def dispatch_triggered_workflow(
                     )
                 continue
             except Exception:
+                quota_charge.refund()
+
                 logger.exception(
                     "Failed to invoke trigger event for app %s",
                     plugin_trigger.app_id,
@@ -333,6 +356,8 @@ def dispatch_triggered_workflow(
                 continue
 
             if invoke_response is not None and invoke_response.cancelled:
+                quota_charge.refund()
+
                 logger.info(
                     "Trigger ignored for app %s with trigger event %s",
                     plugin_trigger.app_id,
@@ -366,6 +391,8 @@ def dispatch_triggered_workflow(
                     event_name,
                 )
             except Exception:
+                quota_charge.refund()
+
                 logger.exception(
                     "Failed to trigger workflow for app %s",
                     plugin_trigger.app_id,

api/tasks/trigger_subscription_refresh_tasks.py

@@ -6,6 +6,7 @@ from typing import Any
 from celery import shared_task
 from sqlalchemy.orm import Session
 
+from configs import dify_config
 from core.plugin.entities.plugin_daemon import CredentialType
 from core.trigger.utils.locks import build_trigger_refresh_lock_key
 from extensions.ext_database import db
@@ -25,9 +26,10 @@ def _load_subscription(session: Session, tenant_id: str, subscription_id: str) -
 
 
 def _refresh_oauth_if_expired(tenant_id: str, subscription: TriggerSubscription, now: int) -> None:
+    threshold_seconds: int = int(dify_config.TRIGGER_PROVIDER_CREDENTIAL_THRESHOLD_SECONDS)
     if (
         subscription.credential_expires_at != -1
-        and int(subscription.credential_expires_at) <= now
+        and int(subscription.credential_expires_at) <= now + threshold_seconds
         and CredentialType.of(subscription.credential_type) == CredentialType.OAUTH2
     ):
         logger.info(
@@ -53,13 +55,15 @@ def _refresh_subscription_if_expired(
     subscription: TriggerSubscription,
     now: int,
 ) -> None:
-    if subscription.expires_at == -1 or int(subscription.expires_at) > now:
+    threshold_seconds: int = int(dify_config.TRIGGER_PROVIDER_SUBSCRIPTION_THRESHOLD_SECONDS)
+    if subscription.expires_at == -1 or int(subscription.expires_at) > now + threshold_seconds:
         logger.debug(
-            "Subscription not due: tenant=%s subscription_id=%s expires_at=%s now=%s",
+            "Subscription not due: tenant=%s subscription_id=%s expires_at=%s now=%s threshold=%s",
             tenant_id,
             subscription.id,
             subscription.expires_at,
             now,
+            threshold_seconds,
         )
         return
 

api/tasks/workflow_schedule_tasks.py

@@ -8,9 +8,12 @@ from core.workflow.nodes.trigger_schedule.exc import (
     ScheduleNotFoundError,
     TenantOwnerNotFoundError,
 )
+from enums.quota_type import QuotaType, unlimited
 from extensions.ext_database import db
 from models.trigger import WorkflowSchedulePlan
 from services.async_workflow_service import AsyncWorkflowService
+from services.errors.app import QuotaExceededError
+from services.trigger.app_trigger_service import AppTriggerService
 from services.trigger.schedule_service import ScheduleService
 from services.workflow.entities import ScheduleTriggerData
 
@@ -30,6 +33,7 @@ def run_schedule_trigger(schedule_id: str) -> None:
         TenantOwnerNotFoundError: If no owner/admin for tenant
         ScheduleExecutionError: If workflow trigger fails
     """
+
     session_factory = sessionmaker(bind=db.engine, expire_on_commit=False)
 
     with session_factory() as session:
@@ -41,6 +45,14 @@ def run_schedule_trigger(schedule_id: str) -> None:
         if not tenant_owner:
             raise TenantOwnerNotFoundError(f"No owner or admin found for tenant {schedule.tenant_id}")
 
+        quota_charge = unlimited()
+        try:
+            quota_charge = QuotaType.TRIGGER.consume(schedule.tenant_id)
+        except QuotaExceededError:
+            AppTriggerService.mark_tenant_triggers_rate_limited(schedule.tenant_id)
+            logger.info("Tenant %s rate limited, skipping schedule trigger %s", schedule.tenant_id, schedule_id)
+            return
+
         try:
             # Production dispatch: Trigger the workflow normally
             response = AsyncWorkflowService.trigger_workflow_async(
@@ -55,6 +67,7 @@ def run_schedule_trigger(schedule_id: str) -> None:
             )
             logger.info("Schedule %s triggered workflow: %s", schedule_id, response.workflow_trigger_log_id)
         except Exception as e:
+            quota_charge.refund()
             raise ScheduleExecutionError(
                 f"Failed to trigger workflow for schedule {schedule_id}, app {schedule.app_id}"
             ) from e