From e8484fe1506c867199b23fac0f5c5c7aec985f62 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:43:35 +0000 Subject: [PATCH] fix(rag): keep config file team deployments first and try JWT teams in sorted order A team's config file OpenAI deployments now win in file order, as on the base, before its DB ones oldest first, so upgrading never moves a team that has both. A JWT in several teams tries them in sorted order, so every worker resolves the same team. --- litellm/proxy/auth/handle_jwt.py | 2 +- .../openai_files_endpoints/common_utils.py | 34 ++++++++++----- tests/unit/proxy/auth/test_handle_jwt.py | 43 +++++++++++++++++++ .../test_files_common_utils.py | 25 ++++++++++- 4 files changed, 90 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index fec09d38026..18b20d24bd7 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -1705,7 +1705,7 @@ class JWTAuthManager: return None, None any_claim_team_resolved = False - for team_id in team_ids: + for team_id in sorted(team_ids): try: team_object = await get_team_object( team_id=team_id, diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 85688196cbc..5d19f5a04b5 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -486,13 +486,21 @@ def _key_allowlist_covers(public_model_name: str, key_model_allowlist: Sequence[ ) +_CONFIG_FILE_RANK: Final = datetime.min.replace(tzinfo=timezone.utc) _UNKNOWN_CREATION_TIME: Final = datetime.max.replace(tzinfo=timezone.utc) -def _creation_time(created_at: object) -> datetime: +def _team_deployment_rank(model_info: Mapping[str, object], deployment_id: str) -> tuple[datetime, str]: + if model_info.get("db_model") is not True: + return _CONFIG_FILE_RANK, "" + created_at: Final = model_info.get("created_at") if not isinstance(created_at, datetime): - return _UNKNOWN_CREATION_TIME - return created_at if created_at.tzinfo else created_at.replace(tzinfo=timezone.utc) + return _UNKNOWN_CREATION_TIME, deployment_id + return (created_at if created_at.tzinfo else created_at.replace(tzinfo=timezone.utc)), deployment_id + + +def _by_team_deployment_rank(order: tuple[tuple[datetime, str], str]) -> tuple[datetime, str]: + return order[0] def get_team_provider_credentials( @@ -569,7 +577,7 @@ def get_team_provider_credentials( return {key: value for key, value in credentials.items() if key != "model"} return None - def _usable_team_deployment_order(deployment: dict) -> tuple[datetime, str] | None: + def _usable_team_deployment_order(deployment: dict) -> tuple[tuple[datetime, str], str] | None: model_info: Final[Mapping[str, object]] = deployment["model_info"] if model_info.get("team_id") != team_id: return None @@ -578,16 +586,20 @@ def get_team_provider_credentials( deployment_id: Final = model_info.get("id") if not isinstance(deployment_id, str): return None - return _creation_time(model_info.get("created_at")), deployment_id + return _team_deployment_rank(model_info, deployment_id), deployment_id - # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. The oldest one - # wins, then the lowest id, so adding a deployment never moves the team off the one it - # already uses, and every worker picks the same one in any load or edit order. + # 1. Prefer the team's own BYOK deployment, matched by model_info.team_id. Config file + # deployments win, kept in file order by the stable sort, then DB ones oldest first, then + # by lowest id, so adding a deployment never moves the team off the one it already uses, + # and every worker picks the same one in any DB load or edit order. if team_id is not None: team_deployment_orders: Final = sorted( - order - for deployment in llm_router.model_list or [] - if (order := _usable_team_deployment_order(deployment)) is not None + ( + order + for deployment in llm_router.model_list or [] + if (order := _usable_team_deployment_order(deployment)) is not None + ), + key=_by_team_deployment_rank, ) team_deployment_ids: Final = tuple(deployment_id for _, deployment_id in team_deployment_orders) team_credentials: Final = next( diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index e51e1dead73..2b9013a4efa 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -295,6 +295,49 @@ async def test_find_team_with_model_access_uses_request_method_for_passthrough_a assert "allowed_passthrough_routes" in exc_info.value.detail + +class _TeamClaims(set[str]): + def __init__(self, claim_order: tuple[str, ...]): + super().__init__(claim_order) + self.claim_order = claim_order + + def __iter__(self): + return iter(self.claim_order) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("claim_order", [("team-a", "team-b"), ("team-b", "team-a")]) +async def test_a_jwt_in_two_teams_resolves_the_same_team_in_any_claim_order(claim_order: tuple[str, ...]): + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + teams = {team_id: LiteLLM_TeamTable(team_id=team_id, models=["gpt-4"], metadata={}) for team_id in claim_order} + + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + side_effect=lambda team_id, **_: teams[team_id], + ), + patch( + "litellm.proxy.auth.handle_jwt.can_team_access_model", + new_callable=AsyncMock, + return_value=True, + ), + ): + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids=_TeamClaims(claim_order), + requested_model="gpt-4", + route="/v1/rag/query", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + ) + + assert team_id == "team-a" + assert team_obj == teams["team-a"] + _AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = { "test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": { "endpoint_id": "test-uuid-1", diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index 243d960cad4..ff3d39b6c8d 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -715,11 +715,13 @@ def test_a_restricted_key_outside_the_team_wildcard_never_resolves_its_credentia assert _team_only_openai_credentials("openai/*", "openai/*", key_models) is None -def _team_openai_deployment(deployment_id: str, api_key: str, created_at: datetime | None = None) -> dict[str, object]: +def _team_openai_deployment( + deployment_id: str, api_key: str, created_at: datetime | None = None, db_model: bool = True +) -> dict[str, object]: return { "model_name": f"team-gpt-{deployment_id}", "litellm_params": {"model": "openai/gpt-4o", "api_key": api_key}, - "model_info": {"id": deployment_id, "team_id": "team-a"} + "model_info": {"id": deployment_id, "team_id": "team-a", "db_model": db_model} | ({} if created_at is None else {"created_at": created_at}), } @@ -758,3 +760,22 @@ def test_a_team_keeps_its_older_openai_deployment_when_a_newer_one_has_a_lower_i assert credentials is not None assert credentials["api_key"] == "older-key" + + +def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones(): + db_deployment: Final = _team_openai_deployment("deployment-a", "db-key", datetime(2020, 1, 1, tzinfo=timezone.utc)) + config_deployments: Final = tuple( + _team_openai_deployment(deployment_id, f"{deployment_id}-key", db_model=False) + for deployment_id in ("deployment-z", "deployment-b") + ) + router: Final = Router(model_list=[db_deployment, *config_deployments]) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="team-a"), + custom_llm_provider="openai", + include_accessible_deployments=False, + ) + + assert credentials is not None + assert credentials["api_key"] == "deployment-z-key"