From 478de14bb519c2a9d31b5be2dad1963d093672fe Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 23 Sep 2026 18:23:31 -0700 Subject: [PATCH] fix(jwt): scope auto_register_map_existing_key reuse to the JWT-resolved team Only reuse a key whose team_id matches the team auth_builder resolved for the JWT (no team matches no team), so a personal key can no longer bypass the resolved team's model and budget limits. With the flag on, the first JWT request now falls through to the same virtual-key checks later mapped requests get, instead of returning early, so a reused key's own limits apply from request one rather than 200 then 403. Flag off keeps the early return unchanged. --- litellm/proxy/_types.py | 7 +- litellm/proxy/auth/user_api_key_auth.py | 138 +++++++++--------- .../verification_token_repository.py | 5 +- .../proxy/auth/test_user_api_key_auth.py | 131 +++++++++++++++++ 4 files changed, 211 insertions(+), 70 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 29bdc350446..8ad7b99c26b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5303,9 +5303,10 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): description=( "Only used with unregistered_jwt_client_behavior='auto_register'. When True, the JWT claim is " "mapped to a virtual key the resolved internal user already owns instead of minting a new one. " - "If the user owns several, the most recently created key that can call LLM routes is chosen: " - "not blocked, not expired, not an Admin UI session key, and with no route restriction other than " - "llm_api_routes. A new key is only minted when the user has no such key." + "If the user owns several, the most recently created key in the JWT-resolved team (or with no " + "team when the JWT resolves none) that can call LLM routes is chosen: not blocked, not expired, " + "not an Admin UI session key, and with no route restriction other than llm_api_routes. A new key " + "is only minted when the user has no such key." ), ) routing_overrides: list[JWTRoutingOverride] | None = Field( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f537f048d1e..7813578d385 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -930,8 +930,10 @@ class _PendingAutoRegister(NamedTuple): jwt_issuer: str | None = None -async def _latest_active_key_hash_for_user(prisma_client: PrismaClient, user_id: str) -> str | None: - row: Final = await VerificationTokenRepository(prisma_client).find_latest_llm_api_row_by_user_id(user_id) +async def _latest_active_key_hash_for_user( + prisma_client: PrismaClient, user_id: str, team_id: str | None +) -> str | None: + row: Final = await VerificationTokenRepository(prisma_client).find_latest_llm_api_row_by_user_id(user_id, team_id) return None if row is None else row.token @@ -973,7 +975,7 @@ async def _auto_register_jwt_mapping( ) existing_token_hash: Final = ( - await _latest_active_key_hash_for_user(prisma_client, user_id) + await _latest_active_key_hash_for_user(prisma_client, user_id, team_id) if jwt_handler.litellm_jwtauth.auto_register_map_existing_key and user_id is not None else None ) @@ -1772,8 +1774,8 @@ async def _user_api_key_auth_builder( # mapping + virtual key from the *validated* identity, then # replace valid_token with the new key so downstream checks # use the key-scoped path. - if pending_auto_register is not None and prisma_client is not None: - auto_registered: Final = await _auto_register_jwt_mapping( + auto_registered: Final = ( + await _auto_register_jwt_mapping( virtual_key_claim_field=pending_auto_register.claim_field, claim_value=pending_auto_register.claim_value, jwt_handler=jwt_handler, @@ -1789,72 +1791,76 @@ async def _user_api_key_auth_builder( end_user_id=end_user_id, agent_id=agent_id, ) - if auto_registered is not None: - auto_registered.jwt_claims = jwt_claims - auto_registered.user_email = user_email - # The auto-registered token is built from the new key's - # columns, which carry no user budget. Carry over the - # already-loaded user row rather than re-reading it, or - # the budget check below has nothing to enforce. - auto_registered.user_model_max_budget = ( - user_object.model_max_budget if user_object is not None else None - ) - valid_token = auto_registered - api_key = valid_token.token or "" - - # Check if model has zero cost - if so, skip all budget checks - model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, + if pending_auto_register is not None and prisma_client is not None + else None ) - skip_budget_checks = False - if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero - - skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) - if skip_budget_checks: - verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) - - # Fetch project object for JWT path if project_id is set - _jwt_project_obj = None - if valid_token.project_id is not None: - _jwt_project_obj = await get_project_object( - project_id=valid_token.project_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + if auto_registered is not None: + auto_registered.jwt_claims = jwt_claims + auto_registered.user_email = user_email + # The auto-registered token is built from the new key's + # columns, which carry no user budget. Carry over the + # already-loaded user row rather than re-reading it, or + # the budget check below has nothing to enforce. + auto_registered.user_model_max_budget = ( + user_object.model_max_budget if user_object is not None else None ) - if _jwt_project_obj is not None: - valid_token.project_metadata = _jwt_project_obj.metadata - valid_token.project_alias = _jwt_project_obj.project_alias + valid_token = auto_registered + api_key = valid_token.token or "" - # JWT auth returns here rather than falling through to the - # virtual-key checks below, so the user's per-model budget - # has to be enforced on this path too. Without it the - # post-call increment still charges the counter and nothing - # ever reads it, which is worse than not tracking at all. - # Guarded by the same flag the virtual-key path uses, or a - # zero-cost model would be refused here and allowed there, - # while the log above claims all budget checks were skipped. - if not skip_budget_checks: - await _check_user_model_budget( - valid_token=cast(UserAPIKeyAuth, valid_token), - model_max_budget_limiter=model_max_budget_limiter, - models=_get_model_names_for_budget_checks( - model=_get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - ), + if auto_registered is None or not jwt_handler.litellm_jwtauth.auto_register_map_existing_key: + # Check if model has zero cost - if so, skip all budget checks + model = _get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=valid_token.team_id, ) + skip_budget_checks = False + if model is not None and llm_router is not None: + from litellm.proxy.auth.auth_checks import _is_model_cost_zero - return cast(UserAPIKeyAuth, valid_token) + skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router) + if skip_budget_checks: + verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model) + + # Fetch project object for JWT path if project_id is set + _jwt_project_obj = None + if valid_token.project_id is not None: + _jwt_project_obj = await get_project_object( + project_id=valid_token.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if _jwt_project_obj is not None: + valid_token.project_metadata = _jwt_project_obj.metadata + valid_token.project_alias = _jwt_project_obj.project_alias + + # JWT auth returns here rather than falling through to the + # virtual-key checks below, so the user's per-model budget + # has to be enforced on this path too. Without it the + # post-call increment still charges the counter and nothing + # ever reads it, which is worse than not tracking at all. + # Guarded by the same flag the virtual-key path uses, or a + # zero-cost model would be refused here and allowed there, + # while the log above claims all budget checks were skipped. + if not skip_budget_checks: + await _check_user_model_budget( + valid_token=cast(UserAPIKeyAuth, valid_token), + model_max_budget_limiter=model_max_budget_limiter, + models=_get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + ), + ) + + return cast(UserAPIKeyAuth, valid_token) #### ELSE #### ## CHECK PASS-THROUGH ENDPOINTS ## diff --git a/litellm/repositories/verification_token_repository.py b/litellm/repositories/verification_token_repository.py index a01737cfd1e..914dcba90f9 100644 --- a/litellm/repositories/verification_token_repository.py +++ b/litellm/repositories/verification_token_repository.py @@ -124,10 +124,13 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]): records: Final[Sequence[PrismaVerificationToken]] = await self.table.find_many(where={"user_id": user_id}) return self._to_model_list(records) - async def find_latest_llm_api_row_by_user_id(self, user_id: str) -> "PrismaVerificationToken | None": + async def find_latest_llm_api_row_by_user_id( + self, user_id: str, team_id: str | None + ) -> "PrismaVerificationToken | None": row: Final = await self.table.find_first( where={ # mutable-ok: the prisma where clause contract is a plain dict "user_id": user_id, + "team_id": team_id, "AND": [ # mutable-ok: prisma filter literal {"OR": [{"blocked": False}, {"blocked": None}]}, # mutable-ok: prisma filter literal { # mutable-ok: prisma filter literal 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 a012e5b20b4..a3b2badb5a5 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 @@ -2236,6 +2236,39 @@ async def test_auto_register_map_existing_key_skips_keys_that_cannot_call_llm_ro ], f"only unrestricted or llm_api_routes keys may be reused: {where}" +@pytest.mark.asyncio +@pytest.mark.parametrize("resolved_team_id", ["validated-team", None]) +async def test_auto_register_map_existing_key_only_reuses_keys_in_the_jwt_resolved_team( + resolved_team_id: str | None, +) -> None: + from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping + + prisma_client = MagicMock() + prisma_client.db.litellm_verificationtoken.find_first = AsyncMock( + return_value=SimpleNamespace(token="existing-hash") + ) + prisma_client.db.litellm_jwtkeymapping.create = AsyncMock() + + user_api_key_cache = MagicMock() + user_api_key_cache.async_set_cache = AsyncMock() + + jwt_handler = MagicMock() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + auto_register_map_existing_key=True, + virtual_key_mapping_cache_ttl=300, + ) + + generate_patch, resolve_patch = _auto_register_patches(plaintext_key=None) + with generate_patch, resolve_patch: + await _auto_register_jwt_mapping( + **_auto_register_kwargs(prisma_client, user_api_key_cache, jwt_handler, team_id=resolved_team_id) + ) + + where = prisma_client.db.litellm_verificationtoken.find_first.await_args.kwargs["where"] + assert "team_id" in where, f"reuse must be scoped to the JWT-resolved team: {where}" + assert where["team_id"] == resolved_team_id + + @pytest.mark.asyncio async def test_auto_register_map_existing_key_mints_when_user_has_no_key(): """The flag must not leave a keyless user unmapped: with no existing key it @@ -2364,6 +2397,104 @@ async def test_auto_register_map_existing_key_user_id_none_mints(): generate_key.assert_awaited_once() +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reused_key_models", "expect_denied"), + [(["some-other-model"], True), ([], False)], +) +async def test_auto_register_map_existing_key_first_request_runs_key_checks( + reused_key_models: list[str], expect_denied: bool +) -> None: + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature" + user_api_key_cache = DualCache() + prisma_client = MagicMock() + jwt_handler = MagicMock() + jwt_handler.is_jwt.return_value = True + jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1"}) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + virtual_key_claim_field="sub", + virtual_key_mapping_cache_ttl=300, + auto_register_map_existing_key=True, + ) + reused_key = UserAPIKeyAuth( + token="hashed-existing-key", + api_key="hashed-existing-key", + user_id="validated-user", + models=reused_key_models, + ) + mock_jwt_result = { + "is_proxy_admin": False, + "team_object": None, + "user_object": LiteLLM_UserTable(user_id="validated-user", user_role="internal_user"), + "end_user_object": None, + "org_object": None, + "token": jwt_token, + "team_id": None, + "user_id": "validated-user", + "user_email": None, + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": "user1"}, + } + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + with ( + patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(post_call_failure_hook=AsyncMock(return_value=None)), + ), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch( + "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", + new_callable=AsyncMock, + return_value=_PendingAutoRegister( + claim_field="sub", + claim_value="user1", + cache_key="jwt_key_mapping:sub:user1", + ), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=mock_jwt_result, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping", + new_callable=AsyncMock, + return_value=reused_key, + ), + ): + call = _user_api_key_auth_builder( + request=mock_request, + api_key=jwt_token, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "gpt-4o-mini"}, + ) + if expect_denied: + with pytest.raises(ProxyException, match="not available for this API key"): + await call + return + result = await call + + assert result.api_key == "hashed-existing-key" + assert result.user_id == "validated-user" + + @pytest.mark.asyncio @pytest.mark.parametrize("active", [True, False]) async def test_auto_register_first_request_propagates_user_email(active: bool) -> None: