From 8b0441a628c01f0cd6caa10176ae887c06d75fa7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:29:25 -0700 Subject: [PATCH] fix(vector_stores): block caller-supplied embedding selection params on query surfaces --- litellm/proxy/rag_endpoints/endpoints.py | 2 ++ .../proxy/vector_store_endpoints/endpoints.py | 24 ++++++++++++++ .../proxy/rag_endpoints/test_rag_endpoints.py | 24 ++++++++++++++ .../test_vector_store_endpoints.py | 32 +++++++++++++++++++ 4 files changed, 82 insertions(+) diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index d2c7d6f93ee..0ab7d99e4e4 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( ) from litellm.proxy.vector_store_endpoints.endpoints import ( build_request_data_from_managed_vector_store, + reject_caller_embedding_selection_params, ) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, @@ -716,6 +717,7 @@ async def rag_query( status_code=400, detail={"error": "retrieval_config must contain 'vector_store_id'"}, ) + reject_caller_embedding_selection_params(payload=retrieval_config, source="retrieval_config") resolved_stores: Final = await _authorize_nested_vector_store_ids( payload=retrieval_config, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 3fc6749f18a..7d64e648e08 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -29,6 +29,29 @@ from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry router: Final = APIRouter() + +BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS: Final = frozenset( + { + "embedding_model", + "litellm_embedding_model", + "litellm_embedding_config", + "litellm_credential_name", + } +) + + +def reject_caller_embedding_selection_params(payload: Mapping[str, object], source: str) -> None: + blocked: Final = sorted(BLOCKED_QUERY_EMBEDDING_SELECTION_PARAMS & payload.keys()) + if blocked: + raise HTTPException( + status_code=400, + detail={ + "error": f"'{blocked[0]}' cannot be set in {source}. " + "Embedding configuration comes from the vector store's server-side registration." + }, + ) + + ######################################################## # OpenAI Compatible Endpoints ######################################################## @@ -134,6 +157,7 @@ async def vector_store_search( ) data = await _read_request_body(request=request) + reject_caller_embedding_selection_params(payload=data, source="the search request body") data["vector_store_id"] = vector_store_id # Check for legacy vector store registry (non-managed vector stores) diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index 342a4535b21..0085b6ebd36 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -421,6 +421,30 @@ def test_rag_query_store_params_win_over_user_retrieval_config(client_internal_u assert forwarded_config["aws_region_name"] == "eu-west-1" +@pytest.mark.parametrize( + "blocked_key", + ["embedding_model", "litellm_embedding_model", "litellm_embedding_config", "litellm_credential_name"], +) +def test_rag_query_rejects_caller_embedding_selection_params(client_internal_user, blocked_key): + """ + Regression: a caller must not pick the embedding model or credential used at + search time. Those resolve through the Router with the proxy's credentials, + bypassing the key's model permissions, so they may only come from the + managed store's server-side registration. + """ + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": {"vector_store_id": "s3-store", blocked_key: "attacker-choice"}, + }, + ) + + assert response.status_code == 400, response.json() + assert blocked_key in str(response.json()) + + EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index eae6f90863a..45a0221c8a6 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3158,3 +3158,35 @@ class TestAzureAIAnalyzeNamedIndexClassification: user_api_key_dict=self._team_member("analyze", ["read"]), ) assert result is True + + +@pytest.mark.parametrize( + "blocked_key", + ["embedding_model", "litellm_embedding_model", "litellm_embedding_config", "litellm_credential_name"], +) +def test_vector_store_search_rejects_caller_embedding_selection_params(blocked_key): + """ + Regression: the search request body must not pick the embedding model or + credential used to embed the query. Those resolve through the Router with + the proxy's credentials, bypassing the key's model permissions, so they may + only come from the managed store's server-side registration. + """ + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + + mock_auth = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER.value) + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + try: + client = TestClient(app) + response = client.post( + "/v1/vector_stores/s3-store/search", + json={"query": "hello", blocked_key: "attacker-choice"}, + ) + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == 400, response.json() + assert blocked_key in str(response.json())