From 49ccb3369c0e3b203063dd655eae7a9f6df308ab Mon Sep 17 00:00:00 2001
From: user <70670632+stuxf@users.noreply.github.com>
Date: Thu, 30 Apr 2026 16:33:42 -0700
Subject: [PATCH] test(vector stores): pin rag scan depth boundary

---
 litellm/proxy/rag_endpoints/endpoints.py      | 26 +++++++++++++------
 .../test_vector_store_tenant_guard.py         | 16 ++++++++++++
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py
index ecb13638648..df774c1d321 100644
--- a/litellm/proxy/rag_endpoints/endpoints.py
+++ b/litellm/proxy/rag_endpoints/endpoints.py
@@ -30,6 +30,15 @@ from litellm.proxy.vector_store_endpoints.utils import (
 router = APIRouter()
 
 
+def _raise_vector_store_scan_depth_exceeded() -> None:
+    raise HTTPException(
+        status_code=400,
+        detail={
+            "error": f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while scanning vector_store_id values"
+        },
+    )
+
+
 def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]:
     vector_store_ids: set[str] = set()
     payload_stack = [(payload, 0)]
@@ -37,12 +46,7 @@ def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]:
     while payload_stack:
         current_payload, depth = payload_stack.pop()
         if depth > DEFAULT_MAX_RECURSE_DEPTH:
-            raise HTTPException(
-                status_code=400,
-                detail={
-                    "error": f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while scanning vector_store_id values"
-                },
-            )
+            _raise_vector_store_scan_depth_exceeded()
 
         if isinstance(current_payload, dict):
             for key, value in current_payload.items():
@@ -57,9 +61,15 @@ def _collect_vector_store_ids_from_payload(payload: Any) -> set[str]:
                     vector_store_ids.add(value)
                     continue
                 if isinstance(value, (dict, list)):
-                    payload_stack.append((value, depth + 1))
+                    next_depth = depth + 1
+                    if next_depth > DEFAULT_MAX_RECURSE_DEPTH:
+                        _raise_vector_store_scan_depth_exceeded()
+                    payload_stack.append((value, next_depth))
         elif isinstance(current_payload, list):
-            payload_stack.extend((item, depth + 1) for item in current_payload)
+            next_depth = depth + 1
+            if current_payload and next_depth > DEFAULT_MAX_RECURSE_DEPTH:
+                _raise_vector_store_scan_depth_exceeded()
+            payload_stack.extend((item, next_depth) for item in current_payload)
 
     return vector_store_ids
 
diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py
index 380b3963c94..ecde853b0af 100644
--- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py
+++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_tenant_guard.py
@@ -247,6 +247,22 @@ def test_rag_payload_scan_rejects_excessive_nesting():
     assert exc_info.value.status_code == 400
 
 
+def test_rag_payload_scan_accepts_vector_store_id_at_depth_limit():
+    from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+    from litellm.proxy.rag_endpoints.endpoints import (
+        _collect_vector_store_ids_from_payload,
+    )
+
+    payload = {}
+    current = payload
+    for _ in range(DEFAULT_MAX_RECURSE_DEPTH):
+        current["nested"] = {}
+        current = current["nested"]
+    current["vector_store_id"] = "vs_at_limit"
+
+    assert _collect_vector_store_ids_from_payload(payload) == {"vs_at_limit"}
+
+
 @pytest.mark.asyncio
 async def test_responses_file_search_denies_other_team_vector_store():
     from litellm.proxy.common_request_processing import (