From 6d0146a3f34e009420444f9262a2489b2909459b Mon Sep 17 00:00:00 2001 From: lzhan011 <35493221+lzhan011@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:17:30 -0500 Subject: [PATCH 1/3] fix(guardrails): wire _unpack_sequence_ into the custom-code sandbox RestrictedPython rewrites every tuple-unpacking target that is not a for-loop target into a call to _unpack_sequence_, and ships no default for it -- the same way it ships no _inplacevar_, which this module already supplies. The sandbox globals never defined it, so ordinary guardrail code compiled cleanly and then raised NameError on its first run: a, b = pair a, (b, c) = nested a, *rest = seq with ctx as (a, b): RestrictedPython.Guards.guarded_unpack_sequence is the helper the rewritten bytecode expects, and it applies the same _getiter_ guard to each element that guarded_iter_unpack_sequence already applies for `for a, b in x`. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0147HaHohHPnpFBUa4tUkvPg --- .../guardrail_hooks/custom_code/sandbox.py | 7 +++++ .../guardrails/test_custom_code_security.py | 29 +++++++++++++++++++ 2 files changed, 36 insertions(+) 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 f93ecfc3010..d245096764d 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -228,3 +228,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 From 9dcaca25223573a7fcbc937cc0dca2da0de561a8 Mon Sep 17 00:00:00 2001 From: lzhan011 <35493221+lzhan011@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:18:27 -0500 Subject: [PATCH 2/3] fix(guardrails): apply the `with` guards to `async with` in the sandbox AsyncAwareTransformer re-permits the async nodes the RestrictedPython policy rejects outright. AsyncFunctionDef delegates to visit_FunctionDef so it inherits that visitor's checks, but AsyncWith went to node_contents_visit, which only recurses into children and performs no rewriting. The result was that the async spelling enforced less than the sync one: `with x as (a, b)` is rewritten to guard the unpacking, `async with x as (a, b)` was not. AsyncWith has the same _fields and the same security semantics as With, so delegate to visit_With the way AsyncFunctionDef delegates to visit_FunctionDef. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0147HaHohHPnpFBUa4tUkvPg --- .../guardrail_hooks/custom_code/sandbox.py | 7 ++- .../guardrails/test_custom_code_security.py | 54 +++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index 1588836f681..d5707894a9d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -45,7 +45,10 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): has the same ``_fields`` as ``FunctionDef`` and the same security semantics, so we delegate to ``visit_FunctionDef`` — name check, argument check, print-scope wrapping, and any future additions to that method are - inherited automatically. ``AsyncFor``/``AsyncWith``/``Await`` delegate to + inherited automatically. ``AsyncWith`` gets the same treatment for the same + reason: ``node_contents_visit`` only recurses into children, so routing it + there left ``async with x as (a, b)`` without the unpack guard that + ``with x as (a, b)`` gets. ``AsyncFor``/``Await`` delegate to ``node_contents_visit`` so their children still get transformed. """ @@ -56,7 +59,7 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): return self.node_contents_visit(node) def visit_AsyncWith(self, node: ast.AsyncWith) -> ast.AST: - return self.node_contents_visit(node) + return self.visit_With(node) def visit_Await(self, node: ast.Await) -> ast.AST: return self.node_contents_visit(node) 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 d245096764d..679c0966828 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -257,3 +257,57 @@ def test_tuple_unpacking_runs(body: str, expected): assert fn is not None inputs = {"pair": (1, 2), "nested": (1, (2, 3)), "seq": [1, 2, 3], "ctx": _Ctx()} assert fn(inputs, {}, "request") == expected + + +def _guard_names(source: str) -> set[str]: + """Names of the RestrictedPython guards the compiled bytecode calls.""" + import types + + from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( + compile_sandboxed, + ) + + found: set[str] = set() + stack = [compile_sandboxed(source)] + while stack: + code = stack.pop() + found |= {n for n in code.co_names + code.co_varnames if n.startswith("_") and n.endswith("_")} + stack += [c for c in code.co_consts if isinstance(c, types.CodeType)] + return found + + +def test_async_with_gets_the_same_guards_as_with(): + """`async with` is the async spelling of `with`; it must not enforce less.""" + sync_guards = _guard_names("def f(x):\n with x as (a, b):\n pass\n") + async_guards = _guard_names("async def f(x):\n async with x as (a, b):\n pass\n") + + assert "_unpack_sequence_" in sync_guards, "precondition: `with` unpacking is guarded" + assert sync_guards <= async_guards + + +@pytest.mark.asyncio +async def test_async_with_still_executes(): + """Guarding `async with` must not break it.""" + from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( + build_sandbox_globals, + compile_sandboxed, + ) + + class _ACtx: + def __init__(self, value): + self.value = value + + async def __aenter__(self): + return self.value + + async def __aexit__(self, *args): + return False + + sandbox_globals = build_sandbox_globals() + exec( # noqa: S102 + compile_sandboxed( + "async def f(ctx):\n async with ctx as (a, b):\n return a + b\n" + ), + sandbox_globals, + ) + assert await sandbox_globals["f"](_ACtx((5, 6))) == 11 From 03f6d79c43e332a571a7a18037b0fe49617e07e7 Mon Sep 17 00:00:00 2001 From: lzhan011 <35493221+lzhan011@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:24:04 -0500 Subject: [PATCH 3/3] fix(guardrails): apply the `for` guards to `async for` in the sandbox AsyncFor was the last async node still routed to node_contents_visit, which only recurses into children. `for a in x` is rewritten to `for a in _getiter_(x)` and `for a, b in x` to iterate through the unpack guard; the async spellings got neither, so they enforced strictly less than the sync forms they mirror. Delegating to visit_For fixes the iterator guard, but its tuple-target rewrite emits _iter_unpack_sequence_, a plain generator that `async for` cannot consume ("requires an object with __aiter__ method, got generator"). So the call is retargeted to _aiter_unpack_sequence_, an async-generator helper with the same contract: guard the iteration, then guard each element's unpacking. Await keeps node_contents_visit -- it has no synchronous counterpart and only wraps an expression. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_0147HaHohHPnpFBUa4tUkvPg --- .../guardrail_hooks/custom_code/sandbox.py | 46 ++++++++++++++-- .../guardrails/test_custom_code_security.py | 53 +++++++++++++++++++ 2 files changed, 95 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index d5707894a9d..a1f291c0bbc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -16,7 +16,7 @@ restriction intact. import ast import operator -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping from types import CodeType from typing import Final @@ -48,15 +48,19 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): inherited automatically. ``AsyncWith`` gets the same treatment for the same reason: ``node_contents_visit`` only recurses into children, so routing it there left ``async with x as (a, b)`` without the unpack guard that - ``with x as (a, b)`` gets. ``AsyncFor``/``Await`` delegate to - ``node_contents_visit`` so their children still get transformed. + ``with x as (a, b)`` gets. ``AsyncFor`` likewise delegates to ``visit_For``, + which is what wraps the loop iterator in ``_getiter_``. ``Await`` has no + synchronous counterpart and only wraps an expression, so it stays on + ``node_contents_visit``. """ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AST: return self.visit_FunctionDef(node) def visit_AsyncFor(self, node: ast.AsyncFor) -> ast.AST: - return self.node_contents_visit(node) + transformed: Final = self.visit_For(node) + _use_async_iter_unpack(transformed) + return transformed def visit_AsyncWith(self, node: ast.AsyncWith) -> ast.AST: return self.visit_With(node) @@ -65,6 +69,39 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): return self.node_contents_visit(node) +_ITER_UNPACK_NAME: Final = "_iter_unpack_sequence_" +_ASYNC_ITER_UNPACK_NAME: Final = "_aiter_unpack_sequence_" + + +def _use_async_iter_unpack(node: ast.AST) -> None: + """Point a transformed ``async for`` at the async unpack guard. + + ``visit_For`` rewrites ``for a, b in x`` into + ``for (a, b) in _iter_unpack_sequence_(x, spec, _getiter_)``. That helper is + a plain generator, which ``async for`` cannot consume, so the async form + needs the async-generator equivalent under its own name. + """ + iter_node: Final = getattr(node, "iter", None) + if ( + isinstance(iter_node, ast.Call) + and isinstance(iter_node.func, ast.Name) + and iter_node.func.id == _ITER_UNPACK_NAME + ): + iter_node.func.id = _ASYNC_ITER_UNPACK_NAME + + +async def _aiter_unpack_sequence_( + it: object, spec: object, _getiter_: Callable[[object], AsyncIterable[object]] +) -> AsyncIterator[object]: + """``guarded_iter_unpack_sequence`` for ``async for`` targets. + + Same contract as the RestrictedPython helper — guard the iteration, then + guard each element's sequence unpacking — over an async iterator. + """ + async for ob in _getiter_(it): + yield guarded_unpack_sequence(ob, spec, _getiter_) + + _INPLACE_OPS: Final[Mapping[str, Callable[[object, object], object]]] = { "+=": operator.iadd, "-=": operator.isub, @@ -125,6 +162,7 @@ def build_sandbox_globals() -> dict[str, object]: # ships no default. Without it those statements compile and then raise # NameError the first time the guardrail runs. "_unpack_sequence_": guarded_unpack_sequence, + _ASYNC_ITER_UNPACK_NAME: _aiter_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 679c0966828..91bd61f8686 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -311,3 +311,56 @@ async def test_async_with_still_executes(): sandbox_globals, ) assert await sandbox_globals["f"](_ACtx((5, 6))) == 11 + + +def test_async_for_gets_the_same_guards_as_for(): + """`async for` is the async spelling of `for`; it must not enforce less.""" + sync_guards = _guard_names("def f(x):\n for a in x:\n pass\n") + async_guards = _guard_names("async def f(x):\n async for a in x:\n pass\n") + + assert "_getiter_" in sync_guards, "precondition: `for` iteration is guarded" + assert sync_guards <= async_guards + + +def test_async_for_unpacking_uses_the_async_unpack_guard(): + """Tuple targets are guarded too, via the async-iterable variant of the helper.""" + guards = _guard_names("async def f(x):\n async for a, b in x:\n pass\n") + + assert "_aiter_unpack_sequence_" in guards + # The sync generator would raise "requires an object with __aiter__". + assert "_iter_unpack_sequence_" not in guards + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "items", "expected"), + [ + (" async for i in src:\n out.append(i)\n", [1, 2, 3], [1, 2, 3]), + (" async for a, b in src:\n out.append(a + b)\n", [(1, 2), (3, 4)], [3, 7]), + ], +) +async def test_async_for_still_executes(body: str, items: list, expected: list): + """Guarding `async for` must not break it, with or without a tuple target.""" + from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import ( + build_sandbox_globals, + compile_sandboxed, + ) + + class _AIter: + def __init__(self, values): + self.values = list(values) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self.values: + raise StopAsyncIteration + return self.values.pop(0) + + sandbox_globals = build_sandbox_globals() + exec( # noqa: S102 + compile_sandboxed("async def f(src):\n out = []\n" + body + " return out\n"), + sandbox_globals, + ) + assert await sandbox_globals["f"](_AIter(items)) == expected