From 3dfd4420c60be34c1015a0f99ced7c4f6cd0b7ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 14:27:53 -0700 Subject: [PATCH] fix(vector-stores): restore redacted placeholders on update and drop forged admin marker on REST stores Echoing a GET /vector_store/info body back into /vector_store/update stored the literal REDACTED placeholder as the credential. The update now restores every placeholder from the existing row, nested fields included, and drops placeholders for keys the row never had. The non-gRPC persistence branch also strips the internal admin marker so a REST store cannot carry a forged approval into a later transport switch. --- .../llms/milvus/vector_stores/connection.py | 8 +- .../management_endpoints.py | 19 ++++- .../code_coverage_tests/recursive_detector.py | 1 + .../test_vector_store_endpoints.py | 73 +++++++++++++++++++ 4 files changed, 97 insertions(+), 4 deletions(-) diff --git a/litellm/llms/milvus/vector_stores/connection.py b/litellm/llms/milvus/vector_stores/connection.py index 1c7750cda1d..3eed8783362 100644 --- a/litellm/llms/milvus/vector_stores/connection.py +++ b/litellm/llms/milvus/vector_stores/connection.py @@ -94,9 +94,11 @@ def prepare_connection_for_persistence( previous_is_grpc: Final = _is_grpc_connection(existing_custom_llm_provider, existing) effective_is_grpc: Final = _is_grpc_connection(custom_llm_provider, effective) if not previous_is_grpc and not effective_is_grpc: - return ( # mutable-ok: persistence requires an isolated JSON-serializable dict - dict(supplied) if isinstance(litellm_params, dict) else dict(existing) - ) + return { # mutable-ok: persistence requires an isolated JSON-serializable dict + key: value + for key, value in (supplied if isinstance(litellm_params, dict) else existing).items() + if key != MILVUS_ADMIN_CONFIGURED_CONNECTION + } is_create: Final = existing_custom_llm_provider is None provider_changed: Final = not is_create and custom_llm_provider != existing_custom_llm_provider diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 97ff84adeac..ac10d9f0735 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -112,6 +112,20 @@ def _redact_sensitive_litellm_params(litellm_params: object, _depth: int = 0) -> return out +def _restore_redacted_litellm_params(supplied: object, existing: object, _depth: int = 0) -> object: + if supplied == REDACTED_BY_LITELM_STRING: + return existing + if _depth >= _REDACT_LITELLM_PARAMS_MAX_DEPTH or not isinstance(supplied, dict) or not isinstance(existing, dict): + return supplied + supplied_params: Final = deserialize_litellm_params(supplied) + existing_params: Final = deserialize_litellm_params(existing) + return { + key: _restore_redacted_litellm_params(value, existing_params.get(key), _depth + 1) + for key, value in supplied_params.items() + if value != REDACTED_BY_LITELM_STRING or key in existing_params + } + + def _validated_litellm_params( litellm_params: Mapping[str, object], ) -> Mapping[str, object]: @@ -625,12 +639,15 @@ async def update_vector_store( ) existing_litellm_params: Final = deserialize_litellm_params(existing_vector_store.get("litellm_params")) + supplied_litellm_params: Final = _restore_redacted_litellm_params( + update_data.get("litellm_params"), existing_litellm_params + ) effective_provider: Final = update_data.get("custom_llm_provider") or existing_vector_store.get( "custom_llm_provider" ) effective_litellm_params: Final = prepare_vector_store_connection_for_persistence( custom_llm_provider=effective_provider, - litellm_params=update_data.get("litellm_params"), + litellm_params=supplied_litellm_params, user_api_key_dict=user_api_key_dict, existing_custom_llm_provider=existing_vector_store.get("custom_llm_provider"), existing_litellm_params=existing_litellm_params, diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index e9f87ba6cae..cfe634f0eee 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -48,6 +48,7 @@ IGNORE_FUNCTIONS = [ "_read_image_bytes", # max depth set. "_get_masked_values", # max depth set (default 20) to prevent infinite recursion while masking nested sensitive config dicts. "_redact_sensitive_litellm_params", # max depth set (default 10). + "_restore_redacted_litellm_params", # max depth set (default 10), mirrors _redact_sensitive_litellm_params on the write side. "_redact_secret_values_in_obj", # max depth set (default 10, _REDACT_SECRET_MAX_DEPTH); fails closed by returning "REDACTED" at the cap. "_resolve", # OCI: $ref resolver bounded by `resolving_stack` cycle guard. "resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input). 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 53d92e4f713..b2a8fed0685 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 @@ -1090,6 +1090,18 @@ def test_non_grpc_connection_update_keeps_replacement_semantics(): assert params == {"api_key": "new-key"} +def test_non_grpc_connection_update_drops_forged_admin_marker(): + params = prepare_vector_store_connection_for_persistence( + custom_llm_provider="openai", + litellm_params={"api_key": "new-key", MILVUS_ADMIN_CONFIGURED_CONNECTION: True}, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + existing_custom_llm_provider="openai", + existing_litellm_params={"api_key": "old-key"}, + ) + + assert params == {"api_key": "new-key"} + + class TestCheckVectorStorePermission: """Test suite for check_vector_store_permission function.""" @@ -3334,6 +3346,67 @@ class TestUpdateVectorStoreAccessControlAndRedaction: assert params["api_key"] == REDACTED_BY_LITELM_STRING assert params["api_base"] == "https://api.openai.com/v1" + @pytest.mark.asyncio + async def test_update_restores_redacted_placeholder_from_existing_row(self): + from litellm.constants import REDACTED_BY_LITELM_STRING + from litellm.proxy.vector_store_endpoints.management_endpoints import ( + update_vector_store, + ) + from litellm.types.vector_stores import VectorStoreUpdateRequest + + existing_row = MagicMock() + existing_row.model_dump = MagicMock( + return_value={ + "vector_store_id": "vs_owned", + "team_id": "team-A", + "custom_llm_provider": "openai", + "litellm_params": { + "api_key": "sk-real-openai-key-123", + "api_base": "https://old.example", + "litellm_embedding_config": {"api_key": "sk-real-embedding-key", "model": "text-embedding-3-small"}, + }, + } + ) + updated_row = MagicMock() + updated_row.model_dump = MagicMock(return_value={"vector_store_id": "vs_owned", "team_id": "team-A"}) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=updated_row) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.vector_store_registry", None), + ): + await update_vector_store( + data=VectorStoreUpdateRequest( + vector_store_id="vs_owned", + litellm_params={ + "api_key": REDACTED_BY_LITELM_STRING, + "api_base": "https://new.example", + "litellm_embedding_config": { + "api_key": REDACTED_BY_LITELM_STRING, + "model": "text-embedding-3-small", + }, + "never_stored_secret": REDACTED_BY_LITELM_STRING, + }, + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin", team_id="team-A" + ), + ) + + update_call = mock_prisma_client.db.litellm_managedvectorstorestable.update.call_args + persisted = json.loads(update_call.kwargs["data"]["litellm_params"]) + assert persisted["api_key"] == "sk-real-openai-key-123" + assert persisted["api_base"] == "https://new.example" + assert persisted["litellm_embedding_config"] == { + "api_key": "sk-real-embedding-key", + "model": "text-embedding-3-small", + } + assert "never_stored_secret" not in persisted + assert REDACTED_BY_LITELM_STRING not in update_call.kwargs["data"]["litellm_params"] + @pytest.mark.asyncio async def test_update_row_deleted_mid_update_returns_404(self): """A concurrent delete between the authorization read and the write makes Prisma's