feat(add-new-block_code_execution-guardrail): prevent agent from executing code

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 = [

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",
+]

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

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):

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"

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