diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9cb9edc926..ce2fc2e475 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -580,4 +580,12 @@ model LiteLLM_SearchToolsTable { search_tool_info Json? created_at DateTime @default(now()) updated_at DateTime @updatedAt +} + +// SSO configuration table +model LiteLLM_SSOConfig { + id String @id @default("sso_config") + sso_settings Json + created_at DateTime @default(now()) + updated_at DateTime @updatedAt } \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 938c8979d2..78632628ae 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1071,6 +1071,7 @@ celery_fn = None # Redis Queue for handling requests scheduler = None last_model_cost_map_reload = None + ### DB WRITER ### db_writer_client: Optional[AsyncHTTPHandler] = None ### logger ### @@ -3335,6 +3336,28 @@ class ProxyConfig: if self._should_load_db_object(object_type="model_cost_map"): await self._check_and_reload_model_cost_map(prisma_client=prisma_client) + if self._should_load_db_object(object_type="sso_settings"): + await self._init_sso_settings_in_db(prisma_client=prisma_client) + + async def _init_sso_settings_in_db(self, prisma_client: PrismaClient): + """ + Initialize SSO settings from database into the router on startup. + """ + + try: + sso_settings = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) + if sso_settings is not None: + # Capitalize all keys in sso_settings dictionary + uppercase_sso_settings = {key.upper(): value for key, value in sso_settings.sso_settings.items()} + self._decrypt_and_set_db_env_variables(environment_variables=uppercase_sso_settings) + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_sso_settings_in_db - {}".format( + str(e) + ) + ) async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient): """ diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9cb9edc926..ce2fc2e475 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -580,4 +580,12 @@ model LiteLLM_SearchToolsTable { search_tool_info Json? created_at DateTime @default(now()) updated_at DateTime @updatedAt +} + +// SSO configuration table +model LiteLLM_SSOConfig { + id String @id @default("sso_config") + sso_settings Json + created_at DateTime @default(now()) + updated_at DateTime @updatedAt } \ No newline at end of file diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 1bb937f0b4..af68100910 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -390,40 +390,47 @@ async def update_default_team_settings(settings: DefaultTeamSSOParams): ) async def get_sso_settings(): """ - Get all SSO configuration settings from the environment variables. + Get all SSO configuration settings from the dedicated SSO table. Returns a structured object with values and descriptions for UI display. """ - import os - from litellm.proxy.proxy_server import proxy_config + from litellm.proxy.proxy_server import prisma_client, proxy_config - # Load existing config to get both environment variables and general settings - config = await proxy_config.get_config() - general_settings = config.get("general_settings", {}) or {} - environment_variables = config.get("environment_variables", {}) or {} + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected. Please connect a database."}, + ) - # Get user_email from general_settings - proxy_admin_email = general_settings.get("proxy_admin_email", None) + # Get SSO config from dedicated table + sso_db_record = await prisma_client.db.litellm_ssoconfig.find_unique( + where={"id": "sso_config"} + ) - # Helper function to get env var value (first from config, then from environment) - def get_env_value(env_var_name: str): - return environment_variables.get(env_var_name) or os.getenv(env_var_name) + # Initialize with defaults + sso_settings_dict = {} + + if sso_db_record and sso_db_record.sso_settings: + # Load settings from database + sso_settings_dict = dict(sso_db_record.sso_settings) + + decrypted_sso_settings_dict = proxy_config._decrypt_and_set_db_env_variables(environment_variables=sso_settings_dict) - # Get current environment variables for SSO + # Build SSO config with database values or environment fallback sso_config = SSOConfig( - google_client_id=get_env_value("GOOGLE_CLIENT_ID"), - google_client_secret=get_env_value("GOOGLE_CLIENT_SECRET"), - microsoft_client_id=get_env_value("MICROSOFT_CLIENT_ID"), - microsoft_client_secret=get_env_value("MICROSOFT_CLIENT_SECRET"), - microsoft_tenant=get_env_value("MICROSOFT_TENANT"), - generic_client_id=get_env_value("GENERIC_CLIENT_ID"), - generic_client_secret=get_env_value("GENERIC_CLIENT_SECRET"), - generic_authorization_endpoint=get_env_value("GENERIC_AUTHORIZATION_ENDPOINT"), - generic_token_endpoint=get_env_value("GENERIC_TOKEN_ENDPOINT"), - generic_userinfo_endpoint=get_env_value("GENERIC_USERINFO_ENDPOINT"), - proxy_base_url=get_env_value("PROXY_BASE_URL"), - user_email=proxy_admin_email, # Get from config instead of environment - ui_access_mode=general_settings.get("ui_access_mode", None), + google_client_id=decrypted_sso_settings_dict.get("google_client_id", None), + google_client_secret=decrypted_sso_settings_dict.get("google_client_secret", None), + microsoft_client_id=decrypted_sso_settings_dict.get("microsoft_client_id", None), + microsoft_client_secret=decrypted_sso_settings_dict.get("microsoft_client_secret", None), + microsoft_tenant=decrypted_sso_settings_dict.get("microsoft_tenant", None), + generic_client_id=decrypted_sso_settings_dict.get("generic_client_id", None), + generic_client_secret=decrypted_sso_settings_dict.get("generic_client_secret", None), + generic_authorization_endpoint=decrypted_sso_settings_dict.get("generic_authorization_endpoint", None), + generic_token_endpoint=decrypted_sso_settings_dict.get("generic_token_endpoint", None), + generic_userinfo_endpoint=decrypted_sso_settings_dict.get("generic_userinfo_endpoint", None), + proxy_base_url=decrypted_sso_settings_dict.get("proxy_base_url", None), + user_email=decrypted_sso_settings_dict.get("user_email"), + ui_access_mode=decrypted_sso_settings_dict.get("ui_access_mode"), ) # Get the schema for UI display @@ -460,11 +467,26 @@ async def get_sso_settings(): ) async def update_sso_settings(sso_config: SSOConfig): """ - Update SSO configuration by saving to both environment variables and config file. + Update SSO configuration by saving to the dedicated SSO table. """ import os + import json - from litellm.proxy.proxy_server import proxy_config + from litellm.proxy.proxy_server import prisma_client, store_model_in_db, proxy_config + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": "Database not connected. Please connect a database."}, + ) + + if store_model_in_db is not True: + raise HTTPException( + status_code=500, + detail={ + "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." + }, + ) # Update environment variables env_var_mapping = { @@ -495,39 +517,29 @@ async def update_sso_settings(sso_config: SSOConfig): # Update environment variables in config and in memory sso_data = sso_config.model_dump() for field_name, value in sso_data.items(): - if field_name == "user_email": - if value: - # Store user_email in general_settings instead of environment variables - config["general_settings"]["proxy_admin_email"] = value - else: - # Clear user_email if null/empty - config["general_settings"].pop("proxy_admin_email", None) - elif field_name == "ui_access_mode": - if value: - config["general_settings"]["ui_access_mode"] = value - else: - # Clear ui_access_mode if null/empty - config["general_settings"].pop("ui_access_mode", None) - elif field_name in env_var_mapping and value: + if field_name in env_var_mapping: env_var_name = env_var_mapping[field_name] - # Update in config - config["environment_variables"][env_var_name] = value - # Update in runtime environment - os.environ[env_var_name] = value - elif field_name in env_var_mapping: - # Clear environment variable if value is null/empty - env_var_name = env_var_mapping[field_name] - config["environment_variables"].pop(env_var_name, None) - os.environ.pop(env_var_name, None) + if value: + os.environ[env_var_name] = value + else: + # Clear environment variable if value is null/empty + os.environ.pop(env_var_name, None) - stored_config = config - if len(config["environment_variables"]) > 0: + encrypted_sso_data = proxy_config._encrypt_env_variables(environment_variables=sso_data) - stored_config["environment_variables"] = proxy_config._encrypt_env_variables( - environment_variables=config["environment_variables"] - ) - # Save the updated config - await proxy_config.save_config(new_config=stored_config) + # Save to dedicated SSO table + await prisma_client.db.litellm_ssoconfig.upsert( + where={"id": "sso_config"}, + data={ + "create": { + "id": "sso_config", + "sso_settings": json.dumps(encrypted_sso_data), + }, + "update": { + "sso_settings": json.dumps(encrypted_sso_data), + }, + }, + ) return { "message": "SSO settings updated successfully", diff --git a/schema.prisma b/schema.prisma index 9cb9edc926..ce2fc2e475 100644 --- a/schema.prisma +++ b/schema.prisma @@ -580,4 +580,12 @@ model LiteLLM_SearchToolsTable { search_tool_info Json? created_at DateTime @default(now()) updated_at DateTime @updatedAt +} + +// SSO configuration table +model LiteLLM_SSOConfig { + id String @id @default("sso_config") + sso_settings Json + created_at DateTime @default(now()) + updated_at DateTime @updatedAt } \ No newline at end of file diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5af7ac5f3a..e78a9689e8 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2288,3 +2288,158 @@ async def test_tag_cache_update_multiple_tags(): assert tag_updates["tag:tag2"]["spend"] == 25.0 +@pytest.mark.asyncio +async def test_init_sso_settings_in_db(): + """ + Test that _init_sso_settings_in_db properly loads SSO settings from database, + uppercases keys, and calls _decrypt_and_set_db_env_variables. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Test Case 1: SSO settings exist in database + mock_sso_config = MagicMock() + mock_sso_config.sso_settings = { + "google_client_id": "test-client-id", + "google_client_secret": "test-client-secret", + "microsoft_client_id": "ms-client-id", + "microsoft_client_secret": "ms-client-secret", + } + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=mock_sso_config + ) + + # Mock _decrypt_and_set_db_env_variables + with patch.object( + proxy_config, "_decrypt_and_set_db_env_variables" + ) as mock_decrypt_and_set: + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + + # Verify find_unique was called with correct parameters + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( + where={"id": "sso_config"} + ) + + # Verify _decrypt_and_set_db_env_variables was called with uppercased keys + mock_decrypt_and_set.assert_called_once() + call_args = mock_decrypt_and_set.call_args + uppercased_settings = call_args.kwargs["environment_variables"] + + # Verify all keys are uppercased + assert "GOOGLE_CLIENT_ID" in uppercased_settings + assert "GOOGLE_CLIENT_SECRET" in uppercased_settings + assert "MICROSOFT_CLIENT_ID" in uppercased_settings + assert "MICROSOFT_CLIENT_SECRET" in uppercased_settings + + # Verify values are preserved + assert uppercased_settings["GOOGLE_CLIENT_ID"] == "test-client-id" + assert uppercased_settings["GOOGLE_CLIENT_SECRET"] == "test-client-secret" + assert uppercased_settings["MICROSOFT_CLIENT_ID"] == "ms-client-id" + assert uppercased_settings["MICROSOFT_CLIENT_SECRET"] == "ms-client-secret" + + # Verify original lowercase keys are not present + assert "google_client_id" not in uppercased_settings + assert "microsoft_client_id" not in uppercased_settings + + +@pytest.mark.asyncio +async def test_init_sso_settings_in_db_no_settings(): + """ + Test that _init_sso_settings_in_db handles the case when no SSO settings exist in database. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Mock prisma client to return None (no SSO settings) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + + # Mock _decrypt_and_set_db_env_variables + with patch.object( + proxy_config, "_decrypt_and_set_db_env_variables" + ) as mock_decrypt_and_set: + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + + # Verify find_unique was called + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( + where={"id": "sso_config"} + ) + + # Verify _decrypt_and_set_db_env_variables was NOT called when no settings exist + mock_decrypt_and_set.assert_not_called() + + +@pytest.mark.asyncio +async def test_init_sso_settings_in_db_error_handling(): + """ + Test that _init_sso_settings_in_db handles database errors gracefully. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Mock prisma client to raise an exception + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + side_effect=Exception("Database connection error") + ) + + # The method should not raise an exception, it should log it instead + try: + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + # If we get here, the exception was handled properly + assert True + except Exception as e: + # The exception should be caught and logged, not propagated + pytest.fail(f"Exception should have been caught and logged, but was raised: {e}") + + +@pytest.mark.asyncio +async def test_init_sso_settings_in_db_empty_settings(): + """ + Test that _init_sso_settings_in_db handles empty SSO settings dictionary. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + # Mock SSO config with empty settings dictionary + mock_sso_config = MagicMock() + mock_sso_config.sso_settings = {} + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( + return_value=mock_sso_config + ) + + # Mock _decrypt_and_set_db_env_variables + with patch.object( + proxy_config, "_decrypt_and_set_db_env_variables" + ) as mock_decrypt_and_set: + await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) + + # Verify find_unique was called + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( + where={"id": "sso_config"} + ) + + # Verify _decrypt_and_set_db_env_variables was called with empty dict + mock_decrypt_and_set.assert_called_once() + call_args = mock_decrypt_and_set.call_args + uppercased_settings = call_args.kwargs["environment_variables"] + + # Verify empty dictionary + assert uppercased_settings == {} + diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 4b571691e4..eff59729fe 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -234,8 +234,31 @@ class TestProxySettingEndpoints: # Verify save_config was called exactly once assert mock_proxy_config["save_call_count"]() == 1 - def test_get_sso_settings(self, mock_proxy_config, mock_auth): - """Test getting the SSO settings""" + def test_get_sso_settings(self, mock_proxy_config, mock_auth, monkeypatch): + """Test getting the SSO settings from the dedicated database table""" + from unittest.mock import AsyncMock, MagicMock + + # Mock the prisma client with database record + # Note: Prisma returns Json fields as dicts (auto-parsed) + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.sso_settings = { + "google_client_id": "test_google_client_id", + "google_client_secret": "test_google_client_secret", + "microsoft_client_id": "test_microsoft_client_id", + "microsoft_client_secret": "test_microsoft_client_secret", + "proxy_base_url": "https://example.com", + "user_email": "admin@example.com", + } + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock decryption to return the values as-is (simulating decryption) + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", lambda environment_variables: environment_variables + ) + response = client.get("/get/sso_settings") assert response.status_code == 200 @@ -266,10 +289,29 @@ class TestProxySettingEndpoints: assert "properties" in data["field_schema"] assert "google_client_id" in data["field_schema"]["properties"] assert "description" in data["field_schema"]["properties"]["google_client_id"] + + # Verify find_unique was called with correct parameters + mock_prisma.db.litellm_ssoconfig.find_unique.assert_called_once() + call_args = mock_prisma.db.litellm_ssoconfig.find_unique.call_args + assert call_args.kwargs["where"]["id"] == "sso_config" def test_update_sso_settings(self, mock_proxy_config, mock_auth, monkeypatch): + """Test updating the SSO settings to the dedicated database table""" + import json + from unittest.mock import AsyncMock, MagicMock + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") - """Test updating the SSO settings""" + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock the prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock encryption to return values as-is + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables) + # New SSO settings to update new_sso_settings = { "google_client_id": "new_google_client_id", @@ -305,34 +347,44 @@ class TestProxySettingEndpoints: assert settings["proxy_base_url"] == new_sso_settings["proxy_base_url"] assert settings["user_email"] == new_sso_settings["user_email"] - # Verify the config was updated - updated_config = mock_proxy_config["config"] - assert ( - updated_config["environment_variables"]["GOOGLE_CLIENT_ID"] - != new_sso_settings["google_client_id"] - ) - assert ( - updated_config["environment_variables"]["GOOGLE_CLIENT_SECRET"] - != new_sso_settings["google_client_secret"] - ) - assert ( - updated_config["general_settings"]["proxy_admin_email"] - == new_sso_settings["user_email"] - ) - - # Verify save_config was called exactly once - assert mock_proxy_config["save_call_count"]() == 1 + # Verify upsert was called with correct parameters + assert mock_prisma.db.litellm_ssoconfig.upsert.called + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + + # Verify the upsert is using the correct ID + assert call_args.kwargs["where"]["id"] == "sso_config" + + # Verify the data structure for create and update + create_data = call_args.kwargs["data"]["create"] + update_data = call_args.kwargs["data"]["update"] + + assert create_data["id"] == "sso_config" + assert "sso_settings" in create_data + assert "sso_settings" in update_data + + # Verify the data is stored as JSON string (as per implementation) + # The encryption mock returns data as-is, so we verify structure + create_sso_settings = json.loads(create_data["sso_settings"]) + assert create_sso_settings["google_client_id"] == "new_google_client_id" def test_update_sso_settings_with_null_values_clears_env_vars( self, mock_proxy_config, mock_auth, monkeypatch ): - """Test that updating SSO settings with null values clears environment variables""" - monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + """Test that updating SSO settings with null values clears environment variables and updates database""" + import json + from unittest.mock import AsyncMock, MagicMock - # First, verify we have existing environment variables - initial_config = mock_proxy_config["config"] - assert "GOOGLE_CLIENT_ID" in initial_config["environment_variables"] - assert "MICROSOFT_CLIENT_ID" in initial_config["environment_variables"] + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock the prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock encryption to return values as-is + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables) # Set some initial environment variables for runtime testing monkeypatch.setenv("GOOGLE_CLIENT_ID", "test_existing_google_id") @@ -361,32 +413,38 @@ class TestProxySettingEndpoints: data = response.json() assert data["status"] == "success" - # Verify that environment variables were cleared from config - updated_config = mock_proxy_config["config"] - - # These should be removed from environment_variables - assert "GOOGLE_CLIENT_ID" not in updated_config["environment_variables"] - assert "GOOGLE_CLIENT_SECRET" not in updated_config["environment_variables"] - assert "MICROSOFT_CLIENT_ID" not in updated_config["environment_variables"] - assert "MICROSOFT_CLIENT_SECRET" not in updated_config["environment_variables"] - assert "MICROSOFT_TENANT" not in updated_config["environment_variables"] - assert "PROXY_BASE_URL" not in updated_config["environment_variables"] - # Verify that runtime environment variables were cleared assert "GOOGLE_CLIENT_ID" not in os.environ assert "MICROSOFT_CLIENT_ID" not in os.environ - # Verify user_email was cleared from general_settings - assert updated_config["general_settings"].get("proxy_admin_email") is None - - # Verify save_config was called - assert mock_proxy_config["save_call_count"]() == 1 + # Verify upsert was called with correct parameters + assert mock_prisma.db.litellm_ssoconfig.upsert.called + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + + # Verify null values are stored in database + create_data = call_args.kwargs["data"]["create"] + create_sso_settings = json.loads(create_data["sso_settings"]) + assert create_sso_settings["google_client_id"] is None + assert create_sso_settings["microsoft_client_id"] is None def test_update_sso_settings_with_empty_strings_clears_env_vars( self, mock_proxy_config, mock_auth, monkeypatch ): - """Test that updating SSO settings with empty strings also clears environment variables""" + """Test that updating SSO settings with empty strings also clears environment variables and updates database""" + import json + from unittest.mock import AsyncMock, MagicMock + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock the prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock encryption to return values as-is + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables) # Set some initial environment variables for runtime testing monkeypatch.setenv("GOOGLE_CLIENT_ID", "test_existing_google_id") @@ -407,28 +465,38 @@ class TestProxySettingEndpoints: data = response.json() assert data["status"] == "success" - # Verify that environment variables with empty strings were cleared from config - updated_config = mock_proxy_config["config"] - assert "GOOGLE_CLIENT_ID" not in updated_config["environment_variables"] - assert "GOOGLE_CLIENT_SECRET" not in updated_config["environment_variables"] - assert "MICROSOFT_CLIENT_SECRET" not in updated_config["environment_variables"] - assert "PROXY_BASE_URL" not in updated_config["environment_variables"] - # Verify that runtime environment variables were cleared assert "GOOGLE_CLIENT_ID" not in os.environ assert "MICROSOFT_CLIENT_SECRET" not in os.environ - # Verify user_email was cleared from general_settings - assert updated_config["general_settings"].get("proxy_admin_email") is None - - # Verify save_config was called - assert mock_proxy_config["save_call_count"]() == 1 + # Verify upsert was called with correct parameters + assert mock_prisma.db.litellm_ssoconfig.upsert.called + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + + # Verify empty strings are stored in database + create_data = call_args.kwargs["data"]["create"] + create_sso_settings = json.loads(create_data["sso_settings"]) + assert create_sso_settings["google_client_id"] == "" + assert create_sso_settings["microsoft_client_secret"] == "" def test_update_sso_settings_mixed_null_and_valid_values( self, mock_proxy_config, mock_auth, monkeypatch ): - """Test updating SSO settings with mix of null and valid values""" + """Test updating SSO settings with mix of null and valid values - verifies both env vars and database""" + import json + from unittest.mock import AsyncMock, MagicMock + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock the prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock encryption to return values as-is + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables) # Set some initial environment variables monkeypatch.setenv("GOOGLE_CLIENT_ID", "old_google_id") @@ -451,44 +519,43 @@ class TestProxySettingEndpoints: data = response.json() assert data["status"] == "success" - # Verify the config was updated correctly - updated_config = mock_proxy_config["config"] - - # Valid values should be set - assert ( - updated_config["environment_variables"]["GOOGLE_CLIENT_ID"] - != "new_google_client_id" - ) # Encrypted - assert ( - updated_config["environment_variables"]["MICROSOFT_CLIENT_SECRET"] - != "new_microsoft_secret" - ) # Encrypted - assert ( - updated_config["environment_variables"]["PROXY_BASE_URL"] - != "https://newproxy.com" - ) # Encrypted - - # Null values should be cleared - assert "GOOGLE_CLIENT_SECRET" not in updated_config["environment_variables"] - assert "MICROSOFT_CLIENT_ID" not in updated_config["environment_variables"] - # Verify runtime environment variables assert os.environ.get("GOOGLE_CLIENT_ID") == "new_google_client_id" assert os.environ.get("MICROSOFT_CLIENT_SECRET") == "new_microsoft_secret" assert "GOOGLE_CLIENT_SECRET" not in os.environ assert "MICROSOFT_CLIENT_ID" not in os.environ - # Verify user_email was cleared from general_settings - assert updated_config["general_settings"].get("proxy_admin_email") is None - - # Verify save_config was called - assert mock_proxy_config["save_call_count"]() == 1 + # Verify upsert was called with correct parameters + assert mock_prisma.db.litellm_ssoconfig.upsert.called + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + + # Verify the mixed values are stored correctly in database + create_data = call_args.kwargs["data"]["create"] + create_sso_settings = json.loads(create_data["sso_settings"]) + assert create_sso_settings["google_client_id"] == "new_google_client_id" + assert create_sso_settings["google_client_secret"] is None + assert create_sso_settings["microsoft_client_id"] is None + assert create_sso_settings["microsoft_client_secret"] == "new_microsoft_secret" + assert create_sso_settings["proxy_base_url"] == "https://newproxy.com" def test_update_sso_settings_ui_access_mode_handling( self, mock_proxy_config, mock_auth, monkeypatch ): - """Test that ui_access_mode is handled correctly in general_settings""" + """Test that ui_access_mode is handled correctly and stored in database""" + import json + from unittest.mock import AsyncMock, MagicMock + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock the prisma client + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock encryption to return values as-is + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", lambda environment_variables: environment_variables) # Test setting ui_access_mode sso_settings_with_ui_mode = { @@ -502,15 +569,13 @@ class TestProxySettingEndpoints: data = response.json() assert data["status"] == "success" - # Verify ui_access_mode was set in general_settings (not environment_variables) - updated_config = mock_proxy_config["config"] - assert updated_config["general_settings"]["ui_access_mode"] == "admin_only" - assert ( - updated_config["general_settings"]["proxy_admin_email"] == "admin@test.com" - ) - - # Verify ui_access_mode is NOT in environment_variables - assert "ui_access_mode" not in updated_config["environment_variables"] + # Verify upsert was called with correct data + assert mock_prisma.db.litellm_ssoconfig.upsert.called + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + create_data = call_args.kwargs["data"]["create"] + create_sso_settings = json.loads(create_data["sso_settings"]) + assert create_sso_settings["ui_access_mode"] == "admin_only" + assert create_sso_settings["user_email"] == "admin@test.com" # Test clearing ui_access_mode clear_ui_mode = {"ui_access_mode": None, "user_email": None} @@ -519,13 +584,13 @@ class TestProxySettingEndpoints: assert response.status_code == 200 - # Verify ui_access_mode and user_email were cleared - updated_config = mock_proxy_config["config"] - assert updated_config["general_settings"].get("ui_access_mode") is None - assert updated_config["general_settings"].get("proxy_admin_email") is None - - # Verify save_config was called twice (once for each update) - assert mock_proxy_config["save_call_count"]() == 2 + # Verify upsert was called again with null values + assert mock_prisma.db.litellm_ssoconfig.upsert.call_count == 2 + call_args = mock_prisma.db.litellm_ssoconfig.upsert.call_args + create_data = call_args.kwargs["data"]["create"] + create_sso_settings = json.loads(create_data["sso_settings"]) + assert create_sso_settings["ui_access_mode"] is None + assert create_sso_settings["user_email"] is None def test_get_ui_theme_settings(self, mock_proxy_config): """Test getting UI theme settings without authentication""" @@ -556,3 +621,184 @@ class TestProxySettingEndpoints: updated_config = mock_proxy_config["config"] assert "UI_LOGO_PATH" in updated_config["environment_variables"] assert mock_proxy_config["save_call_count"]() == 1 + + def test_get_sso_settings_from_database(self, mock_proxy_config, mock_auth, monkeypatch): + """Test getting SSO settings from the dedicated database table""" + import json + from unittest.mock import AsyncMock, MagicMock + + # Mock the prisma client + mock_prisma = MagicMock() + mock_db_record = MagicMock() + + # Simulate encrypted data from database + mock_sso_settings = { + "google_client_id": "encrypted_google_id", + "google_client_secret": "encrypted_google_secret", + "microsoft_client_id": "encrypted_microsoft_id", + "proxy_base_url": "encrypted_proxy_url", + } + + mock_db_record.sso_settings = mock_sso_settings + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_db_record) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock the decryption method to return decrypted values + def mock_decrypt_and_set(environment_variables): + return { + "google_client_id": "decrypted_google_id", + "google_client_secret": "decrypted_google_secret", + "microsoft_client_id": "decrypted_microsoft_id", + "proxy_base_url": "https://decrypted.example.com", + } + + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt_and_set + ) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + data = response.json() + + # Verify structure + assert "values" in data + assert "field_schema" in data + + # Verify decrypted values are returned + values = data["values"] + assert values["google_client_id"] == "decrypted_google_id" + assert values["google_client_secret"] == "decrypted_google_secret" + assert values["microsoft_client_id"] == "decrypted_microsoft_id" + assert values["proxy_base_url"] == "https://decrypted.example.com" + + def test_update_sso_settings_to_database(self, mock_proxy_config, mock_auth, monkeypatch): + """Test updating SSO settings saves to the dedicated database table""" + import json + from unittest.mock import AsyncMock, MagicMock + + monkeypatch.setenv("LITELLM_SALT_KEY", "test_salt_key") + + # Mock the prisma client + mock_prisma = MagicMock() + upsert_mock = AsyncMock() + mock_prisma.db.litellm_ssoconfig.upsert = upsert_mock + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Track what was encrypted + encrypted_data = {} + + def mock_encrypt(environment_variables): + # Simulate encryption by adding prefix + encrypted = { + k: f"encrypted_{v}" if v else v + for k, v in environment_variables.items() + } + encrypted_data.update(encrypted) + return encrypted + + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr(proxy_config, "_encrypt_env_variables", mock_encrypt) + + # New SSO settings to save + new_sso_settings = { + "google_client_id": "new_google_id", + "google_client_secret": "new_google_secret", + "microsoft_client_id": "new_microsoft_id", + "proxy_base_url": "https://new.example.com", + } + + response = client.patch("/update/sso_settings", json=new_sso_settings) + + assert response.status_code == 200 + data = response.json() + + assert data["status"] == "success" + assert data["settings"]["google_client_id"] == "new_google_id" + + # Verify upsert was called + assert upsert_mock.called + call_args = upsert_mock.call_args + + # Verify it's using the correct ID + assert call_args.kwargs["where"]["id"] == "sso_config" + + # Verify encrypted data was saved + create_data = call_args.kwargs["data"]["create"] + update_data = call_args.kwargs["data"]["update"] + + assert create_data["id"] == "sso_config" + # The sso_settings should be JSON string of encrypted data + assert "sso_settings" in create_data + assert "sso_settings" in update_data + + # Verify the encrypted data is correctly stored + create_sso_settings = json.loads(create_data["sso_settings"]) + assert create_sso_settings["google_client_id"] == "encrypted_new_google_id" + assert create_sso_settings["google_client_secret"] == "encrypted_new_google_secret" + assert create_sso_settings["proxy_base_url"] == "encrypted_https://new.example.com" + + def test_get_sso_settings_empty_database(self, mock_proxy_config, mock_auth, monkeypatch): + """Test getting SSO settings when database table is empty""" + from unittest.mock import AsyncMock, MagicMock + + # Mock the prisma client to return None (no record found) + mock_prisma = MagicMock() + mock_prisma.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + # Mock the decryption method + def mock_decrypt_and_set(environment_variables): + # Should receive empty dict + return environment_variables + + from litellm.proxy.proxy_server import proxy_config + monkeypatch.setattr( + proxy_config, "_decrypt_and_set_db_env_variables", mock_decrypt_and_set + ) + + response = client.get("/get/sso_settings") + + assert response.status_code == 200 + data = response.json() + + # Verify structure is still correct with empty values + assert "values" in data + assert "field_schema" in data + + # All values should be None + values = data["values"] + assert values.get("google_client_id") is None + assert values.get("google_client_secret") is None + assert values.get("microsoft_client_id") is None + + def test_update_sso_settings_no_database_connection(self, mock_proxy_config, mock_auth, monkeypatch): + """Test updating SSO settings when database is not connected""" + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + new_sso_settings = { + "google_client_id": "new_google_id", + } + + response = client.patch("/update/sso_settings", json=new_sso_settings) + + assert response.status_code == 500 + data = response.json() + assert "error" in data["detail"] + assert "Database not connected" in data["detail"]["error"] + + def test_get_sso_settings_no_database_connection(self, mock_proxy_config, mock_auth, monkeypatch): + """Test getting SSO settings when database is not connected""" + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + response = client.get("/get/sso_settings") + + assert response.status_code == 500 + data = response.json() + assert "error" in data["detail"] + assert "Database not connected" in data["detail"]["error"]