Custom Code Guardrails UI Playground - #20377
Conversation
first step in allowing teams to submit custom code for guardrails
support passing custom code for guardrails
allows users to write guardrails based on custom code
allows ui testing playground to sanity check if guardrail is working as expected
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile OverviewGreptile SummaryThis PR adds a "Custom Code Guardrails" feature that allows users to write Python-like code to implement custom guardrail logic. The code runs in a sandboxed environment using Key Changes
Critical Issues Found
Recommendations
Confidence Score: 2/5
|
| Filename | Overview |
|---|---|
| litellm/proxy/guardrails/guardrail_endpoints.py | New endpoint for testing custom code guardrails with security validations, includes regex-based forbidden pattern checking and timeout protection |
| litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py | Core custom code guardrail implementation using exec() with restricted global namespace, sandboxing through primitives-only environment |
| ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodePlayground.tsx | React UI for testing custom code guardrails with test runner and primitives reference |
| ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx | Modal UI for creating custom code guardrails with code templates and inline testing |
Sequence Diagram
sequenceDiagram
participant User
participant UI as LiteLLM Dashboard
participant API as Guardrail Endpoints
participant Guardrail as CustomCodeGuardrail
participant Sandbox as exec() Sandbox
participant Primitives as Safe Primitives
User->>UI: Create Custom Code Guardrail
UI->>UI: Edit code in CustomCodeModal
User->>UI: Click "Test"
UI->>API: POST /guardrails/test_custom_code
API->>API: Validate forbidden patterns (regex)
API->>API: Set __builtins__ = {}
API->>Sandbox: exec(custom_code, restricted_globals)
Sandbox->>Sandbox: Compile apply_guardrail function
API->>Sandbox: call apply_guardrail(inputs, request_data, input_type)
Sandbox->>Primitives: Use regex_match(), block(), etc.
Primitives-->>Sandbox: Return results
Sandbox-->>API: Return action (allow/block/modify)
API-->>UI: TestCustomCodeGuardrailResponse
UI-->>User: Display test result
User->>UI: Click "Save Guardrail"
UI->>API: POST /guardrails
API->>API: Store guardrail in database
API->>Guardrail: Initialize CustomCodeGuardrail
Guardrail->>Sandbox: Compile custom code on init
API-->>UI: Success response
Note over User,Primitives: Production Request Flow
User->>API: POST /chat/completions (with guardrail)
API->>Guardrail: apply_guardrail (pre_call)
Guardrail->>Sandbox: Execute apply_guardrail function
Sandbox->>Primitives: Call primitives
Primitives-->>Sandbox: Results
Sandbox-->>Guardrail: allow/block/modify
alt Block
Guardrail-->>API: Raise HTTPException
API-->>User: 400 Error with reason
else Allow/Modify
Guardrail-->>API: Modified inputs
API->>API: Continue to LLM
end
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| exec_globals = get_custom_code_primitives().copy() | ||
|
|
||
| # Remove access to builtins to prevent escape | ||
| exec_globals["__builtins__"] = {} |
There was a problem hiding this comment.
setting __builtins__ to an empty dict does not provide complete sandbox protection - users can still access builtins through introspection like ().__class__.__bases__[0].__subclasses__() or other dunder methods. the regex-based forbidden pattern check on lines 1357-1406 is good first line of defense, but Python sandboxing via exec() is notoriously difficult to secure
| exec_globals["__builtins__"] = {} | |
| # Remove most builtins but keep safe ones needed by primitives | |
| exec_globals["__builtins__"] = { | |
| "__build_class__": __builtins__["__build_class__"], | |
| "__name__": __builtins__["__name__"], | |
| } |
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/proxy/guardrails/guardrail_endpoints.py
Line: 1412:1412
Comment:
setting `__builtins__` to an empty dict does not provide complete sandbox protection - users can still access builtins through introspection like `().__class__.__bases__[0].__subclasses__()` or other dunder methods. the regex-based forbidden pattern check on lines 1357-1406 is good first line of defense, but Python sandboxing via `exec()` is notoriously difficult to secure
```suggestion
# Remove most builtins but keep safe ones needed by primitives
exec_globals["__builtins__"] = {
"__build_class__": __builtins__["__build_class__"],
"__name__": __builtins__["__name__"],
}
```
How can I resolve this? If you propose a fix, please make it concise.| exec_globals = get_custom_code_primitives().copy() | ||
|
|
||
| # Execute the user code in the restricted environment | ||
| exec(compile(self.custom_code, "<guardrail>", "exec"), exec_globals) |
There was a problem hiding this comment.
using exec() with user-provided code is inherently risky. while you've restricted the globals namespace, Python sandboxing is very difficult - attackers can escape through object introspection (e.g., accessing __subclasses__, __bases__, etc.). the test endpoint adds regex-based pattern blocking, but this implementation in production doesn't have those checks when loading from config.yaml
consider using a more secure sandboxing approach like:
- running code in a separate process with resource limits
- using a proper sandboxing library like RestrictedPython
- or using a language designed for safe user code execution (like Lua, Starlark, or CEL)
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
Line: 136:136
Comment:
using `exec()` with user-provided code is inherently risky. while you've restricted the globals namespace, Python sandboxing is very difficult - attackers can escape through object introspection (e.g., accessing `__subclasses__`, `__bases__`, etc.). the test endpoint adds regex-based pattern blocking, but this implementation in production doesn't have those checks when loading from config.yaml
consider using a more secure sandboxing approach like:
- running code in a separate process with resource limits
- using a proper sandboxing library like RestrictedPython
- or using a language designed for safe user code execution (like Lua, Starlark, or CEL)
How can I resolve this? If you propose a fix, please make it concise.| (r"\bsubprocess\.", "subprocess module access is not allowed"), | ||
| ] | ||
|
|
||
| EXECUTION_TIMEOUT_SECONDS = 5 |
There was a problem hiding this comment.
the timeout is applied via ThreadPoolExecutor, but malicious code could still consume excessive CPU/memory within the 5-second window. threads also share memory space, so a memory exhaustion attack could impact the entire proxy process
consider adding resource limits (CPU, memory) using OS-level controls or running in a separate process
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/proxy/guardrails/guardrail_endpoints.py
Line: 1395:1395
Comment:
the timeout is applied via `ThreadPoolExecutor`, but malicious code could still consume excessive CPU/memory within the 5-second window. threads also share memory space, so a memory exhaustion attack could impact the entire proxy process
consider adding resource limits (CPU, memory) using OS-level controls or running in a separate process
How can I resolve this? If you propose a fix, please make it concise.| # Security validation patterns | ||
| FORBIDDEN_PATTERNS = [ | ||
| # Import statements | ||
| (r"\bimport\s+", "import statements are not allowed"), | ||
| (r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"), | ||
| (r"__import__\s*\(", "__import__() is not allowed"), | ||
| # Dangerous builtins | ||
| (r"\bexec\s*\(", "exec() is not allowed"), | ||
| (r"\beval\s*\(", "eval() is not allowed"), | ||
| (r"\bcompile\s*\(", "compile() is not allowed"), | ||
| (r"\bopen\s*\(", "open() is not allowed"), | ||
| (r"\bgetattr\s*\(", "getattr() is not allowed"), | ||
| (r"\bsetattr\s*\(", "setattr() is not allowed"), | ||
| (r"\bdelattr\s*\(", "delattr() is not allowed"), | ||
| (r"\bglobals\s*\(", "globals() is not allowed"), | ||
| (r"\blocals\s*\(", "locals() is not allowed"), | ||
| (r"\bvars\s*\(", "vars() is not allowed"), | ||
| (r"\bdir\s*\(", "dir() is not allowed"), | ||
| (r"\bbreakpoint\s*\(", "breakpoint() is not allowed"), | ||
| (r"\binput\s*\(", "input() is not allowed"), | ||
| # Dangerous dunder access | ||
| (r"__builtins__", "__builtins__ access is not allowed"), | ||
| (r"__globals__", "__globals__ access is not allowed"), | ||
| (r"__code__", "__code__ access is not allowed"), | ||
| (r"__subclasses__", "__subclasses__ access is not allowed"), | ||
| (r"__bases__", "__bases__ access is not allowed"), | ||
| (r"__mro__", "__mro__ access is not allowed"), | ||
| (r"__class__", "__class__ access is not allowed"), | ||
| (r"__dict__", "__dict__ access is not allowed"), | ||
| (r"__getattribute__", "__getattribute__ access is not allowed"), | ||
| (r"__reduce__", "__reduce__ access is not allowed"), | ||
| (r"__reduce_ex__", "__reduce_ex__ access is not allowed"), | ||
| # OS/system access | ||
| (r"\bos\.", "os module access is not allowed"), | ||
| (r"\bsys\.", "sys module access is not allowed"), | ||
| (r"\bsubprocess\.", "subprocess module access is not allowed"), | ||
| ] |
There was a problem hiding this comment.
security validation via regex patterns is good, but these patterns are only applied in the test endpoint. the production code path (when loading guardrails from config.yaml) in custom_code_guardrail.py:136 doesn't apply these checks
move this validation logic to a shared function that's called both by the test endpoint and during guardrail initialization
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/proxy/guardrails/guardrail_endpoints.py
Line: 1357:1393
Comment:
security validation via regex patterns is good, but these patterns are only applied in the test endpoint. the production code path (when loading guardrails from config.yaml) in `custom_code_guardrail.py:136` doesn't apply these checks
move this validation logic to a shared function that's called both by the test endpoint and during guardrail initialization
How can I resolve this? If you propose a fix, please make it concise.| const response = await testCustomCodeGuardrail(accessToken, { | ||
| custom_code: code, | ||
| test_input: parsedInput, | ||
| input_type: mode as "request" | "response", | ||
| request_data: { | ||
| model: "test-model", | ||
| metadata: {}, | ||
| }, | ||
| }); |
There was a problem hiding this comment.
the test endpoint is called with mode as input_type, but mode can be "pre_call", "post_call", "during_call", or "logging_only" while input_type expects "request" or "response"
| const response = await testCustomCodeGuardrail(accessToken, { | |
| custom_code: code, | |
| test_input: parsedInput, | |
| input_type: mode as "request" | "response", | |
| request_data: { | |
| model: "test-model", | |
| metadata: {}, | |
| }, | |
| }); | |
| const response = await testCustomCodeGuardrail(accessToken, { | |
| custom_code: code, | |
| test_input: parsedInput, | |
| input_type: mode === "post_call" ? "response" : "request", | |
| request_data: { | |
| model: "test-model", | |
| metadata: {}, | |
| }, | |
| }); |
Prompt To Fix With AI
This is a comment left during a code review.
Path: ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx
Line: 255:263
Comment:
the test endpoint is called with `mode` as `input_type`, but `mode` can be "pre_call", "post_call", "during_call", or "logging_only" while `input_type` expects "request" or "response"
```suggestion
const response = await testCustomCodeGuardrail(accessToken, {
custom_code: code,
test_input: parsedInput,
input_type: mode === "post_call" ? "response" : "request",
request_data: {
model: "test-model",
metadata: {},
},
});
```
How can I resolve this? If you propose a fix, please make it concise.| export const testCustomCodeGuardrail = async ( | ||
| accessToken: string, | ||
| request: TestCustomCodeGuardrailRequest |
There was a problem hiding this comment.
the function signature doesn't match the actual usage in CustomCodeModal.tsx:255 and CustomCodePlayground.tsx:113. in those files, the function is called with 5 separate parameters (accessToken, customCode, parsedInputs, parsedRequestData, inputType) but this signature expects 2 parameters (accessToken, request)
check which call sites are using the old signature and update them
Prompt To Fix With AI
This is a comment left during a code review.
Path: ui/litellm-dashboard/src/components/networking.tsx
Line: 7645:7647
Comment:
the function signature doesn't match the actual usage in `CustomCodeModal.tsx:255` and `CustomCodePlayground.tsx:113`. in those files, the function is called with 5 separate parameters `(accessToken, customCode, parsedInputs, parsedRequestData, inputType)` but this signature expects 2 parameters `(accessToken, request)`
check which call sites are using the old signature and update them
How can I resolve this? If you propose a fix, please make it concise.| safe_request_data = self._prepare_safe_request_data(request_data) | ||
|
|
||
| # Execute the custom function | ||
| result = self._compiled_function(inputs, safe_request_data, input_type) |
There was a problem hiding this comment.
the custom function is called synchronously but apply_guardrail is an async method. if the custom code performs any long-running operations (loops, regex on large strings), it will block the async event loop and degrade proxy performance
consider running the custom code execution in a thread pool executor (like the test endpoint does) to avoid blocking
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py
Line: 205:205
Comment:
the custom function is called synchronously but `apply_guardrail` is an async method. if the custom code performs any long-running operations (loops, regex on large strings), it will block the async event loop and degrade proxy performance
consider running the custom code execution in a thread pool executor (like the test endpoint does) to avoid blocking
How can I resolve this? If you propose a fix, please make it concise.| # Python builtins (safe subset) | ||
| "len": len, | ||
| "str": str, | ||
| "int": int, | ||
| "float": float, | ||
| "bool": bool, | ||
| "list": list, | ||
| "dict": dict, | ||
| "True": True, | ||
| "False": False, | ||
| "None": None, | ||
| } |
There was a problem hiding this comment.
exposing Python builtins like list, dict, str, int, float, and bool is necessary for basic operations, but these objects have dunder methods that could potentially be used for sandbox escape (e.g., ().__class__.__bases__[0].__subclasses__())
this is a known limitation of Python sandboxing - consider documenting this security caveat in the code comments and documentation
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py
Line: 576:587
Comment:
exposing Python builtins like `list`, `dict`, `str`, `int`, `float`, and `bool` is necessary for basic operations, but these objects have dunder methods that could potentially be used for sandbox escape (e.g., `().__class__.__bases__[0].__subclasses__()`)
this is a known limitation of Python sandboxing - consider documenting this security caveat in the code comments and documentation
How can I resolve this? If you propose a fix, please make it concise.| const result = await testCustomCodeGuardrail( | ||
| accessToken, | ||
| customCode, | ||
| parsedInputs, | ||
| parsedRequestData, | ||
| inputType | ||
| ); |
There was a problem hiding this comment.
incorrect function signature - testCustomCodeGuardrail expects 2 parameters (accessToken, request) but this is calling it with 5 separate parameters
| const result = await testCustomCodeGuardrail( | |
| accessToken, | |
| customCode, | |
| parsedInputs, | |
| parsedRequestData, | |
| inputType | |
| ); | |
| const result = await testCustomCodeGuardrail( | |
| accessToken, | |
| { | |
| custom_code: customCode, | |
| test_input: parsedInputs, | |
| request_data: parsedRequestData, | |
| input_type: inputType | |
| } | |
| ); |
Prompt To Fix With AI
This is a comment left during a code review.
Path: ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodePlayground.tsx
Line: 113:119
Comment:
incorrect function signature - `testCustomCodeGuardrail` expects 2 parameters `(accessToken, request)` but this is calling it with 5 separate parameters
```suggestion
const result = await testCustomCodeGuardrail(
accessToken,
{
custom_code: customCode,
test_input: parsedInputs,
request_data: parsedRequestData,
input_type: inputType
}
);
```
How can I resolve this? If you propose a fix, please make it concise.| "'output' runs on model → user traffic, and 'both' applies to both." | ||
| ), | ||
| ) | ||
|
|
There was a problem hiding this comment.
removed presidio_filter_scope field - this appears to be an unrelated change that removes functionality. was this intentional? if users were relying on this field to control whether Presidio runs on input/output/both, this is a breaking change
verify this removal is intentional and not an accidental deletion during development
Prompt To Fix With AI
This is a comment left during a code review.
Path: litellm/types/guardrails.py
Line: 300:300
Comment:
removed `presidio_filter_scope` field - this appears to be an unrelated change that removes functionality. was this intentional? if users were relying on this field to control whether Presidio runs on input/output/both, this is a breaking change
verify this removal is intentional and not an accidental deletion during development
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
PresidioPresidioConfigModelUserInterface already has this param
There was a problem hiding this comment.
confirmed - PresidioPresidioConfigModelUserInterface (litellm/types/guardrails.py:281-285) already defines this field with the same signature. removing the duplicate definition from the child class is correct - it now properly inherits the field from the parent.
* feat(guardrails/): allow custom code execution for guardrails first step in allowing teams to submit custom code for guardrails * feat: custom_code_guardrail.md support passing custom code for guardrails * feat: initial commit adding ui for custom code guardrails allows users to write guardrails based on custom code * feat: expose new test custom code guardrail endpoint allows ui testing playground to sanity check if guardrail is working as expected * fix: fix linting errors * fix: fix max recursion check * fix: fix linting error
Relevant issues
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
tests/litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unitCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🆕 New Feature
🐛 Bug Fix
🧹 Refactoring
📖 Documentation
🚄 Infrastructure
✅ Test
Changes