tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 10:24:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

SSO - Allow passing additional headers + Spend Tags - automatically track spend by user agent (allows cost tracking for claude code) (#11781)

Browse Source 
* 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
This commit is contained in:
Krish Dholakia
2025-06-16 21:53:40 -07:00
committed by GitHub
parent 80501b8268
commit 87ae2cf3d7
6 changed files with 250 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/proxy/admin_ui_sso.md
+1
View File
@@ -50,6 +50,7 @@ GENERIC_AUTHORIZATION_ENDPOINT = "<your-okta-domain>/authorize" # https://dev-2k
GENERIC_TOKEN_ENDPOINT = "<your-okta-domain>/token" # https://dev-2kqkcd6lx6kdkuzt.us.auth0.com/oauth/token
GENERIC_USERINFO_ENDPOINT = "<your-okta-domain>/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 `<YOUR-OKTA-DOMAIN>/.well-known/openid-configuration`
docs/my-website/docs/proxy/config_settings.md
+1
View File
@@ -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 <token>, 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
litellm/proxy/litellm_pre_call_utils.py
+17
View File
@@ -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"]
litellm/proxy/management_endpoints/ui_sso.py
+23 -3
View File
@@ -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 <token>, 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 {}
tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+127
View File
@@ -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
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
+81
View File
@@ -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