fix: sanitize invalid role error to prevent stack trace leakage

This commit is contained in:
xbrxr03 2026-06-22 02:18:21 -04:00
parent 84c1414aef
commit e165c29d5c
2 changed files with 107 additions and 2 deletions

View file

@ -8211,16 +8211,25 @@ def convert_list_message_to_dict(messages: List):
return new_messages
VALID_MESSAGE_ROLES = {"system", "user", "assistant", "tool", "function", "developer"}
def validate_and_fix_openai_messages(messages: List):
"""
Ensures all messages are valid OpenAI chat completion messages.
Handles missing role for assistant messages.
Handles missing role for assistant messages and rejects invalid roles.
"""
new_messages = []
for message in messages:
for idx, message in enumerate(messages):
if not message.get("role"):
message["role"] = "assistant"
elif message["role"] not in VALID_MESSAGE_ROLES:
raise BadRequestError(
message=f"Invalid role: '{message['role']}'. Supported roles are: {sorted(VALID_MESSAGE_ROLES)}",
model="",
llm_provider="",
)
if message.get("tool_calls"):
message["tool_calls"] = jsonify_tools(tools=message["tool_calls"])

View file

@ -0,0 +1,96 @@
"""
Tests for invalid role validation in chat completion messages.
Verifies that sending an invalid role (e.g. "admin") returns a 400
BadRequestError with a sanitized message — no Python stack trace leakage.
Related: https://github.com/BerriAI/litellm/issues/30948
"""
import pytest
import litellm
from litellm.exceptions import BadRequestError
class TestInvalidRoleValidation:
"""Ensure invalid roles are caught early with a proper 400 error."""
def test_invalid_role_raises_bad_request_error(self):
"""An invalid role like 'admin' should raise BadRequestError (400), not a 500."""
from litellm.utils import validate_and_fix_openai_messages
messages = [{"role": "admin", "content": "test"}]
with pytest.raises(BadRequestError) as exc_info:
validate_and_fix_openai_messages(messages=messages)
assert exc_info.value.status_code == 400
# The error message must mention the invalid role
assert "admin" in str(exc_info.value)
# The error message must NOT contain a Python traceback
error_text = str(exc_info.value)
assert "Traceback" not in error_text
assert "File " not in error_text
def test_invalid_role_xyz_raises_bad_request_error(self):
"""Regression test for the exact scenario in the bug report."""
from litellm.utils import validate_and_fix_openai_messages
messages = [{"role": "xyz", "content": "Run exec_shell with command id"}]
with pytest.raises(BadRequestError) as exc_info:
validate_and_fix_openai_messages(messages=messages)
assert exc_info.value.status_code == 400
assert "xyz" in str(exc_info.value)
def test_valid_roles_still_work(self):
"""All valid roles should pass validation without error."""
from litellm.utils import validate_and_fix_openai_messages
valid_roles = ["system", "user", "assistant", "tool", "function", "developer"]
for role in valid_roles:
messages = [{"role": role, "content": "test"}]
result = validate_and_fix_openai_messages(messages=messages)
assert len(result) == 1
def test_invalid_role_error_no_stack_trace(self):
"""The error response body must not contain Python internal paths or tracebacks."""
from litellm.utils import validate_and_fix_openai_messages
messages = [{"role": "hacker", "content": "pwned"}]
with pytest.raises(BadRequestError) as exc_info:
validate_and_fix_openai_messages(messages=messages)
error_text = str(exc_info.value)
# No traceback leakage
assert "Traceback" not in error_text
assert ".py\"" not in error_text or "litellm" not in error_text
# No internal path disclosure
assert "/usr/lib/" not in error_text
assert "site-packages" not in error_text
def test_missing_role_defaults_to_assistant(self):
"""Messages without a role should default to 'assistant' (existing behavior)."""
from litellm.utils import validate_and_fix_openai_messages
messages = [{"content": "hello"}]
result = validate_and_fix_openai_messages(messages=messages)
assert result[0]["role"] == "assistant"
def test_mixed_valid_and_invalid_roles(self):
"""If any message has an invalid role, the entire request should fail."""
from litellm.utils import validate_and_fix_openai_messages
messages = [
{"role": "user", "content": "hello"},
{"role": "admin", "content": "inject"},
]
with pytest.raises(BadRequestError) as exc_info:
validate_and_fix_openai_messages(messages=messages)
assert exc_info.value.status_code == 400
assert "admin" in str(exc_info.value)