From dfb5d905eae8397cd977754ff1e8827b08e47029 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:57:41 -0700 Subject: [PATCH] fix(guardrails): block private destinations in custom code http_request and bound guardrail execution time (#43280) * fix(guardrails): block private destinations in custom code http_request and bound guardrail execution time * fix(guardrails): keep startup fail-closed on a custom code compile error and report a load timeout on the test endpoint A compile failure is no longer a ValueError, so a config-file custom code guardrail that does not compile stops the proxy at startup as it did before, while POST /guardrails catches it by name and still rolls back. The admin test endpoint reports a module-level timeout as an execution timeout instead of a compile error, a caller-supplied Host header is stripped from http_* requests while validation is on, and GET keeps the shared client's connect timeout. * test(guardrails): cover the http_request methods, header passthrough and cancellation paths --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../proxy/guardrails/guardrail_endpoints.py | 49 +- .../guardrail_hooks/custom_code/__init__.py | 1 + .../custom_code/bounded_execution.py | 231 ++++++++ .../custom_code/custom_code_guardrail.py | 92 ++- .../guardrail_hooks/custom_code/primitives.py | 78 ++- .../guardrail_hooks/custom_code/sandbox.py | 27 +- .../test_custom_code_bounded_execution.py | 174 ++++++ .../guardrails/test_custom_code_security.py | 535 +++++++++++++++++- .../guardrails/test_guardrail_endpoints.py | 457 +++++++-------- .../proxy/guardrails/test_init_guardrails.py | 22 + 10 files changed, 1334 insertions(+), 332 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/bounded_execution.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 6874d7aa73e..6053ab26726 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2,7 +2,7 @@ CRUD ENDPOINTS FOR GUARDRAILS """ -import concurrent.futures +import asyncio import inspect import json import os @@ -22,6 +22,12 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.path_utils import safe_join +from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution import ( + ExecutionTimeoutError, + await_with_timeout, + call_off_loop_with_timeout, +) +from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeCompilationError from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( build_sandbox_globals, compile_sandboxed, @@ -401,7 +407,7 @@ async def create_guardrail( verbose_proxy_logger.info( "Immediate sync: Successfully initialized guardrail '%s' (ID: %s)", guardrail_name, guardrail_id ) - except (ValueError, TypeError) as init_error: + except (ValueError, TypeError, CustomCodeCompilationError) as init_error: # Configuration error — roll back the DB write so the guardrail isn't orphaned if prisma_client is not None: try: @@ -421,6 +427,8 @@ async def create_guardrail( ) return result + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception("Error adding guardrail to db: %s", e) raise HTTPException(status_code=500, detail=str(e)) @@ -2124,15 +2132,20 @@ async def test_custom_code_guardrail( try: exec_globals: Final = build_sandbox_globals() - try: + def load_module() -> None: compiled: Final[CodeType] = compile_sandboxed(request.custom_code) exec(compiled, exec_globals) # noqa: S102 + + try: + await call_off_loop_with_timeout(load_module, EXECUTION_TIMEOUT_SECONDS, label="test:load") except SyntaxError as e: return TestCustomCodeGuardrailResponse( success=False, error=f"Syntax error in custom code: {e}", error_type="compilation", ) + except ExecutionTimeoutError: + return _execution_timeout_response(EXECUTION_TIMEOUT_SECONDS) except Exception as e: return TestCustomCodeGuardrailResponse( success=False, @@ -2178,16 +2191,9 @@ async def test_custom_code_guardrail( return apply_fn(test_inputs, safe_request_data, request.input_type) try: - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: - future: Final = executor.submit(execute_guardrail) - try: - result: Final = future.result(timeout=EXECUTION_TIMEOUT_SECONDS) - except concurrent.futures.TimeoutError: - return TestCustomCodeGuardrailResponse( - success=False, - error=f"Execution timeout: code took longer than {EXECUTION_TIMEOUT_SECONDS} seconds", - error_type="execution", - ) + result: Final = await _run_test_guardrail(execute_guardrail, EXECUTION_TIMEOUT_SECONDS) + except ExecutionTimeoutError: + return _execution_timeout_response(EXECUTION_TIMEOUT_SECONDS) except Exception as e: return TestCustomCodeGuardrailResponse( success=False, @@ -2219,6 +2225,23 @@ async def test_custom_code_guardrail( ) +def _execution_timeout_response(timeout: float) -> TestCustomCodeGuardrailResponse: + return TestCustomCodeGuardrailResponse( + success=False, + error=f"Execution timeout: code took longer than {timeout:g} seconds", + error_type="execution", + ) + + +async def _run_test_guardrail(execute_guardrail: Callable[[], object], timeout: float) -> object: + deadline: Final = asyncio.get_running_loop().time() + timeout + raw_result: Final = await call_off_loop_with_timeout(execute_guardrail, timeout, label="test") + if not inspect.iscoroutine(raw_result): + return raw_result + remaining: Final = max(deadline - asyncio.get_running_loop().time(), 0.0) + return await await_with_timeout(raw_result, remaining, label="test") + + def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: str) -> Literal["request", "response"]: """Return the effective input_type, auto-upgrading to 'response' for post_call guardrails.""" if input_type == "request": diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py index e187ba7430a..13d7178dfe9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -46,6 +46,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" custom_code_guardrail: Final = CustomCodeGuardrail( guardrail_name=guardrail_name, custom_code=custom_code, + execution_timeout=litellm_params.timeout, event_hook=litellm_params.mode, default_on=litellm_params.default_on, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/bounded_execution.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/bounded_execution.py new file mode 100644 index 00000000000..d926b40da6a --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/bounded_execution.py @@ -0,0 +1,231 @@ +"""Wall-clock bounds for sandboxed guardrail code. + +Sync guardrail code runs on a dedicated daemon thread so a runaway loop never stalls the event loop; async +guardrail code is awaited as its own task. Either way the code's deadline is published through a context +variable, and the sandbox compiler routes every ``while`` test, ``for`` iteration and comprehension through +:func:`budget_ok`, which raises ``ExecutionInterrupted`` once that deadline has passed, whatever the code +catches around the loop body. As a backstop, a worker thread still running at the deadline has +``ExecutionInterrupted`` injected with ``PyThreadState_SetAsyncExc`` and a task still running is cancelled +repeatedly. A long-running C call (a catastrophic regex, for one) only sees any of this once it returns, so +the caller still gets its timeout on schedule while the worker keeps burning CPU until that call ends. +""" + +import asyncio +import concurrent.futures +import contextvars +import ctypes +import threading +import time +from collections.abc import Awaitable, Callable, Iterable, Iterator +from dataclasses import dataclass +from typing import Final, Generic, TypeVar + +from litellm._logging import verbose_proxy_logger + +T: Final = TypeVar("T") + +_INTERRUPT_GRACE_SECONDS: Final = 1.0 +_INTERRUPT_POLL_SECONDS: Final = 0.05 +_deadline: Final[contextvars.ContextVar[float | None]] = contextvars.ContextVar("guardrail_code_deadline", default=None) + + +class ExecutionInterrupted(BaseException): + """Raised inside guardrail code once its budget is spent; a BaseException so sandboxed + ``except Exception`` clauses cannot swallow it.""" + + +class ExecutionTimeoutError(Exception): + """The guardrail code did not finish within its wall-clock budget.""" + + def __init__(self, timeout: float) -> None: + super().__init__(f"exceeded the {timeout:g}s execution timeout") + self.timeout: Final = timeout + + +class SandboxExit(Exception): + """Sandboxed code raised something outside the ``Exception`` tree (``SystemExit``, ``KeyboardInterrupt``, + a bare ``BaseException``). It is delivered as an ordinary exception so it can neither stop the event loop + nor pass for a timeout.""" + + def __init__(self, cause: BaseException) -> None: + super().__init__(f"{type(cause).__name__}: {cause}") + + +def _past_deadline() -> bool: + deadline: Final = _deadline.get() + return deadline is not None and time.monotonic() > deadline + + +def budget_ok() -> bool: + """Bound to ``_budget_ok_`` in the sandbox, where every ``while`` test starts with a call to it.""" + if _past_deadline(): + raise ExecutionInterrupted + return True + + +def budgeted_iter(iterable: Iterable[T]) -> Iterator[T]: + """Bound to ``_getiter_`` in the sandbox, so every ``for`` loop and comprehension checks the budget per item.""" + for item in iterable: + budget_ok() + yield item + + +class _InterruptGate: + """Aims the interrupt at the worker thread only while it is inside the sandboxed call, so a thread id the + OS recycles after the worker exits is never hit.""" + + def __init__(self) -> None: + self._lock: Final = threading.Lock() + self._thread_id: int | None = None + + def open(self) -> None: + self._thread_id = threading.get_ident() + + def close(self) -> None: + with self._lock: + self._thread_id = None + + def is_open(self) -> bool: + return self._thread_id is not None + + def interrupt(self) -> bool: + with self._lock: + if self._thread_id is None: + return False + ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_ulong(self._thread_id), ctypes.py_object(ExecutionInterrupted) + ) + return True + + +@dataclass(frozen=True, slots=True) +class _Worker(Generic[T]): + thread: threading.Thread + outcome: concurrent.futures.Future[T] + gate: _InterruptGate + + +def _run(fn: Callable[[], T], timeout: float, gate: _InterruptGate) -> tuple[T | None, Exception | None]: + _deadline.set(time.monotonic() + timeout) + gate.open() + try: + result: Final = fn() + except Exception as e: # noqa: BLE001 # every failure is handed to the waiting caller through the future + return None, e + except ExecutionInterrupted: + return None, ExecutionTimeoutError(timeout) + except BaseException as e: # noqa: BLE001 # a SystemExit must reach the caller as a failure, not end the worker silently + return None, SandboxExit(e) + finally: + gate.close() + if _past_deadline(): + return None, ExecutionTimeoutError(timeout) + return result, None + + +def _deliver(fn: Callable[[], T], timeout: float, outcome: concurrent.futures.Future[T], gate: _InterruptGate) -> None: + try: + _settle(outcome, *_run(fn, timeout, gate)) + except ExecutionInterrupted: + _settle(outcome, exception=ExecutionTimeoutError(timeout)) + + +def _settle(outcome: concurrent.futures.Future[T], result: T | None = None, exception: Exception | None = None) -> None: + try: + if exception is not None: + outcome.set_exception(exception) + else: + outcome.set_result(result) # pyright: ignore[reportArgumentType] # result is T whenever exception is None + except concurrent.futures.InvalidStateError: + return + + +def _start_worker(fn: Callable[[], T], timeout: float, label: str) -> _Worker[T]: + outcome: Final[concurrent.futures.Future[T]] = concurrent.futures.Future() + gate: Final = _InterruptGate() + thread: Final = threading.Thread( + target=_deliver, args=(fn, timeout, outcome, gate), name=f"guardrail-code:{label}", daemon=True + ) + thread.start() + return _Worker(thread, outcome, gate) + + +def _interrupt(worker: _Worker[T]) -> None: + deadline: Final = time.monotonic() + _INTERRUPT_GRACE_SECONDS + while worker.gate.interrupt() and time.monotonic() < deadline: + worker.thread.join(_INTERRUPT_POLL_SECONDS) + if worker.gate.is_open(): + verbose_proxy_logger.error( + "%s is still running after its timeout; it is stuck in a call Python cannot interrupt", worker.thread.name + ) + + +def call_with_timeout(fn: Callable[[], T], timeout: float, label: str) -> T: + """Run ``fn`` on a worker thread and wait for it, from sync code.""" + worker: Final = _start_worker(fn, timeout, label) + try: + return worker.outcome.result(timeout=timeout) + except concurrent.futures.TimeoutError: + worker.outcome.cancel() + _interrupt(worker) + raise ExecutionTimeoutError(timeout) from None + + +async def call_off_loop_with_timeout(fn: Callable[[], T], timeout: float, label: str) -> T: + """Run ``fn`` on a worker thread and await it without blocking the event loop.""" + worker: Final = _start_worker(fn, timeout, label) + try: + return await asyncio.wait_for(asyncio.wrap_future(worker.outcome), timeout) + except asyncio.TimeoutError: + await asyncio.to_thread(_interrupt, worker) + raise ExecutionTimeoutError(timeout) from None + except asyncio.CancelledError: + threading.Thread(target=_interrupt, args=(worker,), name=f"guardrail-interrupt:{label}", daemon=True).start() + raise + + +def _discard_outcome(task: asyncio.Future[T]) -> None: + if not task.cancelled(): + task.exception() + + +async def _cancel(task: asyncio.Task[T], label: str) -> None: + deadline: Final = time.monotonic() + _INTERRUPT_GRACE_SECONDS + while not task.done() and time.monotonic() < deadline: + task.cancel() + await asyncio.wait((task,), timeout=_INTERRUPT_POLL_SECONDS) + if not task.done(): + verbose_proxy_logger.error( + "guardrail-code:%s is still running after its timeout; it keeps swallowing cancellation", label + ) + + +async def _contain(pending: Awaitable[T], timeout: float) -> T: + try: + result: Final = await pending + except (Exception, asyncio.CancelledError): + raise + except ExecutionInterrupted: + raise ExecutionTimeoutError(timeout) from None + except BaseException as e: # noqa: BLE001 # a SystemExit escaping a task stops the whole event loop + raise SandboxExit(e) from e + if _past_deadline(): + raise ExecutionTimeoutError(timeout) + return result + + +async def await_with_timeout(pending: Awaitable[object], timeout: float, label: str) -> object: + """Await ``pending`` on the event loop and give it up at the deadline, even if it swallows cancellation.""" + context: Final = contextvars.copy_context() + context.run(_deadline.set, time.monotonic() + timeout) + task: Final = context.run(asyncio.ensure_future, _contain(pending, timeout)) + task.add_done_callback(_discard_outcome) + try: + await asyncio.wait((task,), timeout=timeout) + except asyncio.CancelledError: + task.cancel() + raise + if task.done(): + return task.result() + await _cancel(task, label) + raise ExecutionTimeoutError(timeout) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index ea26eafccae..8505ceeb54a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -35,12 +35,16 @@ Example: block when response rejects the user (input_type response only): """ import asyncio +import functools +import inspect import threading import time from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Optional, cast from fastapi import HTTPException +from pydantic import Field from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger @@ -53,11 +57,19 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs +from .bounded_execution import ( + ExecutionTimeoutError, + await_with_timeout, + call_off_loop_with_timeout, + call_with_timeout, +) from .sandbox import build_sandbox_globals, compile_sandboxed if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +DEFAULT_EXECUTION_TIMEOUT_SECONDS: Final = 30.0 + def _metadata_bucket(request_data: Mapping[str, object], key: str) -> Mapping[str, object]: bucket: Final = request_data.get(key) @@ -73,7 +85,8 @@ class CustomCodeGuardrailError(Exception): class CustomCodeCompilationError(CustomCodeGuardrailError): - """Raised when custom code fails to compile.""" + """Raised when custom code fails to compile. Deliberately not a ValueError: a config-file guardrail whose + code does not compile must stop startup instead of being skipped, so the guardrail endpoints catch it by name.""" class CustomCodeExecutionError(CustomCodeGuardrailError): @@ -90,6 +103,15 @@ class CustomCodeGuardrailConfigModel(GuardrailConfigModel): custom_code: str """The Python-like code containing the apply_guardrail function.""" + timeout: float | None = Field( + default=DEFAULT_EXECUTION_TIMEOUT_SECONDS, + gt=0.0, + description=( + "Wall-clock limit in seconds for one run of apply_guardrail, module-level code included. " + "A run that exceeds it fails the request instead of stalling the proxy." + ), + ) + class CustomCodeGuardrail(CustomGuardrail): """ @@ -97,7 +119,8 @@ class CustomCodeGuardrail(CustomGuardrail): The code runs in a sandboxed environment that provides: - Access to LiteLLM primitives (regex_match, json_parse, etc.) - - No file I/O or network access + - No file I/O; network access only through `http_get`/`http_post`/`http_request`, which refuse + private, link-local and loopback destinations unless the host is allowlisted - No imports allowed Users write an `apply_guardrail(inputs, request_data, input_type)` function @@ -119,6 +142,7 @@ class CustomCodeGuardrail(CustomGuardrail): self, custom_code: str, guardrail_name: str | None = "custom_code", + execution_timeout: float | None = None, **kwargs: Unpack[_CustomGuardrailOptions], ) -> None: """ @@ -127,9 +151,15 @@ class CustomCodeGuardrail(CustomGuardrail): Args: custom_code: The source code containing apply_guardrail function guardrail_name: Name of this guardrail instance + execution_timeout: Wall-clock budget in seconds for one run of the code **kwargs: Additional arguments passed to CustomGuardrail """ + if execution_timeout is not None and not execution_timeout > 0: + raise ValueError(f"execution_timeout must be positive, got {execution_timeout}") self.custom_code: str = custom_code + self.execution_timeout: float = ( + DEFAULT_EXECUTION_TIMEOUT_SECONDS if execution_timeout is None else execution_timeout + ) self._compiled_function: Callable[..., object] | None = None self._compile_lock = threading.Lock() self._compile_error: str | None = None @@ -163,7 +193,11 @@ class CustomCodeGuardrail(CustomGuardrail): """Internal compilation method without lock. Expected to run inside _compile_lock.""" exec_globals: Final = build_sandbox_globals() compiled: Final = compile_sandboxed(self.custom_code) - exec(compiled, exec_globals) # noqa: S102 + + def load_module() -> None: + exec(compiled, exec_globals) # noqa: S102 + + call_with_timeout(load_module, self.execution_timeout, label=f"{self.guardrail_name}:load") if "apply_guardrail" not in exec_globals: raise CustomCodeCompilationError( @@ -241,18 +275,10 @@ class CustomCodeGuardrail(CustomGuardrail): start_time: Final = time.time() try: - # Prepare inputs dict for the function - - # Prepare request_data with safe subset of information safe_request_data: Final = self._prepare_safe_request_data(request_data) - - # Execute the custom function - handle both sync and async functions - raw_result: Final = self._compiled_function(inputs, safe_request_data, input_type) - - # If the function is async (returns a coroutine), await it - resolved_result: Final[object] = await raw_result if asyncio.iscoroutine(raw_result) else raw_result - - # Process the result + resolved_result: Final = await self._call_compiled( + self._compiled_function, inputs, safe_request_data, input_type + ) return self._process_result( result=resolved_result, inputs=inputs, @@ -267,6 +293,19 @@ class CustomCodeGuardrail(CustomGuardrail): except ModifyResponseException: # Pre-call block uses passthrough; must not wrap as execution error (500) raise + except ExecutionTimeoutError: + verbose_proxy_logger.error( + "Custom code guardrail '%s' exceeded its %gs execution timeout", + self.guardrail_name, + self.execution_timeout, + ) + raise CustomCodeExecutionError( + f"Custom code guardrail '{self.guardrail_name}' exceeded its " + f"{self.execution_timeout:g}s execution timeout", + details=MappingProxyType( + {"guardrail_name": self.guardrail_name, "input_type": input_type, "timeout": self.execution_timeout} + ), + ) from None except Exception as e: verbose_proxy_logger.error("Custom code guardrail '%s' execution error: %s", self.guardrail_name, e) raise CustomCodeExecutionError( @@ -277,6 +316,31 @@ class CustomCodeGuardrail(CustomGuardrail): }, ) from e + async def _call_compiled( + self, + compiled_function: Callable[..., object], + inputs: GenericGuardrailAPIInputs, + safe_request_data: Mapping[str, object], + input_type: Literal["request", "response"], + ) -> object: + """Run the user's function under the execution budget. + + A coroutine function is awaited on the event loop, so the budget bounds it at its + await points and it is given up at the deadline even if it swallows cancellation. A + plain function runs on a worker thread, which keeps a busy loop from stalling every + other request and lets the runner interrupt it at the deadline. + """ + label: Final = str(self.guardrail_name) + if inspect.iscoroutinefunction(compiled_function): + pending: Final = compiled_function(inputs, safe_request_data, input_type) + return await await_with_timeout(pending, self.execution_timeout, label) + call: Final = functools.partial(compiled_function, inputs, safe_request_data, input_type) + deadline: Final = time.monotonic() + self.execution_timeout + raw_result: Final = await call_off_loop_with_timeout(call, self.execution_timeout, label) + if not asyncio.iscoroutine(raw_result): + return raw_result + return await await_with_timeout(raw_result, max(deadline - time.monotonic(), 0.0), label) + def _prepare_safe_request_data(self, request_data: Mapping[str, object]) -> dict[str, object]: """ Prepare a safe subset of request_data for code execution. diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index d5dbfaeb84b..55f7abd40a8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -5,6 +5,7 @@ These functions are injected into the custom code execution environment and provide safe, sandboxed functionality for common guardrail operations. """ +import asyncio import json import re from collections.abc import Mapping, Sequence @@ -15,7 +16,9 @@ import httpx from pydantic import JsonValue from typing_extensions import ReadOnly, TypedDict +import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, validate_url from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -393,6 +396,8 @@ _HTTP_DEFAULT_TIMEOUT: Final = 30.0 # Maximum allowed timeout (in seconds) _HTTP_MAX_TIMEOUT: Final = 60.0 +_HTTP_ALLOWED_METHODS: Final = ("GET", "POST", "PUT", "DELETE", "PATCH") + class HttpResponseResult(TypedDict): """Outcome of an HTTP primitive call, as handed back to custom code.""" @@ -463,6 +468,11 @@ async def http_request( Uses LiteLLM's global cached AsyncHTTPHandler for connection pooling and better performance. + Destinations go through LiteLLM's SSRF validation: private, link-local, + loopback and cloud-metadata addresses are refused (every redirect hop + included) unless the host is listed in ``litellm_settings.user_url_allowed_hosts`` + or ``litellm_settings.user_url_validation`` is turned off. + Args: url: The URL to request method: HTTP method (GET, POST, PUT, DELETE, PATCH). Defaults to GET. @@ -492,35 +502,35 @@ async def http_request( body={"text": "content to check"} ) """ - # Validate URL if not is_valid_url(url): return _http_error_response(f"Invalid URL: {url}") - # Validate and normalize method - method = method.upper() - allowed_methods: Final = {"GET", "POST", "PUT", "DELETE", "PATCH"} - if method not in allowed_methods: - return _http_error_response(f"Invalid HTTP method: {method}. Allowed: {', '.join(allowed_methods)}") + normalized_method: Final = method.upper() + if normalized_method not in _HTTP_ALLOWED_METHODS: + return _http_error_response( + f"Invalid HTTP method: {normalized_method}. Allowed: {', '.join(_HTTP_ALLOWED_METHODS)}" + ) - # Apply timeout limits - if timeout is None: - timeout = _HTTP_DEFAULT_TIMEOUT - else: - timeout = min(max(0.1, timeout), _HTTP_MAX_TIMEOUT) + effective_timeout: Final = _HTTP_DEFAULT_TIMEOUT if timeout is None else min(max(0.1, timeout), _HTTP_MAX_TIMEOUT) - # Get the global cached async HTTP client client: Final = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, - params={"timeout": httpx.Timeout(timeout=timeout, connect=5.0)}, + params={ + "timeout": httpx.Timeout(timeout=effective_timeout, connect=5.0), + "follow_redirects": not litellm.user_url_validation, + }, ) try: - response: Final = await _execute_http_request(client, method, url, headers, body, timeout) + response: Final = await _execute_http_request(client, normalized_method, url, headers, body, effective_timeout) return _http_success_response(response) + except SSRFError as e: + verbose_proxy_logger.warning("Custom code http_request blocked: %s", e) + return _http_error_response(f"Blocked URL: {e}") except httpx.TimeoutException as e: verbose_proxy_logger.warning("Custom code http_request timeout: %s", e) - return _http_error_response(f"Request timeout after {timeout}s") + return _http_error_response(f"Request timeout after {effective_timeout}s") except httpx.HTTPStatusError as e: # Return the response even for non-2xx status codes return _http_success_response(e.response) @@ -542,21 +552,47 @@ async def _execute_http_request( ) -> httpx.Response: """Execute the HTTP request using the appropriate client method.""" json_body, data_body = _prepare_http_body(body) + outbound_headers: Final = _caller_headers(headers) if method == "GET": - return await client.get(url=url, headers=headers) - elif method == "POST": - return await client.post(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) + return await async_safe_get(client, url, headers=outbound_headers) + + destination_url, destination_headers = await _validated_destination(url, outbound_headers) + if method == "POST": + return await client.post( + url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout + ) elif method == "PUT": - return await client.put(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) + return await client.put( + url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout + ) elif method == "DELETE": - return await client.delete(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) + return await client.delete( + url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout + ) elif method == "PATCH": - return await client.patch(url=url, headers=headers, json=json_body, data=data_body, timeout=timeout) + return await client.patch( + url=destination_url, headers=destination_headers, json=json_body, data=data_body, timeout=timeout + ) else: raise ValueError(f"Unsupported HTTP method: {method}") +def _caller_headers(headers: dict[str, str] | None) -> dict[str, str]: + if headers is None: + return {} + if not litellm.user_url_validation: + return headers + return {name: value for name, value in headers.items() if name.lower() != "host"} + + +async def _validated_destination(url: str, headers: dict[str, str]) -> tuple[str, dict[str, str]]: + if not litellm.user_url_validation: + return url, headers + destination_url, host_header = await asyncio.to_thread(validate_url, url) + return destination_url, {**headers, "Host": host_header} + + async def http_get( url: str, headers: dict[str, str] | None = None, diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index 35f1e6e6515..582ca44f19e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -27,13 +27,15 @@ from RestrictedPython import ( safe_builtins, utility_builtins, ) -from RestrictedPython.Eval import default_guarded_getitem, default_guarded_getiter +from RestrictedPython.Eval import default_guarded_getitem from RestrictedPython.Guards import ( full_write_guard, guarded_iter_unpack_sequence, safer_getattr, ) +from RestrictedPython.transformer import copy_locations +from .bounded_execution import budget_ok, budgeted_iter from .primitives import get_custom_code_primitives @@ -46,11 +48,31 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): check, print-scope wrapping, and any future additions to that method are inherited automatically. ``AsyncFor``/``AsyncWith``/``Await`` delegate to ``node_contents_visit`` so their children still get transformed. + + ``visit_While`` rewrites ``while test:`` to ``while _budget_ok_() and test:`` + so a loop that never yields is still stopped at the execution deadline; + ``for`` loops and comprehensions get the same check through ``_getiter_``. """ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AST: return self.visit_FunctionDef(node) + def visit_While(self, node: ast.While) -> ast.AST: + visited: Final = self.node_contents_visit(node) + budget_check: Final = ast.Call( + func=ast.Name(id="_budget_ok_", ctx=ast.Load()), + args=[], # mutable-ok: ast accepts list fields only + keywords=[], # mutable-ok: ast accepts list fields only + ) + test: Final = ast.BoolOp( + op=ast.And(), + values=[budget_check, visited.test], # mutable-ok: ast accepts list fields only + ) + copy_locations(test, visited.test) + bounded: Final = ast.While(test=test, body=visited.body, orelse=visited.orelse) + copy_locations(bounded, visited) + return bounded + def visit_AsyncFor(self, node: ast.AsyncFor) -> ast.AST: return self.node_contents_visit(node) @@ -113,10 +135,11 @@ def build_sandbox_globals() -> dict[str, object]: "__builtins__": _build_sandbox_builtins(), "_getattr_": safer_getattr, "_getitem_": default_guarded_getitem, - "_getiter_": default_guarded_getiter, + "_getiter_": budgeted_iter, "_iter_unpack_sequence_": guarded_iter_unpack_sequence, "_write_": full_write_guard, "_inplacevar_": _inplacevar_, + "_budget_ok_": budget_ok, } diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py new file mode 100644 index 00000000000..dc3d3c8cf91 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_custom_code_bounded_execution.py @@ -0,0 +1,174 @@ +import asyncio +import threading +import time + +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.custom_code.bounded_execution import ( + ExecutionTimeoutError, + SandboxExit, + await_with_timeout, + call_off_loop_with_timeout, + call_with_timeout, +) + + +def _worker_threads() -> list[str]: + return [t.name for t in threading.enumerate() if t.name.startswith("guardrail-code:")] + + +def _spin_forever() -> None: + n = 0 + while True: + n += 1 + + +def _spin_swallowing_exceptions() -> None: + while True: + try: + _spin_forever() + except Exception: + continue + + +async def _swallow_cancellations_for(seconds: float) -> str: + deadline = time.monotonic() + seconds + while time.monotonic() < deadline: + try: + await asyncio.sleep(deadline - time.monotonic()) + except asyncio.CancelledError: + continue + return "survived" + + +def _exit_now() -> None: + raise SystemExit("bye") + + +def test_call_with_timeout_returns_the_result_and_reraises_failures(): + assert call_with_timeout(lambda: 42, 1.0, label="ok") == 42 + with pytest.raises(ZeroDivisionError): + call_with_timeout(lambda: 1 // 0, 1.0, label="boom") + + +def test_call_with_timeout_delivers_a_system_exit_at_once(): + started = time.monotonic() + + with pytest.raises(SandboxExit, match="SystemExit: bye"): + call_with_timeout(_exit_now, 5.0, label="exit") + + assert time.monotonic() - started < 1.0 + + +async def _exit_later() -> None: + await asyncio.sleep(0) + raise SystemExit("bye") + + +@pytest.mark.asyncio +async def test_await_with_timeout_contains_a_system_exit_instead_of_stopping_the_loop(): + with pytest.raises(SandboxExit, match="SystemExit: bye"): + await await_with_timeout(_exit_later(), 1.0, label="exit") + + assert await asyncio.sleep(0, result="loop still running") == "loop still running" + + +@pytest.mark.parametrize("fn", [_spin_forever, _spin_swallowing_exceptions]) +def test_call_with_timeout_interrupts_a_busy_loop_and_reclaims_the_thread(fn): + started = time.monotonic() + + with pytest.raises(ExecutionTimeoutError, match=r"exceeded the 0\.2s execution timeout") as exc_info: + call_with_timeout(fn, 0.2, label="spin") + + assert exc_info.value.timeout == 0.2 + assert time.monotonic() - started < 1.5 + time.sleep(0.2) + assert _worker_threads() == [] + + +@pytest.mark.asyncio +async def test_call_off_loop_with_timeout_keeps_the_loop_running_and_stops_the_worker(): + ticks = 0 + + async def tick_forever() -> None: + nonlocal ticks + while True: + await asyncio.sleep(0.02) + ticks += 1 + + ticker = asyncio.create_task(tick_forever()) + try: + assert await call_off_loop_with_timeout(lambda: "done", 1.0, label="ok") == "done" + with pytest.raises(ExecutionTimeoutError): + await call_off_loop_with_timeout(_spin_forever, 0.3, label="spin") + finally: + ticker.cancel() + + assert ticks >= 5 + await asyncio.sleep(0.2) + assert _worker_threads() == [] + + +def _stragglers() -> list[asyncio.Task[object]]: + return [task for task in asyncio.all_tasks() if task is not asyncio.current_task()] + + +@pytest.mark.asyncio +async def test_await_with_timeout_keeps_cancelling_a_coroutine_that_swallows_cancellation(): + assert await await_with_timeout(_swallow_cancellations_for(0.0), 1.0, label="ok") == "survived" + started = time.monotonic() + + with pytest.raises(ExecutionTimeoutError, match=r"exceeded the 0\.1s execution timeout"): + await await_with_timeout(_swallow_cancellations_for(0.4), 0.1, label="stubborn") + + assert time.monotonic() - started < 1.0 + assert _stragglers() == [] + + +@pytest.mark.asyncio +async def test_await_with_timeout_abandons_a_coroutine_that_never_stops_swallowing_cancellation(): + started = time.monotonic() + + with pytest.raises(ExecutionTimeoutError): + await await_with_timeout(_swallow_cancellations_for(2.0), 0.1, label="stubborn") + + elapsed = time.monotonic() - started + assert 1.0 <= elapsed < 1.8 + stragglers = _stragglers() + assert len(stragglers) == 1 + with pytest.raises(ExecutionTimeoutError): + await asyncio.gather(*stragglers) + + +@pytest.mark.asyncio +async def test_call_off_loop_with_timeout_stops_the_worker_when_the_caller_is_cancelled(): + waiting = asyncio.create_task(call_off_loop_with_timeout(_spin_forever, 30.0, label="spin")) + await asyncio.sleep(0.1) + + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + + await asyncio.sleep(0.3) + assert _worker_threads() == [] + + +@pytest.mark.asyncio +async def test_await_with_timeout_cancels_the_code_when_the_caller_is_cancelled(): + interrupted = asyncio.Event() + + async def sleep_until_cancelled() -> None: + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + interrupted.set() + raise + + waiting = asyncio.create_task(await_with_timeout(sleep_until_cancelled(), 30.0, label="sleep")) + await asyncio.sleep(0.1) + + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + + await asyncio.wait_for(interrupted.wait(), timeout=1.0) diff --git a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py index 068cd0d8ed7..5532d2c9809 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -1,11 +1,22 @@ +import asyncio +import http.server +import threading +import time +from http.server import ThreadingHTTPServer + import pytest from fastapi import HTTPException +import litellm from litellm.exceptions import ModifyResponseException from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( + DEFAULT_EXECUTION_TIMEOUT_SECONDS, CustomCodeCompilationError, + CustomCodeExecutionError, CustomCodeGuardrail, ) +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler +from litellm.types.guardrails import SupportedGuardrailIntegrations # str.mro() + generator gi_code + code.replace(co_names=...) + __setattr__ # to swap a function's bytecode and read http_get's real builtins dict. @@ -77,18 +88,14 @@ def test_nfkc_homoglyph_rejected_at_compile(): [ # Literal dunder attribute access. "def apply_guardrail(i, r, t):\n return str.__class__\n", - "def apply_guardrail(i, r, t):\n" - " return ().__class__.__bases__[0].__subclasses__()\n", + "def apply_guardrail(i, r, t):\n return ().__class__.__bases__[0].__subclasses__()\n", # gi_code — on the transformer's restricted-names list. - "def apply_guardrail(i, r, t):\n" - " def g():\n yield 1\n" - " return g().gi_code\n", + "def apply_guardrail(i, r, t):\n def g():\n yield 1\n return g().gi_code\n", # Import forms. "import os\ndef apply_guardrail(i, r, t):\n return allow()\n", - "from subprocess import call\n" - "def apply_guardrail(i, r, t):\n return allow()\n", + "from subprocess import call\ndef apply_guardrail(i, r, t):\n return allow()\n", # __import__ is rejected as an underscore-prefixed name. - "def apply_guardrail(i, r, t):\n" ' return __import__("os")\n', + 'def apply_guardrail(i, r, t):\n return __import__("os")\n', ], ) def test_compile_time_rejections(snippet: str): @@ -100,8 +107,7 @@ def test_compile_time_rejections(snippet: str): "snippet", [ # getattr is not in the sandbox builtins — NameError at call time. - "def apply_guardrail(i, r, t):\n" - ' return getattr(str, "_"+"_class_"+"_")\n', + 'def apply_guardrail(i, r, t):\n return getattr(str, "_"+"_class_"+"_")\n', # setattr is guarded_setattr + full_write_guard — setting any attribute # on a user-defined object raises TypeError, whether the name is a # dunder or not. @@ -139,10 +145,7 @@ def test_documented_ssn_example_compiles_and_runs(): @pytest.mark.asyncio async def test_async_guardrail_compiles_and_runs(): - code = ( - "async def apply_guardrail(inputs, request_data, input_type):\n" - " return allow()\n" - ) + code = "async def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" guardrail = _compile(code) from litellm.types.utils import GenericGuardrailAPIInputs @@ -156,10 +159,7 @@ async def test_async_guardrail_compiles_and_runs(): @pytest.mark.asyncio async def test_custom_code_pre_call_block_uses_passthrough(): - code = ( - "def apply_guardrail(inputs, request_data, input_type):\n" - ' return block("blocked by test")\n' - ) + code = 'def apply_guardrail(inputs, request_data, input_type):\n return block("blocked by test")\n' guardrail = _compile(code) with pytest.raises(ModifyResponseException) as exc_info: @@ -176,10 +176,7 @@ async def test_custom_code_pre_call_block_uses_passthrough(): @pytest.mark.asyncio async def test_custom_code_post_call_block_raises_http_400(): - code = ( - "def apply_guardrail(inputs, request_data, input_type):\n" - ' return block("blocked by test")\n' - ) + code = 'def apply_guardrail(inputs, request_data, input_type):\n return block("blocked by test")\n' guardrail = _compile(code) with pytest.raises(HTTPException) as exc_info: @@ -333,10 +330,7 @@ async def test_custom_code_allow_still_records_success_not_flagged(): def test_typical_sync_guardrail_still_works(): - code = ( - "def apply_guardrail(inputs, request_data, input_type):\n" - " return allow()\n" - ) + code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" guardrail = _compile(code) assert guardrail._compiled_function is not None @@ -363,3 +357,492 @@ def test_augmented_assignment_works(): def test_missing_apply_guardrail_raises(): with pytest.raises(CustomCodeCompilationError, match="apply_guardrail"): _compile("x = 1\n") + + +class _QuietServer(ThreadingHTTPServer): + def handle_error(self, request: object, client_address: object) -> None: + return + + +def _guardrail_worker_threads() -> list[str]: + return [t.name for t in threading.enumerate() if t.name.startswith("guardrail-code:")] + + +class _LocalServer: + """Loopback HTTP server that records every request it receives.""" + + def __init__(self) -> None: + self.hits: list[tuple[str, str]] = [] + self.received_headers: list[list[tuple[str, str]]] = [] + server = self + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self) -> None: + server.hits.append(("GET", self.path)) + server.received_headers.append(list(self.headers.items())) + if self.path.startswith("/redirect-to/"): + self._redirect() + return + if self.path == "/slow": + time.sleep(2) + self._reply(b"marker") + + def do_POST(self) -> None: + server.hits.append(("POST", self.path)) + server.received_headers.append(list(self.headers.items())) + if self.path.startswith("/redirect-to/"): + self._redirect() + return + self._reply(b"posted") + + def do_PUT(self) -> None: + self._record_and_reply(b"put") + + def do_DELETE(self) -> None: + self._record_and_reply(b"deleted") + + def do_PATCH(self) -> None: + self._record_and_reply(b"patched") + + def _record_and_reply(self, body: bytes) -> None: + server.hits.append((self.command, self.path)) + server.received_headers.append(list(self.headers.items())) + self._reply(body) + + def _redirect(self) -> None: + target_port = self.path.rsplit("/", 1)[1] + self.send_response(302) + self.send_header("Location", f"http://127.0.0.1:{target_port}/marker") + self.send_header("Content-Length", "0") + self.end_headers() + + def _reply(self, body: bytes) -> None: + self.send_response(200) + self.send_header("Content-Type", "text/plain") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args: object) -> None: + return + + self.httpd = _QuietServer(("127.0.0.1", 0), Handler) + self.port = self.httpd.server_address[1] + threading.Thread(target=self.httpd.serve_forever, daemon=True).start() + + def close(self) -> None: + self.httpd.shutdown() + self.httpd.server_close() + + +@pytest.fixture +def local_server(): + server = _LocalServer() + yield server + server.close() + + +@pytest.fixture +def second_server(): + server = _LocalServer() + yield server + server.close() + + +@pytest.fixture(autouse=True) +def _fresh_http_client_and_url_policy(monkeypatch): + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setattr(litellm, "user_url_validation", True) + monkeypatch.setattr(litellm, "user_url_allowed_hosts", []) + + +def _reporting_guardrail(call: str) -> CustomCodeGuardrail: + code = ( + "async def apply_guardrail(inputs, request_data, input_type):\n" + f" r = await {call}\n" + ' return block("status=" + str(r["status_code"]) + " body=" + str(r["body"])' + ' + " error=" + str(r["error"]))\n' + ) + return _compile(code) + + +async def _block_reason(guardrail: CustomCodeGuardrail) -> str: + with pytest.raises(ModifyResponseException) as exc_info: + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request") + return exc_info.value.message + + +@pytest.mark.asyncio +async def test_http_get_refuses_loopback_by_default(local_server): + guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/marker")') + + reason = await _block_reason(guardrail) + + assert "status=0" in reason + assert "error=Blocked URL" in reason + assert "user_url_allowed_hosts" in reason + assert local_server.hits == [] + + +@pytest.mark.asyncio +async def test_http_post_refuses_loopback_by_default(local_server): + guardrail = _reporting_guardrail(f'http_post("http://127.0.0.1:{local_server.port}/hook", body={{"a": 1}})') + + reason = await _block_reason(guardrail) + + assert "error=Blocked URL" in reason + assert local_server.hits == [] + + +@pytest.mark.asyncio +async def test_http_get_reaches_an_allowlisted_host(local_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/marker")') + + reason = await _block_reason(guardrail) + + assert "status=200 body=marker error=None" in reason + assert local_server.hits == [("GET", "/marker")] + + +@pytest.mark.asyncio +async def test_http_get_refuses_a_redirect_into_a_blocked_host(local_server, second_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + guardrail = _reporting_guardrail( + f'http_get("http://127.0.0.1:{local_server.port}/redirect-to/{second_server.port}")' + ) + + reason = await _block_reason(guardrail) + + assert "error=Blocked URL" in reason + assert local_server.hits == [("GET", f"/redirect-to/{second_server.port}")] + assert second_server.hits == [] + + +@pytest.mark.asyncio +async def test_http_post_does_not_follow_redirects(local_server, second_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + guardrail = _reporting_guardrail( + f'http_post("http://127.0.0.1:{local_server.port}/redirect-to/{second_server.port}", body={{"a": 1}})' + ) + + reason = await _block_reason(guardrail) + + assert "status=302" in reason + assert second_server.hits == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call", ["http_post", "http_get"]) +async def test_caller_host_header_never_reaches_the_validated_destination(local_server, monkeypatch, call): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + guardrail = _reporting_guardrail( + f'{call}("http://127.0.0.1:{local_server.port}/marker", headers={{"host": "spoofed", "X-Extra": "kept"}})' + ) + + reason = await _block_reason(guardrail) + + assert "status=200" in reason + (received,) = local_server.received_headers + assert [value for name, value in received if name.lower() == "host"] == [f"127.0.0.1:{local_server.port}"] + assert ("x-extra", "kept") in received + + +@pytest.mark.asyncio +async def test_caller_headers_pass_through_untouched_when_url_validation_is_disabled(local_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_validation", False) + guardrail = _reporting_guardrail( + f'http_post("http://127.0.0.1:{local_server.port}/marker", headers={{"host": "spoofed", "X-Extra": "kept"}})' + ) + + reason = await _block_reason(guardrail) + + assert "status=200" in reason + (received,) = local_server.received_headers + assert [value for name, value in received if name.lower() == "host"] == ["spoofed"] + assert ("x-extra", "kept") in received + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("method", "body"), [("PUT", "put"), ("DELETE", "deleted"), ("PATCH", "patched")]) +async def test_http_request_other_methods_reach_an_allowlisted_host(local_server, monkeypatch, method, body): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + guardrail = _reporting_guardrail( + f'http_request("http://127.0.0.1:{local_server.port}/marker", method="{method}")' + ) + + reason = await _block_reason(guardrail) + + assert f"status=200 body={body}" in reason + assert local_server.hits == [(method, "/marker")] + + +@pytest.mark.asyncio +async def test_http_request_refuses_a_method_outside_the_allowlist(local_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + guardrail = _reporting_guardrail(f'http_request("http://127.0.0.1:{local_server.port}/marker", method="TRACE")') + + reason = await _block_reason(guardrail) + + assert "error=Invalid HTTP method: TRACE" in reason + assert local_server.hits == [] + + +@pytest.mark.asyncio +async def test_http_get_gives_up_at_its_own_timeout(local_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/slow", timeout=0.5)') + + started = time.monotonic() + reason = await _block_reason(guardrail) + + assert time.monotonic() - started < 1.5 + assert "error=Request timeout after 0.5s" in reason + + +@pytest.mark.asyncio +async def test_sync_guardrail_returning_a_coroutine_has_it_awaited(): + code = ( + "async def decide():\n" + ' return block("decided late")\n' + "def apply_guardrail(inputs, request_data, input_type):\n" + " return decide()\n" + ) + guardrail = _compile(code) + + reason = await _block_reason(guardrail) + + assert "decided late" in reason + + +@pytest.mark.asyncio +async def test_http_get_is_unvalidated_when_url_validation_is_disabled(local_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_validation", False) + guardrail = _reporting_guardrail(f'http_get("http://127.0.0.1:{local_server.port}/marker")') + + reason = await _block_reason(guardrail) + + assert "status=200 body=marker" in reason + assert local_server.hits == [("GET", "/marker")] + + +BUSY_LOOP_GUARDRAIL = ( + "def apply_guardrail(inputs, request_data, input_type):\n n = 0\n while True:\n n += 1\n" +) + +SWALLOWING_BUSY_LOOP_GUARDRAIL = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " n = 0\n" + " while True:\n" + " try:\n" + " n += 1\n" + " except Exception:\n" + " n = 0\n" +) + + +async def _expect_execution_timeout(guardrail: CustomCodeGuardrail) -> float: + started = time.monotonic() + with pytest.raises(CustomCodeExecutionError, match=r"exceeded its 0\.3s execution timeout"): + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request") + return time.monotonic() - started + + +@pytest.mark.asyncio +@pytest.mark.parametrize("code", [BUSY_LOOP_GUARDRAIL, SWALLOWING_BUSY_LOOP_GUARDRAIL]) +async def test_sync_busy_loop_is_stopped_at_the_execution_timeout(code): + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="busy", execution_timeout=0.3) + + elapsed = await _expect_execution_timeout(guardrail) + + assert elapsed < 2.0 + await asyncio.sleep(0.2) + assert _guardrail_worker_threads() == [] + + +@pytest.mark.asyncio +async def test_sync_busy_loop_does_not_stall_the_event_loop(): + guardrail = CustomCodeGuardrail(custom_code=BUSY_LOOP_GUARDRAIL, guardrail_name="busy", execution_timeout=0.3) + ticks = 0 + + async def tick_forever() -> None: + nonlocal ticks + while True: + await asyncio.sleep(0.02) + ticks += 1 + + ticker = asyncio.create_task(tick_forever()) + try: + await _expect_execution_timeout(guardrail) + finally: + ticker.cancel() + + assert ticks >= 5 + + +@pytest.mark.asyncio +async def test_async_guardrail_is_stopped_at_the_execution_timeout(local_server, monkeypatch): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + code = ( + "async def apply_guardrail(inputs, request_data, input_type):\n" + f' await http_get("http://127.0.0.1:{local_server.port}/slow")\n' + " return allow()\n" + ) + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="busy", execution_timeout=0.3) + + elapsed = await _expect_execution_timeout(guardrail) + + assert elapsed < 1.5 + + +@pytest.mark.asyncio +async def test_async_guardrail_that_swallows_cancellation_is_stopped_at_the_execution_timeout( + local_server, monkeypatch +): + monkeypatch.setattr(litellm, "user_url_allowed_hosts", [f"127.0.0.1:{local_server.port}"]) + code = ( + "async def apply_guardrail(inputs, request_data, input_type):\n" + " attempts = 0\n" + " while attempts < 3:\n" + " try:\n" + f' await http_get("http://127.0.0.1:{local_server.port}/slow")\n' + " except BaseException:\n" + " attempts += 1\n" + " return allow()\n" + ) + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="stubborn", execution_timeout=0.3) + + elapsed = await _expect_execution_timeout(guardrail) + + assert elapsed < 1.5 + + +LOOP_SHAPES_GUARDRAIL = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + " n = 0\n" + " while n < 3:\n" + " n += 1\n" + " else:\n" + " n += 10\n" + " pairs = [(k, v) for k, v in request_data['metadata'].items()]\n" + " for k, v in pairs:\n" + " n += v\n" + " for i, (k, v) in zip(range(len(pairs)), pairs):\n" + " n += i\n" + " keys = sorted(k for k, v in pairs)\n" + " return block(reason=str(n) + ' ' + ' '.join(keys))\n" +) + + +@pytest.mark.asyncio +async def test_budget_checks_keep_every_loop_shape_working(): + guardrail = CustomCodeGuardrail(custom_code=LOOP_SHAPES_GUARDRAIL, guardrail_name="loops") + + with pytest.raises(ModifyResponseException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["x"]}, request_data={"model": "m", "metadata": {"b": 2, "a": 5}}, input_type="request" + ) + + assert exc_info.value.message == "21 a b" + + +@pytest.mark.asyncio +@pytest.mark.timeout(10) +@pytest.mark.parametrize( + "code", + [ + "async def apply_guardrail(inputs, request_data, input_type):\n while True:\n pass\n", + ( + "async def apply_guardrail(inputs, request_data, input_type):\n" + " for a in range(500):\n" + " for b in range(500):\n" + " for c in range(500):\n" + " pass\n" + " return allow()\n" + ), + ( + "async def apply_guardrail(inputs, request_data, input_type):\n" + " try:\n" + " while True:\n" + " pass\n" + " except BaseException:\n" + " pass\n" + " return allow()\n" + ), + ], +) +async def test_async_loop_that_never_yields_is_stopped_at_the_execution_timeout(code): + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="spin", execution_timeout=0.3) + + elapsed = await _expect_execution_timeout(guardrail) + + assert elapsed < 1.5 + assert await asyncio.sleep(0, result="loop still running") == "loop still running" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "code", + [ + "def apply_guardrail(inputs, request_data, input_type):\n raise SystemExit('bye')\n", + "async def apply_guardrail(inputs, request_data, input_type):\n raise SystemExit('bye')\n", + ], +) +async def test_system_exit_from_guardrail_code_is_an_execution_error_not_a_timeout(code): + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="exit", execution_timeout=5.0) + started = time.monotonic() + + with pytest.raises(CustomCodeExecutionError, match="execution failed: SystemExit: bye"): + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data={"model": "m"}, input_type="request") + + assert time.monotonic() - started < 1.0 + + +def test_module_level_busy_loop_fails_compilation_at_the_execution_timeout(): + code = "n = 0\nwhile True:\n n += 1\n" + BUSY_LOOP_GUARDRAIL + started = time.monotonic() + + with pytest.raises(CustomCodeCompilationError, match=r"exceeded the 0\.3s execution timeout"): + CustomCodeGuardrail(custom_code=code, guardrail_name="busy", execution_timeout=0.3) + + assert time.monotonic() - started < 2.0 + + +@pytest.mark.parametrize("execution_timeout", [0, -1.0]) +def test_execution_timeout_must_be_positive(execution_timeout): + with pytest.raises(ValueError, match="execution_timeout must be positive"): + CustomCodeGuardrail( + custom_code="def apply_guardrail(i, r, t):\n return allow()\n", execution_timeout=execution_timeout + ) + + +def _initialize_from_config(guardrail_name: str, litellm_params: dict[str, object]) -> CustomCodeGuardrail: + InMemoryGuardrailHandler().initialize_guardrail( + guardrail={ + "guardrail_name": guardrail_name, + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.CUSTOM_CODE.value, + "mode": "pre_call", + "custom_code": "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n", + **litellm_params, + }, + } + ) + initialized = [ + callback + for callback in litellm.callbacks + if isinstance(callback, CustomCodeGuardrail) and callback.guardrail_name == guardrail_name + ] + assert initialized, f"{guardrail_name} was not registered as a callback" + return initialized[-1] + + +def test_config_timeout_reaches_the_guardrail(): + assert _initialize_from_config("custom-code-timeout", {"timeout": 0.2}).execution_timeout == 0.2 + + +def test_config_without_timeout_uses_the_default(): + assert ( + _initialize_from_config("custom-code-default-timeout", {}).execution_timeout + == DEFAULT_EXECUTION_TIMEOUT_SECONDS + ) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index bf641fd6cd0..508736fb78e 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -1,4 +1,5 @@ import json +import time from datetime import datetime from typing import Dict, List, Optional from unittest.mock import AsyncMock @@ -13,6 +14,7 @@ from litellm.proxy.guardrails.guardrail_endpoints import ( CreateGuardrailRequest, PatchGuardrailRequest, RegisterGuardrailRequest, + TestCustomCodeGuardrailRequest, UpdateGuardrailRequest, apply_guardrail, approve_guardrail_submission, @@ -28,6 +30,9 @@ from litellm.proxy.guardrails.guardrail_endpoints import ( reject_guardrail_submission, update_guardrail, ) +from litellm.proxy.guardrails.guardrail_endpoints import ( + test_custom_code_guardrail as run_custom_code_test_endpoint, +) MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) from litellm.proxy.guardrails.guardrail_registry import ( @@ -87,12 +92,8 @@ def mock_prisma_client(mocker): # Create async mocks for the database methods mock_client.db = mocker.Mock() mock_client.db.litellm_guardrailstable = mocker.Mock() - mock_client.db.litellm_guardrailstable.find_many = AsyncMock( - return_value=[MOCK_DB_GUARDRAIL] - ) - mock_client.db.litellm_guardrailstable.find_unique = AsyncMock( - return_value=MOCK_DB_GUARDRAIL - ) + mock_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[MOCK_DB_GUARDRAIL]) + mock_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=MOCK_DB_GUARDRAIL) return mock_client @@ -118,17 +119,13 @@ def mock_guardrail_registry(mocker): return_value={**MOCK_DB_GUARDRAIL, "guardrail_id": "new-test-guardrail-id"} ) mock_registry.delete_guardrail_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) - mock_registry.get_guardrail_by_id_from_db = AsyncMock( - return_value=MOCK_DB_GUARDRAIL - ) + mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) mock_registry.update_guardrail_in_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) return mock_registry @pytest.mark.asyncio -async def test_list_guardrails_v2_with_db_and_config( - mocker, mock_prisma_client, mock_in_memory_handler -): +async def test_list_guardrails_v2_with_db_and_config(mocker, mock_prisma_client, mock_in_memory_handler): """Test listing guardrails from both DB and config""" # Mock the prisma client mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -144,17 +141,13 @@ async def test_list_guardrails_v2_with_db_and_config( assert len(response.guardrails) == 2 # Check DB guardrail - db_guardrail = next( - g for g in response.guardrails if g.guardrail_id == "test-db-guardrail" - ) + db_guardrail = next(g for g in response.guardrails if g.guardrail_id == "test-db-guardrail") assert db_guardrail.guardrail_name == "Test DB Guardrail" assert db_guardrail.guardrail_definition_location == "db" assert isinstance(db_guardrail.litellm_params, BaseLitellmParams) # Check config guardrail - config_guardrail = next( - g for g in response.guardrails if g.guardrail_id == "test-config-guardrail" - ) + config_guardrail = next(g for g in response.guardrails if g.guardrail_id == "test-config-guardrail") assert config_guardrail.guardrail_name == "Test Config Guardrail" assert config_guardrail.guardrail_definition_location == "config" assert isinstance(config_guardrail.litellm_params, BaseLitellmParams) @@ -196,9 +189,7 @@ async def test_list_guardrails_v2_skips_stale_db_backed_in_memory_entries(mocker @pytest.mark.asyncio -async def test_get_guardrail_info_404s_stale_db_backed_entry( - mocker, mock_prisma_client, mock_in_memory_handler -): +async def test_get_guardrail_info_404s_stale_db_backed_entry(mocker, mock_prisma_client, mock_in_memory_handler): """ Stale DB-backed entry (in-memory but not in DB) must 404 instead of being returned as if it were a config-loaded guardrail. @@ -208,9 +199,7 @@ async def test_get_guardrail_info_404s_stale_db_backed_entry( "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler, ) - mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None) # In-memory still has it, but it's tagged as 'db' (stale, awaiting reconcile) mock_in_memory_handler.get_source.return_value = "db" @@ -241,9 +230,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker): mock_prisma_client = mocker.Mock() mock_prisma_client.db = mocker.Mock() mock_prisma_client.db.litellm_guardrailstable = mocker.Mock() - mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( - return_value=[db_guardrail_with_secrets] - ) + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[db_guardrail_with_secrets]) mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.list_in_memory_guardrails.return_value = [] @@ -263,11 +250,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker): if isinstance(litellm_params, dict): params = litellm_params else: - params = ( - litellm_params.model_dump() - if hasattr(litellm_params, "model_dump") - else dict(litellm_params) - ) + params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) # Sensitive keys (containing "key", "secret", "token", etc.) should be masked assert params["api_key"] != "sk-1234567890abcdef" @@ -299,9 +282,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) mock_in_memory_handler = mocker.Mock() - mock_in_memory_handler.list_in_memory_guardrails.return_value = [ - config_guardrail_with_secrets - ] + mock_in_memory_handler.list_in_memory_guardrails.return_value = [config_guardrail_with_secrets] mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -318,11 +299,7 @@ async def test_list_guardrails_v2_masks_sensitive_data_in_config_guardrails(mock if isinstance(litellm_params, dict): params = litellm_params else: - params = ( - litellm_params.model_dump() - if hasattr(litellm_params, "model_dump") - else dict(litellm_params) - ) + params = litellm_params.model_dump() if hasattr(litellm_params, "model_dump") else dict(litellm_params) # Sensitive keys should be masked assert params["api_key"] != "my-secret-bedrock-key" @@ -355,9 +332,7 @@ async def test_list_guardrails_v2_admin_viewer_sees_guardrails_of_teams_they_are mock_prisma_client = mocker.Mock() mock_prisma_client.db = mocker.Mock() mock_prisma_client.db.litellm_guardrailstable = mocker.Mock() - mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( - return_value=[other_team_guardrail] - ) + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[other_team_guardrail]) mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.list_in_memory_guardrails.return_value = [] @@ -372,9 +347,7 @@ async def test_list_guardrails_v2_admin_viewer_sees_guardrails_of_teams_they_are AsyncMock(return_value=[]), ) - viewer_auth = UserAPIKeyAuth( - user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ) + viewer_auth = UserAPIKeyAuth(user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) response = await list_guardrails_v2(user_api_key_dict=viewer_auth) assert [g.guardrail_id for g in response.guardrails] == ["other-team-guardrail"] @@ -421,16 +394,10 @@ async def test_list_guardrails_v2_masks_sensitive_data_for_admin_viewer(mocker): AsyncMock(return_value=[]), ) - viewer_auth = UserAPIKeyAuth( - user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ) + viewer_auth = UserAPIKeyAuth(user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) response = await list_guardrails_v2(user_api_key_dict=viewer_auth) - guardrail = next( - g - for g in response.guardrails - if g.guardrail_id == "other-team-secret-guardrail" - ) + guardrail = next(g for g in response.guardrails if g.guardrail_id == "other-team-secret-guardrail") params = guardrail.litellm_params.model_dump() assert params["api_key"] != "sk-viewer-must-not-see-this" assert "****" in str(params["api_key"]) @@ -451,9 +418,7 @@ async def test_get_guardrail_info_from_db(mocker, mock_prisma_client): @pytest.mark.asyncio -async def test_get_guardrail_info_from_config( - mocker, mock_prisma_client, mock_in_memory_handler -): +async def test_get_guardrail_info_from_config(mocker, mock_prisma_client, mock_in_memory_handler): """Test getting guardrail info from config when not found in DB""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -462,9 +427,7 @@ async def test_get_guardrail_info_from_config( ) # Mock DB to return None - mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None) response = await get_guardrail_info("test-config-guardrail") @@ -475,9 +438,7 @@ async def test_get_guardrail_info_from_config( @pytest.mark.asyncio -async def test_get_guardrail_info_not_found( - mocker, mock_prisma_client, mock_in_memory_handler -): +async def test_get_guardrail_info_not_found(mocker, mock_prisma_client, mock_in_memory_handler): """Test getting guardrail info when not found in either DB or config""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -486,9 +447,7 @@ async def test_get_guardrail_info_not_found( ) # Mock both DB and in-memory handler to return None - mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=None) mock_in_memory_handler.get_guardrail_by_id.return_value = None with pytest.raises(HTTPException) as exc_info: @@ -499,9 +458,7 @@ async def test_get_guardrail_info_not_found( @pytest.mark.asyncio -async def test_list_guardrails_v2_without_prisma_returns_config_guardrails( - mocker, mock_in_memory_handler -): +async def test_list_guardrails_v2_without_prisma_returns_config_guardrails(mocker, mock_in_memory_handler): """ A proxy without a DB must still list config-defined guardrails instead of raising 500 'Prisma client not initialized'. @@ -535,18 +492,14 @@ async def test_list_guardrails_v2_without_prisma_non_admin_sees_unrestricted_con mock_in_memory_handler, ) - non_admin_auth = UserAPIKeyAuth( - user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal-user-1" - ) + non_admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal-user-1") response = await list_guardrails_v2(user_api_key_dict=non_admin_auth) assert [g.guardrail_id for g in response.guardrails] == ["test-config-guardrail"] @pytest.mark.asyncio -async def test_get_guardrail_info_without_prisma_returns_config_guardrail( - mocker, mock_in_memory_handler -): +async def test_get_guardrail_info_without_prisma_returns_config_guardrail(mocker, mock_in_memory_handler): """ The info endpoint must serve config-defined guardrails from the in-memory registry when no DB is attached instead of raising 500. @@ -565,9 +518,7 @@ async def test_get_guardrail_info_without_prisma_returns_config_guardrail( @pytest.mark.asyncio -async def test_get_guardrail_info_without_prisma_404s_unknown_id( - mocker, mock_in_memory_handler -): +async def test_get_guardrail_info_without_prisma_404s_unknown_id(mocker, mock_in_memory_handler): mocker.patch("litellm.proxy.proxy_server.prisma_client", None) mocker.patch( "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", @@ -630,10 +581,7 @@ def test_get_provider_specific_params(): assert "optional_params" in fields # Check the structure of a simple field - assert ( - fields["api_key"]["description"] - == "API key for the Azure Content Safety Prompt Shield guardrail" - ) + assert fields["api_key"]["description"] == "API key for the Azure Content Safety Prompt Shield guardrail" assert fields["api_key"]["required"] == False assert fields["api_key"]["type"] == "string" # Should be string, not None @@ -657,17 +605,13 @@ def test_get_provider_specific_params(): == "Severity threshold for the Azure Content Safety Text Moderation guardrail across all categories" ) assert nested_fields["severity_threshold"]["required"] == False - assert ( - nested_fields["severity_threshold"]["type"] == "number" - ) # Should be number, not None + assert nested_fields["severity_threshold"]["type"] == "number" # Should be number, not None # Check other field types assert nested_fields["categories"]["type"] == "multiselect" assert nested_fields["blocklistNames"]["type"] == "array" assert nested_fields["haltOnBlocklistHit"]["type"] == "boolean" - assert ( - nested_fields["outputType"]["type"] == "select" - ) # Literal type should be select + assert nested_fields["outputType"]["type"] == "select" # Literal type should be select @pytest.mark.asyncio @@ -769,17 +713,11 @@ def test_optional_params_returned_when_properly_overridden(): # Create specific optional params model class SpecificOptionalParams(BaseModel): - threshold: Optional[float] = Field( - default=0.5, description="Detection threshold" - ) - categories: Optional[List[str]] = Field( - default=None, description="Categories to check" - ) + threshold: Optional[float] = Field(default=0.5, description="Detection threshold") + categories: Optional[List[str]] = Field(default=None, description="Categories to check") # Create a config model that DOES override optional_params with a specific type - class TestGuardrailConfigWithOptionalParams( - GuardrailConfigModel[SpecificOptionalParams] - ): + class TestGuardrailConfigWithOptionalParams(GuardrailConfigModel[SpecificOptionalParams]): api_key: Optional[str] = Field( default=None, description="Test API key", @@ -806,9 +744,7 @@ async def test_bedrock_guardrail_prepare_request_with_api_key(): ) # Setup guardrail hook - guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" - ) + guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1") mock_credentials = Mock() test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} @@ -839,9 +775,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch): ) # Setup guardrail hook - guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" - ) + guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1") # Mock credentials mock_credentials = Mock() @@ -854,7 +788,6 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch): patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, patch("botocore.awsrequest.AWSRequest") as mock_aws_request, ): - # Mock SigV4Auth mock_sigv4_instance = Mock() mock_sigv4_auth.return_value = mock_sigv4_instance @@ -873,9 +806,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch): ) # Verify SigV4 auth was used - mock_sigv4_auth.assert_called_once_with( - mock_credentials, "bedrock", "us-east-1" - ) + mock_sigv4_auth.assert_called_once_with(mock_credentials, "bedrock", "us-east-1") mock_sigv4_instance.add_auth.assert_called_once() @@ -889,9 +820,7 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(monkeypat ) # Setup guardrail hook - guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" - ) + guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1") # Mock credentials mock_credentials = Mock() @@ -928,9 +857,7 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): BedrockGuardrail, ) - guardrail_hook = BedrockGuardrail( - guardrailIdentifier="test-guardrail-id", guardrailVersion="1" - ) + guardrail_hook = BedrockGuardrail(guardrailIdentifier="test-guardrail-id", guardrailVersion="1") guardrail_hook.async_handler = Mock() mock_response = Mock() @@ -940,20 +867,13 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): test_request_data = {"api_key": "test-api-key-789"} with ( - patch.object( - guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response) - ), + patch.object(guardrail_hook.async_handler, "post", AsyncMock(return_value=mock_response)), patch.object(guardrail_hook, "_load_credentials") as mock_load_creds, patch.object(guardrail_hook, "convert_to_bedrock_format") as mock_convert, - patch.object( - guardrail_hook, "get_guardrail_dynamic_request_body_params" - ) as mock_get_params, - patch.object( - guardrail_hook, "add_standard_logging_guardrail_information_to_request_data" - ), + patch.object(guardrail_hook, "get_guardrail_dynamic_request_body_params") as mock_get_params, + patch.object(guardrail_hook, "add_standard_logging_guardrail_information_to_request_data"), patch("botocore.awsrequest.AWSRequest") as mock_aws_request, ): - mock_load_creds.return_value = (Mock(), "us-east-1") mock_convert.return_value = {"source": "INPUT", "content": [{"text": {"text": "test"}}]} mock_get_params.return_value = {} @@ -965,9 +885,7 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key(): "Content-Type": "application/json", "Authorization": "Bearer test-api-key-789", } - mock_request_instance.prepare.return_value = Mock( - headers=mock_request_instance.headers - ) + mock_request_instance.prepare.return_value = Mock(headers=mock_request_instance.headers) mock_aws_request.return_value = mock_request_instance await guardrail_hook.make_bedrock_api_request( @@ -1025,12 +943,8 @@ async def test_create_guardrail_endpoint( elif scenario == "success_sync_fails": mock_prisma_client = mocker.Mock() - mock_in_memory_handler.initialize_guardrail.side_effect = Exception( - "Sync failed" - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) + mock_in_memory_handler.initialize_guardrail.side_effect = Exception("Sync failed") + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -1044,9 +958,7 @@ async def test_create_guardrail_endpoint( elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception( - "Database error" - ) + mock_guardrail_registry.add_guardrail_to_db.side_effect = Exception("Database error") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -1060,9 +972,7 @@ async def test_create_guardrail_endpoint( # Run the test if expected_exception: with pytest.raises(expected_exception) as exc_info: - await create_guardrail( - MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER - ) + await create_guardrail(MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) if scenario == "database_failure": assert "Database error" in str(exc_info.value.detail) @@ -1070,9 +980,7 @@ async def test_create_guardrail_endpoint( assert "Prisma client not initialized" in str(exc_info.value.detail) else: - result = await create_guardrail( - MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER - ) + result = await create_guardrail(MOCK_CREATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" @@ -1086,9 +994,7 @@ async def test_create_guardrail_endpoint( if scenario == "success_sync_fails": assert mock_logger is not None mock_logger.warning.assert_called_once() - assert "Failed to initialize guardrail" in str( - mock_logger.warning.call_args - ) + assert "Failed to initialize guardrail" in str(mock_logger.warning.call_args) @pytest.mark.parametrize( @@ -1139,12 +1045,8 @@ async def test_update_guardrail_endpoint( # so it keeps the pre-existing swallow-and-warn behavior rather than # rolling back the DB write. mock_prisma_client = mocker.Mock() - mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( - side_effect=Exception("Sync failed") - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(side_effect=Exception("Sync failed")) + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -1177,9 +1079,7 @@ async def test_update_guardrail_endpoint( elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( - "Database error" - ) + mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -1209,15 +1109,10 @@ async def test_update_guardrail_endpoint( # Rolled back: update_guardrail_in_db is called once for the # rejected write and once more to restore the previous config. assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2 - assert ( - mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"] - == MOCK_DB_GUARDRAIL - ) + assert mock_guardrail_registry.update_guardrail_in_db.call_args.kwargs["guardrail"] == MOCK_DB_GUARDRAIL else: - result = await update_guardrail( - "test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER - ) + result = await update_guardrail("test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" @@ -1228,9 +1123,7 @@ async def test_update_guardrail_endpoint( prisma_client=mocker.ANY, ) - mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with( - guardrail=mocker.ANY - ) + mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(guardrail=mocker.ANY) if scenario == "success_sync_fails_unexpected_error": assert mock_logger is not None @@ -1286,12 +1179,8 @@ async def test_patch_guardrail_endpoint( # config-rejection signal, so it keeps the pre-existing swallow-and-warn # behavior rather than rolling back the DB write. mock_prisma_client = mocker.Mock() - mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock( - side_effect=Exception("Sync failed") - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) + mock_in_memory_handler.sync_guardrail_from_db = mocker.Mock(side_effect=Exception("Sync failed")) + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -1324,9 +1213,7 @@ async def test_patch_guardrail_endpoint( elif scenario == "database_failure": mock_prisma_client = mocker.Mock() - mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception( - "Database error" - ) + mock_guardrail_registry.update_guardrail_in_db.side_effect = Exception("Database error") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( @@ -1358,18 +1245,14 @@ async def test_patch_guardrail_endpoint( assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2 else: - result = await patch_guardrail( - "test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER - ) + result = await patch_guardrail("test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER) assert result["guardrail_id"] == expected_result assert result["guardrail_name"] == "Test DB Guardrail" mock_guardrail_registry.update_guardrail_in_db.assert_called_once() - mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with( - guardrail=mocker.ANY - ) + mock_in_memory_handler.sync_guardrail_from_db.assert_called_once_with(guardrail=mocker.ANY) if scenario == "success_sync_fails_unexpected_error": assert mock_logger is not None @@ -1428,12 +1311,8 @@ async def test_delete_guardrail_endpoint( ) elif scenario == "success_sync_fails": - mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception( - "Sync failed" - ) - mock_logger = mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger" - ) + mock_in_memory_handler.delete_in_memory_guardrail.side_effect = Exception("Sync failed") + mock_logger = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.verbose_proxy_logger") mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mocker.patch( "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", @@ -1446,13 +1325,9 @@ async def test_delete_guardrail_endpoint( if expected_exception: with pytest.raises(expected_exception): - await delete_guardrail( - guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER - ) + await delete_guardrail(guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER) else: - result = await delete_guardrail( - guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER - ) + result = await delete_guardrail(guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER) assert result == MOCK_DB_GUARDRAIL @@ -1463,9 +1338,7 @@ async def test_delete_guardrail_endpoint( guardrail_id=expected_result, prisma_client=mock_prisma_client ) - mock_in_memory_handler.delete_in_memory_guardrail.assert_called_once_with( - guardrail_id=expected_result - ) + mock_in_memory_handler.delete_in_memory_guardrail.assert_called_once_with(guardrail_id=expected_result) if scenario == "success_sync_fails": assert mock_logger is not None @@ -1483,9 +1356,7 @@ async def test_apply_guardrail_not_found(mocker): # Mock the GUARDRAIL_REGISTRY to return None (guardrail not found) mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = None - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) mock_proxy_logging = mocker.Mock() mock_proxy_logging.post_call_failure_hook = AsyncMock() @@ -1495,9 +1366,7 @@ async def test_apply_guardrail_not_found(mocker): mocker.patch("litellm.proxy.proxy_server.version", "test") # Create request - request = ApplyGuardrailRequest( - guardrail_name="non-existent-guardrail", text="Test input text" - ) + request = ApplyGuardrailRequest(guardrail_name="non-existent-guardrail", text="Test input text") # Mock user auth mock_user_auth = UserAPIKeyAuth() @@ -1531,9 +1400,7 @@ async def test_apply_guardrail_execution_error(mocker): # Mock the GUARDRAIL_REGISTRY mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) mock_logging_obj = mocker.Mock() mock_logging_obj.async_failure_handler = AsyncMock() @@ -1555,9 +1422,7 @@ async def test_apply_guardrail_execution_error(mocker): mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor") # Create request - request = ApplyGuardrailRequest( - guardrail_name="test-guardrail", text="Test input text with forbidden content" - ) + request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="Test input text with forbidden content") # Mock user auth mock_user_auth = UserAPIKeyAuth() @@ -1581,9 +1446,7 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker): mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) mock_logging_obj = mocker.Mock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1604,13 +1467,9 @@ async def test_apply_guardrail_invokes_logging_pipeline(mocker): mocker.patch("litellm.proxy.proxy_server.proxy_config", mocker.Mock()) mocker.patch("litellm.proxy.proxy_server.version", "test") mock_executor = mocker.Mock() - mocker.patch( - "litellm.litellm_core_utils.thread_pool_executor.executor", mock_executor - ) + mocker.patch("litellm.litellm_core_utils.thread_pool_executor.executor", mock_executor) - request = ApplyGuardrailRequest( - guardrail_name="test-guardrail", text="hello@example.com" - ) + request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello@example.com") response = await apply_guardrail( fastapi_request=mocker.Mock(), request=request, @@ -1634,9 +1493,7 @@ def _patch_apply_guardrail_env(mocker, guardrail_result): mock_registry = mocker.Mock() mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) mock_logging_obj = mocker.Mock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1772,9 +1629,7 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker): # Mock the GUARDRAIL_REGISTRY to return None from DB (so it checks config) mock_registry = mocker.Mock() mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=None) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) # Mock IN_MEMORY_GUARDRAIL_HANDLER at its source to return config guardrail mock_in_memory_handler = mocker.Mock() @@ -1814,12 +1669,8 @@ async def test_get_guardrail_info_endpoint_db_guardrail(mocker): # Mock the GUARDRAIL_REGISTRY to return a guardrail from DB mock_registry = mocker.Mock() - mock_registry.get_guardrail_by_id_from_db = AsyncMock( - return_value=MOCK_DB_GUARDRAIL - ) - mocker.patch( - "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry - ) + mock_registry.get_guardrail_by_id_from_db = AsyncMock(return_value=MOCK_DB_GUARDRAIL) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_registry) # Mock IN_MEMORY_GUARDRAIL_HANDLER to return None mock_in_memory_handler = mocker.Mock() @@ -1978,9 +1829,7 @@ async def test_register_guardrail_non_admin_cross_team_allowed(mocker): team_id="team-beta", litellm_params=MOCK_REGISTER_REQUEST.litellm_params, ) - user = UserAPIKeyAuth( - user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha" - ) + user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha") result = await register_guardrail(req, user) @@ -2000,9 +1849,7 @@ async def test_register_guardrail_non_admin_cross_team_forbidden(mocker): team_id="team-other", litellm_params=MOCK_REGISTER_REQUEST.litellm_params, ) - user = UserAPIKeyAuth( - user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha" - ) + user = UserAPIKeyAuth(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-alpha") with pytest.raises(HTTPException) as exc_info: await register_guardrail(req, user) @@ -2184,9 +2031,7 @@ async def test_list_guardrail_submissions_team_id_filter(mocker): mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) - result = await list_guardrail_submissions( - user_api_key_dict=user, team_id="team-abc" - ) + result = await list_guardrail_submissions(user_api_key_dict=user, team_id="team-abc") assert len(result.submissions) == 1 assert result.submissions[0].guardrail_id == "team-1" @@ -2288,9 +2133,7 @@ async def test_get_guardrail_submission_admin_viewer_other_team_allowed(mocker): "litellm.proxy.guardrails.guardrail_endpoints._get_user_team_ids", AsyncMock(return_value=[]), ) - user = UserAPIKeyAuth( - user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ) + user = UserAPIKeyAuth(user_id="viewer-1", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) result = await get_guardrail_submission("sub-1", user) @@ -2370,9 +2213,7 @@ async def test_reject_guardrail_submission_success(mocker): async def test_reject_guardrail_submission_not_pending(mocker): """Reject returns 400 when status is not pending_review (e.g. already active).""" mock_prisma = mocker.Mock() - row = mocker.Mock( - guardrail_id="already-active", guardrail_name="g", status="active" - ) + row = mocker.Mock(guardrail_id="already-active", guardrail_name="g", status="active") mock_prisma.db.litellm_guardrailstable.find_unique = AsyncMock(return_value=row) mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) @@ -2404,9 +2245,7 @@ async def test_reject_guardrail_submission_not_pending(mocker): "no_hostname", ], ) -async def test_register_guardrail_rejects_bad_api_base( - mocker, api_base, expected_detail -): +async def test_register_guardrail_rejects_bad_api_base(mocker, api_base, expected_detail): """Register returns 400 when api_base has invalid scheme or missing hostname.""" mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) req = RegisterGuardrailRequest( @@ -2474,9 +2313,7 @@ async def test_approve_guardrail_init_failure_returns_warning(mocker): mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma) mock_handler = mocker.Mock() - mock_handler.initialize_guardrail = mocker.Mock( - side_effect=Exception("missing dependency") - ) + mock_handler.initialize_guardrail = mocker.Mock(side_effect=Exception("missing dependency")) mocker.patch( "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_handler, @@ -2572,9 +2409,7 @@ async def test_list_submissions_summary_counts_unaffected_by_filters(mocker): user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) # Filter to only pending, but summary should still show both - result = await list_guardrail_submissions( - status="pending_review", user_api_key_dict=user - ) + result = await list_guardrail_submissions(status="pending_review", user_api_key_dict=user) assert len(result.submissions) == 1 # filtered assert result.summary.total == 2 # unfiltered @@ -2630,15 +2465,13 @@ async def test_ui_settings_map_matches_runtime_supported_event_hooks(): for provider, guardrail_class in guardrail_class_registry.items(): declared = guardrail_class.get_supported_event_hooks() if declared is None: - assert ( - provider not in result.supported_modes_by_provider - ), f"{provider} returned None from classmethod but appears in map" + assert provider not in result.supported_modes_by_provider, ( + f"{provider} returned None from classmethod but appears in map" + ) continue assert provider in result.supported_modes_by_provider, provider - assert result.supported_modes_by_provider[provider] == [ - hook.value for hook in declared - ], provider + assert result.supported_modes_by_provider[provider] == [hook.value for hook in declared], provider def test_content_filter_runtime_rejects_unsupported_mcp_hook(): @@ -2725,3 +2558,115 @@ def test_field_type_inference_handles_pep604_unions(): assert _get_field_type_from_annotation(list[str] | None) == "array" assert _get_field_type_from_annotation(bool | None) == "boolean" assert _unwrap_optional_type(str | None) is str + + +@pytest.mark.asyncio +@pytest.mark.timeout(20) +async def test_test_custom_code_endpoint_returns_a_timeout_for_an_infinite_loop(): + """The endpoint used to join the worker thread after its timeout fired, so an infinite + loop hung the request forever.""" + request = TestCustomCodeGuardrailRequest( + custom_code="def apply_guardrail(inputs, request_data, input_type):\n n = 0\n while True:\n n += 1\n", + test_input={"texts": ["x"]}, + ) + started = time.monotonic() + + response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER) + + assert response.success is False + assert response.error_type == "execution" + assert response.error is not None + assert response.error.startswith("Execution timeout: code took longer than 5 seconds") + assert time.monotonic() - started < 8.0 + + +@pytest.mark.asyncio +@pytest.mark.timeout(20) +async def test_test_custom_code_endpoint_reports_a_module_level_infinite_loop_as_a_timeout(): + """Module-level code that outran the load deadline was reported as a compile failure, as if the + source were invalid.""" + request = TestCustomCodeGuardrailRequest( + custom_code=( + "n = 0\nwhile True:\n n += 1\n\n" + "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" + ), + test_input={"texts": ["x"]}, + ) + started = time.monotonic() + + response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER) + + assert response.success is False + assert response.error_type == "execution" + assert response.error is not None + assert response.error.startswith("Execution timeout: code took longer than 5 seconds") + assert time.monotonic() - started < 8.0 + + +@pytest.mark.asyncio +async def test_test_custom_code_endpoint_awaits_an_async_guardrail(): + request = TestCustomCodeGuardrailRequest( + custom_code=( + 'async def apply_guardrail(inputs, request_data, input_type):\n return block("async said no")\n' + ), + test_input={"texts": ["x"]}, + ) + + response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER) + + assert response.success is True + assert response.result is not None + assert response.result["action"] == "block" + assert response.result["reason"] == "async said no" + + +@pytest.mark.asyncio +async def test_test_custom_code_endpoint_returns_a_sync_guardrails_result(): + request = TestCustomCodeGuardrailRequest( + custom_code='def apply_guardrail(inputs, request_data, input_type):\n return block("sync said no")\n', + test_input={"texts": ["x"]}, + ) + + response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER) + + assert response.success is True + assert response.result is not None + assert response.result["action"] == "block" + assert response.result["reason"] == "sync said no" + + +@pytest.mark.asyncio +async def test_add_guardrail_rolls_back_a_custom_code_guardrail_that_fails_to_compile(mocker, mock_guardrail_registry): + stored = { + "guardrail_id": "custom-code-broken", + "guardrail_name": "custom-code-broken", + "litellm_params": {"guardrail": "custom_code", "mode": "pre_call", "custom_code": "x = 1\n"}, + "guardrail_info": {}, + } + mock_guardrail_registry.add_guardrail_to_db = AsyncMock(return_value=stored) + mocker.patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY", mock_guardrail_registry) + mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.Mock()) + delete_row = mocker.patch("litellm.proxy.guardrails.guardrail_endpoints._delete_guardrail_row", AsyncMock()) + + with pytest.raises(HTTPException) as exc_info: + await create_guardrail(CreateGuardrailRequest(guardrail=stored), user_api_key_dict=MOCK_ADMIN_USER) + + assert exc_info.value.status_code == 400 + assert "apply_guardrail" in exc_info.value.detail + delete_row.assert_awaited_once_with(mocker.ANY, where={"guardrail_id": "custom-code-broken"}) + + +@pytest.mark.asyncio +async def test_test_custom_code_endpoint_reports_a_system_exit_as_an_execution_error(): + request = TestCustomCodeGuardrailRequest( + custom_code="def apply_guardrail(inputs, request_data, input_type):\n raise SystemExit('bye')\n", + test_input={"texts": ["x"]}, + ) + started = time.monotonic() + + response = await run_custom_code_test_endpoint(request=request, user_api_key_dict=MOCK_ADMIN_USER) + + assert response.success is False + assert response.error == "Execution error: SystemExit: bye" + assert response.error_type == "execution" + assert time.monotonic() - started < 2.0 diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index fc2fb949143..39f9f9458b7 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, patch import pytest +from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import CustomCodeCompilationError from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -335,6 +336,27 @@ def test_init_guardrails_v2_skips_invalid_guardrail_instead_of_crashing_boot(): assert "healthy_presidio" in guardrail_names +def test_init_guardrails_v2_stops_boot_when_a_custom_code_guardrail_does_not_compile(): + from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + + IN_MEMORY_GUARDRAIL_HANDLER.IN_MEMORY_GUARDRAILS.clear() + IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.clear() + + all_guardrails = [ + { + "guardrail_name": "custom-code-without-apply-guardrail", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.CUSTOM_CODE.value, + "mode": "pre_call", + "custom_code": "x = 1\n", + }, + }, + ] + + with pytest.raises(CustomCodeCompilationError, match="apply_guardrail"): + init_guardrails_v2(all_guardrails=all_guardrails) + + def test_init_guardrails_v2_accepts_during_call_advisory_mode(): """ Maintainer finding on BerriAI/litellm#34940: on_flagged='inject_system_message'