diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3f8b01cc865..c5676fbb32f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11581,9 +11581,12 @@ async def _get_caller_byok_team_scope( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ): return None + key_team_scope: set[str] = ( + {user_api_key_dict.team_id} if user_api_key_dict.team_id else set() + ) user_id = user_api_key_dict.user_id if user_id is None: - return set() + return key_team_scope try: user_row = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_id} @@ -11591,12 +11594,12 @@ async def _get_caller_byok_team_scope( except Exception: verbose_proxy_logger.exception( "Failed to look up caller teams while scoping BYOK search; " - "defaulting to no team access." + "defaulting to key team scope only." ) - return set() + return key_team_scope if user_row is None: - return set() - return set(user_row.teams or []) + return key_team_scope + return key_team_scope | set(user_row.teams or []) def _byok_row_outside_caller_teams( diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 017f4bd4368..3bf14c08d14 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -9,7 +9,7 @@ Pins (PR2): from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock import pytest @@ -128,6 +128,104 @@ def test_v1_model_info_no_model_list_error(client, auth_as, null_router, path): assert "LLM Model List not loaded" in response.text +# --------------------------------------------------------------------------- +# GET /model/info — team BYOK scoping (issue #30983) +# --------------------------------------------------------------------------- + +_BYOK_TEAM_ID = "team-abc" +_BYOK_PUBLIC_NAME = "my-byok-gpt-4" +_BYOK_INTERNAL_NAME = f"model_name_{_BYOK_TEAM_ID}_0123456789abcdef" + + +@pytest.fixture +def byok_team_router(monkeypatch): + """Router holding one team-scoped BYOK deployment for team `team-abc`. + + Mirrors how a team's own-key BYOK model lives in the router: the routing + key is an internal mangled name while the public name lives in + `model_info.team_public_model_name`. + """ + byok_deployment = { + "model_name": _BYOK_INTERNAL_NAME, + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": { + "id": "byok-deployment-id", + "db_model": True, + "team_id": _BYOK_TEAM_ID, + "team_public_model_name": _BYOK_PUBLIC_NAME, + }, + } + + router = MagicMock() + router.model_list = [byok_deployment] + router.get_model_list_from_model_alias = MagicMock(return_value=[]) + router.get_model_names = MagicMock(return_value=[]) + router.get_model_access_groups = MagicMock(return_value={}) + router.get_model_ids = MagicMock(return_value=[]) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", [byok_deployment]) + monkeypatch.setattr(proxy_server, "user_model", None) + yield router + + +@pytest.mark.parametrize("path", ["/v1/model/info", "/model/info"]) +def test_model_info_team_key_sees_own_byok_model(client, auth_as, byok_team_router, mock_prisma, monkeypatch, path): + """Regression for #30983: a team key (user_id=None) must see its own + team's BYOK model under the public name. + + Before the fix `_get_caller_byok_team_scope` keyed only off the bound + user's team memberships, returned an empty set for a team key, and the + BYOK row was dropped -> `{"data": []}`. + """ + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + mock_prisma.db.litellm_usertable.find_unique.return_value = None + + with auth_as( + role=LitellmUserRoles.INTERNAL_USER, + user_id=None, + team_id=_BYOK_TEAM_ID, + team_models=[_BYOK_PUBLIC_NAME], + ): + response = client.get(path) + + assert response.status_code == 200 + data = response.json()["data"] + surfaced_names = [m.get("model_name") for m in data] + assert _BYOK_PUBLIC_NAME in surfaced_names + assert _BYOK_INTERNAL_NAME not in surfaced_names + + +@pytest.mark.parametrize("path", ["/v1/model/info", "/model/info"]) +def test_model_info_team_key_cannot_see_other_teams_byok_model( + client, auth_as, byok_team_router, mock_prisma, monkeypatch, path +): + """A team key for a different team must NOT see team-abc's BYOK row. + + Guards the fix from over-broadening into a cross-team metadata leak. + """ + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + mock_prisma.db.litellm_usertable.find_unique.return_value = None + + with auth_as( + role=LitellmUserRoles.INTERNAL_USER, + user_id=None, + team_id="other-team", + team_models=[_BYOK_PUBLIC_NAME], + ): + response = client.get(path) + + assert response.status_code == 200 + data = response.json()["data"] + surfaced_names = [m.get("model_name") for m in data] + assert _BYOK_PUBLIC_NAME not in surfaced_names + assert _BYOK_INTERNAL_NAME not in surfaced_names + + # --------------------------------------------------------------------------- # GET /model_group/info # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 40d590132aa..34405f20727 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -316,8 +316,10 @@ async def test_model_info_v1_unrestricted_key_hides_other_team_byok(monkeypatch) @pytest.mark.asyncio async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): - """A key without a resolvable user (e.g. CI/service token) sees only - global deployments, never any team-scoped BYOK rows.""" + """A key with no resolvable user and no team (e.g. a CI/service token + created outside any team) sees only global deployments, never team-scoped + BYOK rows. A team-scoped key does see its own team's rows (issue #30983), + pinned by the /model/info route tests.""" team_row = _team_row() other_team_row = _other_team_row() global_row = { @@ -343,7 +345,7 @@ async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): caller = UserAPIKeyAuth( user_id=None, user_role=LitellmUserRoles.INTERNAL_USER, - team_id="team-abc-123", + team_id=None, models=[], team_models=[], ) @@ -352,6 +354,106 @@ async def test_model_info_v1_service_key_hides_all_team_byok(monkeypatch): assert [m["model_info"]["id"] for m in resp["data"]] == ["global-id-1"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "find_unique", + [ + AsyncMock(return_value=MagicMock(teams=[])), + AsyncMock(return_value=None), + AsyncMock(side_effect=RuntimeError("db down")), + ], + ids=["user-not-in-team", "user-row-missing", "user-lookup-error"], +) +async def test_model_info_v1_team_key_sees_own_byok_regardless_of_user_lookup( + monkeypatch, find_unique +): + """A team-scoped key sees its own team's BYOK rows even when the bound user + is not a member of that team, has no DB row, or the lookup errors; the + key's team_id is authoritative (issue #30983). Other teams' rows stay + hidden.""" + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [_team_row(), _other_team_row(), global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = find_unique + + async def _populate(**kwargs): + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _populate) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team-abc-123", + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_info"]["id"] for m in resp["data"]] == ["byok-id-1", "global-id-1"] + + +@pytest.mark.asyncio +async def test_model_info_v1_user_team_membership_grants_byok(monkeypatch): + """A user's own team memberships still grant that team's BYOK rows, unioned + with any team the key itself is scoped to.""" + global_row = { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "global-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [_team_row(), _other_team_row(), global_row] + router.get_model_names.return_value = ["gpt-4"] + router.get_model_access_groups.return_value = {} + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=MagicMock(teams=["team-other"]) + ) + + async def _populate(**kwargs): + return kwargs["all_models"] + + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "llm_model_list", router.model_list) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "prisma_client", prisma_client) + monkeypatch.setattr(ps, "_populate_team_access_on_models", _populate) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + + caller = UserAPIKeyAuth( + user_id="user-2", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=None, + models=[], + team_models=[], + ) + resp = await ps.model_info_v1(user_api_key_dict=caller, litellm_model_id=None) + + assert [m["model_info"]["id"] for m in resp["data"]] == [ + "byok-id-other", + "global-id-1", + ] + + @pytest.mark.asyncio async def test_model_info_v1_populates_access_via_team_ids(monkeypatch): """`/v1/model/info` must populate access_via_team_ids when the DB is connected.""" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8b10539b188..c85c5ccb39f 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1348,6 +1348,7 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): non_admin = MagicMock(spec=UserAPIKeyAuth) non_admin.user_role = LitellmUserRoles.INTERNAL_USER non_admin.user_id = "user-mine" + non_admin.team_id = None filtered, total_count = await _apply_search_filter_to_models( all_models=[caller_team_byok, other_team_byok, public_model], @@ -1381,6 +1382,7 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): admin = MagicMock(spec=UserAPIKeyAuth) admin.user_role = LitellmUserRoles.PROXY_ADMIN admin.user_id = "admin-1" + admin.team_id = None filtered_admin, _ = await _apply_search_filter_to_models( all_models=[caller_team_byok, other_team_byok, public_model],