From ed3b056bc8fa2cfb5e883de2108acc438310d2c9 Mon Sep 17 00:00:00 2001 From: harish-berri Date: Mon, 4 May 2026 18:17:31 +0000 Subject: [PATCH] Implement normalize_nonempty_secret_str function to trim whitespace from secrets and treat empty values as unset. Update proxy_server to use this function for Grafana credentials. Enhance tests to validate the new normalization behavior. --- litellm/proxy/proxy_server.py | 10 ++- litellm/secret_managers/main.py | 13 +++ tests/test_litellm/proxy/test_pyroscope.py | 83 +++++++++++++++++-- .../test_secret_managers_main.py | 16 +++- 4 files changed, 111 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4b663904b5..a5905765c6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -476,6 +476,7 @@ from litellm.secret_managers.main import ( get_secret, get_secret_bool, get_secret_str, + normalize_nonempty_secret_str, str_to_bool, ) from litellm.types.integrations.slack_alerting import SlackAlertingArgs @@ -7516,11 +7517,12 @@ class ProxyStartupEvent: if env_name: tags["environment"] = env_name sample_rate_env = os.getenv("PYROSCOPE_SAMPLE_RATE") - grafana_pyroscope_user = get_secret_str( - "PYROSCOPE_GRAFANA_USER", default_value=None + + grafana_pyroscope_user = normalize_nonempty_secret_str( + get_secret_str("PYROSCOPE_GRAFANA_USER", default_value=None) ) - grafana_api_token = get_secret_str( - "PYROSCOPE_GRAFANA_API_TOKEN", default_value=None + grafana_api_token = normalize_nonempty_secret_str( + get_secret_str("PYROSCOPE_GRAFANA_API_TOKEN", default_value=None) ) if grafana_api_token and not grafana_pyroscope_user: raise ValueError( diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index a560f5222b..5c31d81f04 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -120,6 +120,19 @@ def get_secret_str( return value +def normalize_nonempty_secret_str(val: Optional[str]) -> Optional[str]: + """ + Strip whitespace and treat None, '', and whitespace-only strings as unset. + + Use when pairing secrets (mutual exclusion, optional auth) so whitespace-only + values do not count as present. + """ + if val is None: + return None + stripped = val.strip() + return stripped if stripped else None + + def get_secret_bool( secret_name: str, default_value: Optional[bool] = None, diff --git a/tests/test_litellm/proxy/test_pyroscope.py b/tests/test_litellm/proxy/test_pyroscope.py index fc2499101c..0ac3e11061 100644 --- a/tests/test_litellm/proxy/test_pyroscope.py +++ b/tests/test_litellm/proxy/test_pyroscope.py @@ -2,11 +2,13 @@ import os import sys +from typing import Optional from unittest.mock import MagicMock, patch import pytest from litellm.proxy.proxy_server import ProxyStartupEvent +from litellm.secret_managers.main import get_secret_str as real_get_secret_str def _mock_pyroscope_module(): @@ -16,6 +18,22 @@ def _mock_pyroscope_module(): return m +def _patch_pyroscope_grafana_secrets(user: Optional[str], token: Optional[str]): + """Patch proxy_server.get_secret_str for Grafana keys; defer other secrets to the real helper.""" + + def side_effect(secret_name: str, default_value=None): + if secret_name == "PYROSCOPE_GRAFANA_USER": + return user + if secret_name == "PYROSCOPE_GRAFANA_API_TOKEN": + return token + return real_get_secret_str(secret_name, default_value) + + return patch( + "litellm.proxy.proxy_server.get_secret_str", + side_effect=side_effect, + ) + + def test_init_pyroscope_returns_cleanly_when_disabled(): """When LITELLM_ENABLE_PYROSCOPE is false, _init_pyroscope returns without error.""" with ( @@ -184,6 +202,7 @@ def test_init_pyroscope_configures_grafana_cloud_basic_auth(): "litellm.proxy.proxy_server.get_secret_bool", return_value=True, ), + _patch_pyroscope_grafana_secrets("123456", "glc_test_token"), patch.dict( sys.modules, {"pyroscope": mock_pyroscope}, @@ -194,8 +213,6 @@ def test_init_pyroscope_configures_grafana_cloud_basic_auth(): "LITELLM_ENABLE_PYROSCOPE": "true", "PYROSCOPE_APP_NAME": "myapp", "PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net", - "PYROSCOPE_GRAFANA_USER": "123456", - "PYROSCOPE_GRAFANA_API_TOKEN": "glc_test_token", }, clear=False, ), @@ -214,6 +231,7 @@ def test_init_pyroscope_raises_when_grafana_token_missing_user(): "litellm.proxy.proxy_server.get_secret_bool", return_value=True, ), + _patch_pyroscope_grafana_secrets("", "glc_test_token"), patch.dict( sys.modules, {"pyroscope": mock_pyroscope}, @@ -224,8 +242,6 @@ def test_init_pyroscope_raises_when_grafana_token_missing_user(): "LITELLM_ENABLE_PYROSCOPE": "true", "PYROSCOPE_APP_NAME": "myapp", "PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net", - "PYROSCOPE_GRAFANA_USER": "", - "PYROSCOPE_GRAFANA_API_TOKEN": "glc_test_token", }, clear=False, ), @@ -242,6 +258,7 @@ def test_init_pyroscope_raises_when_grafana_user_missing_token(): "litellm.proxy.proxy_server.get_secret_bool", return_value=True, ), + _patch_pyroscope_grafana_secrets("123456", ""), patch.dict( sys.modules, {"pyroscope": mock_pyroscope}, @@ -252,11 +269,65 @@ def test_init_pyroscope_raises_when_grafana_user_missing_token(): "LITELLM_ENABLE_PYROSCOPE": "true", "PYROSCOPE_APP_NAME": "myapp", "PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net", - "PYROSCOPE_GRAFANA_USER": "123456", - "PYROSCOPE_GRAFANA_API_TOKEN": "", }, clear=False, ), ): with pytest.raises(ValueError, match="PYROSCOPE_GRAFANA_API_TOKEN"): ProxyStartupEvent._init_pyroscope() + + +def test_init_pyroscope_raises_when_grafana_user_whitespace_only_with_token(): + """Whitespace-only user id does not satisfy Grafana mutual exclusion.""" + mock_pyroscope = _mock_pyroscope_module() + with ( + patch( + "litellm.proxy.proxy_server.get_secret_bool", + return_value=True, + ), + _patch_pyroscope_grafana_secrets(" \t", "glc_test_token"), + patch.dict( + sys.modules, + {"pyroscope": mock_pyroscope}, + ), + patch.dict( + os.environ, + { + "LITELLM_ENABLE_PYROSCOPE": "true", + "PYROSCOPE_APP_NAME": "myapp", + "PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net", + }, + clear=False, + ), + ): + with pytest.raises(ValueError, match="PYROSCOPE_GRAFANA_USER"): + ProxyStartupEvent._init_pyroscope() + + +def test_init_pyroscope_strips_grafana_credentials_for_basic_auth(): + """Leading/trailing whitespace on Grafana secrets is trimmed before configure.""" + mock_pyroscope = _mock_pyroscope_module() + with ( + patch( + "litellm.proxy.proxy_server.get_secret_bool", + return_value=True, + ), + _patch_pyroscope_grafana_secrets(" 123456 ", " glc_test_token\n"), + patch.dict( + sys.modules, + {"pyroscope": mock_pyroscope}, + ), + patch.dict( + os.environ, + { + "LITELLM_ENABLE_PYROSCOPE": "true", + "PYROSCOPE_APP_NAME": "myapp", + "PYROSCOPE_SERVER_ADDRESS": "https://profiles-prod-001.grafana.net", + }, + clear=False, + ), + ): + ProxyStartupEvent._init_pyroscope() + call_kw = mock_pyroscope.configure.call_args[1] + assert call_kw["basic_auth_username"] == "123456" + assert call_kw["basic_auth_password"] == "glc_test_token" diff --git a/tests/test_litellm/secret_managers/test_secret_managers_main.py b/tests/test_litellm/secret_managers/test_secret_managers_main.py index 6d0e33b3e2..b406231804 100644 --- a/tests/test_litellm/secret_managers/test_secret_managers_main.py +++ b/tests/test_litellm/secret_managers/test_secret_managers_main.py @@ -4,7 +4,7 @@ from unittest.mock import Mock, patch import pytest -from litellm.secret_managers.main import get_secret +from litellm.secret_managers.main import get_secret, normalize_nonempty_secret_str # Set up logging for debugging logging.basicConfig(level=logging.DEBUG) @@ -253,3 +253,17 @@ def test_unsupported_oidc_provider(): with pytest.raises(ValueError, match="Unsupported OIDC provider"): get_secret(secret_name) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, None), + ("", None), + (" \t\n", None), + ("abc", "abc"), + (" xyz ", "xyz"), + ], +) +def test_normalize_nonempty_secret_str(raw, expected): + assert normalize_nonempty_secret_str(raw) == expected