From 2880cb45a2110b4707be23c54cd5680d20be1739 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 18 Nov 2025 17:36:38 -0800 Subject: [PATCH] [Fix] AI Gateway Auth - Ensure Team Tags works when using JWT Auth (#16797) * fix setting team_metadata * test_team_metadata_with_tags_flows_through_jwt_auth --- litellm/proxy/auth/user_api_key_auth.py | 1 + .../proxy/auth/test_user_api_key_auth.py | 49 ++++++++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index c9589cd7746..05d0006300b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -584,6 +584,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_membership is not None else None ), + team_metadata=team_object.metadata if team_object is not None else None, ) # run through common checks _ = await common_checks( diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index baf8bcfea35..9905ae5d355 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -13,8 +13,8 @@ from unittest.mock import MagicMock import pytest -from litellm.proxy.auth.user_api_key_auth import get_api_key from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.auth.user_api_key_auth import get_api_key def test_get_api_key(): @@ -59,6 +59,53 @@ def test_get_api_key_with_custom_litellm_key_header( ) == (api_key, passed_in_key) +def test_team_metadata_with_tags_flows_through_jwt_auth(): + """ + Test that team_metadata (specifically tags) flows through JWT authentication. + + This is a regression test for the issue where JWT auth was not populating + team_metadata, causing team-level tags to be missing in litellm_pre_call_utils.py + """ + from litellm.proxy._types import LiteLLM_TeamTable, UserAPIKeyAuth + + # Create a team object with metadata containing tags + team_object = LiteLLM_TeamTable( + team_id="test-team-id", + team_alias="test-team-alias", + metadata={"tags": ["production", "high-priority"], "department": "engineering"}, + tpm_limit=1000, + rpm_limit=100, + models=["gpt-4", "gpt-3.5-turbo"], + ) + + # Simulate constructing UserAPIKeyAuth like we do in JWT auth + # This is the pattern from user_api_key_auth.py lines 552-587 + user_api_key_auth = UserAPIKeyAuth( + api_key=None, + team_id=team_object.team_id, + team_tpm_limit=team_object.tpm_limit if team_object is not None else None, + team_rpm_limit=team_object.rpm_limit if team_object is not None else None, + team_models=team_object.models if team_object is not None else [], + team_metadata=team_object.metadata if team_object is not None else None, + user_role="internal_user", + user_id="test-user", + ) + + # Verify team_metadata is set + assert user_api_key_auth.team_metadata is not None, "team_metadata should be populated" + assert user_api_key_auth.team_metadata == team_object.metadata, ( + f"team_metadata not correctly mapped. " + f"Expected: {team_object.metadata}, Got: {user_api_key_auth.team_metadata}" + ) + + # Specifically verify tags are present + assert "tags" in user_api_key_auth.team_metadata, "tags should be in team_metadata" + assert user_api_key_auth.team_metadata["tags"] == ["production", "high-priority"], ( + f"tags not correctly mapped. " + f"Expected: ['production', 'high-priority'], Got: {user_api_key_auth.team_metadata.get('tags')}" + ) + + def test_route_checks_is_llm_api_route(): """Test RouteChecks.is_llm_api_route() correctly identifies LLM API routes including passthrough endpoints"""