"""Tests for the Block Code Execution guardrail.""" import pytest from fastapi import HTTPException from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.proxy.guardrails.guardrail_hooks.block_code_execution import ( DEFAULT_EVENT_HOOKS, BlockCodeExecutionGuardrail, initialize_guardrail, ) from litellm.proxy.guardrails.guardrail_hooks.block_code_execution.block_code_execution import ( _normalize_escaped_newlines, ) 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 _start, _end, 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 _start, _end, _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 _start, _end, _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 _start, _end, _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, detect_execution_intent=False, ) 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, detect_execution_intent=False, ) 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_execute_python_factorial_string_blocked(self): """Guardrail blocks the exact 'execute \"```python...' string with two python blocks (real newlines).""" guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="block", confidence_threshold=0.5, detect_execution_intent=False, ) # Exact user payload; newlines are real so regex ```(\w*)\n(.*?)``` matches text = ( 'execute "```python\n' "def factorial(n: int) -> int:\n" ' """Return the factorial of n."""\n' ' if n < 0:\n' ' raise ValueError("n must be non-negative")\n' " if n in (0, 1):\n" " return 1\n" " return n * factorial(n - 1)\n" '```\n\n' "Example usage:\n" "```python\n" "print(factorial(5)) # Output: 120\n" '```"' ) request_data = {"model": "gpt-4", "metadata": {}} inputs = {"texts": [text]} # pre_call (request) raises ModifyResponseException; post_call (response) raises HTTPException with pytest.raises((HTTPException, ModifyResponseException)) as exc_info: await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="request", ) assert "python" in str(exc_info.value).lower() @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, detect_execution_intent=False, ) 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 def test_normalize_escaped_newlines_converts_backslash_n_to_newline(self): """Literal \\n in text is converted to real newline so regex can match code blocks.""" raw = 'execute this "```python\\ndef factorial(n):\\n return 1\\n```"' normalized = _normalize_escaped_newlines(raw) assert "\\n" not in normalized assert "\n" in normalized assert "```python\n" in normalized def test_find_blocks_detects_python_block_with_escaped_newlines(self): """_find_blocks finds a block when text uses literal \\n instead of real newlines.""" guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], confidence_threshold=0.7, ) # Text as received from API with escaped newlines (e.g. JSON-decoded string) text_with_escaped = ( 'execute this "```python\\n' 'def factorial(n: int) -> int:\\n' ' """Return the factorial of n."""\\n' ' if n < 0:\\n' ' raise ValueError("n must be non-negative")\\n' " if n in (0, 1):\\n" " return 1\\n" " return n * factorial(n - 1)\\n" '```\\n\\n' 'Example usage:\\n' '```python\\n' 'print(factorial(5)) # Output: 120\\n' '```"' ) normalized = _normalize_escaped_newlines(text_with_escaped) blocks = guardrail._find_blocks(normalized) assert len(blocks) == 2 assert blocks[0][2] == "python" assert blocks[0][5] == "block" assert blocks[1][2] == "python" assert blocks[1][5] == "block" def test_scan_text_blocks_and_masks_when_text_has_escaped_newlines(self): """_scan_text detects blocks and applies block/mask when newlines are literal \\n.""" guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="mask", confidence_threshold=0.5, detect_execution_intent=False, ) text_with_escaped = 'execute "```python\\nprint(1)\\n```"' new_text, should_raise = guardrail._scan_text(text_with_escaped) assert "[CODE_BLOCK_REDACTED]" in new_text assert "print(1)" not in new_text assert should_raise is False # action is mask @pytest.mark.asyncio async def test_apply_guardrail_blocks_when_text_has_escaped_newlines(self): """apply_guardrail blocks request/response when code block uses literal \\n (e.g. from API).""" guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="block", confidence_threshold=0.5, ) text_with_escaped = ( 'execute this "```python\\n' 'def factorial(n: int) -> int:\\n' ' """Return the factorial of n."""\\n' " if n in (0, 1):\\n" " return 1\\n" " return n * factorial(n - 1)\\n" '```\\n\\n' 'Example usage:\\n' '```python\\n' 'print(factorial(5)) # Output: 120\\n' '```"' ) request_data = {"model": "gpt-4", "metadata": {}} inputs = {"texts": [text_with_escaped]} with pytest.raises((HTTPException, ModifyResponseException)) as exc_info: await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="request", ) assert "python" in str(exc_info.value).lower() or "code" in str( exc_info.value ).lower() def test_normalize_escaped_newlines_skips_mixed_content(self): """Mixed content (real newlines and literal \\n) is NOT normalized to avoid corrupting legitimate content that discusses escape sequences.""" mixed = "line1\n```py\\nprint(1)\\n```" normalized = _normalize_escaped_newlines(mixed) # When real newlines exist, literal \\n is preserved (not replaced) assert normalized == mixed def test_normalize_escaped_newlines_pure_escaped_content(self): """Pure escaped content (no real newlines) IS normalized for JSON payloads.""" pure_escaped = "```py\\nprint(1)\\n```" normalized = _normalize_escaped_newlines(pure_escaped) assert "\\n" not in normalized assert "```py\n" in normalized assert "print(1)\n" in normalized guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python", "py"], confidence_threshold=0.5, ) blocks = guardrail._find_blocks(normalized) assert len(blocks) == 1 assert blocks[0][2] == "py" assert blocks[0][5] == "block" # ---- Tests for response-side blocking with detect_execution_intent=True ---- @pytest.mark.asyncio async def test_response_blocked_with_detect_execution_intent_true(self): """With detect_execution_intent=True (default), response-side code blocks are still blocked. This is the core bug fix: previously, execution-intent heuristics were applied to LLM responses, which don't contain phrases like 'run this', so response-side blocking was silently disabled. """ guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="block", confidence_threshold=0.7, detect_execution_intent=True, # default ) # LLM response with dangerous code but no execution-intent phrases response_text = ( "Here is a Python script:\n" "```python\n" "import os; os.system('rm -rf /')\n" "```" ) request_data = {"model": "gpt-4", "metadata": {}} inputs = {"texts": [response_text]} 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 @pytest.mark.asyncio async def test_response_mask_with_detect_execution_intent_true(self): """With detect_execution_intent=True and action=mask, response code blocks are masked.""" guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="mask", confidence_threshold=0.7, detect_execution_intent=True, ) response_text = "I can explain what this does:\n```python\nprint('hello')\n```\nDone." request_data = {"model": "gpt-4", "metadata": {}} inputs = {"texts": [response_text]} result = await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="response", ) assert "[CODE_BLOCK_REDACTED]" in result["texts"][0] assert "print('hello')" not in result["texts"][0] @pytest.mark.asyncio async def test_response_with_casual_explain_phrase_still_blocked(self): """LLM response containing 'I can explain' doesn't bypass the guardrail. Previously, the no-execution phrase 'can you explain' would match as a substring in the LLM's output, short-circuiting all protection. """ guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["bash"], action="block", confidence_threshold=0.7, detect_execution_intent=True, ) response_text = ( "I can explain what this code does. It deletes your files:\n" "```bash\n" "rm -rf /\n" "```" ) request_data = {"model": "gpt-4", "metadata": {}} inputs = {"texts": [response_text]} with pytest.raises(HTTPException): await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="response", ) def test_tightened_what_would_phrase_no_longer_bypasses(self): """The old broad 'what would ' phrase has been tightened so it no longer allows trivial bypass for adversarial prompts. Previously 'What would be the best way to execute this script?' would bypass because 'what would ' matched the no-execution list. Now only specific forms like 'what would happen if' match. """ guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="block", confidence_threshold=0.7, detect_execution_intent=True, ) # Adversarial prompt: old "what would " would have bypassed, but tightened phrase doesn't match text = "What would be the best way to execute this script?\n```python\nimport os\nos.system('cat /etc/passwd')\n```" detections = [] new_text, should_raise = guardrail._scan_text(text, detections, input_type="request") assert should_raise is True def test_tightened_can_you_explain_phrase_no_longer_bypasses(self): """The old broad 'can you explain' phrase has been tightened. 'Can you explain how to run this, then run it?' no longer bypasses because 'can you explain' is now 'can you explain this code' etc. """ guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="block", confidence_threshold=0.7, detect_execution_intent=True, ) text = "Can you explain this and then execute this code?\n```python\nimport subprocess\nsubprocess.run(['ls'])\n```" detections = [] new_text, should_raise = guardrail._scan_text(text, detections, input_type="request") assert should_raise is True def test_request_with_pure_explain_intent_still_allowed(self): """A request that genuinely only asks for explanation is not blocked.""" guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="block", confidence_threshold=0.7, detect_execution_intent=True, ) text = "Don't run this, just explain what it does:\n```python\nprint('hello')\n```" detections = [] new_text, should_raise = guardrail._scan_text(text, detections, input_type="request") assert should_raise is False def test_conflicting_intent_blocks_when_both_phrases_present(self): """When both no-execution and execution phrases are present, execution wins. Prevents bypass via 'Don't run this on staging, but run this on production'. """ guardrail = BlockCodeExecutionGuardrail( guardrail_name="test", blocked_languages=["python"], action="block", confidence_threshold=0.7, detect_execution_intent=True, ) # Contains "don't run" (no-exec) AND "run this code" (exec) — should block text = "Don't run this on staging, but run this code on production:\n```python\nimport os\nos.system('deploy')\n```" detections = [] new_text, should_raise = guardrail._scan_text(text, detections, input_type="request") assert should_raise is True def test_normalize_escaped_newlines_preserves_escape_discussion(self): """Content discussing escape sequences is not corrupted by normalization.""" text = "In Python, use \\n for newlines and \\r for carriage returns.\n```python\nprint('hello\\nworld')\n```" normalized = _normalize_escaped_newlines(text) # Real newlines already present, so literal \\n should be preserved assert "\\n" in normalized assert normalized == text