From 87ae2cf3d72c6071d8a4f03cdd7f4f98914d2f93 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Mon, 16 Jun 2025 21:53:40 -0700 Subject: [PATCH] SSO - Allow passing additional headers + Spend Tags - automatically track spend by user agent (allows cost tracking for claude code) (#11781) * feat(ui_sso.py): allow admin to specify additional headers for sso provider some sso providers require special headers to return a json response * test(test_ui_sso.py): add unit tests to ensure custom headers are respect3ed * docs(config_settings.md): document new header param * fix(litellm_pre_call_utils.py): add spend tag tracking by user agent allows checking spend for cli tools like claude code * feat(litellm_pre_call_utils.py): track spend by user agent part if user agent contains "/" allows tracking spend across user agent versions Better cost tracking for claude cod * test(test_litellm_pre_call_utils.py): add testing for pre call utils, user agent parsing * fix: fix linting check --- docs/my-website/docs/proxy/admin_ui_sso.md | 1 + docs/my-website/docs/proxy/config_settings.md | 1 + litellm/proxy/litellm_pre_call_utils.py | 17 +++ litellm/proxy/management_endpoints/ui_sso.py | 26 +++- .../proxy/management_endpoints/test_ui_sso.py | 127 ++++++++++++++++++ .../proxy/test_litellm_pre_call_utils.py | 81 +++++++++++ 6 files changed, 250 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/proxy/admin_ui_sso.md b/docs/my-website/docs/proxy/admin_ui_sso.md index b8aa152ed8..c1a641b4ff 100644 --- a/docs/my-website/docs/proxy/admin_ui_sso.md +++ b/docs/my-website/docs/proxy/admin_ui_sso.md @@ -50,6 +50,7 @@ GENERIC_AUTHORIZATION_ENDPOINT = "/authorize" # https://dev-2k GENERIC_TOKEN_ENDPOINT = "/token" # https://dev-2kqkcd6lx6kdkuzt.us.auth0.com/oauth/token GENERIC_USERINFO_ENDPOINT = "/userinfo" # https://dev-2kqkcd6lx6kdkuzt.us.auth0.com/userinfo GENERIC_CLIENT_STATE = "random-string" # [OPTIONAL] REQUIRED BY OKTA, if not set random state value is generated +GENERIC_SSO_HEADERS = "Content-Type=application/json, X-Custom-Header=custom-value" # [OPTIONAL] Comma-separated list of additional headers to add to the request - e.g. Content-Type=application/json, etc. ``` You can get your domain specific auth/token/userinfo endpoints at `/.well-known/openid-configuration` diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index cebb5fa712..48c48708c5 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -446,6 +446,7 @@ router_settings: | GENERIC_CLIENT_ID | Client ID for generic OAuth providers | GENERIC_CLIENT_SECRET | Client secret for generic OAuth providers | GENERIC_CLIENT_STATE | State parameter for generic client authentication +| GENERIC_SSO_HEADERS | Comma-separated list of additional headers to add to the request - e.g. Authorization=Bearer , Content-Type=application/json, etc. | GENERIC_INCLUDE_CLIENT_ID | Include client ID in requests for OAuth | GENERIC_SCOPE | Scope settings for generic OAuth providers | GENERIC_TOKEN_ENDPOINT | Token endpoint for generic OAuth providers diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d1b57d6401..c694e7f328 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -483,6 +483,23 @@ class LiteLLMProxyRequestSetup: tags = [tag.strip() for tag in _tags] elif isinstance(headers["x-litellm-tags"], list): tags = headers["x-litellm-tags"] + if "user-agent" in headers: + """ + Allow tracking spend by cli tools like Claude Code - e.g. "claude-cli/1.0.25 (external, cli)" + """ + # add user-agent to tags + if tags is None: + tags = [] + user_agent = headers["user-agent"] + if user_agent is not None: + user_agent_part: Optional[str] = None + if "/" in user_agent: + user_agent_part = user_agent.split("/")[ + 0 + ] # extract "claude-cli" - enables spend tracking acrosss versions + if user_agent_part is not None: + tags.append(user_agent_part) + tags.append(user_agent) # append full user-agent # Check request body for tags if "tags" in data and isinstance(data["tags"], list): tags = data["tags"] diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 36ab13c693..03f740214a 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -260,9 +260,29 @@ async def get_generic_sso_response( scope=generic_scope, ) verbose_proxy_logger.debug("calling generic_sso.verify_and_process") - result = await generic_sso.verify_and_process( - request, params={"include_client_id": generic_include_client_id} - ) + additional_generic_sso_headers = os.getenv( + "GENERIC_SSO_HEADERS", None + ) # Comma-separated list of headers to add to the request - e.g. Authorization=Bearer , Content-Type=application/json, etc. + additional_generic_sso_headers_dict = {} + if additional_generic_sso_headers is not None: + additional_generic_sso_headers_split = additional_generic_sso_headers.split(",") + for header in additional_generic_sso_headers_split: + header = header.strip() + if header: + key, value = header.split("=") + additional_generic_sso_headers_dict[key] = value + + try: + result = await generic_sso.verify_and_process( + request, + params={"include_client_id": generic_include_client_id}, + headers=additional_generic_sso_headers_dict, + ) + except Exception as e: + verbose_proxy_logger.exception( + f"Error verifying and processing generic SSO: {e}. Passed in headers: {additional_generic_sso_headers_dict}" + ) + raise e verbose_proxy_logger.debug("generic result: %s", result) return result or {} diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 47684e64d2..8430fa6ee3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -690,3 +690,130 @@ async def test_check_and_update_if_proxy_admin_id_already_admin(): # Assert assert updated_role == LitellmUserRoles.PROXY_ADMIN.value mock_prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_generic_sso_response_with_additional_headers(): + """ + Test that GENERIC_SSO_HEADERS environment variable is correctly processed + and passed to generic_sso.verify_and_process + """ + from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response + + # Arrange + mock_request = MagicMock(spec=Request) + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + generic_client_id = "test_client_id" + redirect_url = "http://test.com/callback" + + # Mock response from verify_and_process + mock_sso_response = { + "sub": "test_user_123", + "email": "test@example.com", + "preferred_username": "testuser", + } + + # Set up environment variables including GENERIC_SSO_HEADERS + test_env_vars = { + "GENERIC_CLIENT_SECRET": "test_secret", + "GENERIC_AUTHORIZATION_ENDPOINT": "https://auth.example.com/auth", + "GENERIC_TOKEN_ENDPOINT": "https://auth.example.com/token", + "GENERIC_USERINFO_ENDPOINT": "https://auth.example.com/userinfo", + "GENERIC_SSO_HEADERS": "Authorization=Bearer token123, Content-Type=application/json, X-Custom-Header=custom-value", + } + + # Expected headers dictionary + expected_headers = { + "Authorization": "Bearer token123", + "Content-Type": "application/json", + "X-Custom-Header": "custom-value", + } + + # Mock the SSO provider and its methods + mock_sso_instance = MagicMock() + mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) + + mock_sso_class = MagicMock(return_value=mock_sso_instance) + + with patch.dict(os.environ, test_env_vars): + with patch("fastapi_sso.sso.base.DiscoveryDocument") as mock_discovery: + with patch( + "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class + ) as mock_create_provider: + # Act + result = await get_generic_sso_response( + request=mock_request, + jwt_handler=mock_jwt_handler, + generic_client_id=generic_client_id, + redirect_url=redirect_url, + ) + + # Assert + # Verify verify_and_process was called with the correct headers + mock_sso_instance.verify_and_process.assert_called_once_with( + mock_request, + params={"include_client_id": False}, + headers=expected_headers, + ) + + # Verify the result is returned correctly + assert result == mock_sso_response + + +@pytest.mark.asyncio +async def test_get_generic_sso_response_with_empty_headers(): + """ + Test that when GENERIC_SSO_HEADERS is not set, an empty headers dict is passed + """ + from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response + + # Arrange + mock_request = MagicMock(spec=Request) + mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler.get_team_ids_from_jwt.return_value = [] + + generic_client_id = "test_client_id" + redirect_url = "http://test.com/callback" + + mock_sso_response = { + "sub": "test_user_123", + "email": "test@example.com", + "preferred_username": "testuser", + } + + # Set up environment variables without GENERIC_SSO_HEADERS + test_env_vars = { + "GENERIC_CLIENT_SECRET": "test_secret", + "GENERIC_AUTHORIZATION_ENDPOINT": "https://auth.example.com/auth", + "GENERIC_TOKEN_ENDPOINT": "https://auth.example.com/token", + "GENERIC_USERINFO_ENDPOINT": "https://auth.example.com/userinfo", + } + + # Mock the SSO provider and its methods + mock_sso_instance = MagicMock() + mock_sso_instance.verify_and_process = AsyncMock(return_value=mock_sso_response) + + mock_sso_class = MagicMock(return_value=mock_sso_instance) + + with patch.dict(os.environ, test_env_vars): + with patch("fastapi_sso.sso.base.DiscoveryDocument") as mock_discovery: + with patch( + "fastapi_sso.sso.generic.create_provider", return_value=mock_sso_class + ) as mock_create_provider: + # Act + result = await get_generic_sso_response( + request=mock_request, + jwt_handler=mock_jwt_handler, + generic_client_id=generic_client_id, + redirect_url=redirect_url, + ) + + # Assert + # Verify verify_and_process was called with empty headers dict + mock_sso_instance.verify_and_process.assert_called_once_with( + mock_request, params={"include_client_id": False}, headers={} + ) + + assert result == mock_sso_response diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 19dadf17b0..beeca265fa 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -10,6 +10,7 @@ from fastapi import Request from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, _get_enforced_params, add_litellm_data_to_request, check_if_token_is_service_account, @@ -215,3 +216,83 @@ async def test_add_litellm_data_to_request_audio_transcription_multipart(): "jobID:214590dsff09fds", "taskName:run_page_classification", ] + + +def test_add_request_tag_to_metadata_user_agent_parsing(): + """ + Test that user agent parsing works correctly in add_request_tag_to_metadata + """ + + # Test case 1: User agent with version (contains "/") + headers_with_version = {"user-agent": "claude-cli/1.0.25 (external, cli)"} + data = {} + result = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( + llm_router=None, + headers=headers_with_version, + data=data, + ) + expected_tags = ["claude-cli", "claude-cli/1.0.25 (external, cli)"] + assert result == expected_tags + + # Test case 2: User agent without version (no "/") + headers_without_version = {"user-agent": "my-custom-client"} + data = {} + result = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( + llm_router=None, + headers=headers_without_version, + data=data, + ) + expected_tags = ["my-custom-client"] + assert result == expected_tags + + # Test case 3: No user agent header + headers_no_user_agent = {} + data = {} + result = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( + llm_router=None, + headers=headers_no_user_agent, + data=data, + ) + assert result is None + + # Test case 4: User agent with existing x-litellm-tags + headers_with_existing_tags = { + "user-agent": "postman/7.36.1", + "x-litellm-tags": "existing-tag1, existing-tag2", + } + data = {} + result = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( + llm_router=None, + headers=headers_with_existing_tags, + data=data, + ) + expected_tags = ["existing-tag1", "existing-tag2", "postman", "postman/7.36.1"] + assert result == expected_tags + + # Test case 5: User agent with tags in request body (body tags override header tags) + headers_with_user_agent = {"user-agent": "curl/7.68.0"} + data = {"tags": ["body-tag1", "body-tag2"]} + result = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( + llm_router=None, + headers=headers_with_user_agent, + data=data, + ) + # When tags exist in data, they override everything else + expected_tags = ["body-tag1", "body-tag2"] + assert result == expected_tags + + # Test case 6: Complex user agent string + headers_complex = { + "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36" + } + data = {} + result = LiteLLMProxyRequestSetup.add_request_tag_to_metadata( + llm_router=None, + headers=headers_complex, + data=data, + ) + expected_tags = [ + "Mozilla", + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", + ] + assert result == expected_tags