tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-06-17 12:48:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: fix ci/cd + handle oidc jwt tokens

Browse Source
This commit is contained in:
Krrish Dholakia
2026-03-30 16:12:03 -07:00
parent e8c860d450
commit 4c00a14ce0
6 changed files with 446 additions and 294 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/auth/handle_jwt.py
+6 -2
View File
@@ -1397,8 +1397,12 @@ class JWTAuthManager:
request_headers: Optional[dict] = None,
) -> JWTAuthBuilderResult:
"""Main authentication and authorization builder"""
# Check if OIDC UserInfo endpoint is enabled
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
# Check if OIDC UserInfo endpoint is enabled, but fall back to standard
# JWT auth if the token itself is a well-formed JWT (3-part structure).
if (
jwt_handler.litellm_jwtauth.oidc_userinfo_enabled
and not jwt_handler.is_jwt(token=api_key)
):
verbose_proxy_logger.debug(
"OIDC UserInfo is enabled. Fetching user info from UserInfo endpoint."
)
litellm/proxy/auth/user_api_key_auth.py
+4 -1
View File
@@ -686,7 +686,10 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if jwt_handler.litellm_jwtauth.virtual_key_claim_field is not None:
# Decode JWT to get claims without running full auth_builder
jwt_claims: Optional[dict]
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled:
if (
jwt_handler.litellm_jwtauth.oidc_userinfo_enabled
and not jwt_handler.is_jwt(token=api_key)
):
jwt_claims = await jwt_handler.get_oidc_userinfo(token=api_key)
else:
jwt_claims = await jwt_handler.auth_jwt(token=api_key)
tests/litellm_utils_tests/test_azure_ai_anthropic_token_counter.py
+6 -8
View File
@@ -26,22 +26,20 @@ class TestAzureAIAnthropicTokenCounter(BaseTokenCounterTest):
return AzureAIAnthropicTokenCounter()
def get_test_model(self) -> str:
return "claude-3-5-sonnet"
return "claude-sonnet-4-6"
def get_test_messages(self) -> List[Dict[str, Any]]:
return [
{"role": "user", "content": "Hello, how are you today?"}
]
return [{"role": "user", "content": "Hello, how are you today?"}]
def get_deployment_config(self) -> Dict[str, Any]:
api_key = os.getenv("AZURE_AI_API_KEY")
api_base = os.getenv("AZURE_AI_API_BASE")
api_key = os.getenv("AZURE_ANTHROPIC_API_KEY")
api_base = os.getenv("AZURE_AI_SWEDEN_API_BASE")
if not api_key:
pytest.skip("AZURE_AI_API_KEY not set")
if not api_base:
pytest.skip("AZURE_AI_API_BASE not set")
return {
"litellm_params": {
"api_key": api_key,
tests/local_testing/test_streaming.py
+1 -79
View File
@@ -471,85 +471,6 @@ def test_completion_azure_stream():
# test_completion_azure_stream()
@pytest.mark.skip("Skipping predibase streaming test - ran out of credits")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_completion_predibase_streaming(sync_mode):
try:
litellm.set_verbose = True
litellm._turn_on_debug()
if sync_mode:
response = completion(
model="predibase/llama-3-8b-instruct",
timeout=5,
tenant_id="c4768f95",
max_tokens=10,
api_base="https://serving.app.predibase.com",
api_key=os.getenv("PREDIBASE_API_KEY"),
messages=[{"role": "user", "content": "What is the meaning of life?"}],
stream=True,
)
complete_response = ""
for idx, init_chunk in enumerate(response):
chunk, finished = streaming_format_tests(idx, init_chunk)
complete_response += chunk
custom_llm_provider = init_chunk._hidden_params["custom_llm_provider"]
print(f"custom_llm_provider: {custom_llm_provider}")
assert custom_llm_provider == "predibase"
if finished:
assert isinstance(
init_chunk.choices[0], litellm.utils.StreamingChoices
)
break
if complete_response.strip() == "":
raise Exception("Empty response received")
else:
response = await litellm.acompletion(
model="predibase/llama-3-8b-instruct",
tenant_id="c4768f95",
timeout=5,
max_tokens=10,
api_base="https://serving.app.predibase.com",
api_key=os.getenv("PREDIBASE_API_KEY"),
messages=[{"role": "user", "content": "What is the meaning of life?"}],
stream=True,
)
# await response
complete_response = ""
idx = 0
async for init_chunk in response:
chunk, finished = streaming_format_tests(idx, init_chunk)
complete_response += chunk
custom_llm_provider = init_chunk._hidden_params["custom_llm_provider"]
print(f"custom_llm_provider: {custom_llm_provider}")
assert custom_llm_provider == "predibase"
idx += 1
if finished:
assert isinstance(
init_chunk.choices[0], litellm.utils.StreamingChoices
)
break
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"complete_response: {complete_response}")
except litellm.Timeout:
pass
except litellm.InternalServerError:
pass
except litellm.ServiceUnavailableError:
pass
except litellm.APIConnectionError:
pass
except Exception as e:
print("ERROR class", e.__class__)
print("ERROR message", e)
print("ERROR traceback", traceback.format_exc())
pytest.fail(f"Error occurred: {e}")
def test_completion_azure_function_calling_stream():
@@ -1143,6 +1064,7 @@ def test_vertex_ai_stream(provider):
# test_completion_vertexai_stream_bad_key()
@pytest.mark.skip(reason="Replicate extremely flaky.")
@pytest.mark.parametrize("sync_mode", [False, True])
@pytest.mark.asyncio
async def test_completion_replicate_llama3_streaming(sync_mode):
tests/proxy_unit_tests/test_jwt_key_mapping.py
+75
View File
@@ -135,6 +135,81 @@ async def test_jwt_to_virtual_key_mapping_no_mapping():
prisma_client.db.litellm_jwtkeymapping.find_first.assert_not_called()
# ──────────────────────────────────────────────
# Tests: OIDC / JWT routing in user_api_key_auth
# ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_virtual_key_mapping_oidc_enabled_jwt_token_uses_auth_jwt():
"""
Regression test for the is_jwt routing fix in user_api_key_auth.py.
When oidc_userinfo_enabled=True and virtual_key_claim_field is set, but
the token is a well-formed JWT (3-part header.payload.sig), the virtual-key
claim lookup must call auth_jwt — not get_oidc_userinfo.
"""
# Three-part token: is_jwt() returns True
api_key = "eyJhbGciOiJSUzI1NiJ9.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20ifQ.sig"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
oidc_userinfo_enabled=True,
virtual_key_claim_field="email",
)
# Confirm our fixture token is treated as a JWT
assert jwt_handler.is_jwt(token=api_key) is True
auth_jwt_mock = AsyncMock(return_value={"email": "user@example.com", "sub": "123"})
oidc_userinfo_mock = AsyncMock(return_value={"email": "user@example.com"})
# Simulate the routing condition from user_api_key_auth.py
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(
token=api_key
):
jwt_claims = await oidc_userinfo_mock(token=api_key)
else:
jwt_claims = await auth_jwt_mock(token=api_key)
auth_jwt_mock.assert_called_once_with(token=api_key)
oidc_userinfo_mock.assert_not_called()
assert jwt_claims["email"] == "user@example.com"
@pytest.mark.asyncio
async def test_virtual_key_mapping_oidc_enabled_opaque_token_uses_oidc_userinfo():
"""
Complement of the test above: when oidc_userinfo_enabled=True and the token
is an opaque access token (not a JWT), the virtual-key claim lookup must
call get_oidc_userinfo — not auth_jwt.
"""
# Opaque token: no dots → is_jwt() returns False
api_key = "some_opaque_access_token_with_no_dots"
jwt_handler = JWTHandler()
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
oidc_userinfo_enabled=True,
virtual_key_claim_field="email",
)
assert jwt_handler.is_jwt(token=api_key) is False
auth_jwt_mock = AsyncMock(return_value={"email": "user@example.com"})
oidc_userinfo_mock = AsyncMock(return_value={"email": "user@example.com", "sub": "123"})
if jwt_handler.litellm_jwtauth.oidc_userinfo_enabled and not jwt_handler.is_jwt(
token=api_key
):
jwt_claims = await oidc_userinfo_mock(token=api_key)
else:
jwt_claims = await auth_jwt_mock(token=api_key)
oidc_userinfo_mock.assert_called_once_with(token=api_key)
auth_jwt_mock.assert_not_called()
assert jwt_claims["sub"] == "123"
# ──────────────────────────────────────────────
# Tests: _to_response redacts hashed token
# ──────────────────────────────────────────────
tests/test_litellm/proxy/auth/test_handle_jwt.py
+354 -204
View File
File diff suppressed because it is too large Load Diff