fix(vector-stores): reject configured ID collisions

This commit is contained in:
Yujong Lee 2026-09-04 15:56:07 -07:00
parent 57b08e8f7a
commit 1a0f4aa4b2
3 changed files with 91 additions and 0 deletions

View file

@ -137,6 +137,16 @@ def _litellm_params_dict(
return {} # mutable-ok: absent persisted parameters normalize to an empty mutable mapping
def _reject_config_vector_store_id(vector_store_id: str) -> None:
registry: Final = litellm.vector_store_registry
if registry is None or vector_store_id not in registry.config_vector_store_ids:
return
raise HTTPException(
status_code=400,
detail=f"Vector store ID {vector_store_id} is defined in proxy configuration and cannot be managed through the API",
)
async def _fetch_and_authorize_vector_store(
vector_store_id: str,
user_api_key_dict: UserAPIKeyAuth,
@ -210,6 +220,8 @@ async def create_vector_store_in_db(
if prisma_client is None:
raise HTTPException(status_code=500, detail="Database not connected")
_reject_config_vector_store_id(vector_store_id)
# Check if vector store already exists
existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique(
where={"vector_store_id": vector_store_id}
@ -629,6 +641,7 @@ async def update_vector_store(
try:
update_data: Final = data.model_dump(exclude_unset=True)
vector_store_id: Final[str] = update_data.pop("vector_store_id")
_reject_config_vector_store_id(vector_store_id)
# Per-store access control: anyone authenticated who passes the
# premium-feature gate could otherwise update *any* vector store —

View file

@ -479,6 +479,8 @@ class VectorStoreRegistry:
"""
Delete a vector store from the registry
"""
if vector_store_id in self.config_vector_store_ids:
return
self.vector_stores = [
vector_store
for vector_store in self.vector_stores
@ -487,6 +489,8 @@ class VectorStoreRegistry:
def update_vector_store_in_registry(self, vector_store_id: str, updated_data: LiteLLM_ManagedVectorStore):
"""Update or add a vector store in the registry"""
if vector_store_id in self.config_vector_store_ids:
return
for i, vector_store in enumerate(self.vector_stores):
if vector_store.get("vector_store_id") == vector_store_id:
self.vector_stores[i] = updated_data

View file

@ -877,6 +877,80 @@ async def test_db_fallback_does_not_evict_config_source():
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_awaited()
@pytest.mark.asyncio
async def test_config_vector_store_id_cannot_be_created_in_database():
registry = VectorStoreRegistry()
registry.config_vector_store_ids = frozenset(("configured",))
prisma_client = MagicMock()
with (
patch.object(litellm, "vector_store_registry", registry),
pytest.raises(HTTPException, match="defined in proxy configuration") as exc_info,
):
await create_vector_store_in_db(
vector_store_id="configured",
custom_llm_provider="milvus",
prisma_client=prisma_client,
)
assert exc_info.value.status_code == 400
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_called()
prisma_client.db.litellm_managedvectorstorestable.create.assert_not_called()
@pytest.mark.asyncio
async def test_config_vector_store_id_cannot_be_updated_in_database():
from litellm.proxy.vector_store_endpoints.management_endpoints import update_vector_store
from litellm.types.vector_stores import VectorStoreUpdateRequest
registry = VectorStoreRegistry()
registry.config_vector_store_ids = frozenset(("configured",))
prisma_client = MagicMock()
with (
patch.object(litellm, "vector_store_registry", registry),
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch(
"litellm.proxy.vector_store_endpoints.management_endpoints.check_feature_access_for_user",
new=AsyncMock(),
),
pytest.raises(HTTPException, match="defined in proxy configuration") as exc_info,
):
await update_vector_store(
data=VectorStoreUpdateRequest(
vector_store_id="configured",
vector_store_description="replacement",
),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert exc_info.value.status_code == 400
prisma_client.db.litellm_managedvectorstorestable.find_unique.assert_not_called()
prisma_client.db.litellm_managedvectorstorestable.update.assert_not_called()
def test_config_vector_store_cannot_be_replaced_or_deleted_from_registry():
configured = LiteLLM_ManagedVectorStore(
vector_store_id="configured",
custom_llm_provider="milvus",
litellm_params={"api_base": "https://configured-milvus:19530"},
)
registry = VectorStoreRegistry(vector_stores=[configured])
registry.config_vector_store_ids = frozenset(("configured",))
registry.update_vector_store_in_registry(
"configured",
LiteLLM_ManagedVectorStore(
vector_store_id="configured",
custom_llm_provider="milvus",
litellm_params={"api_base": "https://attacker.example"},
),
)
registry.delete_vector_store_from_registry("configured")
assert registry.get_litellm_managed_vector_store_from_registry("configured") == configured
def test_admin_persistence_strips_forged_marker_and_adds_server_marker():
params = prepare_milvus_connection_for_persistence(
custom_llm_provider="milvus/probe",