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