diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index fa87814bfaf..6d31fa3f226 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -16,6 +16,7 @@ from typing import ( from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile +from prisma.errors import PrismaError from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model, field_validator from pydantic.fields import FieldInfo, PydanticUndefined from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -26,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper from litellm.proxy.common_utils.validation_error_body import public_validation_errors from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError @@ -1247,6 +1249,37 @@ async def get_sso_settings(): } +_ENVIRONMENT_VARIABLES_PARAM: Final = TypeAdapter(dict[str, object]) + + +def _parse_environment_variables_param(param_value: str | Mapping[str, object] | None) -> Mapping[str, object]: + if isinstance(param_value, str): + return _ENVIRONMENT_VARIABLES_PARAM.validate_python(json.loads(param_value)) + return param_value or {} + + +async def _load_environment_variables_param( + repo: ConfigRepository, +) -> Mapping[str, object] | None | PrismaError | ValueError: + try: + entry: Final = await _config_param_db(repo).find_unique(where={"param_name": "environment_variables"}) + param_value: Final = None if entry is None else _parse_environment_variables_param(entry.param_value) + except (PrismaError, ValueError) as e: + return e + else: + return param_value + + +def _decrypted_sso_fields(environment_variables: Mapping[str, object]) -> Mapping[str, str | None]: + return MappingProxyType( + { + field_name: decrypt_value_helper(value=value, key=env_var, return_original_value=True) + for field_name, env_var in SSO_FIELD_ENV_VARS.items() + if isinstance(value := environment_variables.get(env_var), str) + } + ) + + @router.patch( "/update/sso_settings", tags=["SSO Settings"], @@ -1310,11 +1343,21 @@ async def update_sso_settings( if "general_settings" not in config: config["general_settings"] = {} - # PATCH semantics: omitted fields keep their stored value, an explicit null clears it + # PATCH semantics: omitted fields keep their stored value, an explicit null clears it. SSO values that + # still live only in the environment_variables row are carried over, since that row is stripped below + environment_variables_param: Final = await _load_environment_variables_param( + ConfigRepository(prisma_client, use_writer=True) + ) + legacy_env_values: Final[Mapping[str, str | None]] = ( + _decrypted_sso_fields(environment_variables_param) + if isinstance(environment_variables_param, Mapping) + else MappingProxyType({}) + ) + stored_values: Final = {k: v for k, v in (before_sso_data or {}).items() if v is not None} + fallback_values: Final = {**legacy_env_values, **stored_values} submitted: Final = sso_config.model_dump(include=sso_config.model_fields_set) - stored_values: Final = before_sso_data or {} sso_data: Final = { - field_name: submitted[field_name] if field_name in submitted else stored_values.get(field_name) + field_name: submitted[field_name] if field_name in submitted else fallback_values.get(field_name) for field_name in SSOConfig.model_fields } for field_name, value in submitted.items(): @@ -1353,25 +1396,17 @@ async def update_sso_settings( ) ) - # Remove SSO-related env vars from config.environment_variables - try: - env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client, use_writer=True)).find_unique( - where={"param_name": "environment_variables"} + if isinstance(environment_variables_param, (PrismaError, ValueError)): + raise HTTPException( + status_code=500, + detail={"error": f"Error updating environment_variables: {environment_variables_param}"}, ) - # If no environment_variables entry exists, nothing to clean up - if env_var_entry is not None: - if env_var_entry.param_value is not None: - if isinstance(env_var_entry.param_value, str): - environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value) - else: - environment_variables = dict(env_var_entry.param_value) - else: - environment_variables = {} - + if environment_variables_param is not None: + try: env_vars_to_remove: Final = set(SSO_FIELD_ENV_VARS.values()) filtered_env_vars: Final = { - key: value for key, value in environment_variables.items() if key not in env_vars_to_remove + key: value for key, value in environment_variables_param.items() if key not in env_vars_to_remove } await _config_param_db(ConfigRepository(prisma_client)).update( @@ -1381,11 +1416,11 @@ async def update_sso_settings( }, ) await invalidate_config_param("environment_variables") - except Exception as e: - raise HTTPException( - status_code=500, - detail={"error": f"Error updating environment_variables: {e}"}, - ) + except Exception as e: + raise HTTPException( + status_code=500, + detail={"error": f"Error updating environment_variables: {e}"}, + ) return { "message": "SSO settings updated successfully", diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 0c6c1a00c73..e66050877df 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -717,7 +717,7 @@ class TestProxySettingEndpoints: create_sso_settings = json.loads(create_data["sso_settings"]) assert create_sso_settings["google_client_id"] == "new_google_client_id" - def _patch_sso_settings(self, monkeypatch, stored, body): + def _patch_sso_settings(self, monkeypatch, stored, body, environment_variables=None): """PATCH /update/sso_settings against a stored row; returns (response, saved row).""" import json from unittest.mock import AsyncMock, MagicMock @@ -732,7 +732,12 @@ class TestProxySettingEndpoints: return_value=stored_record ) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() - mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=None) + env_var_entry = None + if environment_variables is not None: + env_var_entry = MagicMock() + env_var_entry.param_value = environment_variables + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock(return_value=env_var_entry) + mock_prisma.db.litellm_config.update = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) from litellm.proxy.proxy_server import proxy_config @@ -840,6 +845,35 @@ class TestProxySettingEndpoints: assert returned["generic_client_secret"] != "stored-client-secret" assert "stored-client-secret" not in response.text + def test_update_sso_settings_moves_omitted_env_row_values_into_sso_row( + self, mock_proxy_config, mock_auth, monkeypatch + ): + import json + + environment_variables = { + "MICROSOFT_TENANT": "env-row-tenant", + "MICROSOFT_CLIENT_ID": "env-row-client-id", + "GENERIC_CLIENT_ID": "env-row-generic-id", + "LANGFUSE_PUBLIC_KEY": "unrelated", + } + + response, saved = self._patch_sso_settings( + monkeypatch, + self._STORED_GENERIC_SSO, + {"user_email": "new@example.com", "microsoft_client_id": None}, + environment_variables=environment_variables, + ) + + from litellm.proxy.proxy_server import prisma_client + + assert response.status_code == 200 + assert saved["microsoft_tenant"] == "env-row-tenant" + assert saved["microsoft_client_id"] is None + assert saved["generic_client_id"] == "stored-client-id" + assert saved["user_email"] == "new@example.com" + env_row_update = prisma_client.db.litellm_config.update.call_args.kwargs + assert json.loads(env_row_update["data"]["param_value"]) == {"LANGFUSE_PUBLIC_KEY": "unrelated"} + def test_update_sso_settings_maps_saml_fields_to_env_vars( self, mock_proxy_config, mock_auth, monkeypatch ): @@ -922,7 +956,7 @@ class TestProxySettingEndpoints: mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() mock_prisma.db.litellm_config = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock( + mock_prisma.writer_db.litellm_config.find_unique = AsyncMock( side_effect=ValueError("cleanup failed") ) mock_prisma.db.litellm_config.update = AsyncMock()