Litellm dev 02 23 2026 p2 (#22118)

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

* feat: working block code execution guardrail

* feat(ui/): add new guardrails to guardrail_garden

* fix: fix greptile feedback

* feat: minor cleanup

* docs: cleanup agent instructions

* feat: cleanup ui

* Update litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* fix: remove unused code

* fixing build

* fix: greptile fixes

* fix: fix linting errors

* fix: address greptile comments

* fixing build and tests

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: yuneng-jiang <yuneng.jiang@gmail.com>
This commit is contained in:
Krish Dholakia 2026-02-25 11:34:48 -08:00 committed by GitHub
parent edd00c025c
commit 21ba9bfe89
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1237 additions and 142 deletions

View file

@ -49,6 +49,17 @@ LiteLLM is a unified interface for 100+ LLMs that:
- Test provider-specific functionality thoroughly
- Consider adding load tests for performance-critical changes
4. **Code Style**: Follow [Google's Python Style Guide](https://google.github.io/styleguide/pyguide.html) for Python; the project also uses Black, Ruff, and MyPy.
### Performance and Database
- **Hot paths:** Do not add duplicate DB queries in proxy/MCP auth or other high-frequency request paths. Reuse already-loaded auth/context (key, team, end-user objects) instead of re-querying.
- **List endpoints:** Avoid N+1 queries when implementing list endpoints (e.g. agents, MCP servers). Use batch loads, `include`/joins, or a single query with needed relations so DB round-trips stay constant.
- **Pre-merge:** For proxy and MCP code, verify that new DB usage in request-handling paths does not introduce duplicate or N+1 queries.
### Database Migrations
- **Do not auto-create or write `migration.sql` files.** If schema changes are needed or a migration is missing, surface the error to the user and instruct them to run `ci_cd/run_migration.py` to generate new schemas/migrations.
### MAKING CODE CHANGES FOR THE UI (IGNORE FOR BACKEND)
1. **Tremor is DEPRECATED, do not use Tremor components in new features/changes**
@ -174,8 +185,11 @@ When opening issues or pull requests, follow these templates:
3. **Rate Limits**: Respect provider rate limits in tests
4. **Memory Usage**: Be mindful of memory usage in streaming scenarios
5. **Dependencies**: Keep dependencies minimal and well-justified
6. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections
7. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks
6. **Duplicate DB queries in hot paths**: Adding redundant DB calls in MCP auth, proxy auth, or other high-frequency paths causes performance degradation at scale.
7. **N+1 queries in list endpoints**: Looping over a list and performing a separate DB query per item (e.g. loading agents or related entities) — use batch/joined queries instead.
8. **Auto-creating migration.sql**: Do not generate or edit migration SQL files. If migrations are needed, tell the user to run `ci_cd/run_migration.py` to generate new schemas.
9. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections
10. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks
## HELPFUL RESOURCES
@ -189,4 +203,5 @@ When opening issues or pull requests, follow these templates:
- Check similar provider implementations
- Ensure comprehensive test coverage
- Update documentation appropriately
- Consider backward compatibility impact
- Consider backward compatibility impact
- For proxy/MCP code, confirm no new duplicate or N+1 DB queries in request paths

View file

@ -91,6 +91,13 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
- Async/await patterns throughout
- Type hints required for all public APIs
- **Avoid imports within methods** — place all imports at the top of the file (module-level). Inline imports inside functions/methods make dependencies harder to trace and hurt readability. The only exception is avoiding circular imports where absolutely necessary.
- Follow [Google's Python Style Guide](https://google.github.io/styleguide/pyguide.html) for Python style and structure; the repo uses Black/Ruff/MyPy on top of that.
### Performance and Database
- **Hot paths:** Avoid adding duplicate or redundant DB queries in hot request paths (e.g. proxy auth, MCP request path). Each extra query per request multiplies load at scale. Prefer reusing already-fetched auth/context (e.g. key/team/end-user objects) instead of re-querying.
- **List endpoints and N+1:** When implementing list endpoints (e.g. agents, MCP servers), avoid N+1 patterns: do not loop over a list and perform a separate DB query per item. Use batch loads, `include`/joins, or a single query with the needed relations so the number of DB round-trips is constant.
- **Review before merging:** For proxy and MCP code, check that new DB access in request-handling paths does not introduce duplicate or N+1 queries.
### Testing Strategy
- Unit tests in `tests/test_litellm/`
@ -103,8 +110,7 @@ LiteLLM is a unified interface for 100+ LLM providers with two main components:
- When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select
### Database Migrations
- Prisma handles schema migrations
- Migration files auto-generated with `prisma migrate dev`
- Prisma handles schema migrations. **Do not auto-create or write `migration.sql` files.** If schema changes are needed or a migration is missing, surface the error to the user and instruct them to run `ci_cd/run_migration.py` to generate new schemas/migrations.
- Always test migrations against both PostgreSQL and SQLite
### Enterprise Features

View file

@ -14,35 +14,28 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import LitellmUserRoles, 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.guardrail_hooks.custom_code.code_validator import (
CustomCodeValidationError,
validate_custom_code,
)
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import (
get_custom_code_primitives,
)
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,
)
CustomCodeValidationError, validate_custom_code)
from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import \
get_custom_code_primitives
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)
#### GUARDRAILS ENDPOINTS ####
@ -161,7 +154,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:
@ -303,7 +297,8 @@ async def create_guardrail(
}
```
"""
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 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
@ -401,7 +396,8 @@ async def update_guardrail(
}
```
"""
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 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
@ -477,7 +473,8 @@ async def delete_guardrail(
}
```
"""
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 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
@ -579,7 +576,8 @@ async def patch_guardrail(
}
```
"""
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 user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
@ -707,7 +705,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
@ -782,10 +781,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 = []
@ -1205,11 +1202,22 @@ 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 ...:
field_dict["default_value"] = field.default
# Copy min, max, step from json_schema_extra for number/percentage inputs
if field_json_schema_extra:
for key in ("min", "max", "step", "default_value"):
if key in field_json_schema_extra:
field_dict[key] = field_json_schema_extra[key]
return field_dict
@ -1357,7 +1365,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()
@ -1485,6 +1494,7 @@ async def test_custom_code_guardrail(
```
"""
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,

View file

@ -0,0 +1,92 @@
"""Block Code Execution guardrail: blocks or masks fenced code blocks by language."""
from typing import TYPE_CHECKING, Any, List, Literal, Optional, Union, 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 _get_param(
litellm_params: "LitellmParams",
guardrail: "Guardrail",
key: str,
default: Any = None,
) -> Any:
"""Get a param from litellm_params, with fallback to raw guardrail litellm_params (for extra fields not on LitellmParams)."""
value = getattr(litellm_params, key, default)
if value is not None:
return value
raw = guardrail.get("litellm_params")
if isinstance(raw, dict) and key in raw:
return raw[key]
return default
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: Optional[List[str]] = cast(
Optional[List[str]],
_get_param(litellm_params, guardrail, "blocked_languages"),
)
action = cast(
Literal["block", "mask"],
_get_param(litellm_params, guardrail, "action", "block"),
)
confidence_threshold = float(
cast(
Union[int, float, str],
_get_param(litellm_params, guardrail, "confidence_threshold", 0.5),
)
)
mode = _get_param(litellm_params, guardrail, "mode")
event_hook = cast(
Optional[Union[Literal["pre_call", "post_call", "during_call"], List[str]]],
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=bool(_get_param(litellm_params, guardrail, "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",
]

View file

@ -0,0 +1,363 @@
"""
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,
ModifyResponseException)
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
# 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_escaped_newlines(text: str) -> str:
"""
Replace literal escaped newlines (backslash + n or backslash + r) with real newlines.
API/JSON payloads sometimes deliver newlines as the two-character sequence \\n.
Applied whenever \\n or \\r appear, including in mixed content with real newlines.
"""
if not text:
return text
if "\\n" not in text and "\\r" not in text:
return text
# Order matters: replace \r\n first so we don't produce extra \n from \r then \n
text = text.replace("\\r\\n", "\n")
text = text.replace("\\n", "\n")
text = text.replace("\\r", "\n")
return text
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
# When block_all is False, caller guarantees blocked_languages is non-empty.
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.5,
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[int, int, str, str, float, CodeBlockActionTaken]]:
"""
Find all fenced code blocks in text. Returns list of
(start, end, language_tag, block_content, confidence, action_taken).
"""
results: List[
Tuple[int, int, 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(
(m.start(), m.end(), 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
text = _normalize_escaped_newlines(text)
blocks = self._find_blocks(text)
if not blocks:
return text, False
should_raise = False
last_end = 0
parts: List[str] = []
for start, end, tag, _body, confidence, action_taken in blocks:
if detections is not None:
detections.append(
cast(
CodeBlockDetection,
{
"type": "code_block",
"language": tag,
"confidence": round(confidence, 2),
"action_taken": action_taken,
},
)
)
if action_taken == "block" and self.action == "block":
should_raise = True
parts.append(text[last_end:start])
if action_taken == "block":
parts.append(self.MASK_PLACEHOLDER)
else:
parts.append(text[start:end])
last_end = 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},
)
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 ModifyResponseException:
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
event_type = (
GuardrailEventHooks.pre_call
if input_type == "request"
else GuardrailEventHooks.post_call
)
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(),
event_type=event_type,
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 as soon as a complete fenced code block is detected (before yielding that chunk)."""
accumulated = ""
async for item in response:
if isinstance(item, ModelResponseStream) and item.choices:
delta_content = ""
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
accumulated += delta_content
# Check after every chunk so we block before yielding the chunk that completes a blocked block
normalized = _normalize_escaped_newlines(accumulated)
blocks = self._find_blocks(normalized)
for _start, _end, _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

View file

@ -5,24 +5,20 @@ from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, ConfigDict, Field, field_validator
from typing_extensions import Required, TypedDict
from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
EnkryptAIGuardrailConfigs,
)
from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import (
GraySwanGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.ibm import (
IBMGuardrailsBaseConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import (
ContentFilterCategoryConfig,
)
from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import (
QualifireGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
ToolPermissionGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import \
BlockCodeExecutionGuardrailConfigModel
from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import \
EnkryptAIGuardrailConfigs
from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import \
GraySwanGuardrailConfigModel
from litellm.types.proxy.guardrails.guardrail_hooks.ibm import \
IBMGuardrailsBaseConfigModel
from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import \
ContentFilterCategoryConfig
from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import \
QualifireGuardrailConfigModel
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import \
ToolPermissionGuardrailConfigModel
"""
Pydantic object defining how to set guardrails on litellm proxy
@ -73,6 +69,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):
@ -707,6 +704,7 @@ class LitellmParams(
EnkryptAIGuardrailConfigs,
IBMGuardrailsBaseConfigModel,
QualifireGuardrailConfigModel,
BlockCodeExecutionGuardrailConfigModel,
):
guardrail: str = Field(description="The type of guardrail integration to use")
mode: Union[str, List[str], Mode] = Field(

View file

@ -0,0 +1,75 @@
"""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.5,
ge=0.0,
le=1.0,
description="Only block or mask when detection confidence >= this value; below threshold, allow or log_only.",
json_schema_extra=cast(
Any,
{
"type": "percentage",
"min": 0.0,
"max": 1.0,
"step": 0.1,
"default_value": 0.5,
},
),
)
@staticmethod
def ui_friendly_name() -> str:
return "Block Code Execution"

View file

@ -0,0 +1,429 @@
"""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,
)
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_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,
)
# 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,
)
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,
)
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()
@pytest.mark.asyncio
async def test_streaming_hook_blocks_before_yielding_chunk_that_completes_block(
self,
):
"""Streaming hook runs block check after every chunk and raises before yielding the chunk that completes a blocked fenced block."""
from litellm.types.utils import (Delta, ModelResponseStream,
StreamingChoices)
guardrail = BlockCodeExecutionGuardrail(
guardrail_name="test",
blocked_languages=["python"],
action="block",
confidence_threshold=0.5,
)
# Chunks that form "```python\nprint(1)\n```" when concatenated
async def mock_stream():
yield ModelResponseStream(
choices=[StreamingChoices(delta=Delta(content="```python\n"))],
)
yield ModelResponseStream(
choices=[StreamingChoices(delta=Delta(content="print(1)\n"))],
)
yield ModelResponseStream(
choices=[StreamingChoices(delta=Delta(content="```"))],
)
request_data = {"model": "gpt-4", "metadata": {}}
yielded_chunks = []
with pytest.raises(HTTPException) as exc_info:
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=None,
response=mock_stream(),
request_data=request_data,
):
yielded_chunks.append(chunk)
assert exc_info.value.status_code == 400
assert "code block" in (exc_info.value.detail or {}).get("error", "")
# The chunk that completes the block (third chunk) must not have been yielded
assert len(yielded_chunks) == 2
@pytest.mark.asyncio
async def test_streaming_hook_blocks_when_accumulated_has_literal_backslash_n(
self,
):
"""Streaming hook normalizes escaped newlines before detection; blocks when literal \\n forms a complete code block."""
from litellm.types.utils import (Delta, ModelResponseStream,
StreamingChoices)
guardrail = BlockCodeExecutionGuardrail(
guardrail_name="test",
blocked_languages=["python"],
action="block",
confidence_threshold=0.5,
)
# Chunks that when concatenated form "```python\\nprint(1)\\n```" (literal backslash-n)
async def mock_stream():
yield ModelResponseStream(
choices=[StreamingChoices(delta=Delta(content="```python\\n"))],
)
yield ModelResponseStream(
choices=[StreamingChoices(delta=Delta(content="print(1)\\n"))],
)
yield ModelResponseStream(
choices=[StreamingChoices(delta=Delta(content="```"))],
)
request_data = {"model": "gpt-4", "metadata": {}}
yielded_chunks = []
with pytest.raises(HTTPException) as exc_info:
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=None,
response=mock_stream(),
request_data=request_data,
):
yielded_chunks.append(chunk)
assert exc_info.value.status_code == 400
assert "code block" in (exc_info.value.detail or {}).get("error", "")
# Block is detected after normalization; chunk that completes the block not yielded
assert len(yielded_chunks) == 2
def test_normalize_escaped_newlines_mixed_content_detects_block(self):
"""Mixed content (real newlines and literal \\n) is normalized so code block is detected."""
# Text with real newline then a fence using literal \n after language tag
mixed = "line1\n```py\\nprint(1)\\n```"
normalized = _normalize_escaped_newlines(mixed)
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"

View file

@ -2,14 +2,7 @@
import React, { useCallback, useDeferredValue, useEffect, useState } from "react";
import { Select, Switch, Tooltip } from "antd";
import {
Table,
TableHead,
TableHeaderCell,
TableBody,
TableRow,
TableCell,
} from "@tremor/react";
import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react";
import { TimeCell } from "./view_logs/time_cell";
import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown";
import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown";
@ -17,14 +10,13 @@ import FilterComponent, { FilterOption } from "./molecules/filter";
import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking";
const POLICY_OPTIONS = [
{ value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" },
{ value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" },
{ value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" },
{ value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" },
] as const;
type PolicyValue = "trusted" | "blocked";
const policyStyle = (p: string) =>
POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1];
const policyStyle = (p: string) => POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1];
type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count";
@ -56,18 +48,6 @@ const PolicySelect: React.FC<{
minWidth: 110,
fontWeight: 500,
}}
styles={{
selector: {
backgroundColor: style.bg,
borderColor: style.border,
color: style.color,
borderRadius: 999,
fontSize: 11,
fontWeight: 600,
paddingLeft: 8,
paddingRight: 4,
},
}}
popupMatchSelectWidth={false}
options={POLICY_OPTIONS.map((o) => ({
value: o.value,
@ -133,7 +113,9 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
}
}, [accessToken]);
useEffect(() => { load(); }, [load]);
useEffect(() => {
load();
}, [load]);
useEffect(() => {
if (!isLiveTail) return;
@ -146,9 +128,7 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
setSaving(toolName);
try {
await updateToolPolicy(accessToken, toolName, newPolicy);
setTools((prev) =>
prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t))
);
setTools((prev) => prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t)));
} catch (e: any) {
alert(`Failed to update policy: ${e.message}`);
} finally {
@ -178,12 +158,14 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
};
// Build unique team/key options from loaded data
const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map(
(v) => ({ label: v as string, value: v as string })
);
const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map(
(v) => ({ label: v as string, value: v as string })
);
const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map((v) => ({
label: v as string,
value: v as string,
}));
const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map((v) => ({
label: v as string,
value: v as string,
}));
const filterOptions: FilterOption[] = [
{
@ -245,7 +227,6 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
<div className="p-6 w-full">
<h1 className="text-2xl font-semibold text-gray-900 mb-6">Tool Policies</h1>
<div className="bg-white rounded-lg shadow w-full max-w-full box-border">
{/* Toolbar */}
<div className="border-b px-6 py-4 w-full max-w-full box-border">
<div className="flex flex-col md:flex-row items-start md:items-center justify-between space-y-4 md:space-y-0 w-full max-w-full box-border">
@ -256,16 +237,29 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
placeholder="Search by Tool Name"
className="w-full px-3 py-2 pl-8 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
value={searchTerm}
onChange={(e) => { setSearchTerm(e.target.value); setCurrentPage(1); }}
onChange={(e) => {
setSearchTerm(e.target.value);
setCurrentPage(1);
}}
/>
<svg className="absolute left-2.5 top-2.5 h-4 w-4 text-gray-500" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z" />
<svg
className="absolute left-2.5 top-2.5 h-4 w-4 text-gray-500"
fill="none"
stroke="currentColor"
viewBox="0 0 24 24"
>
<path
strokeLinecap="round"
strokeLinejoin="round"
strokeWidth={2}
d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z"
/>
</svg>
</div>
<div className="flex items-center gap-2">
<span className="text-sm font-medium text-gray-900">Live Tail</span>
<Switch color="green" checked={isLiveTail} onChange={setIsLiveTail} />
<Switch checked={isLiveTail} onChange={setIsLiveTail} />
</div>
<button
@ -273,8 +267,18 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
disabled={isButtonLoading}
className="flex items-center gap-1.5 px-3 py-2 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-60"
>
<svg className={`w-4 h-4 ${isButtonLoading ? "animate-spin" : ""}`} fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15" />
<svg
className={`w-4 h-4 ${isButtonLoading ? "animate-spin" : ""}`}
fill="none"
stroke="currentColor"
viewBox="0 0 24 24"
>
<path
strokeLinecap="round"
strokeLinejoin="round"
strokeWidth={2}
d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15"
/>
</svg>
{isButtonLoading ? "Fetching" : "Fetch"}
</button>
@ -282,14 +286,27 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
<div className="flex items-center gap-4 text-sm text-gray-600 whitespace-nowrap">
<span>
Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results
Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} -{" "}
{Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results
</span>
<span>
Page {currentPage} of {totalPages}
</span>
<span>Page {currentPage} of {totalPages}</span>
<div className="flex gap-1">
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Previous</button>
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Next</button>
<button
onClick={() => setCurrentPage((p) => Math.max(1, p - 1))}
disabled={currentPage === 1}
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40"
>
Previous
</button>
<button
onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))}
disabled={currentPage === totalPages}
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40"
>
Next
</button>
</div>
</div>
</div>
@ -309,7 +326,9 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
{isLiveTail && (
<div className="bg-green-50 border-b border-green-100 px-6 py-2 flex items-center justify-between">
<span className="text-sm text-green-700">Auto-refreshing every 15 seconds</span>
<button onClick={() => setIsLiveTail(false)} className="text-xs text-green-600 underline">Stop</button>
<button onClick={() => setIsLiveTail(false)} className="text-xs text-green-600 underline">
Stop
</button>
</div>
)}
@ -321,20 +340,34 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
<Table className="[&_td]:py-0.5 [&_th]:py-1 w-full">
<TableHead>
<TableRow>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Discovered" field="created_at" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Tool Name" field="tool_name" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Policy" field="call_policy" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="# Calls" field="call_count" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Team Name" field="team_id" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8">
<SortHeader label="Discovered" field="created_at" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">
<SortHeader label="Tool Name" field="tool_name" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">
<SortHeader label="Policy" field="call_policy" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">
<SortHeader label="# Calls" field="call_count" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">
<SortHeader label="Team Name" field="team_id" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">Key Hash</TableHeaderCell>
<TableHeaderCell className="py-1 h-8"><SortHeader label="Key Name" field="key_alias" /></TableHeaderCell>
<TableHeaderCell className="py-1 h-8">
<SortHeader label="Key Name" field="key_alias" />
</TableHeaderCell>
<TableHeaderCell className="py-1 h-8">Origin</TableHeaderCell>
</TableRow>
</TableHead>
<TableBody>
{loading ? (
<TableRow>
<TableCell colSpan={8} className="h-8 text-center text-gray-500">Loading tools</TableCell>
<TableCell colSpan={8} className="h-8 text-center text-gray-500">
Loading tools
</TableCell>
</TableRow>
) : paginated.length === 0 ? (
<TableRow>
@ -397,12 +430,25 @@ export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
{/* Bottom pagination (only when > 1 page) */}
{totalPages > 1 && (
<div className="border-t px-6 py-3 flex items-center justify-between text-sm text-gray-600">
<span>Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length}</span>
<span>
Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of{" "}
{sorted.length}
</span>
<div className="flex gap-1">
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Previous</button>
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Next</button>
<button
onClick={() => setCurrentPage((p) => Math.max(1, p - 1))}
disabled={currentPage === 1}
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40"
>
Previous
</button>
<button
onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))}
disabled={currentPage === totalPages}
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40"
>
Next
</button>
</div>
</div>
)}

View file

@ -713,7 +713,7 @@ describe("UsagePage", () => {
// Admin should see the user selector select element with the placeholder attribute
const userSelects = screen.getAllByRole("combobox");
const userSelect = userSelects.find(
(el) => el.getAttribute("placeholder") === "All Users (Global View)",
(el) => el.getAttribute("placeholder") === "Select user to filter...",
);
expect(userSelect).toBeDefined();
});
@ -828,7 +828,7 @@ describe("UsagePage", () => {
// Non-admin should not see the user selector
const userSelects = screen.getAllByRole("combobox");
const userSelect = userSelects.find(
(el) => el.getAttribute("placeholder") === "All Users (Global View)",
(el) => el.getAttribute("placeholder") === "Select user to filter...",
);
expect(userSelect).toBeUndefined();
});

View file

@ -166,12 +166,16 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
// Set provider
setSelectedProvider(preset.provider);
form.setFieldsValue({
const baseValues: Record<string, any> = {
provider: preset.provider,
guardrail_name: preset.guardrailNameSuggestion,
mode: preset.mode,
default_on: preset.defaultOn,
});
};
if (preset.provider === "BlockCodeExecution") {
baseValues.confidence_threshold = 0.5;
}
form.setFieldsValue(baseValues);
// Pre-select content category if specified
if (preset.categoryName && guardrailSettings.content_filter_settings?.content_categories) {
@ -195,11 +199,15 @@ const AddGuardrailForm: React.FC<AddGuardrailFormProps> = ({ visible, onClose, a
const handleProviderChange = (value: string) => {
setSelectedProvider(value);
// Reset form fields that are provider-specific
form.setFieldsValue({
const resetValues: Record<string, any> = {
config: undefined,
presidio_analyzer_api_base: undefined,
presidio_anonymizer_api_base: undefined,
});
};
if (value === "BlockCodeExecution") {
resetValues.confidence_threshold = 0.5;
}
form.setFieldsValue(resetValues);
// Reset PII selections when changing provider
setSelectedEntities([]);

View file

@ -148,6 +148,18 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
mode: "pre_call",
defaultOn: false,
},
block_code_execution: {
provider: "BlockCodeExecution",
guardrailNameSuggestion: "Block Code Execution",
mode: "pre_call",
defaultOn: false,
},
cf_competitor_intent: {
provider: "LitellmContentFilter",
guardrailNameSuggestion: "Competitor Name Blocking",
mode: "pre_call",
defaultOn: false,
},
// ── Partner Guardrails ──
presidio: {

View file

@ -213,6 +213,24 @@ export const LITELLM_CONTENT_FILTER_CARDS: GuardrailCardInfo[] = [
logo: `${ASSET_PREFIX}litellm_logo.jpg`,
tags: ["Keywords", "Blocklist"],
},
{
id: "block_code_execution",
name: "Block Code Execution",
description: "Detects markdown fenced code blocks in requests and responses. Block or mask executable code (e.g. Python, JavaScript, Bash) by language with configurable confidence.",
category: "litellm",
subcategory: "Code Safety",
logo: `${ASSET_PREFIX}litellm_logo.jpg`,
tags: ["Code", "Safety", "Prompt Injection"],
},
{
id: "cf_competitor_intent",
name: "Competitor Name Blocking",
description: "Block or reframe competitor comparison and ranking intent. Detect when users ask to compare or recommend competitors (airline or generic competitor lists).",
category: "litellm",
subcategory: "Content Category",
logo: `${ASSET_PREFIX}litellm_logo.jpg`,
tags: ["Content Category", "Competitor", "Topic Blocker"],
},
];
export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [

View file

@ -47,6 +47,7 @@ export const guardrail_provider_map: Record<string, string> = {
Lakera: "lakera_v2",
LitellmContentFilter: "litellm_content_filter",
ToolPermission: "tool_permission",
BlockCodeExecution: "block_code_execution",
};
// Function to populate provider map from API response - updates the original map

View file

@ -1,5 +1,5 @@
import React, { useState, useEffect } from "react";
import { Form, Select, Spin, Input } from "antd";
import { Form, Select, Spin, Input, Slider } from "antd";
import {
guardrail_provider_map,
populateGuardrailProviders,
@ -20,12 +20,15 @@ interface ProviderParam {
param: string;
description: string;
required: boolean;
default_value?: string;
default_value?: string | number;
options?: string[];
type?: string;
fields?: { [key: string]: ProviderParam };
dict_key_options?: string[];
dict_value_type?: string;
min?: number;
max?: number;
step?: number;
}
interface ProviderParamsResponse {
@ -154,6 +157,11 @@ const GuardrailProviderFields: React.FC<GuardrailProviderFieldsProps> = ({
);
}
const percentageInitialValue =
field.type === "percentage" && (fieldValue === undefined || fieldValue === null)
? (field.default_value ?? 0.5)
: undefined;
return (
<Form.Item
key={fullFieldKey}
@ -161,6 +169,7 @@ const GuardrailProviderFields: React.FC<GuardrailProviderFieldsProps> = ({
label={fieldKey}
tooltip={field.description}
rules={field.required ? [{ required: true, message: `${fieldKey} is required` }] : undefined}
initialValue={percentageInitialValue}
>
{field.type === "select" && field.options ? (
<Select placeholder={field.description} defaultValue={fieldValue || field.default_value}>
@ -186,6 +195,17 @@ const GuardrailProviderFields: React.FC<GuardrailProviderFieldsProps> = ({
<Select.Option value="true">True</Select.Option>
<Select.Option value="false">False</Select.Option>
</Select>
) : field.type === "percentage" && field.min != null && field.max != null ? (
<Slider
min={field.min}
max={field.max}
step={field.step ?? 0.1}
marks={{
[field.min]: "0%",
[(field.min + field.max) / 2]: "50%",
[field.max]: "100%",
}}
/>
) : field.type === "number" ? (
<NumericalInput
step={1}

View file

@ -27,6 +27,7 @@ vi.mock("../networking", () => ({
soft_budget: null,
}),
fetchMCPAccessGroups: vi.fn().mockResolvedValue([]),
getAgentsList: vi.fn().mockResolvedValue([]),
}));
vi.mock("../molecules/notifications_manager", () => ({

View file

@ -8,10 +8,10 @@ import { Accordion, AccordionBody, AccordionHeader, Button, Col, Grid, Text, Tex
import { Button as Button2, Form, Input, Modal, Radio, Select, Switch, Tag, Tooltip } from "antd";
import debounce from "lodash/debounce";
import React, { useCallback, useEffect, useState } from "react";
import { CopyToClipboard } from "react-copy-to-clipboard";
import { rolesWithWriteAccess } from "../../utils/roles";
import AgentSelector from "../agent_management/AgentSelector";
import { mapDisplayToInternalNames } from "../callback_info_helpers";
import AccessGroupSelector from "../common_components/AccessGroupSelector";
import BudgetDurationDropdown from "../common_components/budget_duration_dropdown";
import SchemaFormFields from "../common_components/check_openapi_schema";
import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings";
@ -20,7 +20,6 @@ import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSel
import PremiumLoggingSettings from "../common_components/PremiumLoggingSettings";
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion";
import AccessGroupSelector from "../common_components/AccessGroupSelector";
import TeamDropdown from "../common_components/team_dropdown";
import { CreateUserButton } from "../CreateUserButton";
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
@ -40,10 +39,10 @@ import {
proxyBaseUrl,
userFilterUICall,
} from "../networking";
import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
import NumericalInput from "../shared/numerical_input";
import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
import { simplifyKeyGenerateError } from "./utils";
import CreatedKeyDisplay from "../shared/CreatedKeyDisplay";
const { Option } = Select;
@ -299,7 +298,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
formValues.user_id = userID;
} else if (keyOwner === "agent") {
if (!selectedAgentId) {
message.error("Please select an agent");
NotificationsManager.error("Please select an agent");
return;
}
formValues.agent_id = selectedAgentId;
@ -559,7 +558,9 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
<Radio value="you">You</Radio>
<Radio value="service_account">Service Account</Radio>
{userRole === "Admin" && <Radio value="another_user">Another User</Radio>}
<Radio value="agent">Agent <Tag color="purple">New</Tag></Radio>
<Radio value="agent">
Agent <Tag color="purple">New</Tag>
</Radio>
</Radio.Group>
</Form.Item>
@ -1005,9 +1006,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
style={{ width: "100%" }}
disabled={!premiumUser}
placeholder={
!premiumUser
? "Premium feature - Upgrade to set policies by key"
: "Select or enter policies"
!premiumUser ? "Premium feature - Upgrade to set policies by key" : "Select or enter policies"
}
options={policiesList.map((name) => ({ value: name, label: name }))}
/>
@ -1059,9 +1058,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
className="mt-4"
help="Select access groups to assign to this key"
>
<AccessGroupSelector
placeholder="Select access groups (optional)"
/>
<AccessGroupSelector placeholder="Select access groups (optional)" />
</Form.Item>
<Form.Item
label={
@ -1297,7 +1294,11 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey }) => {
accessToken={accessToken || ""}
value={routerSettings || undefined}
onChange={setRouterSettings}
modelData={userModels.length > 0 ? { data: userModels.map((model) => ({ model_name: model })) } : undefined}
modelData={
userModels.length > 0
? { data: userModels.map((model) => ({ model_name: model })) }
: undefined
}
/>
</div>
</AccordionBody>