From 20fdedf44fa66390620745c929c166dc573a75f3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 24 Feb 2026 18:28:19 -0800 Subject: [PATCH] feat(add-new-block_code_execution-guardrail): prevent agent from executing code --- .../proxy/guardrails/guardrail_endpoints.py | 75 ++-- .../block_code_execution/__init__.py | 66 ++++ .../block_code_execution.py | 362 ++++++++++++++++++ litellm/types/guardrails.py | 1 + .../guardrail_hooks/block_code_execution.py | 65 ++++ .../test_block_code_execution.py | 200 ++++++++++ 6 files changed, 734 insertions(+), 35 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index c083c60cb4c..27f42bd9526 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -14,27 +14,23 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry -from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router -from litellm.types.guardrails import ( - PII_ENTITY_CATEGORIES_MAP, - ApplyGuardrailRequest, - ApplyGuardrailResponse, - BaseLitellmParams, - BedrockGuardrailConfigModel, - Guardrail, - GuardrailEventHooks, - GuardrailInfoResponse, - GuardrailUIAddGuardrailSettings, - LakeraV2GuardrailConfigModel, - ListGuardrailsResponse, - LitellmParams, - PatchGuardrailRequest, - PiiAction, - PiiEntityType, - PresidioPresidioConfigModelUserInterface, - SupportedGuardrailIntegrations, - ToolPermissionGuardrailConfigModel, -) +from litellm.proxy.guardrails.usage_endpoints import \ + router as guardrails_usage_router +from litellm.types.guardrails import (PII_ENTITY_CATEGORIES_MAP, + ApplyGuardrailRequest, + ApplyGuardrailResponse, + BaseLitellmParams, + BedrockGuardrailConfigModel, Guardrail, + GuardrailEventHooks, + GuardrailInfoResponse, + GuardrailUIAddGuardrailSettings, + LakeraV2GuardrailConfigModel, + ListGuardrailsResponse, LitellmParams, + PatchGuardrailRequest, PiiAction, + PiiEntityType, + PresidioPresidioConfigModelUserInterface, + SupportedGuardrailIntegrations, + ToolPermissionGuardrailConfigModel) #### GUARDRAILS ENDPOINTS #### @@ -153,7 +149,8 @@ async def list_guardrails_v2(): ``` """ from litellm.litellm_core_utils.litellm_logging import _get_masked_values - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -293,7 +290,8 @@ async def create_guardrail(request: CreateGuardrailRequest): } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -382,7 +380,8 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -450,7 +449,8 @@ async def delete_guardrail(guardrail_id: str): } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -543,7 +543,8 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): } ``` """ - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -665,7 +666,8 @@ async def get_guardrail_info(guardrail_id: str): """ from litellm.litellm_core_utils.litellm_logging import _get_masked_values - from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER + from litellm.proxy.guardrails.guardrail_registry import \ + IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client from litellm.types.guardrails import GUARDRAIL_DEFINITION_LOCATION @@ -740,10 +742,8 @@ async def get_guardrail_ui_settings(): - Content filter settings (patterns and categories) """ from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import ( - PATTERN_CATEGORIES, - get_available_content_categories, - get_pattern_metadata, - ) + PATTERN_CATEGORIES, get_available_content_categories, + get_pattern_metadata) # Convert the PII_ENTITY_CATEGORIES_MAP to the format expected by the UI category_maps = [] @@ -1163,6 +1163,11 @@ def _build_field_dict( # Add options if they exist in json_schema_extra (this takes precedence) if field_json_schema_extra and "options" in field_json_schema_extra: field_dict["options"] = field_json_schema_extra["options"] + elif field_type == "select": + # For Literal types, populate options so the UI can render a dropdown + literal_options = _extract_literal_values(field_annotation) + if literal_options: + field_dict["options"] = literal_options # Add default value if it exists if field.default is not None and field.default is not ...: @@ -1315,7 +1320,8 @@ async def get_provider_specific_params(): } ### get the config model for the guardrail - go through the registry and get the config model for the guardrail - from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry + from litellm.proxy.guardrails.guardrail_registry import \ + guardrail_class_registry for guardrail_name, guardrail_class in guardrail_class_registry.items(): guardrail_config_model = guardrail_class.get_config_model() @@ -1443,9 +1449,8 @@ async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest): import concurrent.futures import re - from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import ( - get_custom_code_primitives, - ) + from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import \ + get_custom_code_primitives # Security validation patterns FORBIDDEN_PATTERNS = [ diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py new file mode 100644 index 00000000000..7fa8f4a5a42 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/__init__.py @@ -0,0 +1,66 @@ +"""Block Code Execution guardrail: blocks or masks fenced code blocks by language.""" + +from typing import TYPE_CHECKING, Literal, cast + +from litellm.types.guardrails import (GuardrailEventHooks, + SupportedGuardrailIntegrations) + +from .block_code_execution import BlockCodeExecutionGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + +# Default: run on both request and response (and during_call is supported too) +DEFAULT_EVENT_HOOKS = [ + GuardrailEventHooks.pre_call.value, + GuardrailEventHooks.post_call.value, +] + + +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", +) -> BlockCodeExecutionGuardrail: + """Initialize the Block Code Execution guardrail from config.""" + import litellm + + guardrail_name = guardrail.get("guardrail_name") + if not guardrail_name: + raise ValueError( + "Block Code Execution guardrail requires a guardrail_name" + ) + + blocked_languages = getattr(litellm_params, "blocked_languages", None) + action = cast( + Literal["block", "mask"], getattr(litellm_params, "action", "block") + ) + confidence_threshold = float( + getattr(litellm_params, "confidence_threshold", 0.7) + ) + mode = getattr(litellm_params, "mode", None) + event_hook = mode if mode is not None else DEFAULT_EVENT_HOOKS + + instance = BlockCodeExecutionGuardrail( + guardrail_name=guardrail_name, + blocked_languages=blocked_languages, + action=action, + confidence_threshold=confidence_threshold, + event_hook=event_hook, + default_on=getattr(litellm_params, "default_on", False), + ) + litellm.logging_callback_manager.add_litellm_callback(instance) + return instance + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.BLOCK_CODE_EXECUTION.value: initialize_guardrail, +} + +guardrail_class_registry = { + SupportedGuardrailIntegrations.BLOCK_CODE_EXECUTION.value: BlockCodeExecutionGuardrail, +} + +__all__ = [ + "BlockCodeExecutionGuardrail", + "initialize_guardrail", +] diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py new file mode 100644 index 00000000000..33edf5735f9 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -0,0 +1,362 @@ +""" +Block Code Execution guardrail. + +Detects markdown fenced code blocks in request/response content and blocks or masks them +when the language is in the blocked list (or all blocks when list is empty). Supports +confidence scoring and a tunable threshold (only block when confidence >= threshold). +""" + +import re +from datetime import datetime +from typing import (TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Literal, + Optional, Tuple, Union, cast) + +from fastapi import HTTPException + +from litellm.integrations.custom_guardrail import (CustomGuardrail, + log_guardrail_information) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.base import \ + GuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import ( + CodeBlockActionTaken, CodeBlockDetection) +from litellm.types.utils import (GenericGuardrailAPIInputs, GuardrailStatus, + GuardrailTracingDetail, ModelResponseStream) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import \ + Logging as LiteLLMLoggingObj + +# Default executable languages when blocked_languages is not set (block-all mode uses this for "block all") +DEFAULT_BLOCKED_LANGUAGES: List[str] = [ + "python", + "javascript", + "js", + "bash", + "sh", + "ruby", + "go", + "java", + "csharp", + "php", + "c", + "cpp", + "rust", + "sql", +] + +# Language tag aliases (normalize to canonical for comparison) +LANGUAGE_ALIASES: Dict[str, str] = { + "js": "javascript", + "py": "python", +} + +# Tags that indicate non-executable / plain text (lower confidence when block-all) +NON_EXECUTABLE_TAGS: frozenset = frozenset( + {"text", "plaintext", "plain", "markdown", "md", "output", "result"} +) + +# Regex: fenced code block with optional language tag. Handles ```lang\n...\n``` +# Content between fences; does not handle nested ``` inside body (documented edge case). +FENCED_BLOCK_RE = re.compile(r"```(\w*)\n(.*?)```", re.DOTALL) + + +def _normalize_language(tag: str) -> str: + """Normalize language tag (lowercase, resolve aliases).""" + tag = (tag or "").strip().lower() + return LANGUAGE_ALIASES.get(tag, tag) + + +def _is_blocked_language( + tag: str, + blocked_languages: Optional[List[str]], + block_all: bool, +) -> bool: + """True if this language tag should be considered blocked.""" + normalized = _normalize_language(tag) + if block_all: + # Block all: only allow through if it's explicitly non-executable (we still block but with lower confidence) + return True + if not blocked_languages: + return True + normalized_list = [_normalize_language(t) for t in blocked_languages] + return normalized in normalized_list + + +def _confidence_for_block( + tag: str, + block_all: bool, + tag_in_blocked_list: bool, +) -> float: + """Return confidence in [0, 1] for this code block detection.""" + normalized = _normalize_language(tag) + if tag_in_blocked_list: + return 1.0 + if block_all: + # Explicit non-executable tags (e.g. text, plaintext) get lower confidence + if normalized in NON_EXECUTABLE_TAGS: + return 0.5 + # Untagged or other tags in block-all mode: treat as executable, high confidence + return 1.0 + return 0.0 + + +class BlockCodeExecutionGuardrail(CustomGuardrail): + """ + Guardrail that detects fenced code blocks (markdown ```) and blocks or masks them + when the language is in the blocked list (or all when list is empty/None). + Supports confidence threshold: only block when confidence >= confidence_threshold. + """ + + MASK_PLACEHOLDER = "[CODE_BLOCK_REDACTED]" + + def __init__( + self, + guardrail_name: Optional[str] = None, + blocked_languages: Optional[List[str]] = None, + action: Literal["block", "mask"] = "block", + confidence_threshold: float = 0.7, + event_hook: Optional[ + Union[Literal["pre_call", "post_call", "during_call"], List[str]] + ] = None, + default_on: bool = False, + **kwargs: Any, + ) -> None: + # Normalize to type expected by CustomGuardrail + _event_hook: Optional[ + Union[GuardrailEventHooks, List[GuardrailEventHooks]] + ] = None + if event_hook is not None: + if isinstance(event_hook, list): + _event_hook = [ + GuardrailEventHooks(h) if isinstance(h, str) else h + for h in event_hook + ] + else: + _event_hook = GuardrailEventHooks(event_hook) + super().__init__( + guardrail_name=guardrail_name or "block_code_execution", + supported_event_hooks=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.during_call, + ], + event_hook=_event_hook + or [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + default_on=default_on, + **kwargs, + ) + self.blocked_languages = blocked_languages + self.block_all = blocked_languages is None or len(blocked_languages) == 0 + self.action = action + self.confidence_threshold = max(0.0, min(1.0, confidence_threshold)) + + @staticmethod + def get_config_model() -> Optional[type[GuardrailConfigModel]]: + from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import \ + BlockCodeExecutionGuardrailConfigModel + return BlockCodeExecutionGuardrailConfigModel + + def _find_blocks(self, text: str) -> List[Tuple[str, str, float, CodeBlockActionTaken]]: + """ + Find all fenced code blocks in text. Returns list of + (language_tag, block_content, confidence, action_taken). + """ + results: List[Tuple[str, str, float, CodeBlockActionTaken]] = [] + for m in FENCED_BLOCK_RE.finditer(text): + tag = (m.group(1) or "").strip() + body = m.group(2) + tag_in_list = not self.block_all and _normalize_language(tag) in [ + _normalize_language(t) for t in (self.blocked_languages or []) + ] + is_blocked = _is_blocked_language( + tag, self.blocked_languages, self.block_all + ) + confidence = _confidence_for_block(tag, self.block_all, tag_in_list) + if not is_blocked: + action_taken: CodeBlockActionTaken = "allow" + elif confidence >= self.confidence_threshold: + action_taken = "block" + else: + action_taken = "log_only" + results.append((tag or "(none)", body, confidence, action_taken)) + return results + + def _scan_text( + self, + text: str, + detections: Optional[List[CodeBlockDetection]] = None, + ) -> Tuple[str, bool]: + """ + Scan one text: find blocks, apply block/mask/allow by confidence. + Returns (modified_text, should_raise). + """ + if not text: + return text, False + blocks = self._find_blocks(text) + if not blocks: + return text, False + + should_raise = False + last_end = 0 + parts: List[str] = [] + for m in FENCED_BLOCK_RE.finditer(text): + tag = (m.group(1) or "").strip() + tag_in_list = not self.block_all and _normalize_language(tag) in [ + _normalize_language(t) for t in (self.blocked_languages or []) + ] + is_blocked = _is_blocked_language( + tag, self.blocked_languages, self.block_all + ) + confidence = _confidence_for_block(tag, self.block_all, tag_in_list) + if not is_blocked: + action_taken: CodeBlockActionTaken = "allow" + elif confidence >= self.confidence_threshold: + action_taken = "block" + else: + action_taken = "log_only" + + if detections is not None: + detections.append( + cast( + CodeBlockDetection, + { + "type": "code_block", + "language": tag or "(none)", + "confidence": round(confidence, 2), + "action_taken": action_taken, + }, + ) + ) + + if action_taken == "block" and self.action == "block": + should_raise = True + parts.append(text[last_end : m.start()]) + if action_taken == "block": + parts.append(self.MASK_PLACEHOLDER) + else: + parts.append(text[m.start() : m.end()]) + last_end = m.end() + + parts.append(text[last_end:]) + new_text = "".join(parts) + return new_text, should_raise + + def _raise_block_error(self, language: str, is_output: bool, request_data: dict) -> None: + msg = ( + f"Content blocked: executable code block detected (language: {language})" + ) + if is_output: + raise HTTPException( + status_code=400, + detail={ + "error": msg, + "guardrail": self.guardrail_name, + "language": language, + }, + ) + self.raise_passthrough_exception( + violation_message=msg, + request_data=request_data, + detection_info={"language": language}, + ) + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + start_time = datetime.now() + detections: List[CodeBlockDetection] = [] + status: GuardrailStatus = "success" + exception_str = "" + + try: + texts = inputs.get("texts", []) + if not texts: + return inputs + + is_output = input_type == "response" + processed: List[str] = [] + for text in texts: + new_text, should_raise = self._scan_text(text, detections) + processed.append(new_text) + if should_raise: + # Determine language from first blocking detection + lang = "unknown" + for d in detections: + if d.get("action_taken") == "block": + lang = d.get("language", "unknown") + break + self._raise_block_error(lang, is_output, request_data) + + inputs["texts"] = processed + return inputs + except HTTPException: + status = "guardrail_intervened" + raise + except Exception as e: + status = "guardrail_failed_to_respond" + exception_str = str(e) + raise + finally: + guardrail_response: Union[List[dict], str] = [dict(d) for d in detections] + if status != "success" and not detections: + guardrail_response = exception_str + max_confidence: Optional[float] = None + for d in detections: + c = d.get("confidence") + if c is not None and (max_confidence is None or c > max_confidence): + max_confidence = c + tracing_kw: Dict[str, Any] = { + "guardrail_id": self.guardrail_name, + "detection_method": "fenced_code_block", + "match_details": guardrail_response, + } + if max_confidence is not None: + tracing_kw["confidence_score"] = max_confidence + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider="block_code_execution", + guardrail_json_response=guardrail_response, + request_data=request_data, + guardrail_status=status, + start_time=start_time.timestamp(), + end_time=datetime.now().timestamp(), + duration=(datetime.now() - start_time).total_seconds(), + tracing_detail=GuardrailTracingDetail(**tracing_kw), # type: ignore[typeddict-item] + ) + + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: Any, + response: Any, + request_data: dict, + ) -> AsyncGenerator[ModelResponseStream, None]: + """Accumulate streamed content and block if a complete fenced code block is detected.""" + accumulated = "" + async for item in response: + if isinstance(item, ModelResponseStream) and item.choices: + delta_content = "" + is_final = False + for choice in item.choices: + if hasattr(choice, "delta") and choice.delta: + content = getattr(choice.delta, "content", None) + if content and isinstance(content, str): + delta_content += content + if getattr(choice, "finish_reason", None): + is_final = True + accumulated += delta_content + if is_final: + # Run detection on full accumulated text (streaming: block only, no mask) + blocks = self._find_blocks(accumulated) + for _tag, _body, confidence, action_taken in blocks: + if action_taken == "block" and confidence >= self.confidence_threshold: + lang = _tag or "unknown" + self._raise_block_error(lang, True, request_data) + yield item diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index c87d1076d0e..47b54a8a0fe 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -73,6 +73,7 @@ class SupportedGuardrailIntegrations(Enum): CUSTOM_CODE = "custom_code" SEMANTIC_GUARD = "semantic_guard" MCP_END_USER_PERMISSION = "mcp_end_user_permission" + BLOCK_CODE_EXECUTION = "block_code_execution" class Role(Enum): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py new file mode 100644 index 00000000000..641892cb674 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/block_code_execution.py @@ -0,0 +1,65 @@ +"""Types for the Block Code Execution guardrail.""" + +from typing import Any, List, Literal, Optional, TypedDict, cast + +from pydantic import Field + +from .base import GuardrailConfigModel + +CodeBlockActionTaken = Literal["block", "allow", "log_only"] + +# Supported language tags for the blocked_languages multiselect dropdown +BLOCKED_LANGUAGES_OPTIONS = [ + "python", + "javascript", + "js", + "bash", + "sh", + "ruby", + "go", + "java", + "csharp", + "php", + "c", + "cpp", + "rust", + "sql", +] + + +class CodeBlockDetection(TypedDict, total=False): + """Detection output for a single fenced code block (for tracing/logging).""" + + type: Literal["code_block"] + language: str + confidence: float + action_taken: CodeBlockActionTaken + evidence: Optional[str] + snippet: Optional[str] + + +class BlockCodeExecutionGuardrailConfigModel(GuardrailConfigModel): + """Configuration for the Block Code Execution guardrail.""" + + blocked_languages: Optional[List[str]] = Field( + default=None, + description="Language tags to block (e.g. python, javascript, bash). Empty or None = block all fenced code blocks.", + json_schema_extra=cast( + Any, + {"type": "multiselect", "options": BLOCKED_LANGUAGES_OPTIONS}, + ), + ) + action: Literal["block", "mask"] = Field( + default="block", + description="'block' raises an error; 'mask' replaces the code block with a placeholder.", + ) + confidence_threshold: float = Field( + default=0.7, + ge=0.0, + le=1.0, + description="Only block or mask when detection confidence >= this value; below threshold, allow or log_only.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Block Code Execution" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py new file mode 100644 index 00000000000..f17a8a224c2 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_block_code_execution.py @@ -0,0 +1,200 @@ +"""Tests for the Block Code Execution guardrail.""" + +import pytest +from fastapi import HTTPException + +from litellm.proxy.guardrails.guardrail_hooks.block_code_execution import ( + DEFAULT_EVENT_HOOKS, BlockCodeExecutionGuardrail, initialize_guardrail) +from litellm.types.guardrails import GuardrailEventHooks + + +class TestBlockCodeExecutionGuardrail: + """Test BlockCodeExecutionGuardrail detection and actions.""" + + def test_detects_python_block_when_in_blocked_list(self): + """Text with ```python block is detected when python is in blocked_languages.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + confidence_threshold=0.7, + ) + blocks = guardrail._find_blocks("Here is code:\n```python\nprint(1)\n```\nDone.") + assert len(blocks) == 1 + tag, _body, confidence, action_taken = blocks[0] + assert tag == "python" + assert confidence == 1.0 + assert action_taken == "block" + + def test_block_all_when_blocked_languages_empty(self): + """When blocked_languages is empty, any fenced block is blocked (block all).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=[], + confidence_threshold=0.7, + ) + blocks = guardrail._find_blocks("```\nfoo\n```") + assert len(blocks) == 1 + _tag, _body, confidence, action_taken = blocks[0] + assert action_taken == "block" + assert confidence in (0.5, 1.0) + + def test_no_block_when_language_not_in_list(self): + """When language is not in blocked_languages, block is not triggered.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + confidence_threshold=0.7, + ) + blocks = guardrail._find_blocks("```text\nplain output\n```") + assert len(blocks) == 1 + _tag, _body, confidence, action_taken = blocks[0] + assert action_taken == "allow" + assert confidence == 0.0 + + def test_confidence_below_threshold_allows(self): + """When confidence < confidence_threshold, action_taken is log_only and we do not block.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=[], # block all + confidence_threshold=0.9, + ) + # Block with no tag or plaintext tag gets confidence 0.5 + blocks = guardrail._find_blocks("```text\nx\n```") + assert len(blocks) == 1 + _tag, _body, confidence, action_taken = blocks[0] + assert confidence == 0.5 + assert action_taken == "log_only" + + @pytest.mark.asyncio + async def test_apply_guardrail_block_raises_for_response(self): + """When action=block and detection above threshold, apply_guardrail raises HTTPException (response).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = { + "texts": [ + "Example:\n```python\ndef factorial(n):\n return 1 if n <= 1 else n * factorial(n - 1)\n```" + ] + } + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + assert exc_info.value.status_code == 400 + assert "code block" in (exc_info.value.detail or {}).get("error", "") + + @pytest.mark.asyncio + async def test_apply_guardrail_mask_returns_placeholder(self): + """When action=mask, code block is replaced with placeholder.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="mask", + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = { + "texts": ["Before\n```python\nx=1\n```\nAfter"] + } + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + assert result["texts"] is not None + assert len(result["texts"]) == 1 + assert "[CODE_BLOCK_REDACTED]" in result["texts"][0] + assert "x=1" not in result["texts"][0] + + @pytest.mark.asyncio + async def test_factorial_scenario_blocked(self): + """Exact user scenario: Python factorial snippet in markdown is blocked when python in list.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="block", + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + text = '''```python +def factorial(n: int) -> int: + """Return the factorial of n.""" + if n < 0: + raise ValueError("n must be non-negative") + if n in (0, 1): + return 1 + return n * factorial(n - 1) +``` + +Example usage: +```python +print(factorial(5)) # Output: 120 +```''' + inputs = {"texts": [text]} + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + + @pytest.mark.asyncio + async def test_detection_includes_confidence_and_action_taken(self): + """Detection output includes confidence and action_taken for tracing.""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + action="mask", # don't raise so we can inspect request_data + confidence_threshold=0.7, + ) + request_data = {"model": "gpt-4", "metadata": {}} + inputs = {"texts": ["```python\n1+1\n```"]} + await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + ) + meta = request_data.get("metadata") or request_data.get("litellm_metadata") or {} + guardrail_info = meta.get("standard_logging_guardrail_information") or [] + assert len(guardrail_info) >= 1 + info = guardrail_info[-1] + assert info.get("guardrail_status") == "success" + # tracing_detail may be in the logged structure + assert "guardrail_response" in info or "guardrail_response" in str(info) + + def test_default_runs_on_pre_call_and_post_call(self): + """When mode is not set, guardrail runs on both pre_call and post_call (and during_call is supported).""" + guardrail = BlockCodeExecutionGuardrail( + guardrail_name="test", + blocked_languages=["python"], + ) + event_hook = guardrail.event_hook + if isinstance(event_hook, list): + values = [h.value if hasattr(h, "value") else h for h in event_hook] + else: + values = [event_hook.value if hasattr(event_hook, "value") else event_hook] + assert GuardrailEventHooks.pre_call.value in values + assert GuardrailEventHooks.post_call.value in values + + def test_initialize_guardrail_default_mode_is_both(self): + """initialize_guardrail with no mode uses DEFAULT_EVENT_HOOKS (pre_call + post_call).""" + from unittest.mock import MagicMock + + litellm_params = MagicMock() + litellm_params.guardrail = "block_code_execution" + litellm_params.blocked_languages = ["python"] + litellm_params.action = "block" + litellm_params.confidence_threshold = 0.7 + litellm_params.default_on = False + litellm_params.mode = None # not set + guardrail = {"guardrail_name": "block-code-test"} + instance = initialize_guardrail(litellm_params, guardrail) + assert instance.event_hook == DEFAULT_EVENT_HOOKS + assert GuardrailEventHooks.pre_call.value in instance.event_hook + assert GuardrailEventHooks.post_call.value in instance.event_hook