mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): keep SSO values that only live in the environment_variables row on a partial SSO PATCH
The SSO PATCH strips every SSO env var from the environment_variables config row. Now that the UI only sends changed fields, an omitted value that lived only in that row was dropped from both tables and lost on restart. Carry those values into the SSO row before stripping them
This commit is contained in:
Jesus Velarde
parent
e01c41bc90
commit
41079f400c
2 changed files with 95 additions and 26 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue