diff --git a/api/controllers/trigger/webhook.py b/api/controllers/trigger/webhook.py index 4747ab9c22d..04b2b50bc4a 100644 --- a/api/controllers/trigger/webhook.py +++ b/api/controllers/trigger/webhook.py @@ -7,11 +7,26 @@ from werkzeug.exceptions import NotFound, RequestEntityTooLarge from controllers.trigger import bp from core.trigger.debug.event_bus import TriggerDebugEventBus from core.trigger.debug.events import WebhookDebugEvent, build_webhook_pool_key +from enums.quota_type import QuotaType from services.errors.app import QuotaExceededError from services.trigger.webhook_service import RawWebhookDataDict, WebhookService logger = logging.getLogger(__name__) +_QUOTA_EXCEEDED_MESSAGES = { + QuotaType.TRIGGER: "Trigger event quota exceeded. Please upgrade your plan.", + QuotaType.WORKFLOW: "Workflow execution quota exceeded. Please upgrade your plan.", +} +_DEFAULT_QUOTA_EXCEEDED_MESSAGE = "Quota exceeded. Please upgrade your plan." + + +def _get_quota_exceeded_message(feature: str) -> str: + try: + quota_type = QuotaType(feature) + except ValueError: + return _DEFAULT_QUOTA_EXCEEDED_MESSAGE + return _QUOTA_EXCEEDED_MESSAGES.get(quota_type, _DEFAULT_QUOTA_EXCEEDED_MESSAGE) + def _prepare_webhook_execution(webhook_id: str, is_debug: bool = False): """Fetch trigger context, extract request data, and validate payload using unified processing. @@ -61,11 +76,11 @@ def handle_webhook(webhook_id: str): response_data, status_code = WebhookService.generate_webhook_response(node_config) return jsonify(response_data), status_code - except QuotaExceededError: + except QuotaExceededError as error: return jsonify( { "error": "Too Many Requests", - "message": "Trigger event quota exceeded. Please upgrade your plan.", + "message": _get_quota_exceeded_message(error.feature), } ), 429 except ValueError as error: diff --git a/api/services/async_workflow_service.py b/api/services/async_workflow_service.py index 601cad7557a..012e979aff2 100644 --- a/api/services/async_workflow_service.py +++ b/api/services/async_workflow_service.py @@ -6,6 +6,7 @@ with support for different subscription tiers, rate limiting, and execution trac """ import json +import logging from datetime import UTC, datetime from typing import Any @@ -21,7 +22,7 @@ from models.model import App, EndUser from models.trigger import WorkflowTriggerLog, WorkflowTriggerLogDict from models.workflow import Workflow from repositories.sqlalchemy_workflow_trigger_log_repository import SQLAlchemyWorkflowTriggerLogRepository -from services.errors.app import QuotaExceededError, WorkflowNotFoundError, WorkflowQuotaLimitError +from services.errors.app import QuotaExceededError, WorkflowNotFoundError from services.quota_service import QuotaService, unlimited from services.workflow.entities import AsyncTriggerResponse, TriggerData, WorkflowTaskData from services.workflow.queue_dispatcher import QueueDispatcherManager, QueuePriority @@ -32,6 +33,8 @@ from tasks.async_workflow_tasks import ( execute_workflow_team, ) +logger = logging.getLogger(__name__) + class AsyncWorkflowService: """ @@ -71,7 +74,7 @@ class AsyncWorkflowService: Raises: WorkflowNotFoundError: If app or workflow not found - InvokeDailyRateLimitError: If daily rate limit exceeded + QuotaExceededError: If workflow execution quota is exhausted Behavior: - Non-blocking: Returns immediately after queuing @@ -145,10 +148,15 @@ class AsyncWorkflowService: trigger_log.error = f"Quota limit reached: {e}" trigger_log_repo.update(trigger_log) session.commit() + logger.info( + "Workflow quota exceeded for tenant %s, app %s, workflow %s, trigger log %s", + trigger_data.tenant_id, + trigger_data.app_id, + workflow.id, + trigger_log.id, + ) - raise WorkflowQuotaLimitError( - f"Workflow execution quota limit reached for tenant {trigger_data.tenant_id}" - ) from e + raise # 8. Create task data queue_name = dispatcher.get_queue_name() @@ -206,6 +214,7 @@ class AsyncWorkflowService: Raises: ValueError: If trigger log not found + QuotaExceededError: If workflow execution quota is exhausted Behavior: - Non-blocking: Returns immediately after queuing retry diff --git a/api/services/errors/app.py b/api/services/errors/app.py index 60e59e97dc2..c9e9df97dea 100644 --- a/api/services/errors/app.py +++ b/api/services/errors/app.py @@ -18,12 +18,6 @@ class WorkflowIdFormatError(Exception): pass -class WorkflowQuotaLimitError(Exception): - """Raised when workflow execution quota is exceeded (for async/background workflows).""" - - pass - - class QuotaExceededError(ValueError): """Raised when billing quota is exceeded for a feature.""" diff --git a/api/services/trigger/webhook_service.py b/api/services/trigger/webhook_service.py index bc0dcc0b5d0..a921c91d64a 100644 --- a/api/services/trigger/webhook_service.py +++ b/api/services/trigger/webhook_service.py @@ -802,7 +802,7 @@ class WebhookService: workflow: The workflow to execute Raises: - QuotaExceededError: If the tenant has exhausted its trigger quota + QuotaExceededError: If the tenant has exhausted its trigger or workflow execution quota ValueError: If tenant owner is not found Exception: If workflow execution fails """ @@ -828,11 +828,6 @@ class WebhookService: quota_charge = QuotaService.reserve(QuotaType.TRIGGER, 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 try: @@ -845,6 +840,14 @@ class WebhookService: quota_charge.refund() raise + except QuotaExceededError as e: + logger.info( + "Tenant %s quota exceeded for feature %s, skipping webhook trigger %s", + webhook_trigger.tenant_id, + e.feature, + webhook_trigger.webhook_id, + ) + raise except Exception: logger.exception("Failed to trigger workflow for webhook %s", webhook_trigger.webhook_id) raise diff --git a/api/tests/unit_tests/controllers/trigger/test_webhook.py b/api/tests/unit_tests/controllers/trigger/test_webhook.py index 7e591af6e96..9acf57089fc 100644 --- a/api/tests/unit_tests/controllers/trigger/test_webhook.py +++ b/api/tests/unit_tests/controllers/trigger/test_webhook.py @@ -103,6 +103,25 @@ class TestHandleWebhook: "message": "Trigger event quota exceeded. Please upgrade your plan.", } + @patch.object(module.WebhookService, "get_webhook_trigger_and_workflow") + @patch.object(module.WebhookService, "extract_and_validate_webhook_data") + @patch.object( + module.WebhookService, + "trigger_workflow_execution", + side_effect=QuotaExceededError(feature="workflow", tenant_id="tenant-1", required=1), + ) + def test_workflow_quota_exceeded(self, mock_trigger, mock_extract, mock_get): + mock_get.return_value = (DummyWebhookTrigger(), "workflow", "node_config") + mock_extract.return_value = {"input": "x"} + + response, status = module.handle_webhook("wh-1") + + assert status == 429 + assert response == { + "error": "Too Many Requests", + "message": "Workflow execution quota exceeded. Please upgrade your plan.", + } + @patch.object( module.WebhookService, "get_webhook_trigger_and_workflow", diff --git a/api/tests/unit_tests/services/test_async_workflow_service.py b/api/tests/unit_tests/services/test_async_workflow_service.py index 567066845bf..93599363c4f 100644 --- a/api/tests/unit_tests/services/test_async_workflow_service.py +++ b/api/tests/unit_tests/services/test_async_workflow_service.py @@ -1,4 +1,5 @@ import json +import logging from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -7,7 +8,7 @@ import pytest import services.async_workflow_service as async_workflow_service_module from models.enums import AppTriggerType, CreatorUserRole, WorkflowRunTriggeredFrom, WorkflowTriggerStatus from services.async_workflow_service import AsyncWorkflowService -from services.errors.app import QuotaExceededError, WorkflowNotFoundError, WorkflowQuotaLimitError +from services.errors.app import QuotaExceededError, WorkflowNotFoundError from services.workflow.entities import AsyncTriggerResponse, TriggerData from services.workflow.queue_dispatcher import QueuePriority @@ -234,8 +235,10 @@ class TestAsyncWorkflowService: trigger_data=trigger_data, ) - def test_should_mark_log_rate_limited_and_raise_when_quota_exceeded(self, async_workflow_trigger_mocks): - """Test quota-exceeded path updates trigger log and raises WorkflowQuotaLimitError.""" + def test_should_mark_log_rate_limited_and_reraise_when_quota_exceeded( + self, async_workflow_trigger_mocks, caplog: pytest.LogCaptureFixture + ): + """Test quota-exceeded path updates trigger log and preserves the quota exception.""" # Arrange session = MagicMock() session.commit = MagicMock() @@ -254,22 +257,27 @@ class TestAsyncWorkflowService: tenant_id="tenant-123", required=1, ) + caplog.set_level(logging.INFO, logger=async_workflow_service_module.__name__) # Act / Assert - with pytest.raises( - WorkflowQuotaLimitError, - match="Workflow execution quota limit reached for tenant tenant-123", - ): + with pytest.raises(QuotaExceededError) as exc_info: AsyncWorkflowService.trigger_workflow_async( session=session, user=SimpleNamespace(id="user-123"), trigger_data=trigger_data, ) + assert exc_info.value.feature == "workflow" + assert exc_info.value.tenant_id == "tenant-123" + assert exc_info.value.required == 1 assert session.commit.call_count == 3 updated_log = mocks["repo"].update.call_args[0][0] assert updated_log.status == WorkflowTriggerStatus.RATE_LIMITED assert "Quota limit reached" in updated_log.error + assert ( + "Workflow quota exceeded for tenant tenant-123, app app-123, workflow workflow-123, " + "trigger log trigger-log-123" + ) in caplog.messages mocks["professional_task"].delay.assert_not_called() mocks["team_task"].delay.assert_not_called() mocks["sandbox_task"].delay.assert_not_called() diff --git a/api/tests/unit_tests/services/test_webhook_service.py b/api/tests/unit_tests/services/test_webhook_service.py index d36c45d7777..4babcd049c8 100644 --- a/api/tests/unit_tests/services/test_webhook_service.py +++ b/api/tests/unit_tests/services/test_webhook_service.py @@ -6,12 +6,56 @@ import pytest from flask import Flask from werkzeug.datastructures import FileStorage +from services.errors.app import QuotaExceededError from services.trigger.webhook_service import WebhookService class TestWebhookServiceUnit: """Unit tests for WebhookService focusing on business logic without database dependencies.""" + def test_trigger_workflow_execution_propagates_quota_error_without_error_log(self): + webhook_trigger = MagicMock( + webhook_id="webhook-123", + tenant_id="tenant-123", + app_id="app-123", + node_id="node-123", + ) + workflow = MagicMock(id="workflow-123") + quota_charge = MagicMock() + quota_error = QuotaExceededError(feature="workflow", tenant_id="tenant-123", required=1) + + with ( + patch( + "services.trigger.webhook_service.EndUserService.get_or_create_end_user_by_type", + return_value=MagicMock(id="end-user-123"), + ), + patch("services.trigger.webhook_service.QuotaService.reserve", return_value=quota_charge), + patch("services.trigger.webhook_service.db"), + patch("services.trigger.webhook_service.Session"), + patch( + "services.trigger.webhook_service.AsyncWorkflowService.trigger_workflow_async", + side_effect=quota_error, + ), + patch("services.trigger.webhook_service.logger.info") as mock_log_info, + patch("services.trigger.webhook_service.logger.exception") as mock_log_exception, + ): + with pytest.raises(QuotaExceededError) as exc_info: + WebhookService.trigger_workflow_execution( + webhook_trigger, + {"body": {}, "headers": {}, "query_params": {}, "files": {}, "method": "POST"}, + workflow, + ) + + assert exc_info.value is quota_error + quota_charge.refund.assert_called_once_with() + mock_log_info.assert_called_once_with( + "Tenant %s quota exceeded for feature %s, skipping webhook trigger %s", + webhook_trigger.tenant_id, + quota_error.feature, + webhook_trigger.webhook_id, + ) + mock_log_exception.assert_not_called() + def test_extract_webhook_data_json(self): """Test webhook data extraction from JSON request.""" app = Flask(__name__)