From a4f021cfa066b2bd5c5ae413ee54f783ee8bd421 Mon Sep 17 00:00:00 2001
From: lzhan011 <35493221+lzhan011@users.noreply.github.com>
Date: Mon, 5 Oct 2026 04:59:13 -0500
Subject: [PATCH] fix(guardrails): apply the `with` guards to `async with` in
 the sandbox

`visit_AsyncWith` routed to `node_contents_visit`, which only recurses into
children. The base transformer's `visit_With` is what rewrites the `as`
target, so `async with x as (a, b)` compiled without the `_unpack_sequence_`
guard that `with x as (a, b)` gets -- the async spelling enforced less than
the sync one.

Delegating to `visit_With` gives it the same treatment `AsyncFunctionDef`
already gets from `visit_FunctionDef`, so any future addition to that method
is inherited rather than re-implemented.

Rebuilt on `main` (this PR was opened against the since-deleted
`litellm_internal_staging`). `main` still routes `visit_AsyncWith` to
`node_contents_visit`, so the gap is unchanged; the test file moved upstream
to tests/unit/proxy/guardrails/ and the new cases are appended there.

Two tests: the compiled bytecode of `async with x as (a, b)` must call at
least the guards that `with x as (a, b)` does (with `_unpack_sequence_` in the
sync set asserted as a precondition), and the guarded form must still run
under a real `__aenter__`/`__aexit__`.

tests/unit/proxy/guardrails/test_custom_code_security.py -> 63 passed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016SxQ2jq4B4KLGwPoLgfHye
---
 .../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 63c731fe5ff..5681438d8d9 100644
--- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py
+++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py
@@ -47,7 +47,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.
 
     ``visit_While`` rewrites ``while test:`` to ``while _budget_ok_() and test:``
@@ -78,7 +81,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/unit/proxy/guardrails/test_custom_code_security.py b/tests/unit/proxy/guardrails/test_custom_code_security.py
index cdab46a47eb..03db0098030 100644
--- a/tests/unit/proxy/guardrails/test_custom_code_security.py
+++ b/tests/unit/proxy/guardrails/test_custom_code_security.py
@@ -875,3 +875,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