diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index 35f1e6e6515..1588836f681 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -31,6 +31,7 @@ from RestrictedPython.Eval import default_guarded_getitem, default_guarded_getit from RestrictedPython.Guards import ( full_write_guard, guarded_iter_unpack_sequence, + guarded_unpack_sequence, safer_getattr, ) @@ -115,6 +116,12 @@ def build_sandbox_globals() -> dict[str, object]: "_getitem_": default_guarded_getitem, "_getiter_": default_guarded_getiter, "_iter_unpack_sequence_": guarded_iter_unpack_sequence, + # RestrictedPython emits _unpack_sequence_ for every tuple-unpacking + # target that is not a for-loop target — ``a, b = pair``, + # ``a, *rest = seq``, ``with x as (a, b)`` — and, like _inplacevar_, + # ships no default. Without it those statements compile and then raise + # NameError the first time the guardrail runs. + "_unpack_sequence_": guarded_unpack_sequence, "_write_": full_write_guard, "_inplacevar_": _inplacevar_, } diff --git a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py index 068cd0d8ed7..ac1b454b8b1 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -363,3 +363,32 @@ def test_augmented_assignment_works(): def test_missing_apply_guardrail_raises(): with pytest.raises(CustomCodeCompilationError, match="apply_guardrail"): _compile("x = 1\n") + + +@pytest.mark.parametrize( + ("body", "expected"), + [ + # Every one of these compiles fine and then raises + # "NameError: name '_unpack_sequence_' is not defined" at call time when + # the guard is missing from the sandbox globals. + (" a, b = inputs['pair']\n return a + b\n", 3), + (" a, (b, c) = inputs['nested']\n return a + b + c\n", 6), + (" a, *rest = inputs['seq']\n return rest\n", [2, 3]), + (" with inputs['ctx'] as (a, b):\n return a + b\n", 15), + ], +) +def test_tuple_unpacking_runs(body: str, expected): + """Ordinary tuple unpacking must work inside a guardrail, not NameError.""" + + class _Ctx: + def __enter__(self): + return (7, 8) + + def __exit__(self, *args): + return False + + guardrail = _compile("def apply_guardrail(inputs, request_data, input_type):\n" + body) + fn = guardrail._compiled_function + assert fn is not None + inputs = {"pair": (1, 2), "nested": (1, (2, 3)), "seq": [1, 2, 3], "ctx": _Ctx()} + assert fn(inputs, {}, "request") == expected