Add app-level Redis-based RateLimiter (#4336)

skyvern/forge/forge_app.py

@@ -21,6 +21,7 @@ from skyvern.forge.sdk.artifact.storage.factory import StorageFactory
 from skyvern.forge.sdk.artifact.storage.s3 import S3Storage
 from skyvern.forge.sdk.cache.base import BaseCache
 from skyvern.forge.sdk.cache.factory import CacheFactory
+from skyvern.forge.sdk.core.rate_limiter import NoopRateLimiter, RateLimiter
 from skyvern.forge.sdk.db.agent_db import AgentDB
 from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider
 from skyvern.forge.sdk.schemas.credentials import CredentialVaultType
@@ -50,6 +51,7 @@ class ForgeApp:
     ARTIFACT_MANAGER: ArtifactManager
     BROWSER_MANAGER: BrowserManager
     EXPERIMENTATION_PROVIDER: BaseExperimentationProvider
+    RATE_LIMITER: RateLimiter
     LLM_API_HANDLER: LLMAPIHandler
     OPENAI_CLIENT: AsyncOpenAI | AsyncAzureOpenAI
     ANTHROPIC_CLIENT: AsyncAnthropic | AsyncAnthropicBedrock
@@ -107,6 +109,7 @@ def create_forge_app() -> ForgeApp:
     app.ARTIFACT_MANAGER = ArtifactManager()
     app.BROWSER_MANAGER = RealBrowserManager()
     app.EXPERIMENTATION_PROVIDER = NoOpExperimentationProvider()
+    app.RATE_LIMITER = NoopRateLimiter()
 
     app.LLM_API_HANDLER = LLMAPIHandlerFactory.get_llm_api_handler(settings.LLM_KEY)
     app.OPENAI_CLIENT = AsyncOpenAI(

skyvern/forge/sdk/core/rate_limiter.py

@@ -0,0 +1,33 @@
+from typing import Protocol
+
+
+class RateLimiter(Protocol):
+    """
+    Protocol for rate limiting submit run requests per organization.
+
+    Implementations should be thread-safe and work correctly in distributed environments.
+    """
+
+    async def rate_limit_submit_run(self, organization_id: str) -> None:
+        """
+        Check and enforce rate limit for submitting a new run (task/workflow)
+        raises RateLimitExceeded exception if rate limit is exceeded.
+
+        Args:
+            organization_id: The organization ID to rate limit
+
+        Raises:
+            Exception: If rate limit is exceeded (implementation-specific exception)
+        """
+        ...
+
+
+class NoopRateLimiter(RateLimiter):
+    """
+    No-op rate limiter.
+
+    This implementation does not enforce any rate limits.
+    """
+
+    async def rate_limit_submit_run(self, organization_id: str) -> None:
+        """No-op implementation that never rate limits."""

skyvern/forge/sdk/routes/agent_protocol.py

@@ -162,6 +162,7 @@ async def run_task(
 ) -> TaskRunResponse:
     analytics.capture("skyvern-oss-run-task", data={"url": run_request.url})
     await PermissionCheckerFactory.get_instance().check(current_org, browser_session_id=run_request.browser_session_id)
+    await app.RATE_LIMITER.rate_limit_submit_run(current_org.organization_id)
 
     if run_request.engine in CUA_ENGINES or run_request.engine == RunEngine.skyvern_v1:
         # create task v1
@@ -347,6 +348,7 @@ async def run_workflow(
     await PermissionCheckerFactory.get_instance().check(
         current_org, browser_session_id=workflow_run_request.browser_session_id
     )
+    await app.RATE_LIMITER.rate_limit_submit_run(current_org.organization_id)
     workflow_id = workflow_run_request.workflow_id
     context = skyvern_context.ensure_context()
     request_id = context.request_id
@@ -1623,6 +1625,7 @@ async def run_task_v1(
 ) -> CreateTaskResponse:
     analytics.capture("skyvern-oss-agent-task-create", data={"url": task.url})
     await PermissionCheckerFactory.get_instance().check(current_org, browser_session_id=task.browser_session_id)
+    await app.RATE_LIMITER.rate_limit_submit_run(current_org.organization_id)
 
     created_task = await task_v1_service.run_task(
         task=task,
@@ -2086,6 +2089,7 @@ async def run_workflow_legacy(
         current_org,
         browser_session_id=workflow_request.browser_session_id,
     )
+    await app.RATE_LIMITER.rate_limit_submit_run(current_org.organization_id)
 
     try:
         workflow_run = await workflow_service.run_workflow(
@@ -2667,6 +2671,7 @@ async def run_task_v2(
             max_steps_override=x_max_steps_override,
         )
     await PermissionCheckerFactory.get_instance().check(organization, browser_session_id=data.browser_session_id)
+    await app.RATE_LIMITER.rate_limit_submit_run(organization.organization_id)
 
     try:
         task_v2 = await task_v2_service.initialize_task_v2(