feat(guardrails): bound every guardrail check with litellm_params.timeout

litellm/integrations/custom_guardrail.py

@@ -1,12 +1,14 @@
+import asyncio
 import contextvars
 import copy
 import hashlib
 import os
 import secrets
-from collections.abc import Mapping, Sequence
+from collections.abc import Awaitable, Mapping, Sequence
+from dataclasses import dataclass
 from datetime import datetime
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
+from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeVar, get_args
 
 from litellm._logging import verbose_logger
 from litellm.caching import DualCache
@@ -51,6 +53,7 @@
     GuardrailRaisedException,
     ModifyResponseException,
     SensitiveDataRouteException,
+    Timeout,
 )
 
 # Per-process secret tagging each recorded marker. The deployment hook only
@@ -167,6 +170,60 @@ def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None
     return None
 
 
+_T = TypeVar("_T")
+
+
+@dataclass(frozen=True)
+class GuardrailHookTimedOut:
+    guardrail_name: str | None
+    timeout: float
+    fail_open: bool
+
+    @property
+    def message(self) -> str:
+        return f"Guardrail '{self.guardrail_name}' did not finish within {self.timeout}s"
+
+
+@dataclass(frozen=True)
+class _HookRaisedTimeout:
+    error: TimeoutError | asyncio.TimeoutError
+
+
+async def _capture_hook_timeout(coro: Awaitable[_T]) -> _T | _HookRaisedTimeout:
+    try:
+        return await coro
+    except (TimeoutError, asyncio.TimeoutError) as e:
+        return _HookRaisedTimeout(e)
+
+
+async def await_within_hook_timeout(callback: object, coro: Awaitable[_T]) -> _T | GuardrailHookTimedOut:
+    """Await a guardrail check, bounded by the callback's configured ``hook_timeout``."""
+    if not isinstance(callback, CustomGuardrail) or callback.hook_timeout is None:
+        return await coro
+    timeout: Final = callback.hook_timeout
+    try:
+        outcome: Final = await asyncio.wait_for(_capture_hook_timeout(coro), timeout)
+    except (TimeoutError, asyncio.TimeoutError):
+        return GuardrailHookTimedOut(
+            guardrail_name=callback.guardrail_name,
+            timeout=timeout,
+            fail_open=callback.hook_timeout_fallback == "fail_open",
+        )
+    if isinstance(outcome, _HookRaisedTimeout):
+        raise outcome.error
+    return outcome
+
+
+def resolve_hook_timeout(outcome: _T | GuardrailHookTimedOut, fail_open_result: _T) -> _T:
+    """Map a timed-out check to ``fail_open_result`` or a 408 ``Timeout``, per ``unreachable_fallback``."""
+    if not isinstance(outcome, GuardrailHookTimedOut):
+        return outcome
+    if outcome.fail_open:
+        verbose_logger.critical("%s; unreachable_fallback=fail_open, skipping it", outcome.message)
+        return fail_open_result
+    raise Timeout(message=outcome.message, model="", llm_provider="guardrail")
+
+
 class CustomGuardrail(CustomLogger):
     # If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
     use_native_during_call_hook: ClassVar[bool] = False
@@ -176,6 +233,12 @@ class CustomGuardrail(CustomLogger):
 
     records_own_guardrail_information: ClassVar[bool] = False
 
+    # If True, this guardrail bounds itself by litellm_params.timeout, so the proxy sets no hook_timeout for it.
+    enforces_own_timeout: ClassVar[bool] = False
+
+    hook_timeout: float | None = None
+    hook_timeout_fallback: Literal["fail_closed", "fail_open"] = "fail_closed"
+
     def __init_subclass__(cls, **kwargs: object) -> None:  # kwargs-ok: forwarded to cooperative __init_subclass__ hooks
         super().__init_subclass__(**kwargs)
         own_apply_guardrail: Final[object] = cls.__dict__.get("apply_guardrail")

litellm/integrations/prometheus.py

@@ -3222,7 +3222,7 @@ def _record_guardrail_metrics(
         Args:
             guardrail_name: Name of the guardrail
             latency_seconds: Execution latency in seconds
-            status: "success", "error", or "intervened"
+            status: "success", "error", "intervened", or "timeout"
             error_type: Type of error if any, None otherwise
             hook_type: "pre_call", "during_call", or "post_call"
         """
@@ -3243,7 +3243,7 @@ def _record_guardrail_metrics(
             ).inc()
 
             # Record error count if there was an error
-            if status == "error" and error_type:
+            if status in ("error", "timeout") and error_type:
                 self.litellm_guardrail_errors_total.labels(
                     guardrail_name=guardrail_name,
                     error_type=error_type,

litellm/proxy/_lazy_openapi_snapshot.json

@@ -10377,12 +10377,12 @@
                   "type": "null"
                 }
               ],
-              "description": "Per-request timeout for the guardrail provider API call (seconds). Accepts int, float, or numeric string; coerced to float on load. Each guardrail handler chooses its own default when unset.",
+              "description": "Maximum seconds one non-streaming guardrail check may take, covering every call it makes. On expiry the check fails with a 408 timeout, or is skipped when unreachable_fallback='fail_open'. Streaming response hooks are not bounded by it. Some handlers also apply it to their provider HTTP call; custom_code applies it as its own execution timeout instead. Accepts int, float, or numeric string; coerced to float on load. Unset means no proxy-side limit; provider HTTP client defaults still apply.",
               "title": "Timeout"
             },
             "unreachable_fallback": {
               "default": "fail_closed",
-              "description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed.",
+              "description": "Behavior when a guardrail endpoint is unreachable due to network errors. Implemented by guardrail='generic_guardrail_api', 'agent_365', 'akto', 'vigil_guard', 'repelloai', 'headroom', 'compresr', and 'typesafe'. 'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed. Every guardrail also honors it when its check exceeds timeout.",
               "enum": [
                 "fail_closed",
                 "fail_open"
@@ -13173,7 +13173,7 @@
                   "type": "null"
                 }
               ],
-              "description": "Per-request timeout for the guardrail provider API call (seconds). Accepts int, float, or numeric string; coerced to float on load. Each guardrail handler chooses its own default when unset.",
+              "description": "Maximum seconds one non-streaming guardrail check may take, covering every call it makes. On expiry the check fails with a 408 timeout, or is skipped when unreachable_fallback='fail_open'. Streaming response hooks are not bounded by it. Some handlers also apply it to their provider HTTP call; custom_code applies it as its own execution timeout instead. Accepts int, float, or numeric string; coerced to float on load. Unset means no proxy-side limit; provider HTTP client defaults still apply.",
               "title": "Timeout"
             },
             "tool_selection_quality_check": {

litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py

@@ -41,7 +41,7 @@ async def apply_guardrail(inputs, request_data, input_type):
 import time
 from collections.abc import Callable, Mapping
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Final, Literal, Optional, cast
+from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional, cast
 
 from fastapi import HTTPException
 from pydantic import Field
@@ -138,6 +138,8 @@ def apply_guardrail(inputs, request_data, input_type):
             return allow()
     """
 
+    enforces_own_timeout: ClassVar[bool] = True
+
     def __init__(
         self,
         custom_code: str,

litellm/proxy/guardrails/guardrail_registry.py

@@ -459,6 +459,10 @@ def _configure_callback_scoping(
             "scanning, so no request content would ever be scanned. Remove one of the two."
         )
     _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params)
+    custom_guardrail_callback.hook_timeout = (
+        None if custom_guardrail_callback.enforces_own_timeout else litellm_params.timeout
+    )
+    custom_guardrail_callback.hook_timeout_fallback = litellm_params.unreachable_fallback
 
 
 class InMemoryGuardrailHandler:

litellm/proxy/policy_engine/pipeline_executor.py

@@ -18,6 +18,8 @@
 from litellm.integrations.custom_guardrail import (
     CustomGuardrail,
     ModifyResponseException,
+    await_within_hook_timeout,
+    resolve_hook_timeout,
 )
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.core_helpers import (
@@ -528,11 +530,17 @@ async def _run_step(
 
         try:
             if mode == "pre_call":
-                response = await target.async_pre_call_hook(
-                    user_api_key_dict=user_api_key_dict,
-                    cache=None,
-                    data=hook_input,
-                    call_type=call_type,
+                response = resolve_hook_timeout(
+                    await await_within_hook_timeout(
+                        callback,
+                        target.async_pre_call_hook(
+                            user_api_key_dict=user_api_key_dict,
+                            cache=None,
+                            data=hook_input,
+                            call_type=call_type,
+                        ),
+                    ),
+                    fail_open_result=None,
                 )
                 if isinstance(callback, CustomGuardrail):
                     callback.mark_pre_call_hook_ran(data)
@@ -557,10 +565,16 @@ async def _run_step(
                 )
                 response = None
             elif mode == "post_call":
-                response = await target.async_post_call_success_hook(
-                    user_api_key_dict=user_api_key_dict,
-                    data=data,
-                    response=data.get("response"),
+                response = resolve_hook_timeout(
+                    await await_within_hook_timeout(
+                        callback,
+                        target.async_post_call_success_hook(
+                            user_api_key_dict=user_api_key_dict,
+                            data=data,
+                            response=data.get("response"),
+                        ),
+                    ),
+                    fail_open_result=None,
                 )
             else:
                 return ("error", None, f"Unsupported pipeline mode: {mode}", None)