fix(guardrails): apply the for guards to async for in the sandbox

`visit_AsyncFor` routed to `node_contents_visit`, which only recurses into
children, so `async for` compiled with none of the guards `for` gets:
`visit_For` is what wraps the loop iterator, and the async spelling therefore
enforced less than the sync one.

Delegating to `visit_For` is not quite enough on its own. It inserts
synchronous wrappers -- `_getiter_(y)`, or
`_iter_unpack_sequence_(y, spec, _getiter_)` for a tuple target -- and both are
plain generators that `async for` cannot consume. `_use_async_iter_guards`
therefore swaps just the wrapper name for an async-capable twin:

    async for a in y        ->  async for a in _agetiter_(y)
    async for a, b in y     ->  async for (a, b) in _aiter_unpack_sequence_(y, spec, _getiter_)

`_agetiter_` is `budgeted_iter` over an async iterator, so an `async for` that
never yields is still stopped at the execution deadline. `_aiter_unpack_sequence_`
budgets the iteration and then applies `guarded_unpack_sequence` per element;
unpacking an element is synchronous, so the `_getiter_` handed in is the
synchronous one and is passed through untouched.

Adapted while rebuilding on `main`, worth flagging: this PR was written when
`_getiter_` was `default_guarded_getiter`. `main` now binds it to
`budgeted_iter`, which adds the per-item budget check -- so the async twins
have to carry that check too, which is why `_agetiter_` exists rather than the
async helper simply calling `_getiter_`. (This PR was opened against the
since-deleted `litellm_internal_staging`; `main` still routes
`visit_AsyncFor` to `node_contents_visit`.)

Tests: the compiled bytecode of each `async for` shape must carry an
async-capable wrapper, with the sync forms' guards asserted as preconditions;
both shapes must still run under a real async generator; and a non-sequence
element must still raise `TypeError` through the unpack guard rather than pass
silently. The guard-parity test fails with `visit_AsyncFor` put back on
`node_contents_visit` (`1 failed | 3 passed`); all 67 in the file pass here.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016SxQ2jq4B4KLGwPoLgfHye

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
 
@@ -50,8 +50,12 @@ 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_``; the wrapper names it
+    inserts are then swapped for async-capable ones, since ``budgeted_iter`` and
+    ``_iter_unpack_sequence_`` are plain generators that ``async for`` cannot
+    consume. ``Await`` has no synchronous counterpart and only wraps an
+    expression, so it stays on ``node_contents_visit``.
 
     ``visit_While`` rewrites ``while test:`` to ``while _budget_ok_() and test:``
     so a loop that never yields is still stopped at the execution deadline;
@@ -78,7 +82,9 @@ class AsyncAwareTransformer(RestrictingNodeTransformer):
         return bounded
 
     def visit_AsyncFor(self, node: ast.AsyncFor) -> ast.AST:
-        return self.node_contents_visit(node)
+        transformed: Final = self.visit_For(node)
+        _use_async_iter_guards(transformed)
+        return transformed
 
     def visit_AsyncWith(self, node: ast.AsyncWith) -> ast.AST:
         return self.visit_With(node)
@@ -87,6 +93,60 @@ class AsyncAwareTransformer(RestrictingNodeTransformer):
         return self.node_contents_visit(node)
 
 
+_GETITER_NAME: Final = "_getiter_"
+_ASYNC_GETITER_NAME: Final = "_agetiter_"
+_ITER_UNPACK_NAME: Final = "_iter_unpack_sequence_"
+_ASYNC_ITER_UNPACK_NAME: Final = "_aiter_unpack_sequence_"
+
+
+def _use_async_iter_guards(node: ast.AST) -> None:
+    """Swap the iterator wrapper ``visit_For`` inserted for its async twin.
+
+    ``visit_For`` rewrites ``for x in y`` into ``for x in _getiter_(y)`` and
+    ``for a, b in y`` into
+    ``for (a, b) in _iter_unpack_sequence_(y, spec, _getiter_)``. Both wrappers
+    are synchronous generators, which ``async for`` cannot consume, so the
+    async form needs the async-generator equivalents under their own names.
+    Only the wrapper name changes: the guarding each one performs is the same.
+    """
+    iter_node: Final = getattr(node, "iter", None)
+    if not (isinstance(iter_node, ast.Call) and isinstance(iter_node.func, ast.Name)):
+        return
+    if iter_node.func.id == _ITER_UNPACK_NAME:
+        iter_node.func.id = _ASYNC_ITER_UNPACK_NAME
+    elif iter_node.func.id == _GETITER_NAME:
+        iter_node.func.id = _ASYNC_GETITER_NAME
+
+
+async def _agetiter_(iterable: AsyncIterable[object]) -> AsyncIterator[object]:
+    """``budgeted_iter`` for ``async for``: the same per-item budget check.
+
+    Bound to the name ``_use_async_iter_guards`` substitutes, so an ``async for``
+    that never yields is still stopped at the execution deadline, exactly as a
+    ``for`` loop is.
+    """
+    async for item in iterable:
+        budget_ok()
+        yield item
+
+
+async def _aiter_unpack_sequence_(
+    it: AsyncIterable[object],
+    spec: object,
+    getiter: Callable[[object], object],
+) -> AsyncIterator[object]:
+    """``guarded_iter_unpack_sequence`` for ``async for`` targets.
+
+    Same contract as the RestrictedPython helper -- budget the iteration, then
+    guard each element's sequence unpacking -- over an async iterator. Unpacking
+    an element is synchronous, so the ``_getiter_`` handed in is the synchronous
+    one and is passed through untouched.
+    """
+    async for ob in it:
+        budget_ok()
+        yield guarded_unpack_sequence(ob, spec, getiter)
+
+
 _INPLACE_OPS: Final[Mapping[str, Callable[[object, object], object]]] = {
     "+=": operator.iadd,
     "-=": operator.isub,
@@ -147,6 +207,8 @@ 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_GETITER_NAME: _agetiter_,
+        _ASYNC_ITER_UNPACK_NAME: _aiter_unpack_sequence_,
         "_write_": full_write_guard,
         "_inplacevar_": _inplacevar_,
         "_budget_ok_": budget_ok,

tests/unit/proxy/guardrails/test_custom_code_security.py

@@ -929,3 +929,77 @@ 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_plain = _guard_names("def f(x):\n    for a in x:\n        pass\n")
+    async_plain = _guard_names("async def f(x):\n    async for a in x:\n        pass\n")
+    sync_unpack = _guard_names("def f(x):\n    for a, b in x:\n        pass\n")
+    async_unpack = _guard_names("async def f(x):\n    async for a, b in x:\n        pass\n")
+
+    # Preconditions: the sync forms are guarded at all.
+    assert "_getiter_" in sync_plain
+    assert "_iter_unpack_sequence_" in sync_unpack
+
+    # The async forms are guarded by the async-capable twins, under their own
+    # names — `budgeted_iter` and `_iter_unpack_sequence_` are plain generators
+    # that `async for` cannot consume.
+    assert "_agetiter_" in async_plain
+    assert "_aiter_unpack_sequence_" in async_unpack
+
+    # And neither async form is left with no wrapper at all.
+    assert async_plain & {"_agetiter_", "_aiter_unpack_sequence_"}
+    assert async_unpack & {"_agetiter_", "_aiter_unpack_sequence_"}
+
+
+async def _arange(*values):
+    for v in values:
+        yield v
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("source", "arg", "expected"),
+    [
+        (
+            "async def f(it):\n    total = 0\n    async for a in it:\n        total += a\n    return total\n",
+            (1, 2, 3),
+            6,
+        ),
+        (
+            "async def f(it):\n    total = 0\n    async for a, b in it:\n        total += a * b\n    return total\n",
+            ((1, 2), (3, 4)),
+            14,
+        ),
+    ],
+)
+async def test_async_for_still_executes(source, arg, expected):
+    """Guarding `async for` must not break it, in either target shape."""
+    from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
+        build_sandbox_globals,
+        compile_sandboxed,
+    )
+
+    sandbox_globals = build_sandbox_globals()
+    exec(compile_sandboxed(source), sandbox_globals)  # noqa: S102
+    assert await sandbox_globals["f"](_arange(*arg)) == expected
+
+
+@pytest.mark.asyncio
+async def test_async_for_unpacking_still_rejects_a_non_sequence():
+    """The unpack guard is what makes the element shape an error, not a crash."""
+    from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
+        build_sandbox_globals,
+        compile_sandboxed,
+    )
+
+    sandbox_globals = build_sandbox_globals()
+    exec(  # noqa: S102
+        compile_sandboxed(
+            "async def f(it):\n    async for a, b in it:\n        pass\n"
+        ),
+        sandbox_globals,
+    )
+    with pytest.raises(TypeError):
+        await sandbox_globals["f"](_arange(1, 2))